diff --git a/.github/workflows/CodeQuality.yml b/.github/workflows/CodeQuality.yml new file mode 100644 index 00000000..d2de0130 --- /dev/null +++ b/.github/workflows/CodeQuality.yml @@ -0,0 +1,52 @@ +name: Code Quality + +on: + workflow_call: + +jobs: + check: + name: ${{ matrix.name }} + runs-on: ${{ matrix.runner }} + strategy: + fail-fast: false + matrix: + include: + - name: Format Check + runner: ubuntu-22.04 + pip_packages: cmake-format 'black==24.*' cxxheaderparser pcpp 'clang_format==11.0.1' + pip_flags: '' + tidy_binary: '' + target: format-check + - name: Tidy Check + runner: ubuntu-24.04 + pip_packages: '"pybind11[global]"' + pip_flags: --break-system-packages + tidy_binary: clang-tidy-18 + target: tidy-check + env: + CC: gcc + CXX: g++ + GEN: ninja + TIDY_THREADS: 4 + TIDY_BINARY: ${{ matrix.tidy_binary }} + steps: + - name: Checkout + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + fetch-depth: 0 + submodules: recursive + + - name: Install Python check tools + shell: bash + run: | + # Ninja and clang-tidy-18 are already on the hosted Ubuntu image. + # Avoid redundant apt installs that fail on the image's package mirrors. + sudo pip3 install ${{ matrix.pip_packages }} ${{ matrix.pip_flags }} + + - name: Run check + run: | + if [ '${{ matrix.target }}' = 'tidy-check' ]; then + ninja --version + "$TIDY_BINARY" --version + fi + make ${{ matrix.target }} diff --git a/.github/workflows/MainDistributionPipeline.yml b/.github/workflows/MainDistributionPipeline.yml index ae93e1c6..dd5710b6 100644 --- a/.github/workflows/MainDistributionPipeline.yml +++ b/.github/workflows/MainDistributionPipeline.yml @@ -14,17 +14,12 @@ concurrency: jobs: duckdb-stable-build: name: Build extension binaries - uses: duckdb/extension-ci-tools/.github/workflows/_extension_distribution.yml@v1.5.4 + uses: duckdb/extension-ci-tools/.github/workflows/_extension_distribution.yml@ebdcc1961780b0ce1272a0d7637dc58447c34885 with: - duckdb_version: v1.5.4 - ci_tools_version: v1.5.4 + duckdb_version: v2.0.0 + ci_tools_version: ebdcc1961780b0ce1272a0d7637dc58447c34885 extension_name: privacy code-quality-check: name: Code Quality Check - uses: duckdb/extension-ci-tools/.github/workflows/_extension_code_quality.yml@v1.5.4 - with: - duckdb_version: v1.5.4 - ci_tools_version: v1.5.4 - extension_name: privacy - format_checks: 'format;tidy' + uses: ./.github/workflows/CodeQuality.yml diff --git a/CMakeLists.txt b/CMakeLists.txt index 9c3325ad..6862aad9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -3,10 +3,11 @@ cmake_minimum_required(VERSION 3.5) # Set extension name here set(TARGET_NAME privacy) -# DuckDB's extension distribution supports vcpkg. As such, dependencies can be added in ./vcpkg.json and then -# used in cmake with find_package. Feel free to remove or replace with other dependencies. -# Note that it should also be removed from vcpkg.json to prevent needlessly installing it.. -#find_package(OpensSL REQUIRED) +# DuckDB's extension distribution supports vcpkg. As such, dependencies can be +# added in ./vcpkg.json and then used in cmake with find_package. Feel free to +# remove or replace with other dependencies. Note that it should also be removed +# from vcpkg.json to prevent needlessly installing it.. find_package(OpensSL +# REQUIRED) set(EXTENSION_NAME ${TARGET_NAME}_extension) set(LOADABLE_EXTENSION_NAME ${TARGET_NAME}_loadable_extension) @@ -15,15 +16,16 @@ project(${TARGET_NAME}) include_directories(src/include) include_directories(benchmark/include) -# Option to disable approximate sum algorithm (use exact cascading instead) -# By default, pac_sum/pac_avg use the approx algorithm (16-bit cascading counters) +# Option to disable approximate sum algorithm (use exact cascading instead) By +# default, pac_sum/pac_avg use the approx algorithm (16-bit cascading counters) # Set PAC_NOAPPROX=ON to use exact cascading instead -option(PAC_NOAPPROX "Disable approximate sum algorithm, use exact cascading" OFF) +option(PAC_NOAPPROX "Disable approximate sum algorithm, use exact cascading" + OFF) if(PAC_NOAPPROX) - add_definitions(-DPAC_EXACTSUM) - message(STATUS "privacy: Using exact cascading sum algorithm (PAC_NOAPPROX)") + add_definitions(-DPAC_EXACTSUM) + message(STATUS "privacy: Using exact cascading sum algorithm (PAC_NOAPPROX)") else() - message(STATUS "privacy: Using approximate sum algorithm (default)") + message(STATUS "privacy: Using approximate sum algorithm (default)") endif() # Extension sources @@ -37,6 +39,7 @@ set(EXTENSION_SOURCES src/aggregates/as_count.cpp src/aggregates/as_min_max.cpp src/aggregates/as_sum.cpp + src/aggregates/as_sum_state_export.cpp src/aggregates/as_finalize.cpp src/aggregates/as_clip_sum.cpp src/aggregates/as_clip_min_max.cpp @@ -48,6 +51,7 @@ set(EXTENSION_SOURCES src/compiler/privacy_mechanisms.cpp src/compiler/privacy_compiler_helpers.cpp src/query_processing/pac_avg_rewriter.cpp + src/query_processing/as_correlation_optimizer.cpp src/query_processing/privacy_expression_builder.cpp src/query_processing/privacy_plan_traversal.cpp src/query_processing/pac_projection_propagation.cpp @@ -69,8 +73,8 @@ build_static_extension(${TARGET_NAME} ${EXTENSION_SOURCES}) build_loadable_extension(${TARGET_NAME} " " ${EXTENSION_SOURCES}) if(WIN32) - target_link_libraries(${EXTENSION_NAME} bcrypt) - target_link_libraries(${LOADABLE_EXTENSION_NAME} bcrypt) + target_link_libraries(${EXTENSION_NAME} bcrypt) + target_link_libraries(${LOADABLE_EXTENSION_NAME} bcrypt) endif() install( @@ -79,76 +83,73 @@ install( LIBRARY DESTINATION "${INSTALL_LIB_DIR}" ARCHIVE DESTINATION "${INSTALL_LIB_DIR}") -# Test runners and benchmarks - only build for native platforms, not WASM or Windows -# (Windows lacks unistd.h and other POSIX features used by the benchmark code) -set(PRIVACY_LINK_LIBS - duckdb_static - duckdb_generated_extension_loader - ${EXTENSION_NAME} - core_functions_extension - parquet_extension - duckdb_static) - -# jemalloc is only built on Linux 64-bit; link it when the extension was loaded -# (can't use if(TARGET ...) here because the target isn't defined yet at this point) -if(DUCKDB_EXTENSION_JEMALLOC_SHOULD_LINK) - list(APPEND PRIVACY_LINK_LIBS jemalloc_extension) -endif() +# Test runners and benchmarks - only build for native platforms, not WASM or +# Windows (Windows lacks unistd.h and other POSIX features used by the benchmark +# code) if(NOT EMSCRIPTEN AND NOT WIN32) - # Build a single unified test executable that invokes all tests via src/test/test_runner.cpp - add_executable(pac_test_runner - src/test/test_runner.cpp - src/test/test_plan_traversal.cpp - src/test/test_schema_metadata.cpp - src/test/test_compiler_functions.cpp - src/test/test_privacy_parser.cpp) - target_link_libraries(pac_test_runner ${PRIVACY_LINK_LIBS}) - - # TPCH + AS benchmark runner - add_executable(as_tpch_benchmark benchmark/tpch/as_tpch_benchmark.cpp) - target_link_libraries(as_tpch_benchmark ${PRIVACY_LINK_LIBS}) - - # WIP rebuttal experiment: four-way MEPS balanced repeated replication benchmark - add_executable(meps_brr_poc benchmark/meps_brr/meps_brr_poc.cpp) - target_link_libraries(meps_brr_poc ${PRIVACY_LINK_LIBS}) - - # WIP rebuttal experiment: standard joined SSB queries and 64 sampled worlds - add_executable(as_ssb_benchmark benchmark/ssb/as_ssb_benchmark.cpp) - target_link_libraries(as_ssb_benchmark ${PRIVACY_LINK_LIBS}) - - # ClickBench + AS benchmark runner - add_executable(as_clickbench_benchmark benchmark/clickbench/as_clickbench_benchmark.cpp) - target_link_libraries(as_clickbench_benchmark ${PRIVACY_LINK_LIBS}) - - # TPC-DS automatic-rewriter coverage and end-to-end performance - add_executable(as_tpcds_benchmark benchmark/tpcds/as_tpcds_benchmark.cpp) - target_link_libraries(as_tpcds_benchmark ${PRIVACY_LINK_LIBS}) - - # Unified DP utility benchmark runner - set(TPCH_STOCK_DP_QUERY_NAMES q01 q05 q06 q14 q19) - set(TPCH_STOCK_DP_QUERY_FILES) - foreach(QUERY_NAME IN LISTS TPCH_STOCK_DP_QUERY_NAMES) - string(TOUPPER ${QUERY_NAME} QUERY_NAME_UPPER) - set(QUERY_FILE "${CMAKE_CURRENT_SOURCE_DIR}/benchmark/dp/tpch_queries/${QUERY_NAME}.sql") - list(APPEND TPCH_STOCK_DP_QUERY_FILES "${QUERY_FILE}") - file(READ - "${QUERY_FILE}" - "TPCH_${QUERY_NAME_UPPER}_SQL") - endforeach() - set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS - ${TPCH_STOCK_DP_QUERY_FILES}) - set(TPCH_QUERY_GENERATED_DIR "${CMAKE_CURRENT_BINARY_DIR}/generated") - file(MAKE_DIRECTORY "${TPCH_QUERY_GENERATED_DIR}") - configure_file( - benchmark/include/tpch_stock_queries.hpp.in - "${TPCH_QUERY_GENERATED_DIR}/tpch_stock_queries.hpp" - @ONLY) - add_executable(dp_benchmark_runner benchmark/dp/dp_benchmark_runner.cpp) - target_include_directories(dp_benchmark_runner PRIVATE "${TPCH_QUERY_GENERATED_DIR}") - target_link_libraries(dp_benchmark_runner ${PRIVACY_LINK_LIBS}) - - # IMDb/JOB benchmark runner - add_executable(imdb_benchmark benchmark/imdb/imdb_benchmark.cpp) - target_link_libraries(imdb_benchmark ${PRIVACY_LINK_LIBS}) + # Build a single unified test executable that invokes all tests via + # src/test/test_runner.cpp + add_executable( + pac_test_runner + src/test/test_runner.cpp src/test/test_plan_traversal.cpp + src/test/test_schema_metadata.cpp src/test/test_compiler_functions.cpp + src/test/test_privacy_parser.cpp) + + # TPCH + AS benchmark runner + add_executable(as_tpch_benchmark benchmark/tpch/as_tpch_benchmark.cpp) + + # WIP rebuttal experiment: four-way MEPS balanced repeated replication + # benchmark + add_executable(meps_brr_poc benchmark/meps_brr/meps_brr_poc.cpp) + + # WIP rebuttal experiment: standard joined SSB queries and 64 sampled worlds + add_executable(as_ssb_benchmark benchmark/ssb/as_ssb_benchmark.cpp) + + # ClickBench + AS benchmark runner + add_executable(as_clickbench_benchmark + benchmark/clickbench/as_clickbench_benchmark.cpp) + + # TPC-DS automatic-rewriter coverage and end-to-end performance + add_executable(as_tpcds_benchmark benchmark/tpcds/as_tpcds_benchmark.cpp) + + # Unified DP utility benchmark runner + set(TPCH_STOCK_DP_QUERY_NAMES q01 q05 q06 q14 q19) + set(TPCH_STOCK_DP_QUERY_FILES) + foreach(QUERY_NAME IN LISTS TPCH_STOCK_DP_QUERY_NAMES) + string(TOUPPER ${QUERY_NAME} QUERY_NAME_UPPER) + set(QUERY_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/benchmark/dp/tpch_queries/${QUERY_NAME}.sql" + ) + list(APPEND TPCH_STOCK_DP_QUERY_FILES "${QUERY_FILE}") + file(READ "${QUERY_FILE}" "TPCH_${QUERY_NAME_UPPER}_SQL") + endforeach() + set_property( + DIRECTORY + APPEND + PROPERTY CMAKE_CONFIGURE_DEPENDS ${TPCH_STOCK_DP_QUERY_FILES}) + set(TPCH_QUERY_GENERATED_DIR "${CMAKE_CURRENT_BINARY_DIR}/generated") + file(MAKE_DIRECTORY "${TPCH_QUERY_GENERATED_DIR}") + configure_file(benchmark/include/tpch_stock_queries.hpp.in + "${TPCH_QUERY_GENERATED_DIR}/tpch_stock_queries.hpp" @ONLY) + add_executable(dp_benchmark_runner benchmark/dp/dp_benchmark_runner.cpp) + target_include_directories(dp_benchmark_runner + PRIVATE "${TPCH_QUERY_GENERATED_DIR}") + + # IMDb/JOB benchmark runner + add_executable(imdb_benchmark benchmark/imdb/imdb_benchmark.cpp) + + foreach( + PRIVACY_RUNNER IN + ITEMS pac_test_runner + as_tpch_benchmark + meps_brr_poc + as_ssb_benchmark + as_clickbench_benchmark + as_tpcds_benchmark + dp_benchmark_runner + imdb_benchmark) + target_link_libraries(${PRIVACY_RUNNER} PRIVATE duckdb_static) + link_extension_libraries(${PRIVACY_RUNNER} PRIVATE) + endforeach() endif() diff --git a/Makefile b/Makefile index df0bff53..c9ca9741 100644 --- a/Makefile +++ b/Makefile @@ -3,6 +3,7 @@ PROJ_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST)))) # Configuration of extension EXT_NAME=privacy EXT_CONFIG=${PROJ_DIR}extension_config.cmake +OVERRIDE_GIT_DESCRIBE ?= v2.0.0 # Include the Makefile from extension-ci-tools -include extension-ci-tools/makefiles/duckdb_extension.Makefile \ No newline at end of file +include extension-ci-tools/makefiles/duckdb_extension.Makefile diff --git a/README.md b/README.md index 5a03c8b9..dcd37d8c 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,7 @@ Both modes share the same DDL and require no changes to query syntax — the extension rewrites aggregate query plans transparently. -This works on DuckDB v1.5 and beyond. See https://duckdb.org/install to install. Or if you do not want to install anything: this extension is also distributed in WASM, so you can run the examples also in https://shell.duckdb.org in a browser. +This branch targets DuckDB v2.0.0. See https://duckdb.org/install to install. Or if you do not want to install anything: this extension is also distributed in WASM, so you can run the examples also in https://shell.duckdb.org in a browser. ## Install @@ -218,6 +218,19 @@ This gives a formal (ε,δ)-DP guarantee at the **user level** — neighboring d | `dp_delta` | `1e-6` | dp_elastic | Failure probability δ (0 = pure DP) | | `dp_sum_bound` | required | dp_elastic | Clipping bound for `SUM`/`AVG` inputs: values are clipped to `[-bound, bound]` | +For PAC, a fixed `privacy_seed` repeats the random draws when the query, data, +settings, execution order, and initial p-tracking state are unchanged on the same +build. Stream allocation resets for every query, including prepared executions, +so unrelated queries do not shift later seeded draws. Streams advance across cells: +identical aggregate counters in different groups still receive independent draws. + +For reproducible tests, use `SET threads = 1` and a stable input order. Parallel +hash aggregation can emit groups in a different order even at a fixed thread count, +so bit-identical per-group results are not guaranteed across parallel executions, +row/batch reordering, or builds. P-tracking updates its calibration state after +releases; comparisons must start from the same state (the RNG regression tests use +`SET pac_ptracking = false`). + ## SQL Reference ### Supported Aggregates and Operators diff --git a/benchmark/meps_brr/ahrq2016/as_queries/q01_office.sql b/benchmark/meps_brr/ahrq2016/as_queries/q01_office.sql index 99eaf846..90fb892d 100644 --- a/benchmark/meps_brr/ahrq2016/as_queries/q01_office.sql +++ b/benchmark/meps_brr/ahrq2016/as_queries/q01_office.sql @@ -2,24 +2,24 @@ SELECT as_sum(p.mask_lo, p.weight) AS event_count_lo, as_sum(p.mask_hi, p.weight) AS event_count_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.self), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_self_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_self_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.self), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_self_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_self_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.private), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_private_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_private_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.private), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_private_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_private_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.medicare), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicare_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicare_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.medicare), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicare_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicare_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.medicaid), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicaid_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicaid_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.medicaid), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicaid_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicaid_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.other), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_other_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_other_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.other), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_other_hi + lambda x: x[1] / nullif(x[2], 0)) AS mean_other_hi FROM persons_as p JOIN events e USING (person_id) WHERE e.kind = 'OB'; diff --git a/benchmark/meps_brr/ahrq2016/as_queries/q02_office_physician.sql b/benchmark/meps_brr/ahrq2016/as_queries/q02_office_physician.sql index d46c93bb..3ed1b356 100644 --- a/benchmark/meps_brr/ahrq2016/as_queries/q02_office_physician.sql +++ b/benchmark/meps_brr/ahrq2016/as_queries/q02_office_physician.sql @@ -2,24 +2,24 @@ SELECT as_sum(p.mask_lo, p.weight) AS event_count_lo, as_sum(p.mask_hi, p.weight) AS event_count_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.self), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_self_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_self_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.self), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_self_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_self_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.private), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_private_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_private_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.private), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_private_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_private_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.medicare), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicare_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicare_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.medicare), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicare_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicare_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.medicaid), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicaid_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicaid_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.medicaid), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicaid_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicaid_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.other), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_other_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_other_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.other), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_other_hi + lambda x: x[1] / nullif(x[2], 0)) AS mean_other_hi FROM persons_as p JOIN events e USING (person_id) WHERE e.kind = 'OB' AND e.physician = 1; diff --git a/benchmark/meps_brr/ahrq2016/as_queries/q03_outpatient.sql b/benchmark/meps_brr/ahrq2016/as_queries/q03_outpatient.sql index 0ac89328..d6af9da7 100644 --- a/benchmark/meps_brr/ahrq2016/as_queries/q03_outpatient.sql +++ b/benchmark/meps_brr/ahrq2016/as_queries/q03_outpatient.sql @@ -2,24 +2,24 @@ SELECT as_sum(p.mask_lo, p.weight) AS event_count_lo, as_sum(p.mask_hi, p.weight) AS event_count_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.self), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_self_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_self_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.self), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_self_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_self_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.private), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_private_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_private_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.private), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_private_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_private_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.medicare), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicare_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicare_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.medicare), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicare_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicare_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.medicaid), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicaid_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicaid_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.medicaid), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicaid_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicaid_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.other), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_other_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_other_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.other), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_other_hi + lambda x: x[1] / nullif(x[2], 0)) AS mean_other_hi FROM persons_as p JOIN events e USING (person_id) WHERE e.kind = 'OP'; diff --git a/benchmark/meps_brr/ahrq2016/as_queries/q04_outpatient_physician.sql b/benchmark/meps_brr/ahrq2016/as_queries/q04_outpatient_physician.sql index 9c528f12..41564706 100644 --- a/benchmark/meps_brr/ahrq2016/as_queries/q04_outpatient_physician.sql +++ b/benchmark/meps_brr/ahrq2016/as_queries/q04_outpatient_physician.sql @@ -2,24 +2,24 @@ SELECT as_sum(p.mask_lo, p.weight) AS event_count_lo, as_sum(p.mask_hi, p.weight) AS event_count_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.self), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_self_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_self_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.self), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_self_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_self_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.private), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_private_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_private_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.private), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_private_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_private_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.medicare), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicare_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicare_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.medicare), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicare_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicare_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.medicaid), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicaid_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicaid_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.medicaid), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_medicaid_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_medicaid_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * e.other), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_other_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_other_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * e.other), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_other_hi + lambda x: x[1] / nullif(x[2], 0)) AS mean_other_hi FROM persons_as p JOIN events e USING (person_id) WHERE e.kind = 'OP' AND e.physician = 1; diff --git a/benchmark/meps_brr/ahrq2016/as_queries/q05_visits_per_person.sql b/benchmark/meps_brr/ahrq2016/as_queries/q05_visits_per_person.sql index c53e1d8a..e01bce2e 100644 --- a/benchmark/meps_brr/ahrq2016/as_queries/q05_visits_per_person.sql +++ b/benchmark/meps_brr/ahrq2016/as_queries/q05_visits_per_person.sql @@ -5,12 +5,12 @@ WITH visits AS ( FROM events WHERE kind = 'OB' GROUP BY person_id ) SELECT list_transform(list_zip(as_sum(p.mask_lo, p.weight * coalesce(v.n, 0)), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_office_visits_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_office_visits_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * coalesce(v.n, 0)), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_office_visits_hi, + lambda x: x[1] / nullif(x[2], 0)) AS mean_office_visits_hi, list_transform(list_zip(as_sum(p.mask_lo, p.weight * coalesce(v.n_physician, 0)), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_office_physician_visits_lo, + lambda x: x[1] / nullif(x[2], 0)) AS mean_office_physician_visits_lo, list_transform(list_zip(as_sum(p.mask_hi, p.weight * coalesce(v.n_physician, 0)), as_sum(p.mask_hi, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_office_physician_visits_hi + lambda x: x[1] / nullif(x[2], 0)) AS mean_office_physician_visits_hi FROM persons_as p LEFT JOIN visits v USING (person_id) WHERE true; diff --git a/benchmark/meps_brr/damico2006/as_queries/q01_adults.sql b/benchmark/meps_brr/damico2006/as_queries/q01_adults.sql index 1f519414..e792657e 100644 --- a/benchmark/meps_brr/damico2006/as_queries/q01_adults.sql +++ b/benchmark/meps_brr/damico2006/as_queries/q01_adults.sql @@ -1,5 +1,5 @@ -- Damico (2009), Table 2: explicit AS rewrite; masks constructed in the timed query. SELECT list_transform(list_zip(as_sum(p.mask_lo, p.weight * p.expenditure), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_expenditure_lo + lambda x: x[1] / nullif(x[2], 0)) AS mean_expenditure_lo FROM persons_as p WHERE p.age > 17; diff --git a/benchmark/meps_brr/damico2006/as_queries/q02_nonelderly.sql b/benchmark/meps_brr/damico2006/as_queries/q02_nonelderly.sql index a138f0fd..5661a874 100644 --- a/benchmark/meps_brr/damico2006/as_queries/q02_nonelderly.sql +++ b/benchmark/meps_brr/damico2006/as_queries/q02_nonelderly.sql @@ -1,5 +1,5 @@ -- Damico (2009), Table 2: explicit AS rewrite; masks constructed in the timed query. SELECT list_transform(list_zip(as_sum(p.mask_lo, p.weight * p.expenditure), as_sum(p.mask_lo, p.weight)), - x -> x[1] / nullif(x[2], 0)) AS mean_expenditure_lo + lambda x: x[1] / nullif(x[2], 0)) AS mean_expenditure_lo FROM persons_as p WHERE p.age > 17 AND p.age < 65; diff --git a/benchmark/meps_brr/meps_brr_poc.cpp b/benchmark/meps_brr/meps_brr_poc.cpp index 8fc9f100..2edd7b66 100644 --- a/benchmark/meps_brr/meps_brr_poc.cpp +++ b/benchmark/meps_brr/meps_brr_poc.cpp @@ -38,14 +38,14 @@ static void Save(const string &path, const string &sql) { throw std::runtime_error("Cannot write " + path); f << sql; } -static unique_ptr Query(Connection &con, const string &sql) { +static unique_ptr Query(Connection &con, const string &sql) { auto r = con.Query(sql); if (!r || r->HasError()) throw std::runtime_error(r ? r->GetError() : "No query result"); return r; } static void Require(Connection &con, const string &sql, const string &message) { - if (!Query(con, sql)->GetValue(0, 0).GetValue()) + if (!Query(con, sql)->Collection().GetValue(0, 0).GetValue()) throw std::runtime_error(message); } static string ReplicateSQL(string sql, int rep) { @@ -78,8 +78,8 @@ static Results Run(Connection &con, const string &sql, const string &mode, int w if (r->RowCount() != 1) throw std::runtime_error("Expected one row of named scalar/list estimates"); for (idx_t col = 0; col < r->ColumnCount(); col++) { - string metric = r->names[col]; - auto value = r->GetValue(col, 0); + string metric = r->GetNames()[col].GetIdentifierName(); + auto value = r->Collection().GetValue(col, 0); if (mode == "duckdb" || mode == "naive_as") { out[metric].push_back(Number(value)); continue; diff --git a/benchmark/microbench/common.sh b/benchmark/microbench/common.sh index 5e2a0e07..d73adc7a 100755 --- a/benchmark/microbench/common.sh +++ b/benchmark/microbench/common.sh @@ -3,7 +3,7 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" BINARIES_DIR="$SCRIPT_DIR/binaries" -RESULTS_DIR="$SCRIPT_DIR/results" +RESULTS_DIR="${RESULTS_DIR:-$SCRIPT_DIR/results}" # Default binary (with all optimizations enabled) DEFAULT_BINARY="$BINARIES_DIR/duckdb_default" TEST_DB="$BINARIES_DIR/test_data.duckdb" @@ -14,10 +14,46 @@ BENCH_RUNS=${BENCH_RUNS:-3} THREADS=${THREADS:-1} # Default to single-threaded to avoid partitioning issues # Query setup commands -QUERY_SETUP="SET pac_mi=0.0; SET threads=4; SET max_memory='30G'; SET max_temp_directory_size='90GB';" +QUERY_SETUP="${QUERY_SETUP:-SET pac_mi=0.0; SET threads=4; SET max_memory='30G'; SET max_temp_directory_size='90GB';}" mkdir -p "$RESULTS_DIR" +median_seconds() { + sort -n | awk '{a[NR]=$1} END {if (NR % 2) print a[(NR+1)/2]; else if (NR) print (a[NR/2]+a[NR/2+1])/2}' +} + +aggregate_profile_seconds() { + python3 - "$1" <<'PY' +import json +import sys + +with open(sys.argv[1]) as file: + profile = json.load(file) + +def aggregate_time(node): + kind = node.get("operator_type", node.get("type", "")) + matched = any(name in kind for name in ("UNGROUPED", "GROUP_BY", "AGGREGATE")) + seconds = node.get("operator_timing", node.get("timing", 0.0)) if matched else 0.0 + count = int(matched) + for child in node.get("children", []): + child_seconds, child_count = aggregate_time(child) + seconds += child_seconds + count += child_count + return seconds, count + +seconds = 0.0 +count = 0 +# DuckDB 2.0 places the operator tree under "operator" and renames its fields. +for root in profile.get("operator", [profile]): + root_seconds, root_count = aggregate_time(root) + seconds += root_seconds + count += root_count +if not count: + raise ValueError("No aggregate operator found in the query profile") +print(f"{seconds:.3f}") +PY +} + # Run a benchmark query and extract timing # Args: binary, query, [runs] # Returns: wall_median,agg_median,wall_times,agg_times @@ -33,13 +69,14 @@ run_bench() { # Warmup for ((i=0; i/dev/null 2>&1 + "$binary" -c "$QUERY_SETUP $query" >/dev/null 2>> "$RESULTS_DIR/errors.log" done # Benchmark runs with profiling local wall_times=() local agg_times=() - local profile_file=$(mktemp) + local profile_dir=$(mktemp -d) + local profile_file="$profile_dir/profile.json" local query_failed=0 local timeout_threshold=25.0 @@ -47,10 +84,10 @@ run_bench() { local start=$(python3 -c 'import time; print(time.time())') if ! "$binary" -c " $QUERY_SETUP - PRAGMA enable_profiling='json'; - PRAGMA profiling_output='$profile_file'; + SET enable_profiling='json'; + SET profiling_output='$profile_file'; $query - " >/dev/null 2>&1; then + " >/dev/null 2>> "$RESULTS_DIR/errors.log"; then query_failed=1 break fi @@ -58,29 +95,11 @@ run_bench() { local elapsed=$(python3 -c "print(f'{$end - $start:.3f}')") wall_times+=("$elapsed") - # Extract aggregate operator time from JSON profile - local agg_time=$(python3 -c " -import json -import sys -try: - with open('$profile_file') as f: - data = json.load(f) - - def find_agg_time(node): - total = 0.0 - op_type = node.get('operator_type', '') - # Match aggregate operators - if 'UNGROUPED' in op_type or 'GROUP_BY' in op_type or 'AGGREGATE' in op_type: - total += node.get('operator_timing', 0.0) - for child in node.get('children', []): - total += find_agg_time(child) - return total - - agg_sec = find_agg_time(data) - print(f'{agg_sec:.3f}') -except: - print('0.000') -" 2>/dev/null) + local agg_time + if ! agg_time=$(aggregate_profile_seconds "$profile_file" 2>> "$RESULTS_DIR/errors.log"); then + query_failed=1 + break + fi agg_times+=("$agg_time") # Stop repeating if query took longer than threshold @@ -90,6 +109,7 @@ except: done rm -f "$profile_file" + rmdir "$profile_dir" # If any benchmark run failed, return error indicator if [[ $query_failed -eq 1 ]]; then @@ -98,8 +118,8 @@ except: fi # Calculate medians - local wall_median=$(printf '%s\n' "${wall_times[@]}" | sort -n | awk 'NR==int((NR+1)/2)') - local agg_median=$(printf '%s\n' "${agg_times[@]}" | sort -n | awk 'NR==int((NR+1)/2)') + local wall_median=$(printf '%s\n' "${wall_times[@]}" | median_seconds) + local agg_median=$(printf '%s\n' "${agg_times[@]}" | median_seconds) echo "$wall_median,$agg_median,${wall_times[*]},${agg_times[*]}" } @@ -200,7 +220,7 @@ run_bench_db() { # Warmup - check if query succeeds local warmup_failed=0 for ((i=0; i/dev/null 2>&1; then + if ! "$binary" "$TEST_DB" -c "$QUERY_SETUP $query" >/dev/null 2>> "$RESULTS_DIR/errors.log"; then warmup_failed=1 break fi @@ -215,7 +235,8 @@ run_bench_db() { # Benchmark runs with profiling local wall_times=() local agg_times=() - local profile_file=$(mktemp) + local profile_dir=$(mktemp -d) + local profile_file="$profile_dir/profile.json" local query_failed=0 local timeout_threshold=25.0 @@ -223,10 +244,10 @@ run_bench_db() { local start=$(python3 -c 'import time; print(time.time())') if ! "$binary" "$TEST_DB" -c " $QUERY_SETUP - PRAGMA enable_profiling='json'; - PRAGMA profiling_output='$profile_file'; + SET enable_profiling='json'; + SET profiling_output='$profile_file'; $query - " >/dev/null 2>&1; then + " >/dev/null 2>> "$RESULTS_DIR/errors.log"; then query_failed=1 break fi @@ -234,29 +255,11 @@ run_bench_db() { local elapsed=$(python3 -c "print(f'{$end - $start:.3f}')") wall_times+=("$elapsed") - # Extract aggregate operator time from JSON profile - local agg_time=$(python3 -c " -import json -import sys -try: - with open('$profile_file') as f: - data = json.load(f) - - def find_agg_time(node): - total = 0.0 - op_type = node.get('operator_type', '') - # Match aggregate operators - if 'UNGROUPED' in op_type or 'GROUP_BY' in op_type or 'AGGREGATE' in op_type: - total += node.get('operator_timing', 0.0) - for child in node.get('children', []): - total += find_agg_time(child) - return total - - agg_sec = find_agg_time(data) - print(f'{agg_sec:.3f}') -except: - print('0.000') -" 2>/dev/null) + local agg_time + if ! agg_time=$(aggregate_profile_seconds "$profile_file" 2>> "$RESULTS_DIR/errors.log"); then + query_failed=1 + break + fi agg_times+=("$agg_time") # Stop repeating if query took longer than threshold @@ -266,6 +269,7 @@ except: done rm -f "$profile_file" + rmdir "$profile_dir" # If any benchmark run failed, return error indicator if [[ $query_failed -eq 1 ]]; then @@ -274,8 +278,8 @@ except: fi # Calculate medians - local wall_median=$(printf '%s\n' "${wall_times[@]}" | sort -n | awk 'NR==int((NR+1)/2)') - local agg_median=$(printf '%s\n' "${agg_times[@]}" | sort -n | awk 'NR==int((NR+1)/2)') + local wall_median=$(printf '%s\n' "${wall_times[@]}" | median_seconds) + local agg_median=$(printf '%s\n' "${agg_times[@]}" | median_seconds) echo "$wall_median,$agg_median,${wall_times[*]},${agg_times[*]}" } diff --git a/benchmark/ssb/as_ssb_benchmark.cpp b/benchmark/ssb/as_ssb_benchmark.cpp index 9899859f..aa148ca6 100644 --- a/benchmark/ssb/as_ssb_benchmark.cpp +++ b/benchmark/ssb/as_ssb_benchmark.cpp @@ -37,14 +37,14 @@ static string Csv(const string &s) { out += c == '"' ? "\"\"" : string(1, c); return out + "\""; } -static unique_ptr Query(Connection &con, const string &sql) { +static unique_ptr Query(Connection &con, const string &sql) { auto r = con.Query(sql); if (!r || r->HasError()) throw std::runtime_error(r ? r->GetError() : "No result"); return r; } static void Require(Connection &con, const string &sql, const string &why) { - if (!Query(con, sql)->GetValue(0, 0).GetValue()) + if (!Query(con, sql)->Collection().GetValue(0, 0).GetValue()) throw std::runtime_error(why); } static string InputQuery(string sql, int world) { @@ -66,10 +66,11 @@ static string QuerySQL(const string &sql) { } using Results = std::map; enum class Mode { FULL, REPEATED, ASA }; -static void Collect(MaterializedQueryResult &r, Results &out, int world, bool is_as, bool grouped) { +static void Collect(QueryResult &r, Results &out, int world, bool is_as, bool grouped) { + auto rows = r.Collection().GetRows(); idx_t aggregate = r.ColumnCount(); for (idx_t col = 0; col < r.ColumnCount(); col++) - if (r.names[col] == "revenue" || r.names[col] == "profit") + if (r.GetNames()[col] == "revenue" || r.GetNames()[col] == "profit") aggregate = col; if (aggregate == r.ColumnCount()) throw std::runtime_error("Missing aggregate column"); @@ -78,13 +79,13 @@ static void Collect(MaterializedQueryResult &r, Results &out, int world, bool is for (idx_t col = 0; col < r.ColumnCount() - (is_as ? 1 : 0); col++) { if (col == aggregate) continue; - auto v = r.GetValue(col, row); + auto v = rows.GetValue(col, row); string s = v.IsNull() ? "NULL" : v.type().ToString() + ":" + v.ToString(); key += std::to_string(s.size()) + ":" + s; } - Value value = r.GetValue(aggregate, row); + Value value = rows.GetValue(aggregate, row); if (is_as) { - auto pv = r.GetValue(r.ColumnCount() - 1, row); + auto pv = rows.GetValue(r.ColumnCount() - 1, row); bool present = false; if (!pv.IsNull()) { const auto &presence = ListValue::GetChildren(pv); @@ -227,17 +228,20 @@ int main(int argc, char **argv) { if (const char *memory_limit = std::getenv("SSB_MEMORY_LIMIT")) { Query(con, "SET memory_limit=" + Quote(memory_limit)); } - string memory_limit = Query(con, "SELECT current_setting('memory_limit')")->GetValue(0, 0).ToString(); + string memory_limit = + Query(con, "SELECT current_setting('memory_limit')")->Collection().GetValue(0, 0).ToString(); std::cerr << "DuckDB memory limit: " << memory_limit << '\n'; Query(con, "SET priv_rewrite=false; SET pac_sample_diversity_check=false; SET pac_mi=0; SET privacy_seed=42; " "SET pac_hash_repair=true"); Query(con, Read(ssb + "/schema.sql")); auto manifest = Query(con, "SELECT \"table\", rows::UBIGINT FROM read_csv(" + Quote(data + "/row_counts.csv") + ", header=true)"); + auto manifest_rows = manifest->Collection().GetRows(); std::map row_counts; for (idx_t row = 0; row < manifest->RowCount(); row++) { if (!row_counts - .emplace(manifest->GetValue(0, row).ToString(), manifest->GetValue(1, row).GetValue()) + .emplace(manifest_rows.GetValue(0, row).ToString(), + manifest_rows.GetValue(1, row).GetValue()) .second) throw std::runtime_error("Duplicate table in row-count manifest"); } @@ -263,8 +267,8 @@ int main(int argc, char **argv) { double import_ms = Ms(start); Query(con, "CHECKPOINT"); // Explicit column protection permits grouping by customer city/nation while retaining key protection. - Query(con, "ALTER TABLE customer ADD PRIVACY_KEY(c_custkey); ALTER TABLE customer SET PU; " - "ALTER TABLE customer ADD PROTECTED(c_custkey); " + Query(con, "ALTER TABLE customer ADD PRIVACY_KEY(c_custkey); " + "ALTER TABLE customer ADD PROTECTED(c_custkey); ALTER TABLE customer SET PU; " "ALTER TABLE lineorder ADD PRIVACY_LINK(lo_custkey) REFERENCES customer(c_custkey)"); for (auto name : {"q11", "q12", "q13", "q21", "q22", "q23", "q31", "q32", "q33", "q34", "q41", "q42", "q43"}) { string sql = Read(ssb + "/queries/" + name + ".sql"); @@ -281,9 +285,10 @@ int main(int argc, char **argv) { for (const auto &variant : {std::make_pair("full", aggregate_sql), std::make_pair("asa", as_aggregate_sql), std::make_pair("world0", InputQuery(aggregate_sql, 0))}) { auto explain = Query(con, "EXPLAIN " + variant.second); + auto explain_rows = explain->Collection().GetRows(); plans << "\n--- " << name << " timed " << variant.first << " ---\n"; for (idx_t row = 0; row < explain->RowCount(); row++) { - plans << explain->GetValue(1, row).ToString() << '\n'; + plans << explain_rows.GetValue(1, row).ToString() << '\n'; } } bool rewritten = false, executed = false; @@ -291,9 +296,10 @@ int main(int argc, char **argv) { Query(con, "SET priv_rewrite=true"); try { auto explain = Query(con, "EXPLAIN " + sql); + auto explain_rows = explain->Collection().GetRows(); string plan; for (idx_t row = 0; row < explain->RowCount(); row++) - plan += explain->GetValue(1, row).ToString(); + plan += explain_rows.GetValue(1, row).ToString(); plans << "\n--- " << name << " ---\n" << plan << '\n'; rewritten = plan.find("as_noised_sum") != string::npos || plan.find("as_sum(") != string::npos; if (!rewritten) diff --git a/benchmark/tpcds/as_tpcds_benchmark.cpp b/benchmark/tpcds/as_tpcds_benchmark.cpp index df757b4f..5fe7a50a 100644 --- a/benchmark/tpcds/as_tpcds_benchmark.cpp +++ b/benchmark/tpcds/as_tpcds_benchmark.cpp @@ -39,7 +39,7 @@ static string Csv(const string &value) { return out + "\""; } -static unique_ptr Query(Connection &con, const string &sql) { +static unique_ptr Query(Connection &con, const string &sql) { auto result = con.Query(sql); if (!result || result->HasError()) { throw std::runtime_error(result ? result->GetError() : "No result"); @@ -86,7 +86,7 @@ int main(int argc, char **argv) { Query(con, "INSTALL tpcds; LOAD tpcds;"); auto tables = Query(con, "SELECT count(*) FROM duckdb_tables() WHERE schema_name='main' " "AND table_name IN ('customer','store_sales','catalog_sales','web_sales')"); - auto present = tables->GetValue(0, 0).GetValue(); + auto present = tables->Collection().GetValue(0, 0).GetValue(); auto sf_value = JsonNumber(config, "sf", 1); if (sf_value <= 0) { throw std::runtime_error("Scale factor must be positive"); @@ -103,15 +103,15 @@ int main(int argc, char **argv) { throw std::runtime_error("Expected all 99 TPC-DS queries"); } for (idx_t row = 0; row < workload->RowCount(); row++) { - auto id = workload->GetValue(0, row).GetValue(); + auto id = workload->Collection().GetValue(0, row).GetValue(); std::ofstream file(queries + "/" + Name(id) + ".sql"); if (!file) { throw std::runtime_error("Cannot export original query"); } - file << workload->GetValue(1, row).ToString(); + file << workload->Collection().GetValue(1, row).ToString(); } std::ofstream metadata(output + "/environment.txt"); - metadata << "DuckDB: " << Query(con, "SELECT version()")->GetValue(0, 0).ToString() << '\n' + metadata << "DuckDB: " << Query(con, "SELECT version()")->Collection().GetValue(0, 0).ToString() << '\n' << "Config:\n" << Read(argv[1]) << "\nOwnership:\n" << Read(schema); @@ -157,7 +157,7 @@ int main(int argc, char **argv) { auto explain = Query(con, "EXPLAIN " + sql); string plan; for (idx_t row = 0; row < explain->RowCount(); row++) { - plan += explain->GetValue(1, row).ToString() + "\n"; + plan += explain->Collection().GetValue(1, row).ToString() + "\n"; } std::ofstream plan_file(output + "/" + name + "_" + mode + ".plan.txt"); plan_file << plan; diff --git a/benchmark/tpch/tpch_as_queries/q20.sql b/benchmark/tpch/tpch_as_queries/q20.sql index 0e7aba4a..645b754c 100755 --- a/benchmark/tpch/tpch_as_queries/q20.sql +++ b/benchmark/tpch/tpch_as_queries/q20.sql @@ -6,6 +6,6 @@ SELECT s_name, s_address AND priv_filter_gt(ps_availqty * 2, (SELECT as_sum(priv_hash(hash(o_custkey)), l_quantity) FROM lineitem JOIN orders ON l_orderkey = o_orderkey WHERE l_partkey = ps_partkey AND l_suppkey = ps_suppkey - AND l_shipdate >= DATE '1994-01-01' AND l_shipdate < DATE '1995-01-01')) - AND n_name = 'CANADA') + AND l_shipdate >= DATE '1994-01-01' AND l_shipdate < DATE '1995-01-01'))) + AND n_name = 'CANADA' ORDER BY s_name; diff --git a/docs/build/README.md b/docs/build/README.md index 38492142..aee454c7 100644 --- a/docs/build/README.md +++ b/docs/build/README.md @@ -3,9 +3,9 @@ ## Prerequisites - **git** (with submodule support) -- **cmake** (>= 3.11) +- **cmake** (>= 3.14) - **ninja** (optional but recommended — significantly faster builds) -- **C++ compiler** with C++11 support (clang recommended, especially for benchmarks) +- **C++ compiler** with C++17 support (clang recommended, especially for benchmarks) ## Clone @@ -28,7 +28,22 @@ make # release build without ninja make debug # debug build ``` -The extension is built into `build/release/extension/pac/pac.duckdb_extension`. +The extension is built into `build/release/extension/privacy/privacy.duckdb_extension`. + +The Makefile sets `OVERRIDE_GIT_DESCRIBE=v2.0.0` so the loadable extension +uses the released DuckDB ABI label. When configuring CMake directly, pass +`-DOVERRIDE_GIT_DESCRIBE=v2.0.0` as well. DuckDB 2.0 requires explicit static +extension registration; `extension_config.cmake` enables it for the CLI, +tests, and benchmark runners. + +The compiler uses DuckDB 2.0's default CTE-based decorrelation. No legacy +`delim_join_as_cte` setting is required. CTE readers carry the corresponding +mask bindings and counter types, and intermediate expressions remain per-world +until the output projection finalizes them. + +DuckDB 2.0 can infer wider DECIMAL types for arithmetic than DuckDB 1.5. +The inferred width selects the SUM state (including the INT128 path), so +approximate outputs can differ across engine versions even with the same seed. ## Run Tests @@ -40,7 +55,7 @@ make test ./build/release/test/unittest --test-dir . "test/sql/pac*" # C++ unit tests -./build/release/extension/pac/pac_test_runner +./build/release/extension/privacy/pac_test_runner ``` ## Build Flags @@ -63,11 +78,13 @@ The DuckDB version is pinned via the `duckdb` and `extension-ci-tools` submodule To update to a new release: ```bash -cd duckdb && git fetch --tags && git checkout v1.X.0 -cd ../extension-ci-tools && git fetch && git checkout origin/v1.X.0 +cd duckdb && git fetch --tags && git checkout v2.0.0 +cd ../extension-ci-tools && git fetch origin && git checkout ebdcc1961780b0ce1272a0d7637dc58447c34885 cd .. git add duckdb extension-ci-tools -git commit -m "Bump DuckDB to v1.X.0" +git commit -m "Bump DuckDB to v2.0.0" ``` -Replace `v1.X.0` with the target version tag. +The upgrade pins DuckDB 2.0.0 and the CI-tools commit shown above. +When changing versions, update both submodule pins and the references in +`.github/workflows/MainDistributionPipeline.yml`. diff --git a/docs/pac/internal_operators.md b/docs/pac/internal_operators.md index 196d7570..fa5974fd 100644 --- a/docs/pac/internal_operators.md +++ b/docs/pac/internal_operators.md @@ -175,6 +175,51 @@ No new aggregate functions, no new noise mechanisms, no plan decomposition. --- +## Stochastic aggregate state spilling (DuckDB 2.0) + +The membership-mask-plus-scalar SUM, COUNT, and MIN/MAX implementations expose +DuckDB's native state-export/import callbacks, including DECIMAL specializations. +List-input SUM, COUNT, and MIN/MAX also provide callbacks. Scalar AVG uses the +SUM and COUNT implementations through the existing rewrite. DuckDB can export arena-backed hash-aggregate states to spillable columns +and import them when combining partitions. Every aggregate in a hash table must +expose a compatible state layout for DuckDB to select this path. + +The exported BLOB contains buffered values and masks, update counts, positive and +negative integer cascade levels and overflow budgets, floating-point counters, +COUNT SWAR bytes and overflow totals, MIN/MAX bounds, and list-aggregate values +and counts. Folded correction arguments are retained when rebinding exported states. +It contains no addresses. Import reconstructs owned arena allocations and retains +the inline-level optimization. COUNT and MIN/MAX expand restored buffers using their existing +kernels only on the import path, so a subsequent combine cannot overwrite a +restored destination buffer. SUM retains restored buffers for its buffered combiner. Export does not finalize, resample, or add noise; +the native update kernels and aggregate-state sizes remain unchanged. The binary +representation is internal to this implementation, rather than a portable state +interchange format. + +SUM retains one/two-value buffers during combine. Finalization materializes copies +in temporary storage backed by `ArenaAllocator::GetUnderlyingAllocator()`, releasing +these allocations after the callback. These two optimizations address the Q78 +allocation path: expanding many small groups during final combine/finalization can +exhaust memory after sink-phase spilling decisions. Disk spilling handles stored +states; it does not itself bound these late allocations. COUNT and MIN/MAX retain +their original in-memory update, combine, and finalization paths. + +Regression coverage extends `test/sql/pac_sum.test`: round trips include signed, +unsigned, floating-point and DECIMAL inputs, buffered/empty states, high cascade +levels, combining imported states, and truncated-input rejection. The pressure +test uses 800,000 rows in 100,000 contiguous groups, three SUMs and a 32 MB memory +limit; its complete result checksum must equal the in-memory result. + +```sh +build/release/test/unittest 'test/sql/pac_sum.test' +``` + +For the Q78 investigation, retain the complete query, sampling settings and thread +count when comparing builds. Record peak buffer memory, temporary storage and +result equivalence alongside latency. Passing the state-spill regression does not +establish that Q78 SF30 fits the target machine: joins and active aggregate +partitions still require working memory. + ## TODO - Categorical expressions (`pac_categorical` rewriter) diff --git a/docs/syntax.md b/docs/syntax.md index 7bd9e989..bf3f3494 100644 --- a/docs/syntax.md +++ b/docs/syntax.md @@ -158,3 +158,7 @@ There is some preliminary support for using multiple PU tables inside the same s | `privacy_diffcols` | `NULL` | [Utility diff](pac/utility.md): compare noised vs exact results | | `privacy_min_group_count` | `NULL` (disabled) | Admin-set minimum support threshold for DP groups. In PAC aggregates, this remains the z-score threshold for [utility NULLing](pac/runtime_checks.md#utility-nulling). | | `priv_clip_support` | `NULL` (disabled) | Enable clipping: aggregate per-PU values with outlier elimination at the given support threshold | + +PAC seeded reproducibility requires stable execution order; use `SET threads = 1` +for repeatable tests. Parallel group order and p-tracking state can change results; +see the [seed reproducibility contract](../README.md#settings). diff --git a/duckdb b/duckdb index 08e34c44..d54cbbe3 160000 --- a/duckdb +++ b/duckdb @@ -1 +1 @@ -Subproject commit 08e34c447bae34eaee3723cac61f2878b6bdf787 +Subproject commit d54cbbe3df17e8e5ce68312f483b4155836112aa diff --git a/extension-ci-tools b/extension-ci-tools index 1d3a50b3..ebdcc196 160000 --- a/extension-ci-tools +++ b/extension-ci-tools @@ -1 +1 @@ -Subproject commit 1d3a50b3414534f6787e28c1c5efaafb4beb4e8f +Subproject commit ebdcc1961780b0ce1272a0d7637dc58447c34885 diff --git a/extension_config.cmake b/extension_config.cmake index 2a11adbb..de2d3703 100644 --- a/extension_config.cmake +++ b/extension_config.cmake @@ -5,6 +5,7 @@ duckdb_extension_load(privacy SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR} LOAD_TESTS ) +duckdb_extension_statically_link(privacy) # Any extra extensions that should be built -# e.g.: duckdb_extension_load(json) \ No newline at end of file +# e.g.: duckdb_extension_load(json) diff --git a/src/aggregates/as_aggregate.cpp b/src/aggregates/as_aggregate.cpp index b2f00c8b..123bc969 100644 --- a/src/aggregates/as_aggregate.cpp +++ b/src/aggregates/as_aggregate.cpp @@ -1,3 +1,4 @@ +#include "duckdb/common/vector/list_vector.hpp" #include "aggregates/as_aggregate.hpp" #include "duckdb.hpp" @@ -7,8 +8,10 @@ #include "duckdb/common/exception.hpp" #include "duckdb/function/aggregate_function.hpp" #include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/main/client_context_state.hpp" #include +#include #ifndef M_PI #define M_PI 3.14159265358979323846 #endif @@ -114,7 +117,7 @@ static double ComputeSecondMomentVariance(const vector &values); // PacNoiseInNull: probabilistically returns true based on bit count in key_hash // mi: controls probabilistic (mi>0) vs deterministic (mi<=0) mode // correction: reduces NULL probability by this factor (considers correction times more non-nulls) -bool PacNoiseInNull(uint64_t key_hash, double mi, double correction, std::mt19937_64 &gen) { +bool PacNoiseInNull(uint64_t key_hash, double mi, double correction, PacRandomEngine &gen) { int popcount = pac_popcount64(key_hash); // Effective popcount considering correction factor (correction times more non-nulls) double effective_popcount = popcount * correction; @@ -133,10 +136,34 @@ bool PacNoiseInNull(uint64_t key_hash, double mi, double correction, std::mt1993 return pac_popcount64(~key_hash) > threshold; } +struct PacRandomState : public ClientContextState { + std::atomic next_stream {0}; + + void QueryBegin(ClientContext &) override { + next_stream.store(0, std::memory_order_relaxed); + } +}; + +shared_ptr GetPacRandomState(ClientContext &context) { + return context.registered_state->GetOrCreate("privacy_random_streams"); +} + +uint64_t GetPacSeed(ClientContext &context, double mi) { + Value value; + if (context.TryGetCurrentSetting(Identifier("privacy_seed"), value) && !value.IsNull()) { + return uint64_t(value.GetValue()); + } + return mi == 0.0 ? 42 : 42 ^ (PAC_MAGIC_HASH * uint64_t(context.ActiveTransaction().GetActiveQuery())); +} + +PacRandomEngine PacRandomStream(uint64_t seed, PacRandomState &state) { + return PacRandomEngine(seed, state.next_stream.fetch_add(1, std::memory_order_relaxed)); +} + // ============================================================================ // Utility NULLing: probabilistically NULL low-SNR cells (post-processing) // ============================================================================ -bool PacUtilityNull(double noised_value, double noise_variance, double threshold, std::mt19937_64 &gen) { +bool PacUtilityNull(double noised_value, double noise_variance, double threshold, PacRandomEngine &gen) { if (std::isnan(threshold)) { return false; // disabled (setting is NULL) } @@ -169,7 +196,7 @@ bool PacUtilityNull(double noised_value, double noise_variance, double threshold // correction: factor to multiply values by after compacting but before noising // pstate: optional p-tracking state for persistent secret composition (Bayesian posterior over worlds) // Returns: correction*yJ + noise where yJ is a randomly selected counter -PAC_FLOAT PacNoisySampleFrom64Counters(const PAC_FLOAT counters[64], double mi, double correction, std::mt19937_64 &gen, +PAC_FLOAT PacNoisySampleFrom64Counters(const PAC_FLOAT counters[64], double mi, double correction, PacRandomEngine &gen, uint64_t is_null, uint64_t counter_selector, const std::shared_ptr &pstate, double *out_noise_variance) { D_ASSERT(~is_null != 0); // at least one bit must be valid @@ -350,19 +377,19 @@ uint64_t PacRepairHash32(uint64_t num) { // k : minimum per-subsample support; cells below it are released as NULL // ============================================================================ struct PacAggregateLocalState : public FunctionLocalState { - explicit PacAggregateLocalState(uint64_t seed) : gen(seed) { + explicit PacAggregateLocalState(PacRandomEngine gen_p) : gen(gen_p) { } - std::mt19937_64 gen; + PacRandomEngine gen; }; -static unique_ptr PacAggregateInit(ExpressionState &, const BoundFunctionExpression &, +static unique_ptr PacAggregateInit(ExpressionState &state, const BoundFunctionExpression &, FunctionData *bind_data) { // Single source-of-truth: prefer seed supplied via bind_data (PrivBindData). uint64_t seed = std::random_device {}(); if (bind_data) { seed = bind_data->Cast().seed; } - return make_uniq(seed); + return make_uniq(PacRandomStream(seed, *GetPacRandomState(state.GetContext()))); } // Templated over values element type T and counts element type C so we can accept @@ -376,8 +403,8 @@ static void PacAggregateScalar(DataChunk &args, ExpressionState &state, Vector & idx_t count = args.size(); result.SetVectorType(VectorType::FLAT_VECTOR); - auto res = FlatVector::GetData(result); - FlatVector::Validity(result).SetAllValid(count); + auto res = FlatVector::GetDataMutable(result); + FlatVector::ValidityMutable(result).SetAllValid(count); auto &local = ExecuteFunctionState::GetFunctionState(state)->Cast(); auto &gen = local.gen; @@ -386,18 +413,18 @@ static void PacAggregateScalar(DataChunk &args, ExpressionState &state, Vector & uint64_t query_hash = 0; std::shared_ptr pstate; auto &bound_expr = state.expr.Cast(); - if (bound_expr.bind_info) { - auto &bd = bound_expr.bind_info->Cast(); + if (bound_expr.BindInfo()) { + auto &bd = bound_expr.BindInfo()->Cast(); query_hash = bd.query_hash; pstate = bd.pstate; } UnifiedVectorFormat vvals, vcnts; - vals.ToUnifiedFormat(count, vvals); - cnts.ToUnifiedFormat(count, vcnts); + vals.ToUnifiedFormat(vvals); + cnts.ToUnifiedFormat(vcnts); UnifiedVectorFormat kvals; - k_vec.ToUnifiedFormat(count, kvals); + k_vec.ToUnifiedFormat(kvals); auto *kdata = UnifiedVectorFormat::GetData(kvals); for (idx_t row = 0; row < count; row++) { @@ -426,10 +453,10 @@ static void PacAggregateScalar(DataChunk &args, ExpressionState &state, Vector & idx_t vals_len = ve.length; idx_t cnts_len = ce.length; - auto &vals_child = ListVector::GetEntry(vals); - auto &cnts_child = ListVector::GetEntry(cnts); - vals_child.Flatten(ve.offset + ve.length); - cnts_child.Flatten(ce.offset + ce.length); + auto &vals_child = ListVector::GetChild(vals); + auto &cnts_child = ListVector::GetChild(cnts); + vals_child.Flatten(); + cnts_child.Flatten(); auto *vdata = FlatVector::GetData(vals_child); auto *cdata = FlatVector::GetData(cnts_child); @@ -521,8 +548,8 @@ static void PacAggregateScalar(DataChunk &args, ExpressionState &state, Vector & } } -static unique_ptr PacAggregateBind(ClientContext &ctx, ScalarFunction &, - vector> &) { +static unique_ptr PacAggregateBind(BindScalarFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); // mi is supplied per-row; the bind-time value is unused for noise but seeds the RNG. return make_uniq(ctx, 128.0); } @@ -533,8 +560,9 @@ void RegisterPacAggregateFunctions(ExtensionLoader &loader) { auto fun = ScalarFunction( "pac_aggregate", {LogicalType::LIST(vals_elem), LogicalType::LIST(cnts_elem), LogicalType::DOUBLE, LogicalType::INTEGER}, - LogicalType::DOUBLE, fn, PacAggregateBind, nullptr, nullptr, PacAggregateInit); - fun.null_handling = FunctionNullHandling::SPECIAL_HANDLING; + LogicalType::DOUBLE, fn, PacAggregateBind, nullptr, PacAggregateInit); + fun.SetNullHandling(FunctionNullHandling::SPECIAL_HANDLING); + fun.SetFallible(); loader.RegisterFunction(fun); }; @@ -569,11 +597,12 @@ void RegisterPacAggregateFunctions(ExtensionLoader &loader) { #undef REG_COUNTS_FOR } -static unique_ptr PacHashBind(ClientContext &ctx, ScalarFunction &, vector> &) { +static unique_ptr PacHashBind(BindScalarFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); double mi = GetPacMiFromSetting(ctx); bool hash_repair = true; Value hr_val; - if (ctx.TryGetCurrentSetting("pac_hash_repair", hr_val) && !hr_val.IsNull()) { + if (ctx.TryGetCurrentSetting(Identifier("pac_hash_repair"), hr_val) && !hr_val.IsNull()) { hash_repair = hr_val.GetValue(); } return make_uniq(ctx, mi, 1.0, 1.0, hash_repair); @@ -585,8 +614,8 @@ static void PacHashFunction(DataChunk &args, ExpressionState &state, Vector &res auto &function = state.expr.Cast(); uint64_t query_hash = 0; bool hash_repair = true; - if (function.bind_info) { - auto &bind_data = function.bind_info->Cast(); + if (function.BindInfo()) { + auto &bind_data = function.BindInfo()->Cast(); query_hash = bind_data.query_hash; hash_repair = bind_data.hash_repair; } diff --git a/src/aggregates/as_clip_min_max.cpp b/src/aggregates/as_clip_min_max.cpp index 5e217eba..5f529993 100644 --- a/src/aggregates/as_clip_min_max.cpp +++ b/src/aggregates/as_clip_min_max.cpp @@ -1,3 +1,4 @@ +#include "duckdb/common/vector/list_vector.hpp" #include "aggregates/as_clip_min_max.hpp" #include "categorical/pac_categorical.hpp" #include "duckdb/common/types/decimal.hpp" @@ -114,12 +115,12 @@ template static void PacClipMinMaxUpdate(Vector inputs[], PacClipMinMaxStateWrapper &state, idx_t count, ArenaAllocator &allocator, int sample_lanes = 0) { UnifiedVectorFormat hash_data, value_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); - if (hash_data.validity.AllValid() && value_data.validity.AllValid()) { + if (hash_data.validity.CannotHaveNull() && value_data.validity.CannotHaveNull()) { for (idx_t i = 0; i < count; i++) { auto h_idx = hash_data.sel->get_index(i); auto v_idx = value_data.sel->get_index(i); @@ -143,9 +144,9 @@ template static void PacClipMinMaxScatterUpdate(Vector inputs[], Vector &states, idx_t count, ArenaAllocator &allocator, int sample_lanes = 0) { UnifiedVectorFormat hash_data, value_data, sdata; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - states.ToUnifiedFormat(count, sdata); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + states.ToUnifiedFormat(sdata); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); @@ -180,11 +181,13 @@ static void PacClipMinMaxScatterUpdate(Vector inputs[], Vector &states, idx_t co // Generate for IS_MAX=true (MAX) #define X(NAME, VALUE_T, INPUT_T, SIGNED_VAL) \ - static void PacClipMaxUpdate##NAME(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_p, \ - idx_t count) { \ - auto &state = *reinterpret_cast *>(state_p); \ - PacClipMinMaxUpdate(inputs, state, count, aggr.allocator, \ - GetPrivSampleLanes(aggr)); \ + static void PacClipMaxUpdate##NAME(Vector inputs[], AggregateInputData &aggr, idx_t input_count, \ + const ClusteredAggr &clusters, idx_t count) { \ + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { \ + auto &state = *reinterpret_cast *>(state_p); \ + PacClipMinMaxUpdate(inputs, state, count, aggr.allocator, \ + GetPrivSampleLanes(aggr)); \ + }); \ } \ static void PacClipMaxScatterUpdate##NAME(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, \ idx_t count) { \ @@ -197,11 +200,13 @@ PCMM_INT_TYPES_UNSIGNED // Generate for IS_MAX=false (MIN) #define X(NAME, VALUE_T, INPUT_T, SIGNED_VAL) \ - static void PacClipMinUpdate##NAME(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_p, \ - idx_t count) { \ - auto &state = *reinterpret_cast *>(state_p); \ - PacClipMinMaxUpdate(inputs, state, count, aggr.allocator, \ - GetPrivSampleLanes(aggr)); \ + static void PacClipMinUpdate##NAME(Vector inputs[], AggregateInputData &aggr, idx_t input_count, \ + const ClusteredAggr &clusters, idx_t count) { \ + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { \ + auto &state = *reinterpret_cast *>(state_p); \ + PacClipMinMaxUpdate(inputs, state, count, aggr.allocator, \ + GetPrivSampleLanes(aggr)); \ + }); \ } \ static void PacClipMinScatterUpdate##NAME(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, \ idx_t count) { \ @@ -219,12 +224,12 @@ template static void PacClipMinMaxUpdateFloat(Vector inputs[], PacClipMinMaxStateWrapper &state, idx_t count, ArenaAllocator &allocator, int sample_lanes = 0) { UnifiedVectorFormat hash_data, value_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); - if (hash_data.validity.AllValid() && value_data.validity.AllValid()) { + if (hash_data.validity.CannotHaveNull() && value_data.validity.CannotHaveNull()) { for (idx_t i = 0; i < count; i++) { auto h_idx = hash_data.sel->get_index(i); auto v_idx = value_data.sel->get_index(i); @@ -248,9 +253,9 @@ template static void PacClipMinMaxScatterUpdateFloat(Vector inputs[], Vector &states, idx_t count, ArenaAllocator &allocator, int sample_lanes = 0) { UnifiedVectorFormat hash_data, value_data, sdata; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - states.ToUnifiedFormat(count, sdata); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + states.ToUnifiedFormat(sdata); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); auto state_ptrs = UnifiedVectorFormat::GetData *>(sdata); @@ -273,22 +278,26 @@ static void PacClipMinMaxScatterUpdateFloat(Vector inputs[], Vector &states, idx XF(SingleDouble, double, CLIP_DOUBLE_SHIFT) #define XF(NAME, FLOAT_T, SHIFT_VAL) \ - static void PacClipMaxUpdate##NAME(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_p, \ - idx_t count) { \ - auto &state = *reinterpret_cast *>(state_p); \ - PacClipMinMaxUpdateFloat(inputs, state, count, aggr.allocator, \ - GetPrivSampleLanes(aggr)); \ + static void PacClipMaxUpdate##NAME(Vector inputs[], AggregateInputData &aggr, idx_t input_count, \ + const ClusteredAggr &clusters, idx_t count) { \ + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { \ + auto &state = *reinterpret_cast *>(state_p); \ + PacClipMinMaxUpdateFloat(inputs, state, count, aggr.allocator, \ + GetPrivSampleLanes(aggr)); \ + }); \ } \ static void PacClipMaxScatterUpdate##NAME(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, \ idx_t count) { \ PacClipMinMaxScatterUpdateFloat(inputs, states, count, aggr.allocator, \ GetPrivSampleLanes(aggr)); \ } \ - static void PacClipMinUpdate##NAME(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_p, \ - idx_t count) { \ - auto &state = *reinterpret_cast *>(state_p); \ - PacClipMinMaxUpdateFloat(inputs, state, count, aggr.allocator, \ - GetPrivSampleLanes(aggr)); \ + static void PacClipMinUpdate##NAME(Vector inputs[], AggregateInputData &aggr, idx_t input_count, \ + const ClusteredAggr &clusters, idx_t count) { \ + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { \ + auto &state = *reinterpret_cast *>(state_p); \ + PacClipMinMaxUpdateFloat(inputs, state, count, aggr.allocator, \ + GetPrivSampleLanes(aggr)); \ + }); \ } \ static void PacClipMinScatterUpdate##NAME(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, \ idx_t count) { \ @@ -341,11 +350,11 @@ static void PacClipMinCombine(Vector &src, Vector &dst, AggregateInputData &aggr // Finalize: noised scalar output // ============================================================================ template -static void PacClipMinMaxFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void PacClipMinMaxFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { auto state_ptrs = FlatVector::GetData *>(states); - auto data = FlatVector::GetData(result); - auto &result_mask = FlatVector::Validity(result); + auto data = FlatVector::GetDataMutable(result); + auto &result_mask = FlatVector::ValidityMutable(result); auto &bind = static_cast(*input.bind_data); double mi = bind.mi; double correction = bind.correction; @@ -354,6 +363,7 @@ static void PacClipMinMaxFinalize(Vector &states, AggregateInputData &input, Vec int clip_support = bind.clip_support_threshold; bool clip_scale = bind.clip_scale; + auto gen = PacRandomStream(bind.seed, *bind.random_state); for (idx_t i = 0; i < count; i++) { PacClipMinMaxFlushBuffer(*state_ptrs[i], *state_ptrs[i], input.allocator); @@ -365,7 +375,6 @@ static void PacClipMinMaxFinalize(Vector &states, AggregateInputData &input, Vec continue; } uint64_t key_hash = (pos ? pos->key_hash : 0) | (neg ? neg->key_hash : 0); - std::mt19937_64 gen(bind.seed); if (PacNoiseInNull(key_hash, mi, correction, gen)) { result_mask.SetInvalid(offset + i); continue; @@ -422,10 +431,12 @@ static void PacClipMinMaxFinalize(Vector &states, AggregateInputData &input, Vec XFIN(HugeInt, hugeint_t) #define XFIN(NAME, ACC_T) \ - static void PacClipMaxNoisedFinalize##NAME(Vector &s, AggregateInputData &i, Vector &r, idx_t c, idx_t o) { \ + static void PacClipMaxNoisedFinalize##NAME(Vector &s, AggregateFinalizeInputData &i, Vector &r, idx_t c, \ + idx_t o) { \ PacClipMinMaxFinalize(s, i, r, c, o); \ } \ - static void PacClipMinNoisedFinalize##NAME(Vector &s, AggregateInputData &i, Vector &r, idx_t c, idx_t o) { \ + static void PacClipMinNoisedFinalize##NAME(Vector &s, AggregateFinalizeInputData &i, Vector &r, idx_t c, \ + idx_t o) { \ PacClipMinMaxFinalize(s, i, r, c, o); \ } PCMM_FINALIZE_TYPES @@ -435,8 +446,8 @@ PCMM_FINALIZE_TYPES // Counters finalize (LIST output) // ============================================================================ template -static void PacClipMinMaxFinalizeCounters(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void PacClipMinMaxFinalizeCounters(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { auto state_ptrs = FlatVector::GetData *>(states); auto &bind = static_cast(*input.bind_data); int clip_support = bind.clip_support_threshold; @@ -444,14 +455,14 @@ static void PacClipMinMaxFinalizeCounters(Vector &states, AggregateInputData &in double float_scale = bind.float_scale; bool clip_scale = bind.clip_scale; - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * 64; ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); for (idx_t i = 0; i < count; i++) { PacClipMinMaxFlushBuffer(*state_ptrs[i], *state_ptrs[i], input.allocator); @@ -499,10 +510,10 @@ static void PacClipMinMaxFinalizeCounters(Vector &states, AggregateInputData &in } } -static void PacClipMaxFinalizeCounters(Vector &s, AggregateInputData &i, Vector &r, idx_t c, idx_t o) { +static void PacClipMaxFinalizeCounters(Vector &s, AggregateFinalizeInputData &i, Vector &r, idx_t c, idx_t o) { PacClipMinMaxFinalizeCounters(s, i, r, c, o); } -static void PacClipMinFinalizeCounters(Vector &s, AggregateInputData &i, Vector &r, idx_t c, idx_t o) { +static void PacClipMinFinalizeCounters(Vector &s, AggregateFinalizeInputData &i, Vector &r, idx_t c, idx_t o) { PacClipMinMaxFinalizeCounters(s, i, r, c, o); } @@ -510,13 +521,16 @@ static void PacClipMinFinalizeCounters(Vector &s, AggregateInputData &i, Vector // State size / init / bind // ============================================================================ template -static idx_t PacClipMinMaxStateSize(const AggregateFunction &) { +static idx_t PacClipMinMaxStateSize(AggregateStateInput &) { return sizeof(PacClipMinMaxStateWrapper); } template -static void PacClipMinMaxInitialize(const AggregateFunction &, data_ptr_t state_p) { - memset(state_p, 0, sizeof(PacClipMinMaxStateWrapper)); +static void PacClipMinMaxInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_p = states[i]; + memset(state_p, 0, sizeof(PacClipMinMaxStateWrapper)); + } } // PacClipBind, PacClipBindFloat, PacClipBindDouble are defined in as_clip_aggr.hpp @@ -534,40 +548,42 @@ static AggregateFunction GetPacClipMinMaxNoisedAggregate(PhysicalType type) { switch (type) { case PhysicalType::INT16: - return AggregateFunction(name, {LogicalType::UBIGINT, LogicalType::SMALLINT}, LogicalType::HUGEINT, state_size, - init, IS_MAX ? PacClipMaxScatterUpdateSmallInt : PacClipMinScatterUpdateSmallInt, - combine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, - IS_MAX ? PacClipMaxUpdateSmallInt : PacClipMinUpdateSmallInt); + return AggregateFunction( + Identifier(name), {LogicalType::UBIGINT, LogicalType::SMALLINT}, LogicalType::HUGEINT, state_size, init, + IS_MAX ? PacClipMaxScatterUpdateSmallInt : PacClipMinScatterUpdateSmallInt, combine, finalize, + FunctionNullHandling::DEFAULT_NULL_HANDLING, IS_MAX ? PacClipMaxUpdateSmallInt : PacClipMinUpdateSmallInt); case PhysicalType::INT32: - return AggregateFunction(name, {LogicalType::UBIGINT, LogicalType::INTEGER}, LogicalType::HUGEINT, state_size, - init, IS_MAX ? PacClipMaxScatterUpdateInteger : PacClipMinScatterUpdateInteger, - combine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, - IS_MAX ? PacClipMaxUpdateInteger : PacClipMinUpdateInteger); + return AggregateFunction( + Identifier(name), {LogicalType::UBIGINT, LogicalType::INTEGER}, LogicalType::HUGEINT, state_size, init, + IS_MAX ? PacClipMaxScatterUpdateInteger : PacClipMinScatterUpdateInteger, combine, finalize, + FunctionNullHandling::DEFAULT_NULL_HANDLING, IS_MAX ? PacClipMaxUpdateInteger : PacClipMinUpdateInteger); case PhysicalType::INT64: - return AggregateFunction(name, {LogicalType::UBIGINT, LogicalType::BIGINT}, LogicalType::HUGEINT, state_size, - init, IS_MAX ? PacClipMaxScatterUpdateBigInt : PacClipMinScatterUpdateBigInt, combine, - finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, - IS_MAX ? PacClipMaxUpdateBigInt : PacClipMinUpdateBigInt); + return AggregateFunction( + Identifier(name), {LogicalType::UBIGINT, LogicalType::BIGINT}, LogicalType::HUGEINT, state_size, init, + IS_MAX ? PacClipMaxScatterUpdateBigInt : PacClipMinScatterUpdateBigInt, combine, finalize, + FunctionNullHandling::DEFAULT_NULL_HANDLING, IS_MAX ? PacClipMaxUpdateBigInt : PacClipMinUpdateBigInt); case PhysicalType::INT128: - return AggregateFunction(name, {LogicalType::UBIGINT, LogicalType::HUGEINT}, LogicalType::HUGEINT, state_size, - init, IS_MAX ? PacClipMaxScatterUpdateBigInt : PacClipMinScatterUpdateBigInt, combine, - IS_MAX ? PacClipMaxNoisedFinalizeHugeInt : PacClipMinNoisedFinalizeHugeInt, - FunctionNullHandling::DEFAULT_NULL_HANDLING, - IS_MAX ? PacClipMaxUpdateBigInt : PacClipMinUpdateBigInt); + return AggregateFunction( + Identifier(name), {LogicalType::UBIGINT, LogicalType::HUGEINT}, LogicalType::HUGEINT, state_size, init, + IS_MAX ? PacClipMaxScatterUpdateBigInt : PacClipMinScatterUpdateBigInt, combine, + IS_MAX ? PacClipMaxNoisedFinalizeHugeInt : PacClipMinNoisedFinalizeHugeInt, + FunctionNullHandling::DEFAULT_NULL_HANDLING, IS_MAX ? PacClipMaxUpdateBigInt : PacClipMinUpdateBigInt); default: throw InternalException("as_noised_clip_min/max: unsupported decimal physical type"); } } template -static unique_ptr BindDecimalPacNoisedClipMinMax(ClientContext &ctx, AggregateFunction &function, - vector> &args) { - auto decimal_type = args[1]->return_type; - function = GetPacClipMinMaxNoisedAggregate(decimal_type.InternalType()); - function.name = IS_MAX ? "as_noised_clip_max" : "as_noised_clip_min"; - function.arguments[1] = decimal_type; - function.return_type = LogicalType::DECIMAL(Decimal::MAX_WIDTH_DECIMAL, DecimalType::GetScale(decimal_type)); - return PacClipBind(ctx, function, args); +static unique_ptr BindDecimalPacNoisedClipMinMax(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &function = bind_input.GetBoundFunction(); + auto &args = bind_input.GetArguments(); + auto decimal_type = args[1]->GetReturnType(); + function.ReplaceImplementation(GetPacClipMinMaxNoisedAggregate(decimal_type.InternalType())); + function.SetName(IS_MAX ? "as_noised_clip_max" : "as_noised_clip_min"); + function.GetArguments()[1] = decimal_type; + function.SetReturnType(LogicalType::DECIMAL(Decimal::MAX_WIDTH_DECIMAL, DecimalType::GetScale(decimal_type))); + return PacClipBind(bind_input); } // ============================================================================ @@ -576,29 +592,29 @@ static unique_ptr BindDecimalPacNoisedClipMinMax(ClientContext &ct template static void AddClipMinMaxCountersFcn(AggregateFunctionSet &set, const string &name, const LogicalType &value_type, aggregate_update_t scatter, aggregate_finalize_t finalize, - aggregate_simple_update_t update) { + aggregate_cluster_update_t update) { auto list_type = LogicalType::LIST(PacFloatLogicalType()); - set.AddFunction(AggregateFunction(name, {LogicalType::UBIGINT, value_type}, list_type, + set.AddFunction(AggregateFunction(Identifier(name), {LogicalType::UBIGINT, value_type}, list_type, PacClipMinMaxStateSize, PacClipMinMaxInitialize, scatter, IS_MAX ? PacClipMaxCombine : PacClipMinCombine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacClipBind)); - set.AddFunction(AggregateFunction(name, {LogicalType::UBIGINT, value_type, LogicalType::DOUBLE}, list_type, - PacClipMinMaxStateSize, PacClipMinMaxInitialize, scatter, - IS_MAX ? PacClipMaxCombine : PacClipMinCombine, finalize, + set.AddFunction(AggregateFunction(Identifier(name), {LogicalType::UBIGINT, value_type, LogicalType::DOUBLE}, + list_type, PacClipMinMaxStateSize, PacClipMinMaxInitialize, + scatter, IS_MAX ? PacClipMaxCombine : PacClipMinCombine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacClipBind)); } template static void AddNoisedClipMinMaxFcn(AggregateFunctionSet &set, const string &name, const LogicalType &value_type, const LogicalType &result_type, aggregate_update_t scatter, - aggregate_finalize_t finalize, aggregate_simple_update_t update) { - set.AddFunction(AggregateFunction(name, {LogicalType::UBIGINT, value_type}, result_type, + aggregate_finalize_t finalize, aggregate_cluster_update_t update) { + set.AddFunction(AggregateFunction(Identifier(name), {LogicalType::UBIGINT, value_type}, result_type, PacClipMinMaxStateSize, PacClipMinMaxInitialize, scatter, IS_MAX ? PacClipMaxCombine : PacClipMinCombine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacClipBind)); - set.AddFunction(AggregateFunction(name, {LogicalType::UBIGINT, value_type, LogicalType::DOUBLE}, result_type, - PacClipMinMaxStateSize, PacClipMinMaxInitialize, scatter, - IS_MAX ? PacClipMaxCombine : PacClipMinCombine, finalize, + set.AddFunction(AggregateFunction(Identifier(name), {LogicalType::UBIGINT, value_type, LogicalType::DOUBLE}, + result_type, PacClipMinMaxStateSize, PacClipMinMaxInitialize, + scatter, IS_MAX ? PacClipMaxCombine : PacClipMinCombine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacClipBind)); } @@ -690,25 +706,25 @@ static void AddFloatDoubleOverloads(AggregateFunctionSet &set, const string &nam // FLOAT set.AddFunction( - AggregateFunction(name, {LogicalType::UBIGINT, LogicalType::FLOAT}, list_type, state_size, init, + AggregateFunction(Identifier(name), {LogicalType::UBIGINT, LogicalType::FLOAT}, list_type, state_size, init, IS_MAX ? PacClipMaxScatterUpdateSingleFloat : PacClipMinScatterUpdateSingleFloat, combine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, IS_MAX ? PacClipMaxUpdateSingleFloat : PacClipMinUpdateSingleFloat, PacClipBindFloat)); set.AddFunction(AggregateFunction( - name, {LogicalType::UBIGINT, LogicalType::FLOAT, LogicalType::DOUBLE}, list_type, state_size, init, - IS_MAX ? PacClipMaxScatterUpdateSingleFloat : PacClipMinScatterUpdateSingleFloat, combine, finalize, + Identifier(name), {LogicalType::UBIGINT, LogicalType::FLOAT, LogicalType::DOUBLE}, list_type, state_size, + init, IS_MAX ? PacClipMaxScatterUpdateSingleFloat : PacClipMinScatterUpdateSingleFloat, combine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, IS_MAX ? PacClipMaxUpdateSingleFloat : PacClipMinUpdateSingleFloat, PacClipBindFloat)); // DOUBLE set.AddFunction( - AggregateFunction(name, {LogicalType::UBIGINT, LogicalType::DOUBLE}, list_type, state_size, init, - IS_MAX ? PacClipMaxScatterUpdateSingleDouble : PacClipMinScatterUpdateSingleDouble, + AggregateFunction(Identifier(name), {LogicalType::UBIGINT, LogicalType::DOUBLE}, list_type, state_size, + init, IS_MAX ? PacClipMaxScatterUpdateSingleDouble : PacClipMinScatterUpdateSingleDouble, combine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, IS_MAX ? PacClipMaxUpdateSingleDouble : PacClipMinUpdateSingleDouble, PacClipBindDouble)); set.AddFunction(AggregateFunction( - name, {LogicalType::UBIGINT, LogicalType::DOUBLE, LogicalType::DOUBLE}, list_type, state_size, init, - IS_MAX ? PacClipMaxScatterUpdateSingleDouble : PacClipMinScatterUpdateSingleDouble, combine, finalize, + Identifier(name), {LogicalType::UBIGINT, LogicalType::DOUBLE, LogicalType::DOUBLE}, list_type, state_size, + init, IS_MAX ? PacClipMaxScatterUpdateSingleDouble : PacClipMinScatterUpdateSingleDouble, combine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, IS_MAX ? PacClipMaxUpdateSingleDouble : PacClipMinUpdateSingleDouble, PacClipBindDouble)); } else { @@ -716,28 +732,28 @@ static void AddFloatDoubleOverloads(AggregateFunctionSet &set, const string &nam auto double_finalize = IS_MAX ? PacClipMaxNoisedFinalizeDouble : PacClipMinNoisedFinalizeDouble; // FLOAT → FLOAT - set.AddFunction( - AggregateFunction(name, {LogicalType::UBIGINT, LogicalType::FLOAT}, LogicalType::FLOAT, state_size, init, - IS_MAX ? PacClipMaxScatterUpdateSingleFloat : PacClipMinScatterUpdateSingleFloat, combine, - float_finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, - IS_MAX ? PacClipMaxUpdateSingleFloat : PacClipMinUpdateSingleFloat, PacClipBindFloat)); set.AddFunction(AggregateFunction( - name, {LogicalType::UBIGINT, LogicalType::FLOAT, LogicalType::DOUBLE}, LogicalType::FLOAT, state_size, init, + Identifier(name), {LogicalType::UBIGINT, LogicalType::FLOAT}, LogicalType::FLOAT, state_size, init, IS_MAX ? PacClipMaxScatterUpdateSingleFloat : PacClipMinScatterUpdateSingleFloat, combine, float_finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, IS_MAX ? PacClipMaxUpdateSingleFloat : PacClipMinUpdateSingleFloat, PacClipBindFloat)); + set.AddFunction(AggregateFunction( + Identifier(name), {LogicalType::UBIGINT, LogicalType::FLOAT, LogicalType::DOUBLE}, LogicalType::FLOAT, + state_size, init, IS_MAX ? PacClipMaxScatterUpdateSingleFloat : PacClipMinScatterUpdateSingleFloat, combine, + float_finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, + IS_MAX ? PacClipMaxUpdateSingleFloat : PacClipMinUpdateSingleFloat, PacClipBindFloat)); // DOUBLE → DOUBLE - set.AddFunction( - AggregateFunction(name, {LogicalType::UBIGINT, LogicalType::DOUBLE}, LogicalType::DOUBLE, state_size, init, - IS_MAX ? PacClipMaxScatterUpdateSingleDouble : PacClipMinScatterUpdateSingleDouble, - combine, double_finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, - IS_MAX ? PacClipMaxUpdateSingleDouble : PacClipMinUpdateSingleDouble, PacClipBindDouble)); set.AddFunction(AggregateFunction( - name, {LogicalType::UBIGINT, LogicalType::DOUBLE, LogicalType::DOUBLE}, LogicalType::DOUBLE, state_size, - init, IS_MAX ? PacClipMaxScatterUpdateSingleDouble : PacClipMinScatterUpdateSingleDouble, combine, + Identifier(name), {LogicalType::UBIGINT, LogicalType::DOUBLE}, LogicalType::DOUBLE, state_size, init, + IS_MAX ? PacClipMaxScatterUpdateSingleDouble : PacClipMinScatterUpdateSingleDouble, combine, double_finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, IS_MAX ? PacClipMaxUpdateSingleDouble : PacClipMinUpdateSingleDouble, PacClipBindDouble)); + set.AddFunction(AggregateFunction( + Identifier(name), {LogicalType::UBIGINT, LogicalType::DOUBLE, LogicalType::DOUBLE}, LogicalType::DOUBLE, + state_size, init, IS_MAX ? PacClipMaxScatterUpdateSingleDouble : PacClipMinScatterUpdateSingleDouble, + combine, double_finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, + IS_MAX ? PacClipMaxUpdateSingleDouble : PacClipMinUpdateSingleDouble, PacClipBindDouble)); } } @@ -748,7 +764,7 @@ template static void RegisterPacClipMinMaxCountersFunctions(ExtensionLoader &loader) { const string name = IS_MAX ? "as_clip_max" : "as_clip_min"; const string short_name = IS_MAX ? "clip_max" : "clip_min"; - AggregateFunctionSet fcn_set(name); + AggregateFunctionSet fcn_set {Identifier(name)}; RegisterClipMinMaxTypeOverloads(fcn_set, name, true); // DECIMAL overloads @@ -771,7 +787,7 @@ static void RegisterPacClipMinMaxCountersFunctions(ExtensionLoader &loader) { template static void RegisterPacNoisedClipMinMaxFunctions(ExtensionLoader &loader) { const string name = IS_MAX ? "as_noised_clip_max" : "as_noised_clip_min"; - AggregateFunctionSet fcn_set(name); + AggregateFunctionSet fcn_set {Identifier(name)}; RegisterClipMinMaxTypeOverloads(fcn_set, name, false); // DECIMAL overloads diff --git a/src/aggregates/as_clip_sum.cpp b/src/aggregates/as_clip_sum.cpp index 30b460a1..6a4cc29c 100644 --- a/src/aggregates/as_clip_sum.cpp +++ b/src/aggregates/as_clip_sum.cpp @@ -1,3 +1,4 @@ +#include "duckdb/common/vector/list_vector.hpp" #include "aggregates/as_clip_sum.hpp" #include "categorical/pac_categorical.hpp" #include "duckdb/common/types/decimal.hpp" @@ -200,12 +201,12 @@ template &state, idx_t count, ArenaAllocator &allocator, int sample_lanes = 0) { UnifiedVectorFormat hash_data, value_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); - if (hash_data.validity.AllValid() && value_data.validity.AllValid()) { + if (hash_data.validity.CannotHaveNull() && value_data.validity.CannotHaveNull()) { for (idx_t i = 0; i < count; i++) { auto h_idx = hash_data.sel->get_index(i); auto v_idx = value_data.sel->get_index(i); @@ -230,9 +231,9 @@ template (hash_data); auto values = UnifiedVectorFormat::GetData(value_data); @@ -266,11 +267,13 @@ static void PacClipSumScatterUpdate(Vector inputs[], Vector &states, idx_t count X(UBigInt, uint64_t, uint64_t, false) #define X(NAME, VALUE_T, INPUT_T, SIGNED) \ - static void PacClipSumUpdate##NAME(Vector input[], AggregateInputData &agg, idx_t, data_ptr_t state_p, \ - idx_t cnt) { \ - auto &state = *reinterpret_cast *>(state_p); \ - PacClipSumUpdate(input, state, cnt, agg.allocator, \ - GetPrivSampleLanes(agg)); \ + static void PacClipSumUpdate##NAME(Vector input[], AggregateInputData &agg, idx_t input_count, \ + const ClusteredAggr &clusters, idx_t cnt) { \ + ExecutePrivClusters(input, input_count, clusters, [&](Vector input[], data_ptr_t state_p, idx_t cnt) { \ + auto &state = *reinterpret_cast *>(state_p); \ + PacClipSumUpdate(input, state, cnt, agg.allocator, \ + GetPrivSampleLanes(agg)); \ + }); \ } \ static void PacClipSumScatterUpdate##NAME(Vector input[], AggregateInputData &agg, idx_t, Vector &sts, \ idx_t cnt) { \ @@ -282,30 +285,34 @@ CLIP_SUM_INT_TYPES_UNSIGNED #undef X // HugeInt update (signed, via hugeint routing — 128-bit needs full 62 levels) -static void PacClipSumUpdateHugeInt(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_p, idx_t count) { - auto &state = *reinterpret_cast *>(state_p); - int sample_lanes = GetPrivSampleLanes(aggr); - UnifiedVectorFormat hash_data, value_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto values = UnifiedVectorFormat::GetData(value_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto v_idx = value_data.sel->get_index(i); - if (!hash_data.validity.RowIsValid(h_idx) || !value_data.validity.RowIsValid(v_idx)) { - continue; +static void PacClipSumUpdateHugeInt(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { + auto &state = *reinterpret_cast *>(state_p); + int sample_lanes = GetPrivSampleLanes(aggr); + UnifiedVectorFormat hash_data, value_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto values = UnifiedVectorFormat::GetData(value_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto v_idx = value_data.sel->get_index(i); + if (!hash_data.validity.RowIsValid(h_idx) || !value_data.validity.RowIsValid(v_idx)) { + continue; + } + PacClipSumUpdateOne(state, hashes[h_idx], values[v_idx], aggr.allocator, + sample_lanes); } - PacClipSumUpdateOne(state, hashes[h_idx], values[v_idx], aggr.allocator, sample_lanes); - } + }); } static void PacClipSumScatterUpdateHugeInt(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { int sample_lanes = GetPrivSampleLanes(aggr); UnifiedVectorFormat hash_data, value_data, sdata; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - states.ToUnifiedFormat(count, sdata); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + states.ToUnifiedFormat(sdata); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); auto state_ptrs = UnifiedVectorFormat::GetData *>(sdata); @@ -321,56 +328,58 @@ static void PacClipSumScatterUpdateHugeInt(Vector inputs[], AggregateInputData & } // UHugeInt update (unsigned, convert to hugeint for routing — 128-bit needs full 62 levels) -static void PacClipSumUpdateUHugeInt(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_p, - idx_t count) { - auto &state = *reinterpret_cast *>(state_p); - int sample_lanes = GetPrivSampleLanes(aggr); - UnifiedVectorFormat hash_data, value_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto values = UnifiedVectorFormat::GetData(value_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto v_idx = value_data.sel->get_index(i); - if (!hash_data.validity.RowIsValid(h_idx) || !value_data.validity.RowIsValid(v_idx)) { - continue; - } - // uhugeint_t is always positive; treat as 128-bit unsigned - auto &v = values[v_idx]; - auto *pos_state = state.EnsureState(aggr.allocator); - if (DUCKDB_LIKELY(hashes[h_idx])) { - uint64_t counter_hash = TransformPacUpdateHash(hashes[h_idx], sample_lanes); - uint64_t upper = static_cast(v.upper); - uint64_t lower = v.lower; - int level = PacClipSumIntState::GetLevel128(upper, lower); - uint64_t shift = level << 1; - uint16_t shifted_val; - if (shift >= 64) { - shifted_val = static_cast(upper >> (shift - 64)); - } else if (shift > 0) { - shifted_val = static_cast((lower >> shift) | (upper << (64 - shift))); - } else { - shifted_val = static_cast(lower); +static void PacClipSumUpdateUHugeInt(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { + auto &state = *reinterpret_cast *>(state_p); + int sample_lanes = GetPrivSampleLanes(aggr); + UnifiedVectorFormat hash_data, value_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto values = UnifiedVectorFormat::GetData(value_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto v_idx = value_data.sel->get_index(i); + if (!hash_data.validity.RowIsValid(h_idx) || !value_data.validity.RowIsValid(v_idx)) { + continue; + } + // uhugeint_t is always positive; treat as 128-bit unsigned + auto &v = values[v_idx]; + auto *pos_state = state.EnsureState(aggr.allocator); + if (DUCKDB_LIKELY(hashes[h_idx])) { + uint64_t counter_hash = TransformPacUpdateHash(hashes[h_idx], sample_lanes); + uint64_t upper = static_cast(v.upper); + uint64_t lower = v.lower; + int level = PacClipSumIntState::GetLevel128(upper, lower); + uint64_t shift = level << 1; + uint16_t shifted_val; + if (shift >= 64) { + shifted_val = static_cast(upper >> (shift - 64)); + } else if (shift > 0) { + shifted_val = static_cast((lower >> shift) | (upper << (64 - shift))); + } else { + shifted_val = static_cast(lower); + } + shifted_val &= 0xFF; + pos_state->key_hash |= counter_hash; + pos_state->EnsureLevelAllocated(aggr.allocator, level); + uint64_t *buf = pos_state->levels[level]; + buf[17] |= (1ULL << (hashes[h_idx] >> 58)); + pos_state->AddToExactCount(buf, shifted_val, aggr.allocator); + Pac2AddToTotalsSWAR16(buf, shifted_val, counter_hash); + pos_state->update_count++; } - shifted_val &= 0xFF; - pos_state->key_hash |= counter_hash; - pos_state->EnsureLevelAllocated(aggr.allocator, level); - uint64_t *buf = pos_state->levels[level]; - buf[17] |= (1ULL << (hashes[h_idx] >> 58)); - pos_state->AddToExactCount(buf, shifted_val, aggr.allocator); - Pac2AddToTotalsSWAR16(buf, shifted_val, counter_hash); - pos_state->update_count++; } - } + }); } static void PacClipSumScatterUpdateUHugeInt(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { int sample_lanes = GetPrivSampleLanes(aggr); UnifiedVectorFormat hash_data, value_data, sdata; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - states.ToUnifiedFormat(count, sdata); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + states.ToUnifiedFormat(sdata); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); auto state_ptrs = UnifiedVectorFormat::GetData *>(sdata); @@ -413,35 +422,38 @@ static void PacClipSumScatterUpdateUHugeInt(Vector inputs[], AggregateInputData // Float/Double update: scale to int64, route through signed path // ============================================================================ template -static void PacClipSumUpdateFloat(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_p, idx_t count) { - auto &state = *reinterpret_cast *>(state_p); - int sample_lanes = GetPrivSampleLanes(aggr); - UnifiedVectorFormat hash_data, value_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto values = UnifiedVectorFormat::GetData(value_data); - - if (hash_data.validity.AllValid() && value_data.validity.AllValid()) { - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto v_idx = value_data.sel->get_index(i); - PacClipSumUpdateOne( - state, hashes[h_idx], ScaleFloatToInt64(values[v_idx]), aggr.allocator, - sample_lanes); - } - } else { - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto v_idx = value_data.sel->get_index(i); - if (!hash_data.validity.RowIsValid(h_idx) || !value_data.validity.RowIsValid(v_idx)) { - continue; +static void PacClipSumUpdateFloat(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { + auto &state = *reinterpret_cast *>(state_p); + int sample_lanes = GetPrivSampleLanes(aggr); + UnifiedVectorFormat hash_data, value_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto values = UnifiedVectorFormat::GetData(value_data); + + if (hash_data.validity.CannotHaveNull() && value_data.validity.CannotHaveNull()) { + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto v_idx = value_data.sel->get_index(i); + PacClipSumUpdateOne( + state, hashes[h_idx], ScaleFloatToInt64(values[v_idx]), aggr.allocator, + sample_lanes); + } + } else { + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto v_idx = value_data.sel->get_index(i); + if (!hash_data.validity.RowIsValid(h_idx) || !value_data.validity.RowIsValid(v_idx)) { + continue; + } + PacClipSumUpdateOne( + state, hashes[h_idx], ScaleFloatToInt64(values[v_idx]), aggr.allocator, + sample_lanes); } - PacClipSumUpdateOne( - state, hashes[h_idx], ScaleFloatToInt64(values[v_idx]), aggr.allocator, - sample_lanes); } - } + }); } template @@ -449,9 +461,9 @@ static void PacClipSumScatterUpdateFloat(Vector inputs[], AggregateInputData &ag idx_t count) { int sample_lanes = GetPrivSampleLanes(aggr); UnifiedVectorFormat hash_data, value_data, sdata; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - states.ToUnifiedFormat(count, sdata); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + states.ToUnifiedFormat(sdata); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); auto state_ptrs = UnifiedVectorFormat::GetData *>(sdata); @@ -469,26 +481,26 @@ static void PacClipSumScatterUpdateFloat(Vector inputs[], AggregateInputData &ag } // Instantiate float/double update functions -static void PacClipSumUpdateSingleFloat(Vector inputs[], AggregateInputData &aggr, idx_t n, data_ptr_t state_p, - idx_t count) { - PacClipSumUpdateFloat(inputs, aggr, n, state_p, count); +static void PacClipSumUpdateSingleFloat(Vector inputs[], AggregateInputData &aggr, idx_t n, + const ClusteredAggr &clusters, idx_t count) { + PacClipSumUpdateFloat(inputs, aggr, n, clusters, count); } static void PacClipSumScatterUpdateSingleFloat(Vector inputs[], AggregateInputData &aggr, idx_t n, Vector &states, idx_t count) { PacClipSumScatterUpdateFloat(inputs, aggr, n, states, count); } -static void PacClipSumUpdateSingleDouble(Vector inputs[], AggregateInputData &aggr, idx_t n, data_ptr_t state_p, - idx_t count) { - PacClipSumUpdateFloat(inputs, aggr, n, state_p, count); +static void PacClipSumUpdateSingleDouble(Vector inputs[], AggregateInputData &aggr, idx_t n, + const ClusteredAggr &clusters, idx_t count) { + PacClipSumUpdateFloat(inputs, aggr, n, clusters, count); } static void PacClipSumScatterUpdateSingleDouble(Vector inputs[], AggregateInputData &aggr, idx_t n, Vector &states, idx_t count) { PacClipSumScatterUpdateFloat(inputs, aggr, n, states, count); } -static void DpSampleClipSumUpdateDouble(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_p, - idx_t count) { - PacClipSumUpdateFloat(inputs, aggr, 0, state_p, count); +static void DpSampleClipSumUpdateDouble(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + PacClipSumUpdateFloat(inputs, aggr, input_count, clusters, count); } static void DpSampleClipSumScatterUpdateDouble(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, @@ -545,10 +557,11 @@ static void DpSampleClipSumCombine(Vector &src, Vector &dst, AggregateInputData // Finalize // ============================================================================ template -static void PacClipSumFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, idx_t offset) { +static void PacClipSumFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, + idx_t offset) { auto state_ptrs = FlatVector::GetData *>(states); - auto data = FlatVector::GetData(result); - auto &result_mask = FlatVector::Validity(result); + auto data = FlatVector::GetDataMutable(result); + auto &result_mask = FlatVector::ValidityMutable(result); auto &bind = static_cast(*input.bind_data); double mi = bind.mi; double correction = bind.correction; @@ -557,6 +570,7 @@ static void PacClipSumFinalize(Vector &states, AggregateInputData &input, Vector int clip_support = bind.clip_support_threshold; bool clip_scale = bind.clip_scale; + auto gen = PacRandomStream(bind.seed, *bind.random_state); for (idx_t i = 0; i < count; i++) { PacClipSumFlushBuffer(*state_ptrs[i], *state_ptrs[i], input.allocator, bind.sample_lanes); @@ -567,7 +581,6 @@ static void PacClipSumFinalize(Vector &states, AggregateInputData &input, Vector continue; } uint64_t key_hash = pos->key_hash; - std::mt19937_64 gen(bind.seed); if (PacNoiseInNull(key_hash, mi, correction, gen)) { result_mask.SetInvalid(offset + i); continue; @@ -607,35 +620,35 @@ static void PacClipSumFinalize(Vector &states, AggregateInputData &input, Vector // Instantiate noised finalize (scalar output for as_noised_clip_sum) // 64-bit types -static void PacClipSumNoisedFinalizeSigned(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void PacClipSumNoisedFinalizeSigned(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { PacClipSumFinalize(states, input, result, count, offset); } -static void PacClipSumNoisedFinalizeUnsigned(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void PacClipSumNoisedFinalizeUnsigned(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { PacClipSumFinalize(states, input, result, count, offset); } // 128-bit types -static void PacClipSumNoisedFinalizeSigned128(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void PacClipSumNoisedFinalizeSigned128(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { PacClipSumFinalize(states, input, result, count, offset); } -static void PacClipSumNoisedFinalizeUnsigned128(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void PacClipSumNoisedFinalizeUnsigned128(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { PacClipSumFinalize(states, input, result, count, offset); } // BIGINT output variant — used for count→sum conversion where the original returned BIGINT -static void PacClipSumNoisedFinalizeBigInt(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void PacClipSumNoisedFinalizeBigInt(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { PacClipSumFinalize(states, input, result, count, offset); } // Float/double output variants -static void PacClipSumNoisedFinalizeFloat(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void PacClipSumNoisedFinalizeFloat(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { PacClipSumFinalize(states, input, result, count, offset); } -static void PacClipSumNoisedFinalizeDouble(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void PacClipSumNoisedFinalizeDouble(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { PacClipSumFinalize(states, input, result, count, offset); } @@ -646,7 +659,7 @@ enum class ClipCountersFinalizeMode : uint8_t { PAC, DP_SAMPLE }; // ============================================================================ template -static void PacClipSumFinalizeCounters(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void PacClipSumFinalizeCounters(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { auto state_ptrs = FlatVector::GetData *>(states); auto &bind = static_cast(*input.bind_data); @@ -657,14 +670,14 @@ static void PacClipSumFinalizeCounters(Vector &states, AggregateInputData &input bool dp_sample = MODE == ClipCountersFinalizeMode::DP_SAMPLE; double dp_sample_rescale = dp_sample && bind.sample_rescale ? DpSampleRescale(bind.sample_lanes) : 1.0; - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * 64; ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); for (idx_t i = 0; i < count; i++) { PacClipSumFlushBuffer(*state_ptrs[i], *state_ptrs[i], input.allocator, bind.sample_lanes); @@ -718,25 +731,25 @@ static void PacClipSumFinalizeCounters(Vector &states, AggregateInputData &input } // 64-bit counters finalize -static void PacClipSumFinalizeCountersSigned(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void PacClipSumFinalizeCountersSigned(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { PacClipSumFinalizeCounters(states, input, result, count, offset); } -static void PacClipSumFinalizeCountersUnsigned(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void PacClipSumFinalizeCountersUnsigned(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { PacClipSumFinalizeCounters(states, input, result, count, offset); } // 128-bit counters finalize -static void PacClipSumFinalizeCountersSigned128(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void PacClipSumFinalizeCountersSigned128(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { PacClipSumFinalizeCounters(states, input, result, count, offset); } -static void PacClipSumFinalizeCountersUnsigned128(Vector &states, AggregateInputData &input, Vector &result, +static void PacClipSumFinalizeCountersUnsigned128(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { PacClipSumFinalizeCounters(states, input, result, count, offset); } -static void DpSampleClipSumFinalizeCounters(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void DpSampleClipSumFinalizeCounters(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { PacClipSumFinalizeCounters(states, input, result, count, offset); } @@ -744,18 +757,24 @@ static void DpSampleClipSumFinalizeCounters(Vector &states, AggregateInputData & // ============================================================================ // State size / init / bind // ============================================================================ -static idx_t PacClipSumStateSize(const AggregateFunction &) { +static idx_t PacClipSumStateSize(AggregateStateInput &) { return sizeof(PacClipSumStateWrapper<>); } -static idx_t PacClipSumStateSize128(const AggregateFunction &) { +static idx_t PacClipSumStateSize128(AggregateStateInput &) { return sizeof(PacClipSumStateWrapper); } -static void PacClipSumInitialize(const AggregateFunction &, data_ptr_t state_p) { - memset(state_p, 0, sizeof(PacClipSumStateWrapper<>)); +static void PacClipSumInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_p = states[i]; + memset(state_p, 0, sizeof(PacClipSumStateWrapper<>)); + } } -static void PacClipSumInitialize128(const AggregateFunction &, data_ptr_t state_p) { - memset(state_p, 0, sizeof(PacClipSumStateWrapper)); +static void PacClipSumInitialize128(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_p = states[i]; + memset(state_p, 0, sizeof(PacClipSumStateWrapper)); + } } // PacClipBind, PacClipBindFloat, PacClipBindDouble are defined in as_clip_aggr.hpp @@ -818,31 +837,35 @@ static AggregateFunction GetPacClipSumCountersAggregate(PhysicalType type) { } } -static unique_ptr BindDecimalPacNoisedClipSum(ClientContext &ctx, AggregateFunction &function, - vector> &args) { - auto decimal_type = args[1]->return_type; - function = GetPacClipSumNoisedAggregate(decimal_type.InternalType()); - function.name = "as_noised_clip_sum"; - function.arguments[1] = decimal_type; - function.return_type = LogicalType::DECIMAL(Decimal::MAX_WIDTH_DECIMAL, DecimalType::GetScale(decimal_type)); - return PacClipBind(ctx, function, args); +static unique_ptr BindDecimalPacNoisedClipSum(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &function = bind_input.GetBoundFunction(); + auto &args = bind_input.GetArguments(); + auto decimal_type = args[1]->GetReturnType(); + function.ReplaceImplementation(GetPacClipSumNoisedAggregate(decimal_type.InternalType())); + function.SetName("as_noised_clip_sum"); + function.GetArguments()[1] = decimal_type; + function.SetReturnType(LogicalType::DECIMAL(Decimal::MAX_WIDTH_DECIMAL, DecimalType::GetScale(decimal_type))); + return PacClipBind(bind_input); } -static unique_ptr BindDecimalPacClipSum(ClientContext &ctx, AggregateFunction &function, - vector> &args) { - auto decimal_type = args[1]->return_type; - function = GetPacClipSumCountersAggregate(decimal_type.InternalType()); - function.name = "as_clip_sum"; - function.arguments[1] = decimal_type; +static unique_ptr BindDecimalPacClipSum(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &function = bind_input.GetBoundFunction(); + auto &args = bind_input.GetArguments(); + auto decimal_type = args[1]->GetReturnType(); + function.ReplaceImplementation(GetPacClipSumCountersAggregate(decimal_type.InternalType())); + function.SetName("as_clip_sum"); + function.GetArguments()[1] = decimal_type; // counters always return LIST, no DECIMAL return type needed - return PacClipBind(ctx, function, args); + return PacClipBind(bind_input); } -static unique_ptr DpSampleClipSumBind(ClientContext &ctx, AggregateFunction &, - vector> &) { +static unique_ptr DpSampleClipSumBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); Value dc_val; int clip_support = 0; - if (ctx.TryGetCurrentSetting("priv_clip_support", dc_val) && !dc_val.IsNull()) { + if (ctx.TryGetCurrentSetting(Identifier("priv_clip_support"), dc_val) && !dc_val.IsNull()) { clip_support = static_cast(dc_val.GetValue()); } if (clip_support <= 0) { @@ -857,25 +880,26 @@ static unique_ptr DpSampleClipSumBind(ClientContext &ctx, Aggregat // ============================================================================ static void AddClipSumCountersFcn(AggregateFunctionSet &set, const string &name, const LogicalType &value_type, aggregate_update_t scatter, aggregate_finalize_t finalize, - aggregate_simple_update_t update) { + aggregate_cluster_update_t update) { auto list_type = LogicalType::LIST(PacFloatLogicalType()); - set.AddFunction(AggregateFunction(name, {LogicalType::UBIGINT, value_type}, list_type, PacClipSumStateSize, - PacClipSumInitialize, scatter, PacClipSumCombine, finalize, - FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacClipBind)); - set.AddFunction(AggregateFunction(name, {LogicalType::UBIGINT, value_type, LogicalType::DOUBLE}, list_type, + set.AddFunction(AggregateFunction(Identifier(name), {LogicalType::UBIGINT, value_type}, list_type, PacClipSumStateSize, PacClipSumInitialize, scatter, PacClipSumCombine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacClipBind)); + set.AddFunction(AggregateFunction(Identifier(name), {LogicalType::UBIGINT, value_type, LogicalType::DOUBLE}, + list_type, PacClipSumStateSize, PacClipSumInitialize, scatter, PacClipSumCombine, + finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacClipBind)); } static void AddNoisedClipSumFcn(AggregateFunctionSet &set, const string &name, const LogicalType &value_type, const LogicalType &result_type, aggregate_update_t scatter, - aggregate_finalize_t finalize, aggregate_simple_update_t update) { - set.AddFunction(AggregateFunction(name, {LogicalType::UBIGINT, value_type}, result_type, PacClipSumStateSize, - PacClipSumInitialize, scatter, PacClipSumCombine, finalize, - FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacClipBind)); - set.AddFunction(AggregateFunction(name, {LogicalType::UBIGINT, value_type, LogicalType::DOUBLE}, result_type, + aggregate_finalize_t finalize, aggregate_cluster_update_t update) { + set.AddFunction(AggregateFunction(Identifier(name), {LogicalType::UBIGINT, value_type}, result_type, PacClipSumStateSize, PacClipSumInitialize, scatter, PacClipSumCombine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacClipBind)); + set.AddFunction(AggregateFunction(Identifier(name), {LogicalType::UBIGINT, value_type, LogicalType::DOUBLE}, + result_type, PacClipSumStateSize, PacClipSumInitialize, scatter, + PacClipSumCombine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, update, + PacClipBind)); } // Helper to register all type overloads for a clip sum function set @@ -903,15 +927,16 @@ static void RegisterClipSumTypeOverloads(AggregateFunctionSet &set, const string // HUGEINT/UHUGEINT: use 128-bit state (62 levels) { auto lt = LogicalType::LIST(PacFloatLogicalType()); - set.AddFunction(AggregateFunction( - name, {LogicalType::UBIGINT, LogicalType::HUGEINT}, lt, PacClipSumStateSize128, PacClipSumInitialize128, - PacClipSumScatterUpdateHugeInt, PacClipSumCombine128, PacClipSumFinalizeCountersSigned128, - FunctionNullHandling::DEFAULT_NULL_HANDLING, PacClipSumUpdateHugeInt, PacClipBind)); set.AddFunction( - AggregateFunction(name, {LogicalType::UBIGINT, LogicalType::UHUGEINT}, lt, PacClipSumStateSize128, - PacClipSumInitialize128, PacClipSumScatterUpdateUHugeInt, PacClipSumCombine128, - PacClipSumFinalizeCountersUnsigned128, FunctionNullHandling::DEFAULT_NULL_HANDLING, - PacClipSumUpdateUHugeInt, PacClipBind)); + AggregateFunction(Identifier(name), {LogicalType::UBIGINT, LogicalType::HUGEINT}, lt, + PacClipSumStateSize128, PacClipSumInitialize128, PacClipSumScatterUpdateHugeInt, + PacClipSumCombine128, PacClipSumFinalizeCountersSigned128, + FunctionNullHandling::DEFAULT_NULL_HANDLING, PacClipSumUpdateHugeInt, PacClipBind)); + set.AddFunction( + AggregateFunction(Identifier(name), {LogicalType::UBIGINT, LogicalType::UHUGEINT}, lt, + PacClipSumStateSize128, PacClipSumInitialize128, PacClipSumScatterUpdateUHugeInt, + PacClipSumCombine128, PacClipSumFinalizeCountersUnsigned128, + FunctionNullHandling::DEFAULT_NULL_HANDLING, PacClipSumUpdateUHugeInt, PacClipBind)); } } else { // Noised (scalar HUGEINT) variants @@ -935,12 +960,12 @@ static void RegisterClipSumTypeOverloads(AggregateFunctionSet &set, const string PacClipSumNoisedFinalizeUnsigned, PacClipSumUpdateUBigInt); // HUGEINT/UHUGEINT: use 128-bit state (62 levels) set.AddFunction( - AggregateFunction(name, {LogicalType::UBIGINT, LogicalType::HUGEINT}, LogicalType::HUGEINT, + AggregateFunction(Identifier(name), {LogicalType::UBIGINT, LogicalType::HUGEINT}, LogicalType::HUGEINT, PacClipSumStateSize128, PacClipSumInitialize128, PacClipSumScatterUpdateHugeInt, PacClipSumCombine128, PacClipSumNoisedFinalizeSigned128, FunctionNullHandling::DEFAULT_NULL_HANDLING, PacClipSumUpdateHugeInt, PacClipBind)); set.AddFunction( - AggregateFunction(name, {LogicalType::UBIGINT, LogicalType::UHUGEINT}, LogicalType::HUGEINT, + AggregateFunction(Identifier(name), {LogicalType::UBIGINT, LogicalType::UHUGEINT}, LogicalType::HUGEINT, PacClipSumStateSize128, PacClipSumInitialize128, PacClipSumScatterUpdateUHugeInt, PacClipSumCombine128, PacClipSumNoisedFinalizeUnsigned128, FunctionNullHandling::DEFAULT_NULL_HANDLING, PacClipSumUpdateUHugeInt, PacClipBind)); diff --git a/src/aggregates/as_count.cpp b/src/aggregates/as_count.cpp index 8aa069b5..d796d15e 100644 --- a/src/aggregates/as_count.cpp +++ b/src/aggregates/as_count.cpp @@ -1,18 +1,23 @@ +#include "duckdb/common/vector/list_vector.hpp" #include #include "aggregates/as_count.hpp" +#include "aggregates/aggregate_state_export.hpp" #include "categorical/pac_categorical.hpp" #include "duckdb/common/types/value_map.hpp" #include "duckdb/parser/parsed_data/create_aggregate_function_info.hpp" namespace duckdb { -static unique_ptr PacCountBind(ClientContext &ctx, AggregateFunction &func, - vector> &args) { +static unique_ptr PacCountBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &func = bind_input.GetBoundFunction(); + auto &args = bind_input.GetArguments(); // Correction is always the last arg when it's a foldable numeric constant idx_t corr_idx = args.size(); // sentinel: no correction if (args.size() >= 2) { auto &last = args.back(); - if (last->IsFoldable() && (last->return_type.IsNumeric() || last->return_type.id() == LogicalTypeId::UNKNOWN)) { + if (last->IsFoldable() && + (last->GetReturnType().IsNumeric() || last->GetReturnType().id() == LogicalTypeId::UNKNOWN)) { corr_idx = args.size() - 1; } } @@ -27,60 +32,158 @@ using ScatterState = PacCountState; using ScatterState = PacCountStateWrapper; #endif +// Preserve SWAR bytes and overflow totals separately: exporting must not flush or +// mutate a live aggregate, and the totals pointer is local to the process. +struct CountStateCodec { + static void WriteInner(AggregateStateWriter &out, const PacCountState &state) { +#ifndef PAC_NOCASCADING + out.Write(state.probabilistic_total8); + out.Write(state.swar_fill); +#endif + out.Write(state.key_hash); + out.Write(state.update_count); + out.Write(uint8_t(state.probabilistic_total != nullptr)); + if (state.probabilistic_total) { + out.WriteData(state.probabilistic_total, 64 * sizeof(uint64_t)); + } + } + static void ReadInner(AggregateStateReader &in, PacCountState &state, ArenaAllocator &allocator) { +#ifndef PAC_NOCASCADING + in.Read(state.probabilistic_total8); + in.Read(state.swar_fill); + if (state.swar_fill >= 255) { + throw InvalidInputException("Invalid stochastic COUNT SWAR fill"); + } +#endif + in.Read(state.key_hash); + in.Read(state.update_count); + if (in.ReadFlag()) { + in.ReadData(state.EnsureTotals(allocator), 64 * sizeof(uint64_t)); + } + } + static void Write(AggregateStateWriter &out, const ScatterState &state) { +#if defined(PAC_NOBUFFERING) || defined(PAC_NOCASCADING) + WriteInner(out, state); +#else + auto buffered = state.n_buffered & 7; + out.Write(uint8_t(buffered)); + for (idx_t i = 0; i < buffered; i++) { + out.Write(state.hash_buf[i]); + } + auto *inner = state.GetState(); + out.Write(uint8_t(inner != nullptr)); + if (inner) { + WriteInner(out, *inner); + } +#endif + } + static void Read(AggregateStateReader &in, ScatterState &state, ArenaAllocator &allocator) { + memset(&state, 0, sizeof(state)); +#if defined(PAC_NOBUFFERING) || defined(PAC_NOCASCADING) + ReadInner(in, state, allocator); +#else + uint8_t buffered; + in.Read(buffered); + if (buffered > 3) { + throw InvalidInputException("Invalid stochastic COUNT buffer size"); + } + for (idx_t i = 0; i < buffered; i++) { + in.Read(state.hash_buf[i]); + } + if (in.ReadFlag()) { + ReadInner(in, *state.EnsureState(allocator), allocator); + } + state.n_buffered |= buffered; + if (buffered) { + state.FlushBuffer(state, allocator); + } +#endif + } +}; + // State functions - uses ScatterState for both simple and scatter -static idx_t PacCountStateSize(const AggregateFunction &) { +static idx_t PacCountStateSize(AggregateStateInput &) { return sizeof(ScatterState); } -static void PacCountInitialize(const AggregateFunction &, data_ptr_t state_ptr) { - memset(state_ptr, 0, sizeof(ScatterState)); +static void PacCountInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_ptr = states[i]; + memset(state_ptr, 0, sizeof(ScatterState)); + } } -void PacCountUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, idx_t count) { - ScatterState &agg = *reinterpret_cast(state_ptr); - int sample_lanes = GetPrivSampleLanes(aggr); - UnifiedVectorFormat idata; - inputs[0].ToUnifiedFormat(count, idata); - auto input_data = UnifiedVectorFormat::GetData(idata); - for (idx_t i = 0; i < count; i++) { - auto idx = idata.sel->get_index(i); - if (idata.validity.RowIsValid(idx)) { - auto key_hash = TransformPacUpdateHash(input_data[idx], sample_lanes); +static AggregateStateLayout CountStateLayout(AggregateLayoutInput &input) { + auto layout = AggregateStateExport::Layout(input); + // COUNT's ANY argument affects updates only. Rebinding a finalized state + // needs its folded correction, including the two-argument shorthand. + if (input.function.GetLogicalArguments().size() == 2 && input.bind_data) { + layout.constant_parameters.emplace(1, Value::DOUBLE(input.bind_data->Cast().correction)); + } + return layout; +} + +static void RegisterCountFunctions(ExtensionLoader &loader, CreateAggregateFunctionInfo &info) { + info.functions.ApplyToFunctions([](AggregateFunction &function) { + if (function.GetStateInitCallback() == PacCountInitialize) { + AggregateStateExport::Register(function); + function.GetCallbacks().get_state_type = CountStateLayout; + } + }); + loader.RegisterFunction(std::move(info)); +} + +void PacCountUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, const ClusteredAggr &clusters, + idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + ScatterState &agg = *reinterpret_cast(state_ptr); + int sample_lanes = GetPrivSampleLanes(aggr); + UnifiedVectorFormat idata; + inputs[0].ToUnifiedFormat(idata); + auto input_data = UnifiedVectorFormat::GetData(idata); + for (idx_t i = 0; i < count; i++) { + auto idx = idata.sel->get_index(i); + if (idata.validity.RowIsValid(idx)) { + auto key_hash = TransformPacUpdateHash(input_data[idx], sample_lanes); #if defined(PAC_NOBUFFERING) || defined(PAC_NOCASCADING) - PacCountUpdateOne(agg, key_hash, aggr.allocator); + PacCountUpdateOne(agg, key_hash, aggr.allocator); #else PacCountBufferOrUpdateOne(agg, key_hash, aggr.allocator); #endif + } } - } + }); } -void PacCountColumnUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, idx_t count) { - ScatterState &agg = *reinterpret_cast(state_ptr); - int sample_lanes = GetPrivSampleLanes(aggr); - UnifiedVectorFormat hash_data, col_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, col_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto c_idx = col_data.sel->get_index(i); - if (hash_data.validity.RowIsValid(h_idx) && col_data.validity.RowIsValid(c_idx)) { - auto key_hash = TransformPacUpdateHash(hashes[h_idx], sample_lanes); +void PacCountColumnUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, const ClusteredAggr &clusters, + idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + ScatterState &agg = *reinterpret_cast(state_ptr); + int sample_lanes = GetPrivSampleLanes(aggr); + UnifiedVectorFormat hash_data, col_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(col_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto c_idx = col_data.sel->get_index(i); + if (hash_data.validity.RowIsValid(h_idx) && col_data.validity.RowIsValid(c_idx)) { + auto key_hash = TransformPacUpdateHash(hashes[h_idx], sample_lanes); #if defined(PAC_NOBUFFERING) || defined(PAC_NOCASCADING) - PacCountUpdateOne(agg, key_hash, aggr.allocator); + PacCountUpdateOne(agg, key_hash, aggr.allocator); #else PacCountBufferOrUpdateOne(agg, key_hash, aggr.allocator); #endif + } } - } + }); } void PacCountScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { int sample_lanes = GetPrivSampleLanes(aggr); UnifiedVectorFormat idata, sdata; - inputs[0].ToUnifiedFormat(count, idata); - states.ToUnifiedFormat(count, sdata); + inputs[0].ToUnifiedFormat(idata); + states.ToUnifiedFormat(sdata); auto input_data = UnifiedVectorFormat::GetData(idata); auto state_p = UnifiedVectorFormat::GetData(sdata); for (idx_t i = 0; i < count; i++) { @@ -99,9 +202,9 @@ void PacCountScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vec void PacCountColumnScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { int sample_lanes = GetPrivSampleLanes(aggr); UnifiedVectorFormat hash_data, col_data, sdata; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, col_data); - states.ToUnifiedFormat(count, sdata); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(col_data); + states.ToUnifiedFormat(sdata); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto state_p = UnifiedVectorFormat::GetData(sdata); for (idx_t i = 0; i < count; i++) { @@ -119,29 +222,31 @@ void PacCountColumnScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_ } } -static void DpSampleCountMaskUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, - idx_t count) { - ScatterState &agg = *reinterpret_cast(state_ptr); - UnifiedVectorFormat idata; - inputs[0].ToUnifiedFormat(count, idata); - auto input_data = UnifiedVectorFormat::GetData(idata); - for (idx_t i = 0; i < count; i++) { - auto idx = idata.sel->get_index(i); - if (idata.validity.RowIsValid(idx)) { +static void DpSampleCountMaskUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + ScatterState &agg = *reinterpret_cast(state_ptr); + UnifiedVectorFormat idata; + inputs[0].ToUnifiedFormat(idata); + auto input_data = UnifiedVectorFormat::GetData(idata); + for (idx_t i = 0; i < count; i++) { + auto idx = idata.sel->get_index(i); + if (idata.validity.RowIsValid(idx)) { #if defined(PAC_NOBUFFERING) || defined(PAC_NOCASCADING) - PacCountUpdateOne(agg, input_data[idx], aggr.allocator); + PacCountUpdateOne(agg, input_data[idx], aggr.allocator); #else PacCountBufferOrUpdateOne(agg, input_data[idx], aggr.allocator); #endif + } } - } + }); } static void DpSampleCountMaskScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { UnifiedVectorFormat idata, sdata; - inputs[0].ToUnifiedFormat(count, idata); - states.ToUnifiedFormat(count, sdata); + inputs[0].ToUnifiedFormat(idata); + states.ToUnifiedFormat(sdata); auto input_data = UnifiedVectorFormat::GetData(idata); auto state_p = UnifiedVectorFormat::GetData(sdata); for (idx_t i = 0; i < count; i++) { @@ -217,24 +322,23 @@ void PacCountCombine(Vector &src, Vector &dst, AggregateInputData &aggr, idx_t c } } -void PacCountFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, idx_t offset) { +void PacCountFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { auto aggs = FlatVector::GetData(states); - auto data = FlatVector::GetData(result); - auto &result_mask = FlatVector::Validity(result); + auto data = FlatVector::GetDataMutable(result); + auto &result_mask = FlatVector::ValidityMutable(result); auto &bind = input.bind_data->Cast(); - std::mt19937_64 gen(bind.seed); double mi = bind.mi; double correction = bind.correction; uint64_t query_hash = bind.query_hash; auto pstate = bind.pstate; PAC_FLOAT buf[64]; + auto gen = PacRandomStream(bind.seed, *bind.random_state); for (idx_t i = 0; i < count; i++) { #if !defined(PAC_NOBUFFERING) && !defined(PAC_NOCASCADING) aggs[i]->FlushBuffer(*aggs[i], input.allocator); // flush values into yourself #endif PacCountState *s = aggs[i]->GetState(); uint64_t key_hash = s ? s->key_hash : 0; - // Check if we should return NULL based on key_hash (uses mi and correction) if (PacNoiseInNull(key_hash, mi, correction, gen)) { result_mask.SetInvalid(offset + i); continue; @@ -271,20 +375,21 @@ void PacCountFinalize(Vector &states, AggregateInputData &input, Vector &result, // Returns LIST with exactly 64 elements (no NULLs). // Position j is 0 if key_hash bit j is 0, otherwise value * 2 * correction. -void PacCountFinalizeCounters(Vector &states, AggregateInputData &input, Vector &result, idx_t count, idx_t offset) { +void PacCountFinalizeCounters(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, + idx_t offset) { auto aggs = FlatVector::GetData(states); double correction = input.bind_data ? input.bind_data->Cast().correction : 1.0; // Result is LIST - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); // Reserve space for all lists (64 elements each) idx_t total_elements = (offset + count) * 64; ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); PAC_FLOAT buf[64]; for (idx_t i = 0; i < count; i++) { @@ -351,7 +456,7 @@ void RegisterPacCountFunctions(ExtensionLoader &loader) { desc.examples = { "SELECT c_mktsegment, as_noised_count(priv_hash(hash(c_custkey))) FROM customer GROUP BY c_mktsegment"}; info.descriptions.push_back(std::move(desc)); - loader.RegisterFunction(std::move(info)); + RegisterCountFunctions(loader, info); } // ============================================================================ @@ -360,21 +465,21 @@ void RegisterPacCountFunctions(ExtensionLoader &loader) { // PAC's `PacCountFinalizeCounters` zeroes positions where (key_hash >> j) & 1 == 0; // SASS samples each lane independently via DpSampleHash, so all 64 are meaningful. // ============================================================================ -static void DpSampleCountFinalizeCounters(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void DpSampleCountFinalizeCounters(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { auto aggs = FlatVector::GetData(states); auto *bind = input.bind_data ? &input.bind_data->Cast() : nullptr; int sample_lanes = bind ? bind->sample_lanes : 1; double rescale = bind && bind->sample_rescale ? DpSampleRescale(sample_lanes) : 1.0; - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * 64; ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); PAC_FLOAT buf[64]; for (idx_t i = 0; i < count; i++) { @@ -399,17 +504,17 @@ static void DpSampleCountFinalizeCounters(Vector &states, AggregateInputData &in } } -static void DpSampleCountMaskFinalizeCounters(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void DpSampleCountMaskFinalizeCounters(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { auto aggs = FlatVector::GetData(states); - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * 64; ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); PAC_FLOAT buf[64]; for (idx_t i = 0; i < count; i++) { @@ -434,8 +539,8 @@ static void DpSampleCountMaskFinalizeCounters(Vector &states, AggregateInputData } } -static unique_ptr DpSampleCountBind(ClientContext &ctx, AggregateFunction &, - vector> &) { +static unique_ptr DpSampleCountBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); return MakeDpSampleBindData(ctx); } @@ -454,7 +559,7 @@ void RegisterDpSampleCountFunctions(ExtensionLoader &loader) { FunctionDescription desc; desc.description = "[INTERNAL] Returns 64 sample-median DP counters for COUNT-like values."; info.descriptions.push_back(std::move(desc)); - loader.RegisterFunction(std::move(info)); + RegisterCountFunctions(loader, info); AggregateFunctionSet mask_set("as_sample_count_mask"); mask_set.AddFunction(AggregateFunction( @@ -465,7 +570,7 @@ void RegisterDpSampleCountFunctions(ExtensionLoader &loader) { FunctionDescription mask_desc; mask_desc.description = "[INTERNAL] Returns 64 sample-median DP counters from precomputed sample masks."; mask_info.descriptions.push_back(std::move(mask_desc)); - loader.RegisterFunction(std::move(mask_info)); + RegisterCountFunctions(loader, mask_info); } // Shared storage for variable-m count aggregates. It lazily allocates one counter per @@ -475,12 +580,15 @@ struct DpSampleMCountState { PAC_FLOAT *counts; }; -static idx_t DpSampleMCountStateSize(const AggregateFunction &) { +static idx_t DpSampleMCountStateSize(AggregateStateInput &) { return sizeof(DpSampleMCountState); } -static void DpSampleMCountInitialize(const AggregateFunction &, data_ptr_t state_ptr) { - memset(state_ptr, 0, sizeof(DpSampleMCountState)); +static void DpSampleMCountInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_ptr = states[i]; + memset(state_ptr, 0, sizeof(DpSampleMCountState)); + } } static PAC_FLOAT *EnsureDpSampleMCountArray(DpSampleMCountState &state, ArenaAllocator &allocator, int sample_count) { @@ -497,42 +605,47 @@ static void DpSampleMCountUpdateOne(DpSampleMCountState &state, ArenaAllocator & counts[DpSassLaneIndex(hash, sample_count)] += PAC_FLOAT(1.0); } -static void DpSampleMCountUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - int sample_count = GetPrivSampleCount(aggr); - UnifiedVectorFormat hash_data; - inputs[0].ToUnifiedFormat(count, hash_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - if (hash_data.validity.RowIsValid(h_idx)) { - DpSampleMCountUpdateOne(state, aggr.allocator, hashes[h_idx], sample_count); +static void DpSampleMCountUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + int sample_count = GetPrivSampleCount(aggr); + UnifiedVectorFormat hash_data; + inputs[0].ToUnifiedFormat(hash_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + if (hash_data.validity.RowIsValid(h_idx)) { + DpSampleMCountUpdateOne(state, aggr.allocator, hashes[h_idx], sample_count); + } } - } + }); } -static void DpSampleMCountColumnUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, - idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - int sample_count = GetPrivSampleCount(aggr); - UnifiedVectorFormat hash_data, col_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, col_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto c_idx = col_data.sel->get_index(i); - if (hash_data.validity.RowIsValid(h_idx) && col_data.validity.RowIsValid(c_idx)) { - DpSampleMCountUpdateOne(state, aggr.allocator, hashes[h_idx], sample_count); +static void DpSampleMCountColumnUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + int sample_count = GetPrivSampleCount(aggr); + UnifiedVectorFormat hash_data, col_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(col_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto c_idx = col_data.sel->get_index(i); + if (hash_data.validity.RowIsValid(h_idx) && col_data.validity.RowIsValid(c_idx)) { + DpSampleMCountUpdateOne(state, aggr.allocator, hashes[h_idx], sample_count); + } } - } + }); } static void DpSampleMCountScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { int sample_count = GetPrivSampleCount(aggr); UnifiedVectorFormat hash_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); for (idx_t i = 0; i < count; i++) { @@ -548,9 +661,9 @@ static void DpSampleMCountColumnScatterUpdate(Vector inputs[], AggregateInputDat idx_t count) { int sample_count = GetPrivSampleCount(aggr); UnifiedVectorFormat hash_data, col_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, col_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(col_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); for (idx_t i = 0; i < count; i++) { @@ -563,44 +676,46 @@ static void DpSampleMCountColumnScatterUpdate(Vector inputs[], AggregateInputDat } } -static void DpSampleMCountIfUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, - idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - int sample_count = GetPrivSampleCount(aggr); - UnifiedVectorFormat hash_data, predicate_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, predicate_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto &predicate_child = ListVector::GetEntry(inputs[1]); - UnifiedVectorFormat predicate_child_data; - predicate_child.ToUnifiedFormat(ListVector::GetListSize(inputs[1]), predicate_child_data); - auto predicate_values = UnifiedVectorFormat::GetData(predicate_child_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto p_idx = predicate_data.sel->get_index(i); - if (!hash_data.validity.RowIsValid(h_idx)) { - continue; - } - idx_t lane = DpSassLaneIndex(hashes[h_idx], sample_count); - if (DpSampleMPredicateLaneIsTrue(predicate_data, predicate_child_data, predicate_values, p_idx, lane)) { - auto counts = EnsureDpSampleMCountArray(state, aggr.allocator, sample_count); - counts[lane] += PAC_FLOAT(1.0); +static void DpSampleMCountIfUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + int sample_count = GetPrivSampleCount(aggr); + UnifiedVectorFormat hash_data, predicate_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(predicate_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto &predicate_child = ListVector::GetChild(inputs[1]); + UnifiedVectorFormat predicate_child_data; + predicate_child.ToUnifiedFormat(predicate_child_data); + auto predicate_values = UnifiedVectorFormat::GetData(predicate_child_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto p_idx = predicate_data.sel->get_index(i); + if (!hash_data.validity.RowIsValid(h_idx)) { + continue; + } + idx_t lane = DpSassLaneIndex(hashes[h_idx], sample_count); + if (DpSampleMPredicateLaneIsTrue(predicate_data, predicate_child_data, predicate_values, p_idx, lane)) { + auto counts = EnsureDpSampleMCountArray(state, aggr.allocator, sample_count); + counts[lane] += PAC_FLOAT(1.0); + } } - } + }); } static void DpSampleMCountIfScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { int sample_count = GetPrivSampleCount(aggr); UnifiedVectorFormat hash_data, predicate_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, predicate_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(predicate_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); - auto &predicate_child = ListVector::GetEntry(inputs[1]); + auto &predicate_child = ListVector::GetChild(inputs[1]); UnifiedVectorFormat predicate_child_data; - predicate_child.ToUnifiedFormat(ListVector::GetListSize(inputs[1]), predicate_child_data); + predicate_child.ToUnifiedFormat(predicate_child_data); auto predicate_values = UnifiedVectorFormat::GetData(predicate_child_data); for (idx_t i = 0; i < count; i++) { auto h_idx = hash_data.sel->get_index(i); @@ -634,18 +749,18 @@ static void DpSampleMCountCombine(Vector &src, Vector &dst, AggregateInputData & } } -static void DpSampleMCountFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void DpSampleMCountFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { int sample_count = GetPrivSampleCount(input); auto state_ptrs = FlatVector::GetData(states); - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * static_cast(sample_count); ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); bool sample_rescale = input.bind_data ? input.bind_data->Cast().sample_rescale : true; // A disjoint lane contains about 1/m of the PUs. Scaling COUNT by m turns each lane // answer into an estimator of the full-data answer; disabling it preserves subsample scale. @@ -665,12 +780,12 @@ static void DpSampleMCountFinalize(Vector &states, AggregateInputData &input, Ve } } -static void DpSampleMCountDistinctFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void DpSampleMCountDistinctFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { int sample_count = GetPrivSampleCount(input); auto state_ptrs = FlatVector::GetData(states); - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * static_cast(sample_count); ListVector::Reserve(result, total_elements); @@ -679,7 +794,7 @@ static void DpSampleMCountDistinctFinalize(Vector &states, AggregateInputData &i // The compiler has already produced one row per (group, lane, distinct value). // Keep these as lane-level DISTINCT answers. Generic COUNT rescaling would instead // change the target to a full-data cardinality extrapolation. - auto child_data = FlatVector::GetData(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); for (idx_t i = 0; i < count; i++) { idx_t base = (offset + i) * static_cast(sample_count); list_entries[offset + i].offset = base; @@ -695,32 +810,34 @@ static void DpSampleMCountDistinctFinalize(Vector &states, AggregateInputData &i } } -static void DpSampleMCountDistinctIfUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, - idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - int sample_count = GetPrivSampleCount(aggr); - UnifiedVectorFormat hash_data, include_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, include_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto include_values = UnifiedVectorFormat::GetData(include_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto include_idx = include_data.sel->get_index(i); - if (hash_data.validity.RowIsValid(h_idx) && include_data.validity.RowIsValid(include_idx) && - include_values[include_idx]) { - DpSampleMCountUpdateOne(state, aggr.allocator, hashes[h_idx], sample_count); +static void DpSampleMCountDistinctIfUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + int sample_count = GetPrivSampleCount(aggr); + UnifiedVectorFormat hash_data, include_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(include_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto include_values = UnifiedVectorFormat::GetData(include_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto include_idx = include_data.sel->get_index(i); + if (hash_data.validity.RowIsValid(h_idx) && include_data.validity.RowIsValid(include_idx) && + include_values[include_idx]) { + DpSampleMCountUpdateOne(state, aggr.allocator, hashes[h_idx], sample_count); + } } - } + }); } static void DpSampleMCountDistinctIfScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { int sample_count = GetPrivSampleCount(aggr); UnifiedVectorFormat hash_data, include_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, include_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(include_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto include_values = UnifiedVectorFormat::GetData(include_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); @@ -735,8 +852,8 @@ static void DpSampleMCountDistinctIfScatterUpdate(Vector inputs[], AggregateInpu } } -static unique_ptr DpSampleMCountBind(ClientContext &ctx, AggregateFunction &, - vector> &) { +static unique_ptr DpSampleMCountBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); int sample_lanes = GetDpSampleLanes(ctx); if (sample_lanes != 1) { throw InvalidInputException("as_sample_m_count: dp_sample_lanes must be 1 when dp_sass_m > 64"); @@ -798,12 +915,15 @@ struct DpSampleMCountDistinctValuesState { value_set_t *sets; }; -static idx_t DpSampleMCountDistinctValuesStateSize(const AggregateFunction &) { +static idx_t DpSampleMCountDistinctValuesStateSize(AggregateStateInput &) { return sizeof(DpSampleMCountDistinctValuesState); } -static void DpSampleMCountDistinctValuesInitialize(const AggregateFunction &, data_ptr_t state_ptr) { - memset(state_ptr, 0, sizeof(DpSampleMCountDistinctValuesState)); +static void DpSampleMCountDistinctValuesInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_ptr = states[i]; + memset(state_ptr, 0, sizeof(DpSampleMCountDistinctValuesState)); + } } static value_set_t *EnsureDpSampleMCountDistinctValueSets(DpSampleMCountDistinctValuesState &state, @@ -832,7 +952,7 @@ static void DpSampleMCountDistinctValuesUpdateOne(DpSampleMCountDistinctValuesSt } auto sets = EnsureDpSampleMCountDistinctValueSets(state, allocator, sample_count); auto &lane_set = sets[DpSassLaneIndex(hash, sample_count)]; - auto &child_vec = ListVector::GetEntry(list_input); + auto &child_vec = ListVector::GetChild(list_input); for (idx_t offset = 0; offset < entry.length; offset++) { auto child_idx = child_data.sel->get_index(entry.offset + offset); if (child_data.validity.RowIsValid(child_idx)) { @@ -841,39 +961,41 @@ static void DpSampleMCountDistinctValuesUpdateOne(DpSampleMCountDistinctValuesSt } } -static void DpSampleMCountDistinctValuesUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, - idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - int sample_count = GetPrivSampleCount(aggr); - UnifiedVectorFormat hash_data, list_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, list_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto &child_vec = ListVector::GetEntry(inputs[1]); - UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(inputs[1]), child_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto list_idx = list_data.sel->get_index(i); - if (hash_data.validity.RowIsValid(h_idx)) { - DpSampleMCountDistinctValuesUpdateOne(state, aggr.allocator, hashes[h_idx], inputs[1], list_data, - child_data, list_idx, sample_count); +static void DpSampleMCountDistinctValuesUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + int sample_count = GetPrivSampleCount(aggr); + UnifiedVectorFormat hash_data, list_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(list_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto &child_vec = ListVector::GetChild(inputs[1]); + UnifiedVectorFormat child_data; + child_vec.ToUnifiedFormat(child_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto list_idx = list_data.sel->get_index(i); + if (hash_data.validity.RowIsValid(h_idx)) { + DpSampleMCountDistinctValuesUpdateOne(state, aggr.allocator, hashes[h_idx], inputs[1], list_data, + child_data, list_idx, sample_count); + } } - } + }); } static void DpSampleMCountDistinctValuesScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { int sample_count = GetPrivSampleCount(aggr); UnifiedVectorFormat hash_data, list_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, list_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(list_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); - auto &child_vec = ListVector::GetEntry(inputs[1]); + auto &child_vec = ListVector::GetChild(inputs[1]); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(inputs[1]), child_data); + child_vec.ToUnifiedFormat(child_data); for (idx_t i = 0; i < count; i++) { auto h_idx = hash_data.sel->get_index(i); auto list_idx = list_data.sel->get_index(i); @@ -904,19 +1026,19 @@ static void DpSampleMCountDistinctValuesCombine(Vector &src, Vector &dst, Aggreg } } -static void DpSampleMCountDistinctValuesFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void DpSampleMCountDistinctValuesFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { int sample_count = GetPrivSampleCount(input); auto state_ptrs = FlatVector::GetData(states); - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * static_cast(sample_count); ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); - auto &child_validity = FlatVector::Validity(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); + auto &child_validity = FlatVector::ValidityMutable(child_vec); child_validity.SetAllValid(total_elements); for (idx_t i = 0; i < count; i++) { idx_t base = (offset + i) * static_cast(sample_count); @@ -948,12 +1070,14 @@ static void DpSampleMCountDistinctValuesDestroy(Vector &states, AggregateInputDa } } -static unique_ptr DpSampleMCountDistinctValuesBind(ClientContext &ctx, AggregateFunction &func, - vector> &args) { - if (args.size() != 2 || args[1]->return_type.id() != LogicalTypeId::LIST) { +static unique_ptr DpSampleMCountDistinctValuesBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &func = bind_input.GetBoundFunction(); + auto &args = bind_input.GetArguments(); + if (args.size() != 2 || args[1]->GetReturnType().id() != LogicalTypeId::LIST) { throw InvalidInputException("as_sample_m_count_distinct_values expects a LIST of per-PU distinct values"); } - return DpSampleMCountBind(ctx, func, args); + return DpSampleMCountBind(bind_input); } void RegisterDpSampleMCountDistinctValuesFunctions(ExtensionLoader &loader) { @@ -993,43 +1117,48 @@ static void PacSampleMCountUpdateOne(DpSampleMCountState &state, ArenaAllocator ForEachPacMSelectedLane(hash, bind, [&](idx_t lane) { counts[lane] += PAC_FLOAT(1.0); }); } -static void PacSampleMCountUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - auto &bind = aggr.bind_data->Cast(); - UnifiedVectorFormat hash_data; - inputs[0].ToUnifiedFormat(count, hash_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - if (hash_data.validity.RowIsValid(h_idx)) { - PacSampleMCountUpdateOne(state, aggr.allocator, hashes[h_idx], bind); +static void PacSampleMCountUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + auto &bind = aggr.bind_data->Cast(); + UnifiedVectorFormat hash_data; + inputs[0].ToUnifiedFormat(hash_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + if (hash_data.validity.RowIsValid(h_idx)) { + PacSampleMCountUpdateOne(state, aggr.allocator, hashes[h_idx], bind); + } } - } + }); } -static void PacSampleMCountColumnUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, - idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - auto &bind = aggr.bind_data->Cast(); - UnifiedVectorFormat hash_data, col_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, col_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto c_idx = col_data.sel->get_index(i); - if (hash_data.validity.RowIsValid(h_idx) && col_data.validity.RowIsValid(c_idx)) { - PacSampleMCountUpdateOne(state, aggr.allocator, hashes[h_idx], bind); +static void PacSampleMCountColumnUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + auto &bind = aggr.bind_data->Cast(); + UnifiedVectorFormat hash_data, col_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(col_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto c_idx = col_data.sel->get_index(i); + if (hash_data.validity.RowIsValid(h_idx) && col_data.validity.RowIsValid(c_idx)) { + PacSampleMCountUpdateOne(state, aggr.allocator, hashes[h_idx], bind); + } } - } + }); } static void PacSampleMCountScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { auto &bind = aggr.bind_data->Cast(); UnifiedVectorFormat hash_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); for (idx_t i = 0; i < count; i++) { @@ -1045,9 +1174,9 @@ static void PacSampleMCountColumnScatterUpdate(Vector inputs[], AggregateInputDa idx_t count) { auto &bind = aggr.bind_data->Cast(); UnifiedVectorFormat hash_data, col_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, col_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(col_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); for (idx_t i = 0; i < count; i++) { @@ -1060,45 +1189,47 @@ static void PacSampleMCountColumnScatterUpdate(Vector inputs[], AggregateInputDa } } -static void PacSampleMCountIfUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, - idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - auto &bind = aggr.bind_data->Cast(); - UnifiedVectorFormat hash_data, predicate_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, predicate_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto &predicate_child = ListVector::GetEntry(inputs[1]); - UnifiedVectorFormat predicate_child_data; - predicate_child.ToUnifiedFormat(ListVector::GetListSize(inputs[1]), predicate_child_data); - auto predicate_values = UnifiedVectorFormat::GetData(predicate_child_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto p_idx = predicate_data.sel->get_index(i); - if (!hash_data.validity.RowIsValid(h_idx)) { - continue; - } - auto counts = EnsureDpSampleMCountArray(state, aggr.allocator, bind.sample_count); - ForEachPacMSelectedLane(hashes[h_idx], bind, [&](idx_t lane) { - if (DpSampleMPredicateLaneIsTrue(predicate_data, predicate_child_data, predicate_values, p_idx, lane)) { - counts[lane] += PAC_FLOAT(1.0); +static void PacSampleMCountIfUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + auto &bind = aggr.bind_data->Cast(); + UnifiedVectorFormat hash_data, predicate_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(predicate_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto &predicate_child = ListVector::GetChild(inputs[1]); + UnifiedVectorFormat predicate_child_data; + predicate_child.ToUnifiedFormat(predicate_child_data); + auto predicate_values = UnifiedVectorFormat::GetData(predicate_child_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto p_idx = predicate_data.sel->get_index(i); + if (!hash_data.validity.RowIsValid(h_idx)) { + continue; } - }); - } + auto counts = EnsureDpSampleMCountArray(state, aggr.allocator, bind.sample_count); + ForEachPacMSelectedLane(hashes[h_idx], bind, [&](idx_t lane) { + if (DpSampleMPredicateLaneIsTrue(predicate_data, predicate_child_data, predicate_values, p_idx, lane)) { + counts[lane] += PAC_FLOAT(1.0); + } + }); + } + }); } static void PacSampleMCountIfScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { auto &bind = aggr.bind_data->Cast(); UnifiedVectorFormat hash_data, predicate_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, predicate_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(predicate_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); - auto &predicate_child = ListVector::GetEntry(inputs[1]); + auto &predicate_child = ListVector::GetChild(inputs[1]); UnifiedVectorFormat predicate_child_data; - predicate_child.ToUnifiedFormat(ListVector::GetListSize(inputs[1]), predicate_child_data); + predicate_child.ToUnifiedFormat(predicate_child_data); auto predicate_values = UnifiedVectorFormat::GetData(predicate_child_data); for (idx_t i = 0; i < count; i++) { auto h_idx = hash_data.sel->get_index(i); @@ -1116,18 +1247,18 @@ static void PacSampleMCountIfScatterUpdate(Vector inputs[], AggregateInputData & } } -static void PacSampleMCountFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void PacSampleMCountFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { int sample_count = GetPrivSampleCount(input); auto state_ptrs = FlatVector::GetData(states); - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * static_cast(sample_count); ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); for (idx_t i = 0; i < count; i++) { idx_t base = (offset + i) * static_cast(sample_count); list_entries[offset + i].offset = base; @@ -1143,8 +1274,8 @@ static void PacSampleMCountFinalize(Vector &states, AggregateInputData &input, V } } -static unique_ptr PacSampleMCountBind(ClientContext &ctx, AggregateFunction &, - vector> &) { +static unique_ptr PacSampleMCountBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); return MakePacSampleMBindData(ctx); } @@ -1205,7 +1336,7 @@ void RegisterPacCountCountersFunctions(ExtensionLoader &loader) { FunctionDescription counters_desc; counters_desc.description = "[INTERNAL] Returns 64 PAC subsample counters as LIST for categorical queries."; counters_info.descriptions.push_back(std::move(counters_desc)); - loader.RegisterFunction(std::move(counters_info)); + RegisterCountFunctions(loader, counters_info); } // ============================================================================ @@ -1275,7 +1406,7 @@ void RegisterPacNoisedClipCountFunctions(ExtensionLoader &loader) { FunctionNullHandling::SPECIAL_HANDLING, PacCountColumnUpdate, PacCountBind)); CreateAggregateFunctionInfo info(fcn_set); - loader.RegisterFunction(std::move(info)); + RegisterCountFunctions(loader, info); } void RegisterPacClipCountFunctions(ExtensionLoader &loader) { @@ -1293,7 +1424,7 @@ void RegisterPacClipCountFunctions(ExtensionLoader &loader) { AddPacListAggregateOverload(fcn_set, "clip_count"); CreateAggregateFunctionInfo info(fcn_set); - loader.RegisterFunction(std::move(info)); + RegisterCountFunctions(loader, info); } } // namespace duckdb diff --git a/src/aggregates/as_finalize.cpp b/src/aggregates/as_finalize.cpp index 65915eb1..73357426 100644 --- a/src/aggregates/as_finalize.cpp +++ b/src/aggregates/as_finalize.cpp @@ -1,3 +1,4 @@ +#include "duckdb/common/vector/list_vector.hpp" #include "aggregates/as_aggregate.hpp" #include "duckdb/common/types/vector.hpp" #include "duckdb/function/scalar_function.hpp" @@ -18,27 +19,22 @@ static void PacFinalizeFunction(DataChunk &args, ExpressionState &state, Vector auto &context = state.GetContext(); double mi = GetPacMiFromSetting(context); - Value pac_seed_val; - uint64_t seed = (context.TryGetCurrentSetting("privacy_seed", pac_seed_val) && !pac_seed_val.IsNull()) - ? uint64_t(pac_seed_val.GetValue()) - : 42; - if (mi != 0.0) { - seed ^= PAC_MAGIC_HASH * static_cast(context.ActiveTransaction().GetActiveQuery()); - } + uint64_t seed = GetPacSeed(context, mi); uint64_t query_hash = (seed * PAC_MAGIC_HASH) ^ PAC_MAGIC_HASH; + auto gen = PacRandomStream(seed, *GetPacRandomState(context)); auto &list_vec = args.data[0]; idx_t count = args.size(); UnifiedVectorFormat list_data; - list_vec.ToUnifiedFormat(count, list_data); + list_vec.ToUnifiedFormat(list_data); - auto result_data = FlatVector::GetData(result); - auto &result_validity = FlatVector::Validity(result); + auto result_data = FlatVector::GetDataMutable(result); + auto &result_validity = FlatVector::ValidityMutable(result); - auto &child_vec = ListVector::GetEntry(list_vec); + auto &child_vec = ListVector::GetChild(list_vec); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(list_vec), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); auto list_entries = UnifiedVectorFormat::GetData(list_data); @@ -63,9 +59,6 @@ static void PacFinalizeFunction(DataChunk &args, ExpressionState &state, Vector counters[j] = child_data.validity.RowIsValid(child_idx) ? child_values[child_idx] : 0; } - // RNG seeded per row using seed + list offset (stable across DataChunk batches) - std::mt19937_64 gen(seed ^ (entry.offset * PAC_MAGIC_HASH)); - result_data[i] = PacNoisySampleFrom64Counters(counters, mi, 1.0, gen, 0ULL, query_hash, nullptr); } } @@ -74,6 +67,7 @@ void RegisterPacFinalizeFunction(ExtensionLoader &loader) { auto list_type = LogicalType::LIST(PacFloatLogicalType()); auto scalar_type = PacFloatLogicalType(); ScalarFunction priv_finalize("priv_finalize", {list_type}, scalar_type, PacFinalizeFunction); + priv_finalize.SetStability(FunctionStability::VOLATILE); loader.RegisterFunction(priv_finalize); // Register implicit casts LIST → FLOAT and LIST → DOUBLE so the binder @@ -83,19 +77,20 @@ void RegisterPacFinalizeFunction(ExtensionLoader &loader) { double mi = 0.0; uint64_t seed = 42; uint64_t query_hash = (seed * PAC_MAGIC_HASH) ^ PAC_MAGIC_HASH; + PacRandomEngine gen(seed, 0); UnifiedVectorFormat list_data; - source.ToUnifiedFormat(count, list_data); + source.ToUnifiedFormat(list_data); - auto &child_vec = ListVector::GetEntry(source); + auto &child_vec = ListVector::GetChild(source); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(source), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); auto list_entries = UnifiedVectorFormat::GetData(list_data); // Write to a PAC_FLOAT temp buffer, then cast to the result type if needed vector tmp(count); - auto &result_validity = FlatVector::Validity(result); + auto &result_validity = FlatVector::ValidityMutable(result); for (idx_t i = 0; i < count; i++) { auto list_idx = list_data.sel->get_index(i); @@ -113,18 +108,17 @@ void RegisterPacFinalizeFunction(ExtensionLoader &loader) { auto child_idx = child_data.sel->get_index(entry.offset + j); counters[j] = child_data.validity.RowIsValid(child_idx) ? child_values[child_idx] : 0; } - std::mt19937_64 gen(seed ^ (entry.offset * PAC_MAGIC_HASH)); tmp[i] = PacNoisySampleFrom64Counters(counters, mi, 1.0, gen, 0ULL, query_hash, nullptr); } // Copy to result — handles both FLOAT and DOUBLE result types if (result.GetType().id() == LogicalTypeId::FLOAT) { - auto result_data = FlatVector::GetData(result); + auto result_data = FlatVector::GetDataMutable(result); for (idx_t i = 0; i < count; i++) { result_data[i] = static_cast(tmp[i]); } } else { - auto result_data = FlatVector::GetData(result); + auto result_data = FlatVector::GetDataMutable(result); for (idx_t i = 0; i < count; i++) { result_data[i] = static_cast(tmp[i]); } @@ -140,16 +134,16 @@ void RegisterPacFinalizeFunction(ExtensionLoader &loader) { // The post-optimizer rewrites these comparisons to priv_filter_ calls before // execution, so the cast is never actually run in comparison contexts. auto scalar_to_list_cast = [](Vector &source, Vector &result, idx_t count, CastParameters ¶meters) -> bool { - auto &child_vec = ListVector::GetEntry(result); - auto child_data = FlatVector::GetData(child_vec); - auto list_entries = FlatVector::GetData(result); + auto &child_vec = ListVector::GetChildMutable(result); + auto child_data = FlatVector::GetDataMutable(child_vec); + auto list_entries = FlatVector::GetDataMutable(result); UnifiedVectorFormat source_data; - source.ToUnifiedFormat(count, source_data); + source.ToUnifiedFormat(source_data); idx_t offset = 0; for (idx_t i = 0; i < count; i++) { auto src_idx = source_data.sel->get_index(i); if (!source_data.validity.RowIsValid(src_idx)) { - FlatVector::Validity(result).SetInvalid(i); + FlatVector::ValidityMutable(result).SetInvalid(i); list_entries[i] = {offset, 0}; continue; } diff --git a/src/aggregates/as_min_max.cpp b/src/aggregates/as_min_max.cpp index 3412c7c6..e4d81a94 100644 --- a/src/aggregates/as_min_max.cpp +++ b/src/aggregates/as_min_max.cpp @@ -1,4 +1,6 @@ +#include "duckdb/common/vector/list_vector.hpp" #include "aggregates/as_min_max.hpp" +#include "aggregates/aggregate_state_export.hpp" #include "categorical/pac_categorical.hpp" #include "duckdb/parser/parsed_data/create_aggregate_function_info.hpp" @@ -17,31 +19,94 @@ template using MinMaxState = PacMinMaxStateWrapper; #endif +template +struct MinMaxStateCodec { + using State = MinMaxState; + using Inner = PacMinMaxState; + static void WriteInner(AggregateStateWriter &out, const Inner &state) { + out.Write(uint8_t(state.initialized)); + out.Write(state.update_count); + out.Write(state.key_hash); + out.Write(state.global_bound); + out.Write(state.extremes); + } + static void ReadInner(AggregateStateReader &in, Inner &state) { + state.initialized = in.ReadFlag(); + in.Read(state.update_count); + in.Read(state.key_hash); + in.Read(state.global_bound); + in.Read(state.extremes); + } + static void Write(AggregateStateWriter &out, const State &state) { +#ifdef PAC_NOBUFFERING + WriteInner(out, state); +#else + auto buffered = state.n_buffered & State::BUF_MASK; + out.Write(uint8_t(buffered)); + for (idx_t i = 0; i < buffered; i++) { + out.Write(state.val_buf[i]); + out.Write(state.hash_buf[i]); + } + auto *inner = state.GetState(); + out.Write(uint8_t(inner != nullptr)); + if (inner) { + WriteInner(out, *inner); + } +#endif + } + static void Read(AggregateStateReader &in, State &state, ArenaAllocator &allocator) { + memset(&state, 0, sizeof(state)); +#ifdef PAC_NOBUFFERING + ReadInner(in, state); +#else + uint8_t buffered; + in.Read(buffered); + if (buffered > State::BUF_SIZE) { + throw InvalidInputException("Invalid stochastic MIN/MAX buffer size"); + } + for (idx_t i = 0; i < buffered; i++) { + in.Read(state.val_buf[i]); + in.Read(state.hash_buf[i]); + } + if (in.ReadFlag()) { + ReadInner(in, *state.EnsureState(allocator)); + } + state.n_buffered |= buffered; + if (buffered) { + state.FlushBuffer(state, allocator); + } +#endif + } +}; + // ============================================================================ // Update functions // ============================================================================ // Non-grouped update template -static void PacMinMaxUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_p, idx_t count) { - auto &agg = *reinterpret_cast *>(state_p); - int sample_lanes = GetPrivSampleLanes(aggr); - - UnifiedVectorFormat hash_data, value_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto values = UnifiedVectorFormat::GetData(value_data); - - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto v_idx = value_data.sel->get_index(i); - if (hash_data.validity.RowIsValid(h_idx) && value_data.validity.RowIsValid(v_idx)) { - // min/max do not use buffering for global aggregates (without GROUP BY) because it was slower - PacMinMaxUpdateOne(agg, TransformPacUpdateHash(hashes[h_idx], sample_lanes), values[v_idx], - aggr.allocator); +static void PacMinMaxUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, const ClusteredAggr &clusters, + idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { + auto &agg = *reinterpret_cast *>(state_p); + int sample_lanes = GetPrivSampleLanes(aggr); + + UnifiedVectorFormat hash_data, value_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto values = UnifiedVectorFormat::GetData(value_data); + + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto v_idx = value_data.sel->get_index(i); + if (hash_data.validity.RowIsValid(h_idx) && value_data.validity.RowIsValid(v_idx)) { + // min/max do not use buffering for global aggregates (without GROUP BY) because it was slower + PacMinMaxUpdateOne(agg, TransformPacUpdateHash(hashes[h_idx], sample_lanes), values[v_idx], + aggr.allocator); + } } - } + }); } // Grouped (scatter) update @@ -49,9 +114,9 @@ template static void PacMinMaxScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { int sample_lanes = GetPrivSampleLanes(aggr); UnifiedVectorFormat hash_data, value_data, sdata; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - states.ToUnifiedFormat(count, sdata); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + states.ToUnifiedFormat(sdata); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); @@ -95,15 +160,19 @@ static void PacMinMaxCombine(Vector &src, Vector &dst, AggregateInputData &aggr, } template -static void PacMinMaxFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, idx_t offset) { +static void PacMinMaxFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, + idx_t offset) { auto state_ptrs = FlatVector::GetData *>(states); - auto data = FlatVector::GetData(result); - auto &result_mask = FlatVector::Validity(result); + auto data = FlatVector::GetDataMutable(result); + auto &result_mask = FlatVector::ValidityMutable(result); double mi = input.bind_data ? input.bind_data->Cast().mi : 0.0; double correction = input.bind_data ? input.bind_data->Cast().correction : 1.0; uint64_t query_hash = input.bind_data ? input.bind_data->Cast().query_hash : 0; auto pstate = input.bind_data ? input.bind_data->Cast().pstate : nullptr; + auto gen = input.bind_data ? PacRandomStream(input.bind_data->Cast().seed, + *input.bind_data->Cast().random_state) + : PacRandomEngine(0, 0); for (idx_t i = 0; i < count; i++) { #ifndef PAC_NOBUFFERING @@ -113,8 +182,6 @@ static void PacMinMaxFinalize(Vector &states, AggregateInputData &input, Vector auto *s = state_ptrs[i]->GetState(); // Check if we should return NULL based on key_hash uint64_t key_hash = s ? s->key_hash : 0; - std::mt19937_64 gen(input.bind_data->Cast().seed); - // mi controls noise mode, correction reduces NULL probability (but does NOT scale min/max values) if (PacNoiseInNull(key_hash, mi, correction, gen)) { result_mask.SetInvalid(offset + i); continue; @@ -145,13 +212,16 @@ static void PacMinMaxFinalize(Vector &states, AggregateInputData &input, Vector // ============================================================================ template -static idx_t PacMinMaxStateSize(const AggregateFunction &) { +static idx_t PacMinMaxStateSize(AggregateStateInput &) { return sizeof(MinMaxState); } template -static void PacMinMaxInitialize(const AggregateFunction &, data_ptr_t p) { - memset(p, 0, sizeof(MinMaxState)); +static void PacMinMaxInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto p = states[i]; + memset(p, 0, sizeof(MinMaxState)); + } } // ============================================================================ @@ -159,23 +229,26 @@ static void PacMinMaxInitialize(const AggregateFunction &, data_ptr_t p) { // ============================================================================ template -static unique_ptr PacMinMaxBind(ClientContext &ctx, AggregateFunction &function, - vector> &args) { - auto &value_type = args[1]->return_type; +static unique_ptr PacMinMaxBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &function = bind_input.GetBoundFunction(); + auto &args = bind_input.GetArguments(); + auto &value_type = args[1]->GetReturnType(); auto physical_type = value_type.InternalType(); - function.return_type = value_type; - function.arguments[1] = value_type; + function.SetReturnType(value_type); + function.GetArguments()[1] = value_type; // Select implementation based on physical type #define BIND_TYPE(PHYS_TYPE, CPP_TYPE) \ case PhysicalType::PHYS_TYPE: \ - function.state_size = PacMinMaxStateSize; \ - function.initialize = PacMinMaxInitialize; \ - function.update = PacMinMaxScatterUpdate; \ - function.simple_update = PacMinMaxUpdate; \ - function.combine = PacMinMaxCombine; \ - function.finalize = PacMinMaxFinalize; \ + function.SetStateSizeCallback(PacMinMaxStateSize); \ + function.SetStateInitCallback(PacMinMaxInitialize); \ + function.SetStateUpdateCallback(PacMinMaxScatterUpdate); \ + function.SetStateClusterUpdateCallback(PacMinMaxUpdate); \ + function.SetStateCombineCallback(PacMinMaxCombine); \ + function.SetStateFinalizeCallback(PacMinMaxFinalize); \ + AggregateStateExport, MinMaxStateCodec>::Register(function); \ break switch (physical_type) { @@ -204,20 +277,20 @@ static unique_ptr PacMinMaxBind(ClientContext &ctx, AggregateFunct // PAC_MIN/MAX_COUNTERS: Returns all 64 counters as LIST for categorical queries // ============================================================================ template -static void PacMinMaxFinalizeCounters(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void PacMinMaxFinalizeCounters(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { auto state_ptrs = FlatVector::GetData *>(states); // Result is LIST - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); // Reserve space for all lists (64 elements each) idx_t total_elements = (offset + count) * 64; ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); for (idx_t i = 0; i < count; i++) { #ifndef PAC_NOBUFFERING @@ -263,23 +336,26 @@ static void PacMinMaxFinalizeCounters(Vector &states, AggregateInputData &input, // Bind function for as_min/max_counters - similar to PacMinMaxBind but returns LIST template -static unique_ptr PacMinMaxCountersBind(ClientContext &ctx, AggregateFunction &function, - vector> &args) { - auto &value_type = args[1]->return_type; +static unique_ptr PacMinMaxCountersBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &function = bind_input.GetBoundFunction(); + auto &args = bind_input.GetArguments(); + auto &value_type = args[1]->GetReturnType(); auto physical_type = value_type.InternalType(); - function.return_type = LogicalType::LIST(PacFloatLogicalType()); - function.arguments[1] = value_type; + function.SetReturnType(LogicalType::LIST(PacFloatLogicalType())); + function.GetArguments()[1] = value_type; // Select implementation based on physical type #define BIND_COUNTERS_TYPE(PHYS_TYPE, CPP_TYPE) \ case PhysicalType::PHYS_TYPE: \ - function.state_size = PacMinMaxStateSize; \ - function.initialize = PacMinMaxInitialize; \ - function.update = PacMinMaxScatterUpdate; \ - function.simple_update = PacMinMaxUpdate; \ - function.combine = PacMinMaxCombine; \ - function.finalize = PacMinMaxFinalizeCounters; \ + function.SetStateSizeCallback(PacMinMaxStateSize); \ + function.SetStateInitCallback(PacMinMaxInitialize); \ + function.SetStateUpdateCallback(PacMinMaxScatterUpdate); \ + function.SetStateClusterUpdateCallback(PacMinMaxUpdate); \ + function.SetStateCombineCallback(PacMinMaxCombine); \ + function.SetStateFinalizeCallback(PacMinMaxFinalizeCounters); \ + AggregateStateExport, MinMaxStateCodec>::Register(function); \ break switch (physical_type) { @@ -308,18 +384,18 @@ static unique_ptr PacMinMaxCountersBind(ClientContext &ctx, Aggreg // SASS MIN/MAX counters: Returns all 64 sample values as LIST // ============================================================================ template -static void DpSampleMinMaxFinalizeCounters(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void DpSampleMinMaxFinalizeCounters(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { auto state_ptrs = FlatVector::GetData *>(states); - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * 64; ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); - auto &child_validity = FlatVector::Validity(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); + auto &child_validity = FlatVector::ValidityMutable(child_vec); child_validity.SetAllValid(total_elements); PAC_FLOAT buf[64]; double scale_divisor = input.bind_data ? input.bind_data->Cast().scale_divisor : 1.0; @@ -356,25 +432,27 @@ static void DpSampleMinMaxFinalizeCounters(Vector &states, AggregateInputData &i } template -static unique_ptr DpSampleMinMaxBind(ClientContext &ctx, AggregateFunction &function, - vector> &args) { - auto &value_type = args[1]->return_type; +static unique_ptr DpSampleMinMaxBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &function = bind_input.GetBoundFunction(); + auto &args = bind_input.GetArguments(); + auto &value_type = args[1]->GetReturnType(); auto physical_type = value_type.InternalType(); double scale_divisor = value_type.id() == LogicalTypeId::DECIMAL ? std::pow(10.0, static_cast(DecimalType::GetScale(value_type))) : 1.0; - function.return_type = LogicalType::LIST(PacFloatLogicalType()); - function.arguments[1] = value_type; + function.SetReturnType(LogicalType::LIST(PacFloatLogicalType())); + function.GetArguments()[1] = value_type; #define BIND_DP_SAMPLE_TYPE(PHYS_TYPE, CPP_TYPE) \ case PhysicalType::PHYS_TYPE: \ - function.state_size = PacMinMaxStateSize; \ - function.initialize = PacMinMaxInitialize; \ - function.update = PacMinMaxScatterUpdate; \ - function.simple_update = PacMinMaxUpdate; \ - function.combine = PacMinMaxCombine; \ - function.finalize = DpSampleMinMaxFinalizeCounters; \ + function.SetStateSizeCallback(PacMinMaxStateSize); \ + function.SetStateInitCallback(PacMinMaxInitialize); \ + function.SetStateUpdateCallback(PacMinMaxScatterUpdate); \ + function.SetStateClusterUpdateCallback(PacMinMaxUpdate); \ + function.SetStateCombineCallback(PacMinMaxCombine); \ + function.SetStateFinalizeCallback(DpSampleMinMaxFinalizeCounters); \ break switch (physical_type) { @@ -411,13 +489,16 @@ struct DpSampleMMinMaxState { }; template -static idx_t DpSampleMMinMaxStateSize(const AggregateFunction &) { +static idx_t DpSampleMMinMaxStateSize(AggregateStateInput &) { return sizeof(DpSampleMMinMaxState); } template -static void DpSampleMMinMaxInitialize(const AggregateFunction &, data_ptr_t state_ptr) { - memset(state_ptr, 0, sizeof(DpSampleMMinMaxState)); +static void DpSampleMMinMaxInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_ptr = states[i]; + memset(state_ptr, 0, sizeof(DpSampleMMinMaxState)); + } } template @@ -442,21 +523,24 @@ static void DpSampleMMinMaxUpdateOne(DpSampleMMinMaxState &state, Are } template -static void DpSampleMMinMaxUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, idx_t count) { - auto &state = *reinterpret_cast *>(state_ptr); - int sample_count = GetPrivSampleCount(aggr); - UnifiedVectorFormat hash_data, value_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto values = UnifiedVectorFormat::GetData(value_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto v_idx = value_data.sel->get_index(i); - if (hash_data.validity.RowIsValid(h_idx) && value_data.validity.RowIsValid(v_idx)) { - DpSampleMMinMaxUpdateOne(state, aggr.allocator, hashes[h_idx], values[v_idx], sample_count); +static void DpSampleMMinMaxUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast *>(state_ptr); + int sample_count = GetPrivSampleCount(aggr); + UnifiedVectorFormat hash_data, value_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto values = UnifiedVectorFormat::GetData(value_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto v_idx = value_data.sel->get_index(i); + if (hash_data.validity.RowIsValid(h_idx) && value_data.validity.RowIsValid(v_idx)) { + DpSampleMMinMaxUpdateOne(state, aggr.allocator, hashes[h_idx], values[v_idx], sample_count); + } } - } + }); } template @@ -464,9 +548,9 @@ static void DpSampleMMinMaxScatterUpdate(Vector inputs[], AggregateInputData &ag idx_t count) { int sample_count = GetPrivSampleCount(aggr); UnifiedVectorFormat hash_data, value_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); auto state_ptrs = UnifiedVectorFormat::GetData *>(state_data); @@ -502,19 +586,19 @@ static void DpSampleMMinMaxCombine(Vector &src, Vector &dst, AggregateInputData } template -static void DpSampleMMinMaxFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void DpSampleMMinMaxFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { int sample_count = GetPrivSampleCount(input); auto state_ptrs = FlatVector::GetData *>(states); - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * static_cast(sample_count); ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); - auto &child_validity = FlatVector::Validity(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); + auto &child_validity = FlatVector::ValidityMutable(child_vec); child_validity.SetAllValid(total_elements); double scale_divisor = input.bind_data ? input.bind_data->Cast().scale_divisor : 1.0; @@ -535,30 +619,32 @@ static void DpSampleMMinMaxFinalize(Vector &states, AggregateInputData &input, V } template -static unique_ptr DpSampleMMinMaxBind(ClientContext &ctx, AggregateFunction &function, - vector> &args) { +static unique_ptr DpSampleMMinMaxBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &function = bind_input.GetBoundFunction(); + auto &args = bind_input.GetArguments(); int sample_lanes = GetDpSampleLanes(ctx); if (sample_lanes != 1) { throw InvalidInputException("as_sample_m_%s: dp_sample_lanes must be 1 when dp_sass_m > 64", IS_MAX ? "max" : "min"); } - auto &value_type = args[1]->return_type; + auto &value_type = args[1]->GetReturnType(); auto physical_type = value_type.InternalType(); double scale_divisor = value_type.id() == LogicalTypeId::DECIMAL ? std::pow(10.0, static_cast(DecimalType::GetScale(value_type))) : 1.0; - function.return_type = LogicalType::LIST(PacFloatLogicalType()); - function.arguments[1] = value_type; + function.SetReturnType(LogicalType::LIST(PacFloatLogicalType())); + function.GetArguments()[1] = value_type; #define BIND_DP_SAMPLE_M_TYPE(PHYS_TYPE, CPP_TYPE) \ case PhysicalType::PHYS_TYPE: \ - function.state_size = DpSampleMMinMaxStateSize; \ - function.initialize = DpSampleMMinMaxInitialize; \ - function.update = DpSampleMMinMaxScatterUpdate; \ - function.simple_update = DpSampleMMinMaxUpdate; \ - function.combine = DpSampleMMinMaxCombine; \ - function.finalize = DpSampleMMinMaxFinalize; \ + function.SetStateSizeCallback(DpSampleMMinMaxStateSize); \ + function.SetStateInitCallback(DpSampleMMinMaxInitialize); \ + function.SetStateUpdateCallback(DpSampleMMinMaxScatterUpdate); \ + function.SetStateClusterUpdateCallback(DpSampleMMinMaxUpdate); \ + function.SetStateCombineCallback(DpSampleMMinMaxCombine); \ + function.SetStateFinalizeCallback(DpSampleMMinMaxFinalize); \ break switch (physical_type) { @@ -721,8 +807,8 @@ void RegisterDpSampleMMaxFunctions(ExtensionLoader &loader) { // Explicit template instantiations #define INST_ALL(T) \ - template void PacMinMaxUpdate(Vector[], AggregateInputData &, idx_t, data_ptr_t, idx_t); \ - template void PacMinMaxUpdate(Vector[], AggregateInputData &, idx_t, data_ptr_t, idx_t); \ + template void PacMinMaxUpdate(Vector[], AggregateInputData &, idx_t, const ClusteredAggr &, idx_t); \ + template void PacMinMaxUpdate(Vector[], AggregateInputData &, idx_t, const ClusteredAggr &, idx_t); \ template void PacMinMaxScatterUpdate(Vector[], AggregateInputData &, idx_t, Vector &, idx_t); \ template void PacMinMaxScatterUpdate(Vector[], AggregateInputData &, idx_t, Vector &, idx_t); \ template void PacMinMaxCombine(Vector &, Vector &, AggregateInputData &, idx_t); \ diff --git a/src/aggregates/as_sum.cpp b/src/aggregates/as_sum.cpp index 9653a47c..eb53bab2 100644 --- a/src/aggregates/as_sum.cpp +++ b/src/aggregates/as_sum.cpp @@ -1,4 +1,6 @@ +#include "duckdb/common/vector/list_vector.hpp" #include "aggregates/as_sum.hpp" +#include "aggregates/as_sum_state_export.hpp" #include "categorical/pac_categorical.hpp" #include "duckdb/common/types/decimal.hpp" #include "duckdb/parser/parsed_data/create_aggregate_function_info.hpp" @@ -196,6 +198,48 @@ inline void PacSumFlushBuffer(PacSumDoubleStateWrapper &src, PacSumDoubleStateWr } } +template +static bool SumHasAllocatedState(const WrapperT &state) { + return state.GetState() +#ifndef PAC_SIGNEDSUM + || state.GetNegState() +#endif + ; +} + +template +static bool CombineSumBuffers(const WrapperT &src, WrapperT &dst) { + if (SumHasAllocatedState(src) || SumHasAllocatedState(dst)) { + return false; + } + auto src_count = src.n_buffered & WrapperT::BUF_MASK; + auto dst_count = dst.n_buffered & WrapperT::BUF_MASK; + if (src_count + dst_count > WrapperT::BUF_SIZE) { + return false; + } + for (idx_t i = 0; i < src_count; i++) { + dst.val_buf[dst_count + i] = src.val_buf[i]; + dst.hash_buf[dst_count + i] = src.hash_buf[i]; + } + dst.n_buffered = src_count + dst_count; + return true; +} + +template +static WrapperT *PrepareSumFinalization(WrapperT &source, WrapperT &temporary, ArenaAllocator &scratch, + ArenaAllocator *&allocator) { + auto *state = &source; + if (!SumHasAllocatedState(source)) { + // Do not leave a full state in the source arena for every one/two-row + // group. The source stays buffered and can safely be finalized again. + temporary = source; + state = &temporary; + allocator = &scratch; + } + PacSumFlushBuffer(*state, *state, *allocator); + return state; +} + // Buffering update for int wrappers template AUTOVECTORIZE inline void PacSumUpdateOne(WrapperT &agg, uint64_t key_hash, ValueT value, ArenaAllocator &a) { @@ -245,27 +289,29 @@ inline void PacSumUpdateOne(PacSumIntStateWrapper &agg, uint64_t key_has #endif // PAC_NOBUFFERING -// PacSumUpdate - unified for both int and double states -template -static void PacSumUpdate(Vector inputs[], State &state, idx_t count, ArenaAllocator &allocator, int sample_lanes = 0) { - UnifiedVectorFormat hash_data, value_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); +// Update a single state from the original input vectors and an optional group selection. +template +static void PacSumUpdateValues(const UnifiedVectorFormat &hash_data, const UnifiedVectorFormat &value_data, + State &state, idx_t count, ArenaAllocator &allocator, int sample_lanes, + const sel_t *selection) { auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); - if (hash_data.validity.AllValid() && value_data.validity.AllValid()) { + if (hash_data.validity.CannotHaveNull() && value_data.validity.CannotHaveNull()) { for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto v_idx = value_data.sel->get_index(i); + auto row = SELECTED ? selection[i] : i; + auto h_idx = hash_data.sel->get_index(row); + auto v_idx = value_data.sel->get_index(row); PacSumUpdateOne(state, TransformPacUpdateHash(HASH_TRANSFORM::Transform(hashes[h_idx]), sample_lanes), ConvertValue::convert(values[v_idx]), allocator); } } else { for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto v_idx = value_data.sel->get_index(i); + auto row = SELECTED ? selection[i] : i; + auto h_idx = hash_data.sel->get_index(row); + auto v_idx = value_data.sel->get_index(row); if (!hash_data.validity.RowIsValid(h_idx) || !value_data.validity.RowIsValid(v_idx)) { continue; } @@ -276,13 +322,40 @@ static void PacSumUpdate(Vector inputs[], State &state, idx_t count, ArenaAlloca } } +template +static void PacSumUpdate(Vector inputs[], State &state, idx_t count, ArenaAllocator &allocator, int sample_lanes = 0) { + UnifiedVectorFormat hash_data, value_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + PacSumUpdateValues( + hash_data, value_data, state, count, allocator, sample_lanes, nullptr); +} + +template +static void PacSumUpdateClusters(Vector inputs[], const ClusteredAggr &clusters, ArenaAllocator &allocator, + int sample_lanes) { + UnifiedVectorFormat hash_data, value_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + for (auto &run : clusters.runs()) { + auto &state = *reinterpret_cast(run.state); + if (run.sel) { + PacSumUpdateValues(hash_data, value_data, state, run.count, + allocator, sample_lanes, run.sel); + } else { + PacSumUpdateValues(hash_data, value_data, state, run.count, + allocator, sample_lanes, nullptr); + } + } +} + template static void PacSumScatterUpdate(Vector inputs[], Vector &states, idx_t count, ArenaAllocator &allocator, int sample_lanes = 0) { UnifiedVectorFormat hash_data, value_data, sdata; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - states.ToUnifiedFormat(count, sdata); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + states.ToUnifiedFormat(sdata); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); @@ -324,6 +397,10 @@ AUTOVECTORIZE static void PacSumCombineInt(Vector &src, Vector &dst, idx_t count for (idx_t i = 0; i < count; i++) { #ifndef PAC_NOBUFFERING + if (CombineSumBuffers(*src_wrapper[i], *dst_wrapper[i])) { + continue; + } + PacSumFlushBuffer>(*dst_wrapper[i], *dst_wrapper[i], allocator); // Flush src's buffer into dst's inner state (avoids allocating src inner) PacSumFlushBuffer>(*src_wrapper[i], *dst_wrapper[i], allocator); #endif @@ -407,6 +484,10 @@ AUTOVECTORIZE static void PacSumCombineDouble(Vector &src, Vector &dst, idx_t co auto dst_wrapper = FlatVector::GetData(dst); for (idx_t i = 0; i < count; i++) { #ifndef PAC_NOBUFFERING + if (CombineSumBuffers(*src_wrapper[i], *dst_wrapper[i])) { + continue; + } + PacSumFlushBuffer(*dst_wrapper[i], *dst_wrapper[i], allocator); // Flush src's buffer into dst's inner state (avoids allocating src inner) PacSumFlushBuffer(*src_wrapper[i], *dst_wrapper[i], allocator); #endif @@ -426,41 +507,47 @@ AUTOVECTORIZE static void PacSumCombineDouble(Vector &src, Vector &dst, idx_t co // Unified Finalize for both int and double states template -void PacSumFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, idx_t offset) { +void PacSumFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { auto state_ptrs = FlatVector::GetData(states); - auto data = FlatVector::GetData(result); - auto &result_mask = FlatVector::Validity(result); + auto data = FlatVector::GetDataMutable(result); + auto &result_mask = FlatVector::ValidityMutable(result); double mi = input.bind_data ? input.bind_data->Cast().mi : 0.0; double correction = input.bind_data ? input.bind_data->Cast().correction : 1.0; uint64_t query_hash = input.bind_data ? input.bind_data->Cast().query_hash : 0; auto pstate = input.bind_data ? input.bind_data->Cast().pstate : nullptr; + auto gen = input.bind_data ? PacRandomStream(input.bind_data->Cast().seed, + *input.bind_data->Cast().random_state) + : PacRandomEngine(0, 0); +#ifndef PAC_NOBUFFERING + ArenaAllocator scratch(input.allocator.GetUnderlyingAllocator()); +#endif for (idx_t i = 0; i < count; i++) { + auto *state = state_ptrs[i]; + auto *allocator = &input.allocator; #ifndef PAC_NOBUFFERING - PacSumFlushBuffer(*state_ptrs[i], *state_ptrs[i], input.allocator); + State temporary; + state = PrepareSumFinalization(*state, temporary, scratch, allocator); #endif PAC_FLOAT buf[64] = {0}; - auto *pos = state_ptrs[i]->GetState(); + auto *pos = state->GetState(); if (!pos) { result_mask.SetInvalid(offset + i); // return NULL (no values seen) continue; } - // Use per-group deterministic RNG seeded by both privacy_seed and key_hash - // This ensures each group gets the same noise regardless of processing order uint64_t key_hash = pos->key_hash; - std::mt19937_64 gen(input.bind_data->Cast().seed); if (PacNoiseInNull(key_hash, mi, correction, gen)) { - result_mask.SetInvalid(offset + i); // return NULL (probabilistic decision) + result_mask.SetInvalid(offset + i); continue; } #ifndef PAC_SIGNEDSUM // Double-sided mode: use SFINAE helpers (cast to unsigned for int states only) - FlushPosState(pos, input.allocator); + FlushPosState(pos, *allocator); GetPosStateTotals(pos, buf); uint64_t update_count = pos->update_count; - update_count += SubtractNegStateTotals(state_ptrs[i], buf, key_hash, input.allocator); + update_count += SubtractNegStateTotals(state, buf, key_hash, *allocator); #else - pos->Flush(input.allocator); + pos->Flush(*allocator); pos->GetTotals(buf); uint64_t update_count = pos->update_count; #endif @@ -513,11 +600,10 @@ void PacSumFinalize(Vector &states, AggregateInputData &input, Vector &result, i // Generate exact as_sum Update/ScatterUpdate functions via X-macros // ============================================================================ #define X(NAME, VALUE_T, INPUT_T, SIGNED) \ - void PacSumUpdate##NAME(Vector inputs[], AggregateInputData &aggr_input_data, idx_t, data_ptr_t state_p, \ - idx_t count) { \ - auto &state = *reinterpret_cast *>(state_p); \ - PacSumUpdate, SIGNED, VALUE_T, INPUT_T>( \ - inputs, state, count, aggr_input_data.allocator, GetPrivSampleLanes(aggr_input_data)); \ + void PacSumUpdate##NAME(Vector inputs[], AggregateInputData &aggr_input_data, idx_t input_count, \ + const ClusteredAggr &clusters, idx_t count) { \ + PacSumUpdateClusters, SIGNED, VALUE_T, INPUT_T>( \ + inputs, clusters, aggr_input_data.allocator, GetPrivSampleLanes(aggr_input_data)); \ } \ void PacSumScatterUpdate##NAME(Vector inputs[], AggregateInputData &aggr_input_data, idx_t, Vector &states, \ idx_t count) { \ @@ -529,27 +615,25 @@ PAC_INT_TYPES_UNSIGNED #undef X // Double/Float/[U]HugeInt use double state (not generated by X-macro) -void PacSumUpdateHugeIntDouble(Vector inputs[], AggregateInputData &aggr_input_data, idx_t, data_ptr_t state_p, - idx_t count) { - auto &state = *reinterpret_cast(state_p); - PacSumUpdate(inputs, state, count, aggr_input_data.allocator, - GetPrivSampleLanes(aggr_input_data)); -} -void PacSumUpdateUHugeInt(Vector inputs[], AggregateInputData &aggr_input_data, idx_t, data_ptr_t state_p, - idx_t count) { - auto &state = *reinterpret_cast(state_p); - PacSumUpdate(inputs, state, count, aggr_input_data.allocator, - GetPrivSampleLanes(aggr_input_data)); -} -void PacSumUpdateFloat(Vector inputs[], AggregateInputData &aggr_input_data, idx_t, data_ptr_t state_p, idx_t count) { - auto &state = *reinterpret_cast(state_p); - PacSumUpdate(inputs, state, count, aggr_input_data.allocator, - GetPrivSampleLanes(aggr_input_data)); -} -void PacSumUpdateDouble(Vector inputs[], AggregateInputData &aggr_input_data, idx_t, data_ptr_t state_p, idx_t count) { - auto &state = *reinterpret_cast(state_p); - PacSumUpdate(inputs, state, count, aggr_input_data.allocator, - GetPrivSampleLanes(aggr_input_data)); +void PacSumUpdateHugeIntDouble(Vector inputs[], AggregateInputData &aggr_input_data, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + PacSumUpdateClusters(inputs, clusters, aggr_input_data.allocator, + GetPrivSampleLanes(aggr_input_data)); +} +void PacSumUpdateUHugeInt(Vector inputs[], AggregateInputData &aggr_input_data, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + PacSumUpdateClusters(inputs, clusters, aggr_input_data.allocator, + GetPrivSampleLanes(aggr_input_data)); +} +void PacSumUpdateFloat(Vector inputs[], AggregateInputData &aggr_input_data, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + PacSumUpdateClusters(inputs, clusters, aggr_input_data.allocator, + GetPrivSampleLanes(aggr_input_data)); +} +void PacSumUpdateDouble(Vector inputs[], AggregateInputData &aggr_input_data, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + PacSumUpdateClusters(inputs, clusters, aggr_input_data.allocator, + GetPrivSampleLanes(aggr_input_data)); } void PacSumScatterUpdateHugeIntDouble(Vector inputs[], AggregateInputData &aggr_input_data, idx_t, Vector &states, idx_t count) { @@ -572,11 +656,10 @@ void PacSumScatterUpdateDouble(Vector inputs[], AggregateInputData &aggr_input_d GetPrivSampleLanes(aggr_input_data)); } -static void DpSampleSumUpdateDouble(Vector inputs[], AggregateInputData &aggr_input_data, idx_t, data_ptr_t state_p, - idx_t count) { - auto &state = *reinterpret_cast(state_p); - PacSumUpdate(inputs, state, count, aggr_input_data.allocator, - GetPrivSampleLanes(aggr_input_data)); +static void DpSampleSumUpdateDouble(Vector inputs[], AggregateInputData &aggr_input_data, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + PacSumUpdateClusters(inputs, clusters, aggr_input_data.allocator, + GetPrivSampleLanes(aggr_input_data)); } static void DpSampleSumScatterUpdateDouble(Vector inputs[], AggregateInputData &aggr_input_data, idx_t, Vector &states, @@ -597,16 +680,19 @@ void PacSumCombineDoubleWrapper(Vector &src, Vector &dst, AggregateInputData &ag } // instantiate Finalize methods for as_sum -void PacSumFinalizeSigned(Vector &states, AggregateInputData &input, Vector &result, idx_t count, idx_t offset) { +void PacSumFinalizeSigned(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, + idx_t offset) { PacSumFinalize, hugeint_t, true>(states, input, result, count, offset); } -void PacSumFinalizeUnsigned(Vector &states, AggregateInputData &input, Vector &result, idx_t count, idx_t offset) { +void PacSumFinalizeUnsigned(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, + idx_t offset) { PacSumFinalize, hugeint_t, false>(states, input, result, count, offset); } -void PacSumFinalizeDoubleWrapper(Vector &states, AggregateInputData &input, Vector &result, idx_t count, idx_t offset) { +void PacSumFinalizeDoubleWrapper(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, + idx_t offset) { PacSumFinalize(states, input, result, count, offset); } -void PacSumFinalizeDoubleToHugeint(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +void PacSumFinalizeDoubleToHugeint(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { PacSumFinalize(states, input, result, count, offset); } @@ -618,48 +704,62 @@ static unique_ptr PacBindWithScaleDivisor(ClientContext &ctx, vect } unique_ptr // Bind function for as_sum with optional mi parameter (must be constant) -PacSumBind(ClientContext &ctx, AggregateFunction &, vector> &args) { +PacSumBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &args = bind_input.GetArguments(); return PacBindWithScaleDivisor(ctx, args, 1.0); // scale_divisor=1.0 for sum } -static unique_ptr DpSampleSumBind(ClientContext &ctx, AggregateFunction &, - vector> &) { +static unique_ptr DpSampleSumBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); return MakeDpSampleBindData(ctx); } -idx_t PacSumIntStateSize(const AggregateFunction &) { +idx_t PacSumIntStateSize(AggregateStateInput &) { return sizeof(ScatterIntState); // signed/unsigned have same wrapper size } -void PacSumIntInitialize(const AggregateFunction &, data_ptr_t state_p) { - memset(state_p, 0, sizeof(ScatterIntState)); +void PacSumIntInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_p = states[i]; + memset(state_p, 0, sizeof(ScatterIntState)); #if !defined(PAC_EXACTSUM) && defined(PAC_NOBUFFERING) - // In APPROX mode without buffering, the raw state needs max_level_used=-1 and inline_level_idx=-1 - // to indicate no levels are allocated yet (memset sets them to 0, which is incorrect) - auto *state = reinterpret_cast *>(state_p); - state->max_level_used = -1; - state->inline_level_idx = -1; + // In APPROX mode without buffering, the raw state needs max_level_used=-1 and inline_level_idx=-1 + // to indicate no levels are allocated yet (memset sets them to 0, which is incorrect) + auto *state = reinterpret_cast *>(state_p); + state->max_level_used = -1; + state->inline_level_idx = -1; #endif + } } -idx_t PacSumDoubleStateSize(const AggregateFunction &) { +idx_t PacSumDoubleStateSize(AggregateStateInput &) { return sizeof(ScatterDoubleState); } -void PacSumDoubleInitialize(const AggregateFunction &, data_ptr_t state_ptr) { - memset(state_ptr, 0, sizeof(ScatterDoubleState)); +void PacSumDoubleInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_ptr = states[i]; + memset(state_ptr, 0, sizeof(ScatterDoubleState)); + } +} + +// State export is cold-path metadata; the update kernels and native state sizes are unchanged. +static AggregateFunction SpillableSum(AggregateFunction function) { + RegisterSumStateExport(function); + return function; } -// Helper to register both 2-param and 3-param (with optional mi) versions +// Helper to register both 2-param and 3-param (with optional correction) versions static void AddFcn(AggregateFunctionSet &set, const LogicalType &value_type, const LogicalType &result_type, aggregate_size_t state_size, aggregate_initialize_t init, aggregate_update_t scatter, - aggregate_combine_t combine, aggregate_finalize_t finalize, aggregate_simple_update_t update, + aggregate_combine_t combine, aggregate_finalize_t finalize, aggregate_cluster_update_t update, aggregate_destructor_t destructor = nullptr) { - set.AddFunction(AggregateFunction("as_noised_sum", {LogicalType::UBIGINT, value_type}, result_type, state_size, - init, scatter, combine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, - update, PacSumBind, destructor)); - set.AddFunction(AggregateFunction("as_noised_sum", {LogicalType::UBIGINT, value_type, LogicalType::DOUBLE}, - result_type, state_size, init, scatter, combine, finalize, - FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacSumBind, destructor)); + set.AddFunction(SpillableSum(AggregateFunction( + "as_noised_sum", {LogicalType::UBIGINT, value_type}, result_type, state_size, init, scatter, combine, finalize, + FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacSumBind, destructor))); + set.AddFunction(SpillableSum(AggregateFunction( + "as_noised_sum", {LogicalType::UBIGINT, value_type, LogicalType::DOUBLE}, result_type, state_size, init, + scatter, combine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacSumBind, destructor))); } // Helper to get the right AggregateFunction for as_sum given a physical type (for DECIMAL support) @@ -699,13 +799,20 @@ static AggregateFunction GetPacSumAggregate(PhysicalType type) { } // Dynamic dispatch for DECIMAL as_sum: selects the right integer implementation based on decimal width -unique_ptr BindDecimalPacSum(ClientContext &ctx, AggregateFunction &function, - vector> &args) { - auto decimal_type = args[1]->return_type; // value is arg 1 (arg 0 is hash) - function = GetPacSumAggregate(decimal_type.InternalType()); - function.name = "as_noised_sum"; - function.arguments[1] = decimal_type; - function.return_type = LogicalType::DECIMAL(Decimal::MAX_WIDTH_DECIMAL, DecimalType::GetScale(decimal_type)); +unique_ptr BindDecimalPacSum(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &function = bind_input.GetBoundFunction(); + auto &args = bind_input.GetArguments(); + auto decimal_type = args[1]->GetReturnType(); // value is arg 1 (arg 0 is hash) + function.ReplaceImplementation(SpillableSum(GetPacSumAggregate(decimal_type.InternalType()))); + function.SetName("as_noised_sum"); + function.GetArguments()[1] = decimal_type; + if (args.size() == 3) { + // Specialization uses a two-argument factory; retain correction in the + // exported call signature so FINALIZE/COMBINE can rebind it correctly. + function.GetArguments().push_back(LogicalType::DOUBLE); + } + function.SetReturnType(LogicalType::DECIMAL(Decimal::MAX_WIDTH_DECIMAL, DecimalType::GetScale(decimal_type))); return PacBindWithScaleDivisor(ctx, args, 1.0); } @@ -786,21 +893,21 @@ enum class PacSumCountersFinalizeMode : uint8_t { PAC, DP_SAMPLE, DP_SAMPLE_NULL // Returns LIST with exactly 64 elements (no NULLs). // Position j is 0 if key_hash bit j is 0, otherwise value * 2 (to compensate for 50% sampling). template -static void PacSumAvgFinalizeCountersInternal(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void PacSumAvgFinalizeCountersInternal(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { auto state_ptrs = FlatVector::GetData(states); // Result is LIST - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); // Reserve space for all lists (64 elements each) idx_t total_elements = (offset + count) * 64; ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); - auto &child_validity = FlatVector::Validity(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); + auto &child_validity = FlatVector::ValidityMutable(child_vec); // correction factor for value scaling auto *bind = input.bind_data ? &input.bind_data->Cast() : nullptr; @@ -812,11 +919,17 @@ static void PacSumAvgFinalizeCountersInternal(Vector &states, AggregateInputData child_validity.SetAllValid(total_elements); } +#ifndef PAC_NOBUFFERING + ArenaAllocator scratch(input.allocator.GetUnderlyingAllocator()); +#endif for (idx_t i = 0; i < count; i++) { + auto *state = state_ptrs[i]; + auto *allocator = &input.allocator; #ifndef PAC_NOBUFFERING - PacSumFlushBuffer(*state_ptrs[i], *state_ptrs[i], input.allocator); + State temporary; + state = PrepareSumFinalization(*state, temporary, scratch, allocator); #endif - auto *s = state_ptrs[i]->GetState(); + auto *s = state->GetState(); // Set up the list entry - always needed even for NULL results idx_t base = (offset + i) * 64; @@ -834,13 +947,13 @@ static void PacSumAvgFinalizeCountersInternal(Vector &states, AggregateInputData update_count = s->update_count; #ifndef PAC_SIGNEDSUM // Double-sided mode: use SFINAE helpers - FlushPosState(s, input.allocator); + FlushPosState(s, *allocator); GetPosStateTotals(s, buf); // For signed int states, subtract neg_state counters element-wise. The neg_state's worlds are OR-ed // into key_hash: a world that only received negative values must not be zeroed below. - update_count += SubtractNegStateTotals(state_ptrs[i], buf, key_hash, input.allocator); + update_count += SubtractNegStateTotals(state, buf, key_hash, *allocator); #else - s->Flush(input.allocator); + s->Flush(*allocator); s->GetTotals(buf); #endif } @@ -867,31 +980,32 @@ static void PacSumAvgFinalizeCountersInternal(Vector &states, AggregateInputData } template -void PacSumAvgFinalizeCounters(Vector &states, AggregateInputData &input, Vector &result, idx_t count, idx_t offset) { +void PacSumAvgFinalizeCounters(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, + idx_t offset) { PacSumAvgFinalizeCountersInternal(states, input, result, count, offset); } // Instantiate counter finalize methods for as_sum -static void PacSumFinalizeCountersSigned(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void PacSumFinalizeCountersSigned(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { PacSumAvgFinalizeCounters, true>(states, input, result, count, offset); } -static void PacSumFinalizeCountersUnsigned(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void PacSumFinalizeCountersUnsigned(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { PacSumAvgFinalizeCounters, false>(states, input, result, count, offset); } -static void PacSumFinalizeCountersDouble(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void PacSumFinalizeCountersDouble(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { PacSumAvgFinalizeCounters(states, input, result, count, offset); } -static void DpSampleSumFinalizeCountersDouble(Vector &states, AggregateInputData &input, Vector &result, idx_t count, - idx_t offset) { +static void DpSampleSumFinalizeCountersDouble(Vector &states, AggregateFinalizeInputData &input, Vector &result, + idx_t count, idx_t offset) { PacSumAvgFinalizeCountersInternal( states, input, result, count, offset); } -static void DpSampleNullableSumFinalizeCountersDouble(Vector &states, AggregateInputData &input, Vector &result, +static void DpSampleNullableSumFinalizeCountersDouble(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { PacSumAvgFinalizeCountersInternal( states, input, result, count, offset); @@ -900,14 +1014,14 @@ static void DpSampleNullableSumFinalizeCountersDouble(Vector &states, AggregateI // Helper to register both 2-param and 3-param versions for as_sum static void AddCountersFcn(AggregateFunctionSet &set, const LogicalType &value_type, aggregate_size_t state_size, aggregate_initialize_t init, aggregate_update_t scatter, aggregate_combine_t combine, - aggregate_finalize_t finalize, aggregate_simple_update_t update) { + aggregate_finalize_t finalize, aggregate_cluster_update_t update) { auto list_double_type = LogicalType::LIST(PacFloatLogicalType()); - set.AddFunction(AggregateFunction("as_sum", {LogicalType::UBIGINT, value_type}, list_double_type, state_size, init, - scatter, combine, finalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, update, - PacSumBind)); - set.AddFunction(AggregateFunction("as_sum", {LogicalType::UBIGINT, value_type, LogicalType::DOUBLE}, - list_double_type, state_size, init, scatter, combine, finalize, - FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacSumBind)); + set.AddFunction(SpillableSum(AggregateFunction("as_sum", {LogicalType::UBIGINT, value_type}, list_double_type, + state_size, init, scatter, combine, finalize, + FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacSumBind))); + set.AddFunction(SpillableSum(AggregateFunction("as_sum", {LogicalType::UBIGINT, value_type, LogicalType::DOUBLE}, + list_double_type, state_size, init, scatter, combine, finalize, + FunctionNullHandling::DEFAULT_NULL_HANDLING, update, PacSumBind))); } void RegisterPacSumCountersFunctions(ExtensionLoader &loader) { @@ -1004,12 +1118,15 @@ struct DpSampleAvgState { PAC_FLOAT counts[64]; }; -static idx_t DpSampleAvgStateSize(const AggregateFunction &) { +static idx_t DpSampleAvgStateSize(AggregateStateInput &) { return sizeof(DpSampleAvgState); } -static void DpSampleAvgInitialize(const AggregateFunction &, data_ptr_t state_ptr) { - memset(state_ptr, 0, sizeof(DpSampleAvgState)); +static void DpSampleAvgInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_ptr = states[i]; + memset(state_ptr, 0, sizeof(DpSampleAvgState)); + } } static void DpSampleAvgUpdateOne(DpSampleAvgState &state, uint64_t key_hash, double sum_value, double count_value) { @@ -1021,40 +1138,43 @@ static void DpSampleAvgUpdateOne(DpSampleAvgState &state, uint64_t key_hash, dou } } -static void DpSampleAvgUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - int sample_lanes = GetPrivSampleLanes(aggr); +static void DpSampleAvgUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + int sample_lanes = GetPrivSampleLanes(aggr); - UnifiedVectorFormat hash_data, sum_data, count_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, sum_data); - inputs[2].ToUnifiedFormat(count, count_data); + UnifiedVectorFormat hash_data, sum_data, count_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(sum_data); + inputs[2].ToUnifiedFormat(count_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto sums = UnifiedVectorFormat::GetData(sum_data); - auto counts = UnifiedVectorFormat::GetData(count_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto sums = UnifiedVectorFormat::GetData(sum_data); + auto counts = UnifiedVectorFormat::GetData(count_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto s_idx = sum_data.sel->get_index(i); - auto c_idx = count_data.sel->get_index(i); - if (!hash_data.validity.RowIsValid(h_idx) || !sum_data.validity.RowIsValid(s_idx) || - !count_data.validity.RowIsValid(c_idx)) { - continue; + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto s_idx = sum_data.sel->get_index(i); + auto c_idx = count_data.sel->get_index(i); + if (!hash_data.validity.RowIsValid(h_idx) || !sum_data.validity.RowIsValid(s_idx) || + !count_data.validity.RowIsValid(c_idx)) { + continue; + } + auto key_hash = TransformPacUpdateHash(hashes[h_idx], sample_lanes); + DpSampleAvgUpdateOne(state, key_hash, sums[s_idx], counts[c_idx]); } - auto key_hash = TransformPacUpdateHash(hashes[h_idx], sample_lanes); - DpSampleAvgUpdateOne(state, key_hash, sums[s_idx], counts[c_idx]); - } + }); } static void DpSampleAvgScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { int sample_lanes = GetPrivSampleLanes(aggr); UnifiedVectorFormat hash_data, sum_data, count_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, sum_data); - inputs[2].ToUnifiedFormat(count, count_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(sum_data); + inputs[2].ToUnifiedFormat(count_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto sums = UnifiedVectorFormat::GetData(sum_data); @@ -1088,17 +1208,18 @@ static void DpSampleAvgCombine(Vector &src, Vector &dst, AggregateInputData &, i } } -static void DpSampleAvgFinalize(Vector &states, AggregateInputData &, Vector &result, idx_t count, idx_t offset) { +static void DpSampleAvgFinalize(Vector &states, AggregateFinalizeInputData &, Vector &result, idx_t count, + idx_t offset) { auto state_ptrs = FlatVector::GetData(states); - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * 64; ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); - auto &child_validity = FlatVector::Validity(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); + auto &child_validity = FlatVector::ValidityMutable(child_vec); child_validity.SetAllValid(total_elements); for (idx_t i = 0; i < count; i++) { @@ -1118,8 +1239,8 @@ static void DpSampleAvgFinalize(Vector &states, AggregateInputData &, Vector &re } } -static unique_ptr DpSampleAvgBind(ClientContext &ctx, AggregateFunction &, - vector> &) { +static unique_ptr DpSampleAvgBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); return MakeDpSampleBindData(ctx); } @@ -1149,20 +1270,26 @@ struct DpSampleMNullableSumState { uint8_t *occupied; }; -static idx_t DpSampleMSumStateSize(const AggregateFunction &) { +static idx_t DpSampleMSumStateSize(AggregateStateInput &) { return sizeof(DpSampleMSumState); } -static void DpSampleMSumInitialize(const AggregateFunction &, data_ptr_t state_ptr) { - memset(state_ptr, 0, sizeof(DpSampleMSumState)); +static void DpSampleMSumInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_ptr = states[i]; + memset(state_ptr, 0, sizeof(DpSampleMSumState)); + } } -static idx_t DpSampleMNullableSumStateSize(const AggregateFunction &) { +static idx_t DpSampleMNullableSumStateSize(AggregateStateInput &) { return sizeof(DpSampleMNullableSumState); } -static void DpSampleMNullableSumInitialize(const AggregateFunction &, data_ptr_t state_ptr) { - memset(state_ptr, 0, sizeof(DpSampleMNullableSumState)); +static void DpSampleMNullableSumInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_ptr = states[i]; + memset(state_ptr, 0, sizeof(DpSampleMNullableSumState)); + } } static PAC_FLOAT *EnsureDpSampleMArray(PAC_FLOAT *&values, ArenaAllocator &allocator, int sample_count) { @@ -1187,29 +1314,32 @@ static void DpSampleMSumUpdateOne(DpSampleMSumState &state, ArenaAllocator &allo values[DpSassLaneIndex(hash, sample_count)] += static_cast(value); } -static void DpSampleMSumUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - int sample_count = GetPrivSampleCount(aggr); - UnifiedVectorFormat hash_data, value_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto values = UnifiedVectorFormat::GetData(value_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto v_idx = value_data.sel->get_index(i); - if (hash_data.validity.RowIsValid(h_idx) && value_data.validity.RowIsValid(v_idx)) { - DpSampleMSumUpdateOne(state, aggr.allocator, hashes[h_idx], values[v_idx], sample_count); +static void DpSampleMSumUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + int sample_count = GetPrivSampleCount(aggr); + UnifiedVectorFormat hash_data, value_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto values = UnifiedVectorFormat::GetData(value_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto v_idx = value_data.sel->get_index(i); + if (hash_data.validity.RowIsValid(h_idx) && value_data.validity.RowIsValid(v_idx)) { + DpSampleMSumUpdateOne(state, aggr.allocator, hashes[h_idx], values[v_idx], sample_count); + } } - } + }); } static void DpSampleMSumScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { int sample_count = GetPrivSampleCount(aggr); UnifiedVectorFormat hash_data, value_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); @@ -1232,31 +1362,33 @@ static void DpSampleMNullableSumUpdateOne(DpSampleMNullableSumState &state, Aren occupied[lane] = 1; } -static void DpSampleMNullableSumUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, - idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - int sample_count = GetPrivSampleCount(aggr); - UnifiedVectorFormat hash_data, value_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto values = UnifiedVectorFormat::GetData(value_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto v_idx = value_data.sel->get_index(i); - if (hash_data.validity.RowIsValid(h_idx) && value_data.validity.RowIsValid(v_idx)) { - DpSampleMNullableSumUpdateOne(state, aggr.allocator, hashes[h_idx], values[v_idx], sample_count); +static void DpSampleMNullableSumUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + int sample_count = GetPrivSampleCount(aggr); + UnifiedVectorFormat hash_data, value_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto values = UnifiedVectorFormat::GetData(value_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto v_idx = value_data.sel->get_index(i); + if (hash_data.validity.RowIsValid(h_idx) && value_data.validity.RowIsValid(v_idx)) { + DpSampleMNullableSumUpdateOne(state, aggr.allocator, hashes[h_idx], values[v_idx], sample_count); + } } - } + }); } static void DpSampleMNullableSumScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { int sample_count = GetPrivSampleCount(aggr); UnifiedVectorFormat hash_data, value_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); @@ -1270,47 +1402,50 @@ static void DpSampleMNullableSumScatterUpdate(Vector inputs[], AggregateInputDat } } -static void DpSampleMSumIfUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - int sample_count = GetPrivSampleCount(aggr); - UnifiedVectorFormat hash_data, value_data, predicate_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - inputs[2].ToUnifiedFormat(count, predicate_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto values = UnifiedVectorFormat::GetData(value_data); - auto &predicate_child = ListVector::GetEntry(inputs[2]); - UnifiedVectorFormat predicate_child_data; - predicate_child.ToUnifiedFormat(ListVector::GetListSize(inputs[2]), predicate_child_data); - auto predicate_values = UnifiedVectorFormat::GetData(predicate_child_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto v_idx = value_data.sel->get_index(i); - auto p_idx = predicate_data.sel->get_index(i); - if (!hash_data.validity.RowIsValid(h_idx) || !value_data.validity.RowIsValid(v_idx)) { - continue; - } - idx_t lane = DpSassLaneIndex(hashes[h_idx], sample_count); - if (DpSampleMPredicateLaneIsTrue(predicate_data, predicate_child_data, predicate_values, p_idx, lane)) { - auto values_ptr = EnsureDpSampleMArray(state.values, aggr.allocator, sample_count); - values_ptr[lane] += static_cast(values[v_idx]); +static void DpSampleMSumIfUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + int sample_count = GetPrivSampleCount(aggr); + UnifiedVectorFormat hash_data, value_data, predicate_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + inputs[2].ToUnifiedFormat(predicate_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto values = UnifiedVectorFormat::GetData(value_data); + auto &predicate_child = ListVector::GetChild(inputs[2]); + UnifiedVectorFormat predicate_child_data; + predicate_child.ToUnifiedFormat(predicate_child_data); + auto predicate_values = UnifiedVectorFormat::GetData(predicate_child_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto v_idx = value_data.sel->get_index(i); + auto p_idx = predicate_data.sel->get_index(i); + if (!hash_data.validity.RowIsValid(h_idx) || !value_data.validity.RowIsValid(v_idx)) { + continue; + } + idx_t lane = DpSassLaneIndex(hashes[h_idx], sample_count); + if (DpSampleMPredicateLaneIsTrue(predicate_data, predicate_child_data, predicate_values, p_idx, lane)) { + auto values_ptr = EnsureDpSampleMArray(state.values, aggr.allocator, sample_count); + values_ptr[lane] += static_cast(values[v_idx]); + } } - } + }); } static void DpSampleMSumIfScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { int sample_count = GetPrivSampleCount(aggr); UnifiedVectorFormat hash_data, value_data, predicate_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - inputs[2].ToUnifiedFormat(count, predicate_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + inputs[2].ToUnifiedFormat(predicate_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); - auto &predicate_child = ListVector::GetEntry(inputs[2]); + auto &predicate_child = ListVector::GetChild(inputs[2]); UnifiedVectorFormat predicate_child_data; - predicate_child.ToUnifiedFormat(ListVector::GetListSize(inputs[2]), predicate_child_data); + predicate_child.ToUnifiedFormat(predicate_child_data); auto predicate_values = UnifiedVectorFormat::GetData(predicate_child_data); for (idx_t i = 0; i < count; i++) { auto h_idx = hash_data.sel->get_index(i); @@ -1366,17 +1501,18 @@ static void DpSampleMNullableSumCombine(Vector &src, Vector &dst, AggregateInput } } -static void DpSampleMSumFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, idx_t offset) { +static void DpSampleMSumFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, + idx_t offset) { int sample_count = GetPrivSampleCount(input); auto state_ptrs = FlatVector::GetData(states); - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * static_cast(sample_count); ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); bool sample_rescale = input.bind_data ? input.bind_data->Cast().sample_rescale : true; // SUM is extensive: multiply disjoint 1/m-sample totals by m only when the caller // wants a full-data estimator. AVG and MIN/MAX are not rescaled. @@ -1396,17 +1532,17 @@ static void DpSampleMSumFinalize(Vector &states, AggregateInputData &input, Vect } } -static void DpSampleMNullableSumFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void DpSampleMNullableSumFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { int sample_count = GetPrivSampleCount(input); auto state_ptrs = FlatVector::GetData(states); - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * static_cast(sample_count); ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); - auto &child_validity = FlatVector::Validity(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); + auto &child_validity = FlatVector::ValidityMutable(child_vec); child_validity.SetAllValid(total_elements); bool sample_rescale = input.bind_data ? input.bind_data->Cast().sample_rescale : true; PAC_FLOAT rescale = sample_rescale ? static_cast(sample_count) : PAC_FLOAT(1.0); @@ -1426,8 +1562,8 @@ static void DpSampleMNullableSumFinalize(Vector &states, AggregateInputData &inp } } -static unique_ptr DpSampleMBind(ClientContext &ctx, AggregateFunction &, - vector> &) { +static unique_ptr DpSampleMBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); int sample_lanes = GetDpSampleLanes(ctx); if (sample_lanes != 1) { throw InvalidInputException("as_sample_m_*: dp_sample_lanes must be 1 when dp_sass_m > 64"); @@ -1482,12 +1618,15 @@ struct DpSampleMAvgState { PAC_FLOAT *counts; }; -static idx_t DpSampleMAvgStateSize(const AggregateFunction &) { +static idx_t DpSampleMAvgStateSize(AggregateStateInput &) { return sizeof(DpSampleMAvgState); } -static void DpSampleMAvgInitialize(const AggregateFunction &, data_ptr_t state_ptr) { - memset(state_ptr, 0, sizeof(DpSampleMAvgState)); +static void DpSampleMAvgInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_ptr = states[i]; + memset(state_ptr, 0, sizeof(DpSampleMAvgState)); + } } static void DpSampleMAvgUpdateOne(DpSampleMAvgState &state, ArenaAllocator &allocator, uint64_t hash, double sum_value, @@ -1499,34 +1638,37 @@ static void DpSampleMAvgUpdateOne(DpSampleMAvgState &state, ArenaAllocator &allo counts[lane] += static_cast(count_value); } -static void DpSampleMAvgUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - int sample_count = GetPrivSampleCount(aggr); - UnifiedVectorFormat hash_data, sum_data, count_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, sum_data); - inputs[2].ToUnifiedFormat(count, count_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto sums = UnifiedVectorFormat::GetData(sum_data); - auto counts = UnifiedVectorFormat::GetData(count_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto s_idx = sum_data.sel->get_index(i); - auto c_idx = count_data.sel->get_index(i); - if (hash_data.validity.RowIsValid(h_idx) && sum_data.validity.RowIsValid(s_idx) && - count_data.validity.RowIsValid(c_idx)) { - DpSampleMAvgUpdateOne(state, aggr.allocator, hashes[h_idx], sums[s_idx], counts[c_idx], sample_count); +static void DpSampleMAvgUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + int sample_count = GetPrivSampleCount(aggr); + UnifiedVectorFormat hash_data, sum_data, count_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(sum_data); + inputs[2].ToUnifiedFormat(count_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto sums = UnifiedVectorFormat::GetData(sum_data); + auto counts = UnifiedVectorFormat::GetData(count_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto s_idx = sum_data.sel->get_index(i); + auto c_idx = count_data.sel->get_index(i); + if (hash_data.validity.RowIsValid(h_idx) && sum_data.validity.RowIsValid(s_idx) && + count_data.validity.RowIsValid(c_idx)) { + DpSampleMAvgUpdateOne(state, aggr.allocator, hashes[h_idx], sums[s_idx], counts[c_idx], sample_count); + } } - } + }); } static void DpSampleMAvgScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { int sample_count = GetPrivSampleCount(aggr); UnifiedVectorFormat hash_data, sum_data, count_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, sum_data); - inputs[2].ToUnifiedFormat(count, count_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(sum_data); + inputs[2].ToUnifiedFormat(count_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto sums = UnifiedVectorFormat::GetData(sum_data); auto counts = UnifiedVectorFormat::GetData(count_data); @@ -1562,11 +1704,12 @@ static void DpSampleMAvgCombine(Vector &src, Vector &dst, AggregateInputData &ag } } -static void DpSampleMAvgFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, idx_t offset) { +static void DpSampleMAvgFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, + idx_t offset) { int sample_count = GetPrivSampleCount(input); auto state_ptrs = FlatVector::GetData(states); - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * static_cast(sample_count); ListVector::Reserve(result, total_elements); @@ -1574,8 +1717,8 @@ static void DpSampleMAvgFinalize(Vector &states, AggregateInputData &input, Vect // Form each lane's ratio before the SAA terminal combines lanes. The common m factor // would cancel between numerator and denominator, so AVG deliberately ignores rescaling. - auto child_data = FlatVector::GetData(child_vec); - auto &child_validity = FlatVector::Validity(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); + auto &child_validity = FlatVector::ValidityMutable(child_vec); child_validity.SetAllValid(total_elements); for (idx_t i = 0; i < count; i++) { auto *state = state_ptrs[i]; @@ -1614,29 +1757,32 @@ static void PacSampleMSumUpdateOne(DpSampleMSumState &state, ArenaAllocator &all ForEachPacMSelectedLane(hash, bind, [&](idx_t lane) { values[lane] += static_cast(value); }); } -static void PacSampleMSumUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - auto &bind = aggr.bind_data->Cast(); - UnifiedVectorFormat hash_data, value_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto values = UnifiedVectorFormat::GetData(value_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto v_idx = value_data.sel->get_index(i); - if (hash_data.validity.RowIsValid(h_idx) && value_data.validity.RowIsValid(v_idx)) { - PacSampleMSumUpdateOne(state, aggr.allocator, hashes[h_idx], values[v_idx], bind); +static void PacSampleMSumUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + auto &bind = aggr.bind_data->Cast(); + UnifiedVectorFormat hash_data, value_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto values = UnifiedVectorFormat::GetData(value_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto v_idx = value_data.sel->get_index(i); + if (hash_data.validity.RowIsValid(h_idx) && value_data.validity.RowIsValid(v_idx)) { + PacSampleMSumUpdateOne(state, aggr.allocator, hashes[h_idx], values[v_idx], bind); + } } - } + }); } static void PacSampleMSumScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { auto &bind = aggr.bind_data->Cast(); UnifiedVectorFormat hash_data, value_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); @@ -1650,49 +1796,52 @@ static void PacSampleMSumScatterUpdate(Vector inputs[], AggregateInputData &aggr } } -static void PacSampleMSumIfUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - auto &bind = aggr.bind_data->Cast(); - UnifiedVectorFormat hash_data, value_data, predicate_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - inputs[2].ToUnifiedFormat(count, predicate_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto values = UnifiedVectorFormat::GetData(value_data); - auto &predicate_child = ListVector::GetEntry(inputs[2]); - UnifiedVectorFormat predicate_child_data; - predicate_child.ToUnifiedFormat(ListVector::GetListSize(inputs[2]), predicate_child_data); - auto predicate_values = UnifiedVectorFormat::GetData(predicate_child_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto v_idx = value_data.sel->get_index(i); - auto p_idx = predicate_data.sel->get_index(i); - if (!hash_data.validity.RowIsValid(h_idx) || !value_data.validity.RowIsValid(v_idx)) { - continue; - } - auto values_ptr = EnsureDpSampleMArray(state.values, aggr.allocator, bind.sample_count); - ForEachPacMSelectedLane(hashes[h_idx], bind, [&](idx_t lane) { - if (DpSampleMPredicateLaneIsTrue(predicate_data, predicate_child_data, predicate_values, p_idx, lane)) { - values_ptr[lane] += static_cast(values[v_idx]); +static void PacSampleMSumIfUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + auto &bind = aggr.bind_data->Cast(); + UnifiedVectorFormat hash_data, value_data, predicate_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + inputs[2].ToUnifiedFormat(predicate_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto values = UnifiedVectorFormat::GetData(value_data); + auto &predicate_child = ListVector::GetChild(inputs[2]); + UnifiedVectorFormat predicate_child_data; + predicate_child.ToUnifiedFormat(predicate_child_data); + auto predicate_values = UnifiedVectorFormat::GetData(predicate_child_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto v_idx = value_data.sel->get_index(i); + auto p_idx = predicate_data.sel->get_index(i); + if (!hash_data.validity.RowIsValid(h_idx) || !value_data.validity.RowIsValid(v_idx)) { + continue; } - }); - } + auto values_ptr = EnsureDpSampleMArray(state.values, aggr.allocator, bind.sample_count); + ForEachPacMSelectedLane(hashes[h_idx], bind, [&](idx_t lane) { + if (DpSampleMPredicateLaneIsTrue(predicate_data, predicate_child_data, predicate_values, p_idx, lane)) { + values_ptr[lane] += static_cast(values[v_idx]); + } + }); + } + }); } static void PacSampleMSumIfScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { auto &bind = aggr.bind_data->Cast(); UnifiedVectorFormat hash_data, value_data, predicate_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, value_data); - inputs[2].ToUnifiedFormat(count, predicate_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(value_data); + inputs[2].ToUnifiedFormat(predicate_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto values = UnifiedVectorFormat::GetData(value_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); - auto &predicate_child = ListVector::GetEntry(inputs[2]); + auto &predicate_child = ListVector::GetChild(inputs[2]); UnifiedVectorFormat predicate_child_data; - predicate_child.ToUnifiedFormat(ListVector::GetListSize(inputs[2]), predicate_child_data); + predicate_child.ToUnifiedFormat(predicate_child_data); auto predicate_values = UnifiedVectorFormat::GetData(predicate_child_data); for (idx_t i = 0; i < count; i++) { auto h_idx = hash_data.sel->get_index(i); @@ -1711,18 +1860,18 @@ static void PacSampleMSumIfScatterUpdate(Vector inputs[], AggregateInputData &ag } } -static void PacSampleMSumFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void PacSampleMSumFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { int sample_count = GetPrivSampleCount(input); auto state_ptrs = FlatVector::GetData(states); - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); idx_t total_elements = (offset + count) * static_cast(sample_count); ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); for (idx_t i = 0; i < count; i++) { idx_t base = (offset + i) * static_cast(sample_count); list_entries[offset + i].offset = base; @@ -1748,34 +1897,37 @@ static void PacSampleMAvgUpdateOne(DpSampleMAvgState &state, ArenaAllocator &all }); } -static void PacSampleMAvgUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_ptr, idx_t count) { - auto &state = *reinterpret_cast(state_ptr); - auto &bind = aggr.bind_data->Cast(); - UnifiedVectorFormat hash_data, sum_data, count_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, sum_data); - inputs[2].ToUnifiedFormat(count, count_data); - auto hashes = UnifiedVectorFormat::GetData(hash_data); - auto sums = UnifiedVectorFormat::GetData(sum_data); - auto counts = UnifiedVectorFormat::GetData(count_data); - for (idx_t i = 0; i < count; i++) { - auto h_idx = hash_data.sel->get_index(i); - auto s_idx = sum_data.sel->get_index(i); - auto c_idx = count_data.sel->get_index(i); - if (hash_data.validity.RowIsValid(h_idx) && sum_data.validity.RowIsValid(s_idx) && - count_data.validity.RowIsValid(c_idx)) { - PacSampleMAvgUpdateOne(state, aggr.allocator, hashes[h_idx], sums[s_idx], counts[c_idx], bind); +static void PacSampleMAvgUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_ptr, idx_t count) { + auto &state = *reinterpret_cast(state_ptr); + auto &bind = aggr.bind_data->Cast(); + UnifiedVectorFormat hash_data, sum_data, count_data; + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(sum_data); + inputs[2].ToUnifiedFormat(count_data); + auto hashes = UnifiedVectorFormat::GetData(hash_data); + auto sums = UnifiedVectorFormat::GetData(sum_data); + auto counts = UnifiedVectorFormat::GetData(count_data); + for (idx_t i = 0; i < count; i++) { + auto h_idx = hash_data.sel->get_index(i); + auto s_idx = sum_data.sel->get_index(i); + auto c_idx = count_data.sel->get_index(i); + if (hash_data.validity.RowIsValid(h_idx) && sum_data.validity.RowIsValid(s_idx) && + count_data.validity.RowIsValid(c_idx)) { + PacSampleMAvgUpdateOne(state, aggr.allocator, hashes[h_idx], sums[s_idx], counts[c_idx], bind); + } } - } + }); } static void PacSampleMAvgScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { auto &bind = aggr.bind_data->Cast(); UnifiedVectorFormat hash_data, sum_data, count_data, state_data; - inputs[0].ToUnifiedFormat(count, hash_data); - inputs[1].ToUnifiedFormat(count, sum_data); - inputs[2].ToUnifiedFormat(count, count_data); - states.ToUnifiedFormat(count, state_data); + inputs[0].ToUnifiedFormat(hash_data); + inputs[1].ToUnifiedFormat(sum_data); + inputs[2].ToUnifiedFormat(count_data); + states.ToUnifiedFormat(state_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto sums = UnifiedVectorFormat::GetData(sum_data); auto counts = UnifiedVectorFormat::GetData(count_data); @@ -1792,8 +1944,8 @@ static void PacSampleMAvgScatterUpdate(Vector inputs[], AggregateInputData &aggr } } -static unique_ptr PacSampleMBind(ClientContext &ctx, AggregateFunction &, - vector> &) { +static unique_ptr PacSampleMBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); return MakePacSampleMBindData(ctx); } diff --git a/src/aggregates/as_sum_state_export.cpp b/src/aggregates/as_sum_state_export.cpp new file mode 100644 index 00000000..719bfd36 --- /dev/null +++ b/src/aggregates/as_sum_state_export.cpp @@ -0,0 +1,197 @@ +#include "aggregates/as_sum_state_export.hpp" +#include "aggregates/as_sum.hpp" +#include "aggregates/aggregate_state_export.hpp" +#include "duckdb/common/vector/flat_vector.hpp" +#include "duckdb/common/vector/string_vector.hpp" + +namespace duckdb { +namespace { + +// Spill the native counters, not finalized totals: finalization can round and +// applies sampling correction. No process-local pointers enter the exported BLOB. +template +struct SumStateExport; + +template +struct SumStateExport> { + using State = PacSumIntState; + + template + static void WriteLevel(AggregateStateWriter &out, const T *ptr) { + out.Write(uint8_t(ptr != nullptr)); + if (ptr) { + out.WriteData(ptr, sizeof(T) * N); + } + } + + template + static void ReadLevel(AggregateStateReader &in, T *&ptr, ArenaAllocator &allocator) { + if (!in.ReadFlag()) { + return; + } + ptr = reinterpret_cast(allocator.Allocate(sizeof(T) * N)); + in.ReadData(ptr, sizeof(T) * N); + } + + static void Write(AggregateStateWriter &out, const State &state) { + out.Write(state.key_hash); + out.Write(state.update_count); +#ifndef PAC_EXACTSUM + out.Write(state.max_level_used); + for (int k = 0; k <= state.max_level_used; k++) { + out.Write(state.exact_total[k]); + WriteLevel(out, state.u.levels[k]); + } +#elif defined(PAC_NOCASCADING) + out.Write(state.probabilistic_total128); +#else + out.Write(state.exact_total8); + out.Write(state.exact_total16); + out.Write(state.exact_total32); + out.Write(state.exact_total64); + WriteLevel(out, state.probabilistic_total8); + WriteLevel(out, state.probabilistic_total16); + WriteLevel(out, state.probabilistic_total32); + WriteLevel(out, state.probabilistic_total64); + WriteLevel(out, state.probabilistic_total128); +#endif + } + + static void Read(AggregateStateReader &in, State &state, ArenaAllocator &allocator) { + memset(&state, 0, sizeof(state)); + in.Read(state.key_hash); + in.Read(state.update_count); +#ifndef PAC_EXACTSUM + in.Read(state.max_level_used); + state.inline_level_idx = -1; + if (state.max_level_used < -1 || state.max_level_used >= PAC_APPROX_NUM_LEVELS) { + throw InvalidInputException("Invalid stochastic SUM cascade level"); + } + for (int k = 0; k <= state.max_level_used; k++) { + in.Read(state.exact_total[k]); +#ifdef PAC_SIGNEDSUM + if (state.exact_total[k] < State::UNDERFLOW_THRESHOLD || state.exact_total[k] > State::OVERFLOW_THRESHOLD) { +#else + if (state.exact_total[k] < 0 || state.exact_total[k] > PacSumIntState::OVERFLOW_THRESHOLD) { +#endif + throw InvalidInputException("Invalid stochastic SUM cascade budget"); + } + if (!in.ReadFlag()) { + throw InvalidInputException("Missing stochastic SUM cascade level"); + } + // Retain the inline optimization when all pointer slots fit before it. + if (state.max_level_used < 9) { + state.AllocateLevel(allocator, k); + } else { + state.u.levels[k] = + reinterpret_cast(allocator.Allocate(sizeof(uint64_t) * PAC_APPROX_SWAR_ELEMENTS)); + } + in.ReadData(state.u.levels[k], sizeof(uint64_t) * PAC_APPROX_SWAR_ELEMENTS); + } +#elif defined(PAC_NOCASCADING) + in.Read(state.probabilistic_total128); +#else + in.Read(state.exact_total8); + in.Read(state.exact_total16); + in.Read(state.exact_total32); + in.Read(state.exact_total64); + ReadLevel(in, state.probabilistic_total8, allocator); + ReadLevel(in, state.probabilistic_total16, allocator); + ReadLevel(in, state.probabilistic_total32, allocator); + ReadLevel(in, state.probabilistic_total64, allocator); + ReadLevel(in, state.probabilistic_total128, allocator); +#endif + } +}; + +template <> +struct SumStateExport { + static void Write(AggregateStateWriter &out, const PacSumDoubleState &state) { + out.Write(state.key_hash); + out.Write(state.update_count); + out.Write(state.probabilistic_total); + } + + static void Read(AggregateStateReader &in, PacSumDoubleState &state, ArenaAllocator &) { + in.Read(state.key_hash); + in.Read(state.update_count); + in.Read(state.probabilistic_total); + } +}; + +#ifndef PAC_NOBUFFERING +template +struct SumStateExport> { + using State = PacSumStateWrapper; + + static void Write(AggregateStateWriter &out, const State &state) { + auto buffered = state.n_buffered & State::BUF_MASK; + out.Write(uint8_t(buffered)); + for (idx_t i = 0; i < buffered; i++) { + out.Write(state.val_buf[i]); + out.Write(state.hash_buf[i]); + } + auto *inner = state.GetState(); + out.Write(uint8_t(inner != nullptr)); + if (inner) { + SumStateExport::Write(out, *inner); + } +#ifndef PAC_SIGNEDSUM + out.Write(uint8_t(state.neg_state != nullptr)); + if (state.neg_state) { + SumStateExport>::Write(out, *state.neg_state); + } +#endif + } + + static void Read(AggregateStateReader &in, State &state, ArenaAllocator &allocator) { + memset(&state, 0, sizeof(state)); + uint8_t buffered; + in.Read(buffered); + if (buffered > State::BUF_SIZE) { + throw InvalidInputException("Invalid stochastic SUM buffer size"); + } + for (idx_t i = 0; i < buffered; i++) { + in.Read(state.val_buf[i]); + in.Read(state.hash_buf[i]); + } + if (in.ReadFlag()) { + auto *inner = state.EnsureState(allocator); + SumStateExport::Read(in, *inner, allocator); + } + state.n_buffered |= buffered; +#ifndef PAC_SIGNEDSUM + if (in.ReadFlag()) { + auto *negative = state.EnsureNegState(allocator); + SumStateExport>::Read(in, *negative, allocator); + } +#endif + } +}; +#endif + +} // namespace + +void RegisterSumStateExport(AggregateFunction &function) { + // Linkers may fold identical initializer functions, so callback addresses do not identify the state layout. + auto value_type = function.GetSignature().GetParameter(1).GetType().id(); + if (value_type == LogicalTypeId::FLOAT || value_type == LogicalTypeId::DOUBLE || + value_type == LogicalTypeId::UHUGEINT +#ifndef PAC_EXACTSUM + || value_type == LogicalTypeId::HUGEINT +#endif + ) { + AggregateStateExport>::Register(function); + } else if (value_type == LogicalTypeId::UTINYINT || value_type == LogicalTypeId::USMALLINT || + value_type == LogicalTypeId::UINTEGER || value_type == LogicalTypeId::UBIGINT) { + AggregateStateExport, SumStateExport>>::Register(function); + } else if (value_type == LogicalTypeId::BOOLEAN || value_type == LogicalTypeId::TINYINT || + value_type == LogicalTypeId::SMALLINT || value_type == LogicalTypeId::INTEGER || + value_type == LogicalTypeId::BIGINT || value_type == LogicalTypeId::HUGEINT) { + AggregateStateExport, SumStateExport>>::Register(function); + } else { + throw InternalException("Unsupported stochastic SUM state type %s", LogicalType(value_type).ToString()); + } +} + +} // namespace duckdb diff --git a/src/aggregates/dp_approx_bounds.cpp b/src/aggregates/dp_approx_bounds.cpp index c936a61f..2706be5c 100644 --- a/src/aggregates/dp_approx_bounds.cpp +++ b/src/aggregates/dp_approx_bounds.cpp @@ -1,3 +1,5 @@ +#include "duckdb/common/vector/struct_vector.hpp" +#include "duckdb/common/vector/list_vector.hpp" #include "aggregates/dp_approx_bounds.hpp" #include "aggregates/as_clip_aggr.hpp" @@ -111,7 +113,7 @@ DpApproxBoundsParameters GetDpApproxBoundsParameters(const LogicalType &type) { } static double EvaluateConfig(ClientContext &context, const Expression &expression, const string &function_name, - const string &name) { + const Identifier &name) { if (!expression.IsFoldable()) { throw InvalidInputException(function_name + ": " + name + " must be a constant"); } @@ -128,32 +130,41 @@ static idx_t EvaluateBinCount(ClientContext &context, const Expression &expressi return static_cast(value); } -static unique_ptr BindDpApproxBounds(ClientContext &context, AggregateFunction &function, - vector> &arguments) { +static unique_ptr BindDpApproxBounds(BindAggregateFunctionInput &bind_input) { + auto &context = bind_input.GetClientContext(); + auto &function = bind_input.GetBoundFunction(); + auto &arguments = bind_input.GetArguments(); if (arguments.size() != 5 && arguments.size() != 6) { - throw InternalException(function.name + ": invalid internal argument count"); + throw InternalException(function.GetName().GetIdentifierName() + ": invalid internal argument count"); } - double epsilon = EvaluateConfig(context, *arguments[1], function.name, "epsilon"); - double max_groups = EvaluateConfig(context, *arguments[2], function.name, "max groups"); + double epsilon = EvaluateConfig(context, *arguments[1], function.GetName().GetIdentifierName(), "epsilon"); + double max_groups = EvaluateConfig(context, *arguments[2], function.GetName().GetIdentifierName(), "max groups"); double max_contributions = - arguments.size() == 6 ? EvaluateConfig(context, *arguments[3], function.name, "max contributions") : 1.0; + arguments.size() == 6 + ? EvaluateConfig(context, *arguments[3], function.GetName().GetIdentifierName(), "max contributions") + : 1.0; idx_t profile_offset = arguments.size() == 6 ? 4 : 3; - double scale = EvaluateConfig(context, *arguments[profile_offset], function.name, "scale"); - idx_t num_bins = EvaluateBinCount(context, *arguments[profile_offset + 1], function.name); + double scale = EvaluateConfig(context, *arguments[profile_offset], function.GetName().GetIdentifierName(), "scale"); + idx_t num_bins = EvaluateBinCount(context, *arguments[profile_offset + 1], function.GetName().GetIdentifierName()); if (!std::isfinite(epsilon) || epsilon <= 0.0) { - throw InvalidInputException(function.name + ": epsilon must be a positive finite number"); + throw InvalidInputException(function.GetName().GetIdentifierName() + + ": epsilon must be a positive finite number"); } if (!std::isfinite(max_groups) || max_groups <= 0.0) { - throw InvalidInputException(function.name + ": max groups must be a positive finite number"); + throw InvalidInputException(function.GetName().GetIdentifierName() + + ": max groups must be a positive finite number"); } if (!std::isfinite(max_contributions) || max_contributions <= 0.0) { - throw InvalidInputException(function.name + ": max contributions must be a positive finite number"); + throw InvalidInputException(function.GetName().GetIdentifierName() + + ": max contributions must be a positive finite number"); } if (arguments.size() == 6 && std::round(max_contributions) != max_contributions) { - throw InvalidInputException(function.name + ": max contributions must be integer-valued"); + throw InvalidInputException(function.GetName().GetIdentifierName() + + ": max contributions must be integer-valued"); } if (!std::isfinite(scale) || scale <= 0.0) { - throw InvalidInputException(function.name + ": scale must be a positive finite number"); + throw InvalidInputException(function.GetName().GetIdentifierName() + + ": scale must be a positive finite number"); } auto result = make_uniq(); result->epsilon = epsilon; @@ -165,12 +176,15 @@ static unique_ptr BindDpApproxBounds(ClientContext &context, Aggre return std::move(result); } -static idx_t DpApproxBoundsStateSize(const AggregateFunction &) { +static idx_t DpApproxBoundsStateSize(AggregateStateInput &) { return sizeof(DpApproxBoundsState); } -static void DpApproxBoundsInitialize(const AggregateFunction &, data_ptr_t state_p) { - memset(state_p, 0, sizeof(DpApproxBoundsState)); +static void DpApproxBoundsInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_p = states[i]; + memset(state_p, 0, sizeof(DpApproxBoundsState)); + } } static DpApproxBoundsBin *EnsureBins(DpApproxBoundsBin *&bins, idx_t num_bins, ArenaAllocator &allocator) { @@ -227,7 +241,7 @@ struct DpApproxBoundsInputs { const double *values; DpApproxBoundsInputs(Vector inputs[], idx_t count) { - inputs[0].ToUnifiedFormat(count, value); + inputs[0].ToUnifiedFormat(value); values = UnifiedVectorFormat::GetData(value); } @@ -249,15 +263,18 @@ static void UpdateRows(Vector inputs[], AggregateInputData &input, idx_t count, } } -static void DpApproxBoundsUpdate(Vector inputs[], AggregateInputData &input, idx_t, data_ptr_t state_p, idx_t count) { - auto state = reinterpret_cast(state_p); - UpdateRows(inputs, input, count, [state](idx_t) { return state; }); +static void DpApproxBoundsUpdate(Vector inputs[], AggregateInputData &input, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { + auto state = reinterpret_cast(state_p); + UpdateRows(inputs, input, count, [state](idx_t) { return state; }); + }); } static void DpApproxBoundsScatterUpdate(Vector inputs[], AggregateInputData &input, idx_t, Vector &states, idx_t count) { UnifiedVectorFormat state_data; - states.ToUnifiedFormat(count, state_data); + states.ToUnifiedFormat(state_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); UpdateRows(inputs, input, count, [&](idx_t row) { return state_ptrs[state_data.sel->get_index(row)]; }); } @@ -475,8 +492,9 @@ struct DpBoundedValueListState { idx_t size; }; -static unique_ptr BindDpBoundedValueList(ClientContext &context, AggregateFunction &, - vector> &arguments) { +static unique_ptr BindDpBoundedValueList(BindAggregateFunctionInput &bind_input) { + auto &context = bind_input.GetClientContext(); + auto &arguments = bind_input.GetArguments(); if (arguments.size() != 2 || !arguments[1]->IsFoldable()) { throw InvalidInputException("dp_bounded_value_list: max values must be a constant"); } @@ -493,12 +511,15 @@ static unique_ptr BindDpBoundedValueList(ClientContext &context, A return std::move(result); } -static idx_t DpBoundedValueListStateSize(const AggregateFunction &) { +static idx_t DpBoundedValueListStateSize(AggregateStateInput &) { return sizeof(DpBoundedValueListState); } -static void DpBoundedValueListInitialize(const AggregateFunction &, data_ptr_t state_p) { - memset(state_p, 0, sizeof(DpBoundedValueListState)); +static void DpBoundedValueListInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_p = states[i]; + memset(state_p, 0, sizeof(DpBoundedValueListState)); + } } static uint64_t BoundedValueScore(double value) { @@ -555,7 +576,7 @@ static void InsertBoundedValue(DpBoundedValueListState &state, double value, idx template static void UpdateBoundedValueRows(Vector inputs[], AggregateInputData &input, idx_t count, STATE_GETTER get_state) { UnifiedVectorFormat values; - inputs[0].ToUnifiedFormat(count, values); + inputs[0].ToUnifiedFormat(values); auto value_data = UnifiedVectorFormat::GetData(values); auto max_values = input.bind_data->Cast().max_values; for (idx_t row = 0; row < count; row++) { @@ -567,16 +588,18 @@ static void UpdateBoundedValueRows(Vector inputs[], AggregateInputData &input, i } } -static void DpBoundedValueListUpdate(Vector inputs[], AggregateInputData &input, idx_t, data_ptr_t state_p, - idx_t count) { - auto state = reinterpret_cast(state_p); - UpdateBoundedValueRows(inputs, input, count, [state](idx_t) { return state; }); +static void DpBoundedValueListUpdate(Vector inputs[], AggregateInputData &input, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { + auto state = reinterpret_cast(state_p); + UpdateBoundedValueRows(inputs, input, count, [state](idx_t) { return state; }); + }); } static void DpBoundedValueListScatterUpdate(Vector inputs[], AggregateInputData &input, idx_t, Vector &states, idx_t count) { UnifiedVectorFormat state_data; - states.ToUnifiedFormat(count, state_data); + states.ToUnifiedFormat(state_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); UpdateBoundedValueRows(inputs, input, count, [&](idx_t row) { return state_ptrs[state_data.sel->get_index(row)]; }); } @@ -592,18 +615,18 @@ static void DpBoundedValueListCombine(Vector &source, Vector &target, AggregateI } } -static void DpBoundedValueListFinalize(Vector &states, AggregateInputData &, Vector &result, idx_t count, +static void DpBoundedValueListFinalize(Vector &states, AggregateFinalizeInputData &, Vector &result, idx_t count, idx_t offset) { auto state_ptrs = FlatVector::GetData(states); - auto entries = FlatVector::GetData(result); - auto &child = ListVector::GetEntry(result); + auto entries = FlatVector::GetDataMutable(result); + auto &child = ListVector::GetChildMutable(result); idx_t total = ListVector::GetListSize(result); for (idx_t i = 0; i < count; i++) { total += state_ptrs[i]->size; } ListVector::Reserve(result, total); ListVector::SetListSize(result, total); - auto output = FlatVector::GetData(child); + auto output = FlatVector::GetDataMutable(child); idx_t cursor = total; for (idx_t i = count; i > 0; i--) { auto *state = state_ptrs[i - 1]; @@ -730,24 +753,24 @@ static LogicalType DpApproxBoundsDebugType() { static void WriteDebug(Vector &result, idx_t row, const DpApproxBoundsResult &value) { auto &children = StructVector::GetEntries(result); - FlatVector::GetData(*children[0])[row] = value.success; - FlatVector::GetData(*children[1])[row] = value.lower_bound; - FlatVector::GetData(*children[2])[row] = value.upper_bound; - FlatVector::GetData(*children[3])[row] = value.clipped_value; - FlatVector::GetData(*children[4])[row] = value.noise_scale; - FlatVector::GetData(*children[5])[row] = value.selected_bin; - FlatVector::GetData(*children[6])[row] = value.attempts; - FlatVector::GetData(*children[7])[row] = value.contributions; + FlatVector::GetDataMutable(children[0])[row] = value.success; + FlatVector::GetDataMutable(children[1])[row] = value.lower_bound; + FlatVector::GetDataMutable(children[2])[row] = value.upper_bound; + FlatVector::GetDataMutable(children[3])[row] = value.clipped_value; + FlatVector::GetDataMutable(children[4])[row] = value.noise_scale; + FlatVector::GetDataMutable(children[5])[row] = value.selected_bin; + FlatVector::GetDataMutable(children[6])[row] = value.attempts; + FlatVector::GetDataMutable(children[7])[row] = value.contributions; } template static void UpdateMeanRows(Vector inputs[], AggregateInputData &input, idx_t count, STATE_GETTER get_state) { UnifiedVectorFormat lists; - inputs[0].ToUnifiedFormat(count, lists); + inputs[0].ToUnifiedFormat(lists); auto entries = UnifiedVectorFormat::GetData(lists); - auto &child = ListVector::GetEntry(inputs[0]); + auto &child = ListVector::GetChild(inputs[0]); UnifiedVectorFormat values; - child.ToUnifiedFormat(ListVector::GetListSize(inputs[0]), values); + child.ToUnifiedFormat(values); auto value_data = UnifiedVectorFormat::GetData(values); auto max_contributions = input.bind_data->Cast().max_contributions; for (idx_t row = 0; row < count; row++) { @@ -771,16 +794,18 @@ static void UpdateMeanRows(Vector inputs[], AggregateInputData &input, idx_t cou } } -static void DpApproxBoundsMeanUpdate(Vector inputs[], AggregateInputData &input, idx_t, data_ptr_t state_p, - idx_t count) { - auto state = reinterpret_cast(state_p); - UpdateMeanRows(inputs, input, count, [state](idx_t) { return state; }); +static void DpApproxBoundsMeanUpdate(Vector inputs[], AggregateInputData &input, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { + auto state = reinterpret_cast(state_p); + UpdateMeanRows(inputs, input, count, [state](idx_t) { return state; }); + }); } static void DpApproxBoundsMeanScatterUpdate(Vector inputs[], AggregateInputData &input, idx_t, Vector &states, idx_t count) { UnifiedVectorFormat state_data; - states.ToUnifiedFormat(count, state_data); + states.ToUnifiedFormat(state_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); UpdateMeanRows(inputs, input, count, [&](idx_t row) { return state_ptrs[state_data.sel->get_index(row)]; }); } @@ -789,15 +814,15 @@ static double ClampMean(double value, double lower, double upper) { return std::max(lower, std::min(upper, value)); } -static void DpApproxBoundsMeanFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void DpApproxBoundsMeanFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { auto state_ptrs = FlatVector::GetData(states); auto &bind = input.bind_data->Cast(); - auto result_data = FlatVector::GetData(result); + auto result_data = FlatVector::GetDataMutable(result); for (idx_t i = 0; i < count; i++) { auto bounds = FinalizeState(*state_ptrs[i], bind, false); if (!bounds.success) { - FlatVector::Validity(result).SetInvalid(offset + i); + FlatVector::ValidityMutable(result).SetInvalid(offset + i); continue; } // A point interval has zero sensitivity but cannot define the normalized mean. Widen it by one @@ -838,11 +863,11 @@ static void DpApproxBoundsMeanFinalize(Vector &states, AggregateInputData &input } template -static void DpApproxBoundsFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void DpApproxBoundsFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { auto state_ptrs = FlatVector::GetData(states); auto &bind = input.bind_data->Cast(); - auto result_data = DEBUG ? nullptr : FlatVector::GetData(result); + auto result_data = DEBUG ? nullptr : FlatVector::GetDataMutable(result); for (idx_t i = 0; i < count; i++) { auto value = FinalizeState(*state_ptrs[i], bind, true); if (DEBUG) { @@ -850,7 +875,7 @@ static void DpApproxBoundsFinalize(Vector &states, AggregateInputData &input, Ve continue; } if (!value.success) { - FlatVector::Validity(result).SetInvalid(offset + i); + FlatVector::ValidityMutable(result).SetInvalid(offset + i); continue; } result_data[offset + i] = @@ -858,10 +883,10 @@ static void DpApproxBoundsFinalize(Vector &states, AggregateInputData &input, Ve } } -static AggregateFunction MakeDpApproxBoundsFunction(const string &name, const LogicalType &return_type, +static AggregateFunction MakeDpApproxBoundsFunction(const Identifier &name, const LogicalType &return_type, aggregate_finalize_t finalize) { return AggregateFunction( - name, + Identifier(name), {LogicalType::DOUBLE, LogicalType::DOUBLE, LogicalType::DOUBLE, LogicalType::DOUBLE, LogicalType::UBIGINT}, return_type, DpApproxBoundsStateSize, DpApproxBoundsInitialize, DpApproxBoundsScatterUpdate, DpApproxBoundsCombine, finalize, FunctionNullHandling::SPECIAL_HANDLING, DpApproxBoundsUpdate, diff --git a/src/aggregates/dp_laplace_noise.cpp b/src/aggregates/dp_laplace_noise.cpp index 590c4979..5c25a7e3 100644 --- a/src/aggregates/dp_laplace_noise.cpp +++ b/src/aggregates/dp_laplace_noise.cpp @@ -1,3 +1,5 @@ +#include "duckdb/common/vector/struct_vector.hpp" +#include "duckdb/common/vector/list_vector.hpp" #include "aggregates/dp_laplace_noise.hpp" #include "aggregates/as_aggregate.hpp" #include "privacy_debug.hpp" @@ -128,32 +130,32 @@ static void DpSassStabilityFunction(DataChunk &args, ExpressionState &, Vector & idx_t count = args.size(); auto &list_vec = args.data[0]; UnifiedVectorFormat list_data; - list_vec.ToUnifiedFormat(count, list_data); + list_vec.ToUnifiedFormat(list_data); auto list_entries = UnifiedVectorFormat::GetData(list_data); - auto &child_vec = ListVector::GetEntry(list_vec); + auto &child_vec = ListVector::GetChild(list_vec); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(list_vec), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); auto &entries = StructVector::GetEntries(result); - auto lane_count_data = FlatVector::GetData(*entries[0]); - auto valid_count_data = FlatVector::GetData(*entries[1]); - auto null_count_data = FlatVector::GetData(*entries[2]); - auto mean_data = FlatVector::GetData(*entries[3]); - auto stddev_data = FlatVector::GetData(*entries[4]); - auto cv_data = FlatVector::GetData(*entries[5]); - auto min_data = FlatVector::GetData(*entries[6]); - auto median_data = FlatVector::GetData(*entries[7]); - auto max_data = FlatVector::GetData(*entries[8]); - auto range_data = FlatVector::GetData(*entries[9]); - auto mad_data = FlatVector::GetData(*entries[10]); - auto max_abs_dev_data = FlatVector::GetData(*entries[11]); - auto stable_stddev_1_data = FlatVector::GetData(*entries[12]); - auto stable_mad_1_data = FlatVector::GetData(*entries[13]); - - auto &result_validity = FlatVector::Validity(result); - auto &cv_validity = FlatVector::Validity(*entries[5]); + auto lane_count_data = FlatVector::GetDataMutable(entries[0]); + auto valid_count_data = FlatVector::GetDataMutable(entries[1]); + auto null_count_data = FlatVector::GetDataMutable(entries[2]); + auto mean_data = FlatVector::GetDataMutable(entries[3]); + auto stddev_data = FlatVector::GetDataMutable(entries[4]); + auto cv_data = FlatVector::GetDataMutable(entries[5]); + auto min_data = FlatVector::GetDataMutable(entries[6]); + auto median_data = FlatVector::GetDataMutable(entries[7]); + auto max_data = FlatVector::GetDataMutable(entries[8]); + auto range_data = FlatVector::GetDataMutable(entries[9]); + auto mad_data = FlatVector::GetDataMutable(entries[10]); + auto max_abs_dev_data = FlatVector::GetDataMutable(entries[11]); + auto stable_stddev_1_data = FlatVector::GetDataMutable(entries[12]); + auto stable_mad_1_data = FlatVector::GetDataMutable(entries[13]); + + auto &result_validity = FlatVector::ValidityMutable(result); + auto &cv_validity = FlatVector::ValidityMutable(entries[5]); for (idx_t i = 0; i < count; i++) { auto list_idx = list_data.sel->get_index(i); if (!list_data.validity.RowIsValid(list_idx)) { @@ -272,18 +274,18 @@ static void DpSassRecordStabilityFunction(DataChunk &args, ExpressionState &, Ve idx_t count = args.size(); auto &list_vec = args.data[0]; UnifiedVectorFormat list_data, aggregate_index_data; - list_vec.ToUnifiedFormat(count, list_data); - args.data[1].ToUnifiedFormat(count, aggregate_index_data); + list_vec.ToUnifiedFormat(list_data); + args.data[1].ToUnifiedFormat(aggregate_index_data); auto list_entries = UnifiedVectorFormat::GetData(list_data); auto aggregate_indexes = UnifiedVectorFormat::GetData(aggregate_index_data); - auto &child_vec = ListVector::GetEntry(list_vec); + auto &child_vec = ListVector::GetChild(list_vec); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(list_vec), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); - auto result_data = FlatVector::GetData(result); - auto &result_validity = FlatVector::Validity(result); + auto result_data = FlatVector::GetDataMutable(result); + auto &result_validity = FlatVector::ValidityMutable(result); for (idx_t i = 0; i < count; i++) { auto list_idx = list_data.sel->get_index(i); auto aggregate_index_idx = aggregate_index_data.sel->get_index(i); @@ -494,15 +496,15 @@ static void DpLaplaceNoiseFunction(DataChunk &args, ExpressionState &, Vector &r auto count = args.size(); UnifiedVectorFormat values; UnifiedVectorFormat scales; - args.data[0].ToUnifiedFormat(count, values); - args.data[1].ToUnifiedFormat(count, scales); + args.data[0].ToUnifiedFormat(values); + args.data[1].ToUnifiedFormat(scales); auto value_data = UnifiedVectorFormat::GetData(values); auto scale_data = UnifiedVectorFormat::GetData(scales); // BinaryExecutor preserves a CONSTANT vector by invoking its callback once. Noise is volatile, so // materialize a FLAT result and sample once per output row even when value and scale are constants. result.SetVectorType(VectorType::FLAT_VECTOR); - auto result_data = FlatVector::GetData(result); - auto &result_validity = FlatVector::Validity(result); + auto result_data = FlatVector::GetDataMutable(result); + auto &result_validity = FlatVector::ValidityMutable(result); result_validity.Reset(); unique_ptr shared_distribution; double shared_scale = 0.0; @@ -538,7 +540,8 @@ void RegisterDpLaplaceNoiseFunction(ExtensionLoader &loader) { DpLaplaceNoiseFunction); // Every released cell needs an independent draw. VOLATILE prevents constant folding and common- // subexpression elimination from reusing one draw across aggregate cells. - function.stability = FunctionStability::VOLATILE; + function.SetStability(FunctionStability::VOLATILE); + function.SetFallible(); set.AddFunction(std::move(function)); CreateScalarFunctionInfo info(set); FunctionDescription desc; @@ -758,7 +761,7 @@ static void DpSmoothMedianNoiseFunctionInternal(DataChunk &args, ExpressionState auto &context = state.GetContext(); bool noise_enabled = true; Value noise_val; - if (context.TryGetCurrentSetting("privacy_noise", noise_val) && !noise_val.IsNull()) { + if (context.TryGetCurrentSetting(Identifier("privacy_noise"), noise_val) && !noise_val.IsNull()) { noise_enabled = noise_val.GetValue(); } @@ -772,28 +775,28 @@ static void DpSmoothMedianNoiseFunctionInternal(DataChunk &args, ExpressionState auto &list_vec = args.data[0]; UnifiedVectorFormat list_data, epsilon_data, delta_data, sample_lanes_data, lower_data, upper_data, empty_data, aggregate_index_data, aggregate_count_data, group_identity_data; - list_vec.ToUnifiedFormat(count, list_data); - args.data[1].ToUnifiedFormat(count, epsilon_data); - args.data[2].ToUnifiedFormat(count, delta_data); - args.data[3].ToUnifiedFormat(count, sample_lanes_data); + list_vec.ToUnifiedFormat(list_data); + args.data[1].ToUnifiedFormat(epsilon_data); + args.data[2].ToUnifiedFormat(delta_data); + args.data[3].ToUnifiedFormat(sample_lanes_data); if (bounded) { - args.data[4].ToUnifiedFormat(count, lower_data); - args.data[5].ToUnifiedFormat(count, upper_data); + args.data[4].ToUnifiedFormat(lower_data); + args.data[5].ToUnifiedFormat(upper_data); } if (has_empty_default) { - args.data[6].ToUnifiedFormat(count, empty_data); + args.data[6].ToUnifiedFormat(empty_data); } if (record_scale) { - args.data[args.ColumnCount() - 1].ToUnifiedFormat(count, aggregate_index_data); + args.data[args.ColumnCount() - 1].ToUnifiedFormat(aggregate_index_data); } else if (has_cell_identity) { - args.data[7].ToUnifiedFormat(count, aggregate_index_data); - args.data[8].ToUnifiedFormat(count, aggregate_count_data); - args.data[9].ToUnifiedFormat(count, group_identity_data); + args.data[7].ToUnifiedFormat(aggregate_index_data); + args.data[8].ToUnifiedFormat(aggregate_count_data); + args.data[9].ToUnifiedFormat(group_identity_data); } - auto &child_vec = ListVector::GetEntry(list_vec); + auto &child_vec = ListVector::GetChild(list_vec); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(list_vec), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); auto list_entries = UnifiedVectorFormat::GetData(list_data); auto epsilons = UnifiedVectorFormat::GetData(epsilon_data); @@ -806,8 +809,8 @@ static void DpSmoothMedianNoiseFunctionInternal(DataChunk &args, ExpressionState (record_scale || has_cell_identity) ? UnifiedVectorFormat::GetData(aggregate_index_data) : nullptr; auto aggregate_counts = has_cell_identity ? UnifiedVectorFormat::GetData(aggregate_count_data) : nullptr; auto group_identities = has_cell_identity ? UnifiedVectorFormat::GetData(group_identity_data) : nullptr; - auto result_data = FlatVector::GetData(result); - auto &result_validity = FlatVector::Validity(result); + auto result_data = FlatVector::GetDataMutable(result); + auto &result_validity = FlatVector::ValidityMutable(result); for (idx_t i = 0; i < count; i++) { auto list_idx = list_data.sel->get_index(i); @@ -887,23 +890,23 @@ static void DpAggregateFunction(DataChunk &args, ExpressionState &state, Vector auto &context = state.GetContext(); bool noise_enabled = true; Value noise_val; - if (context.TryGetCurrentSetting("privacy_noise", noise_val) && !noise_val.IsNull()) { + if (context.TryGetCurrentSetting(Identifier("privacy_noise"), noise_val) && !noise_val.IsNull()) { noise_enabled = noise_val.GetValue(); } idx_t count = args.size(); auto &list_vec = args.data[0]; // samples (counts at args.data[1] is intentionally ignored) UnifiedVectorFormat list_data, epsilon_data, delta_data, sample_lanes_data, lower_data, upper_data; - list_vec.ToUnifiedFormat(count, list_data); - args.data[2].ToUnifiedFormat(count, epsilon_data); - args.data[3].ToUnifiedFormat(count, delta_data); - args.data[4].ToUnifiedFormat(count, sample_lanes_data); - args.data[5].ToUnifiedFormat(count, lower_data); - args.data[6].ToUnifiedFormat(count, upper_data); - - auto &child_vec = ListVector::GetEntry(list_vec); + list_vec.ToUnifiedFormat(list_data); + args.data[2].ToUnifiedFormat(epsilon_data); + args.data[3].ToUnifiedFormat(delta_data); + args.data[4].ToUnifiedFormat(sample_lanes_data); + args.data[5].ToUnifiedFormat(lower_data); + args.data[6].ToUnifiedFormat(upper_data); + + auto &child_vec = ListVector::GetChild(list_vec); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(list_vec), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); auto list_entries = UnifiedVectorFormat::GetData(list_data); auto epsilons = UnifiedVectorFormat::GetData(epsilon_data); @@ -911,8 +914,8 @@ static void DpAggregateFunction(DataChunk &args, ExpressionState &state, Vector auto sample_lanes_values = UnifiedVectorFormat::GetData(sample_lanes_data); auto lower_values = UnifiedVectorFormat::GetData(lower_data); auto upper_values = UnifiedVectorFormat::GetData(upper_data); - auto result_data = FlatVector::GetData(result); - auto &result_validity = FlatVector::Validity(result); + auto result_data = FlatVector::GetDataMutable(result); + auto &result_validity = FlatVector::ValidityMutable(result); for (idx_t i = 0; i < count; i++) { auto list_idx = list_data.sel->get_index(i); @@ -988,7 +991,7 @@ static void DpGuptMeanNoiseFunctionInternal(DataChunk &args, ExpressionState &st auto &context = state.GetContext(); bool noise_enabled = true; Value noise_val; - if (context.TryGetCurrentSetting("privacy_noise", noise_val) && !noise_val.IsNull()) { + if (context.TryGetCurrentSetting(Identifier("privacy_noise"), noise_val) && !noise_val.IsNull()) { noise_enabled = noise_val.GetValue(); } @@ -1000,25 +1003,25 @@ static void DpGuptMeanNoiseFunctionInternal(DataChunk &args, ExpressionState &st auto &list_vec = args.data[0]; UnifiedVectorFormat list_data, epsilon_data, sample_lanes_data, lower_data, upper_data, empty_data, aggregate_index_data, aggregate_count_data, group_identity_data; - list_vec.ToUnifiedFormat(count, list_data); - args.data[1].ToUnifiedFormat(count, epsilon_data); - args.data[2].ToUnifiedFormat(count, sample_lanes_data); - args.data[3].ToUnifiedFormat(count, lower_data); - args.data[4].ToUnifiedFormat(count, upper_data); + list_vec.ToUnifiedFormat(list_data); + args.data[1].ToUnifiedFormat(epsilon_data); + args.data[2].ToUnifiedFormat(sample_lanes_data); + args.data[3].ToUnifiedFormat(lower_data); + args.data[4].ToUnifiedFormat(upper_data); if (has_empty_default) { - args.data[5].ToUnifiedFormat(count, empty_data); + args.data[5].ToUnifiedFormat(empty_data); } if (record_scale) { - args.data[args.ColumnCount() - 1].ToUnifiedFormat(count, aggregate_index_data); + args.data[args.ColumnCount() - 1].ToUnifiedFormat(aggregate_index_data); } else if (has_cell_identity) { - args.data[6].ToUnifiedFormat(count, aggregate_index_data); - args.data[7].ToUnifiedFormat(count, aggregate_count_data); - args.data[8].ToUnifiedFormat(count, group_identity_data); + args.data[6].ToUnifiedFormat(aggregate_index_data); + args.data[7].ToUnifiedFormat(aggregate_count_data); + args.data[8].ToUnifiedFormat(group_identity_data); } - auto &child_vec = ListVector::GetEntry(list_vec); + auto &child_vec = ListVector::GetChild(list_vec); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(list_vec), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); auto list_entries = UnifiedVectorFormat::GetData(list_data); auto epsilons = UnifiedVectorFormat::GetData(epsilon_data); @@ -1030,8 +1033,8 @@ static void DpGuptMeanNoiseFunctionInternal(DataChunk &args, ExpressionState &st (record_scale || has_cell_identity) ? UnifiedVectorFormat::GetData(aggregate_index_data) : nullptr; auto aggregate_counts = has_cell_identity ? UnifiedVectorFormat::GetData(aggregate_count_data) : nullptr; auto group_identities = has_cell_identity ? UnifiedVectorFormat::GetData(group_identity_data) : nullptr; - auto result_data = FlatVector::GetData(result); - auto &result_validity = FlatVector::Validity(result); + auto result_data = FlatVector::GetDataMutable(result); + auto &result_validity = FlatVector::ValidityMutable(result); for (idx_t i = 0; i < count; i++) { auto list_idx = list_data.sel->get_index(i); @@ -1212,7 +1215,7 @@ static void DpNonEmptyMeanNoiseFunctionInternal(DataChunk &args, ExpressionState auto &context = state.GetContext(); bool noise_enabled = true; Value noise_val; - if (context.TryGetCurrentSetting("privacy_noise", noise_val) && !noise_val.IsNull()) { + if (context.TryGetCurrentSetting(Identifier("privacy_noise"), noise_val) && !noise_val.IsNull()) { noise_enabled = noise_val.GetValue(); } @@ -1221,22 +1224,22 @@ static void DpNonEmptyMeanNoiseFunctionInternal(DataChunk &args, ExpressionState UnifiedVectorFormat list_data, epsilon_data, sample_lanes_data, lower_data, upper_data, baseline_data, sensitivity_data, aggregate_index_data; UnifiedVectorFormat aggregate_count_data, group_identity_data; - list_vec.ToUnifiedFormat(count, list_data); - args.data[1].ToUnifiedFormat(count, epsilon_data); - args.data[2].ToUnifiedFormat(count, sample_lanes_data); - args.data[3].ToUnifiedFormat(count, lower_data); - args.data[4].ToUnifiedFormat(count, upper_data); - args.data[5].ToUnifiedFormat(count, baseline_data); - args.data[6].ToUnifiedFormat(count, sensitivity_data); - args.data[7].ToUnifiedFormat(count, aggregate_index_data); + list_vec.ToUnifiedFormat(list_data); + args.data[1].ToUnifiedFormat(epsilon_data); + args.data[2].ToUnifiedFormat(sample_lanes_data); + args.data[3].ToUnifiedFormat(lower_data); + args.data[4].ToUnifiedFormat(upper_data); + args.data[5].ToUnifiedFormat(baseline_data); + args.data[6].ToUnifiedFormat(sensitivity_data); + args.data[7].ToUnifiedFormat(aggregate_index_data); if (!record_scales) { - args.data[8].ToUnifiedFormat(count, aggregate_count_data); - args.data[9].ToUnifiedFormat(count, group_identity_data); + args.data[8].ToUnifiedFormat(aggregate_count_data); + args.data[9].ToUnifiedFormat(group_identity_data); } - auto &child_vec = ListVector::GetEntry(list_vec); + auto &child_vec = ListVector::GetChild(list_vec); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(list_vec), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); auto list_entries = UnifiedVectorFormat::GetData(list_data); auto epsilons = UnifiedVectorFormat::GetData(epsilon_data); @@ -1248,8 +1251,8 @@ static void DpNonEmptyMeanNoiseFunctionInternal(DataChunk &args, ExpressionState auto aggregate_indexes = UnifiedVectorFormat::GetData(aggregate_index_data); auto aggregate_counts = record_scales ? nullptr : UnifiedVectorFormat::GetData(aggregate_count_data); auto group_identities = record_scales ? nullptr : UnifiedVectorFormat::GetData(group_identity_data); - auto result_data = FlatVector::GetData(result); - auto &result_validity = FlatVector::Validity(result); + auto result_data = FlatVector::GetDataMutable(result); + auto &result_validity = FlatVector::ValidityMutable(result); for (idx_t i = 0; i < count; i++) { auto list_idx = list_data.sel->get_index(i); diff --git a/src/aggregates/filterless_aggregate.cpp b/src/aggregates/filterless_aggregate.cpp index 66af87f7..cade04e6 100644 --- a/src/aggregates/filterless_aggregate.cpp +++ b/src/aggregates/filterless_aggregate.cpp @@ -1,3 +1,4 @@ +#include "duckdb/common/vector/struct_vector.hpp" #include "aggregates/filterless_aggregate.hpp" #include "aggregates/as_clip_aggr.hpp" @@ -54,20 +55,20 @@ void ValidateFilterlessBoundsEpsilonFraction(double fraction) { FilterlessSettings GetFilterlessSettings(ClientContext &context) { Value value; int64_t sample_bits = 0; - if (context.TryGetCurrentSetting("dp_filterless_sample_bits", value) && !value.IsNull()) { + if (context.TryGetCurrentSetting(Identifier("dp_filterless_sample_bits"), value) && !value.IsNull()) { sample_bits = value.GetValue(); } ValidateFilterlessSampleBits(sample_bits); double clip_support = std::numeric_limits::quiet_NaN(); - if (!context.TryGetCurrentSetting("dp_filterless_clip_support", value) || value.IsNull()) { + if (!context.TryGetCurrentSetting(Identifier("dp_filterless_clip_support"), value) || value.IsNull()) { throw InvalidInputException("dp_filterless_clip_support must be set to a positive finite number"); } clip_support = value.GetValue(); ValidateFilterlessClipSupport(clip_support); double bounds_fraction = 0.25; - if (context.TryGetCurrentSetting("dp_filterless_bounds_epsilon_fraction", value) && !value.IsNull()) { + if (context.TryGetCurrentSetting(Identifier("dp_filterless_bounds_epsilon_fraction"), value) && !value.IsNull()) { bounds_fraction = value.GetValue(); } ValidateFilterlessBoundsEpsilonFraction(bounds_fraction); @@ -191,7 +192,7 @@ static double SaturatingNoiseScale(long double sensitivity, long double epsilon) : static_cast(scale); } -static double EvaluateConstantDouble(ClientContext &context, const Expression &expression, const string &name) { +static double EvaluateConstantDouble(ClientContext &context, const Expression &expression, const Identifier &name) { if (!expression.IsFoldable()) { throw InvalidInputException("filterless: %s must be a constant", name); } @@ -236,20 +237,24 @@ static unique_ptr BindFilterless(ClientContext &context, vector BindFilterlessSum(ClientContext &context, AggregateFunction &, - vector> &arguments) { - idx_t bin_count = arguments[2]->return_type.InternalType() == PhysicalType::INT64 ? FILTERLESS_COUNT_BIN_COUNT - : FILTERLESS_HUGEINT_BIN_COUNT; +static unique_ptr BindFilterlessSum(BindAggregateFunctionInput &bind_input) { + auto &context = bind_input.GetClientContext(); + auto &arguments = bind_input.GetArguments(); + idx_t bin_count = arguments[2]->GetReturnType().InternalType() == PhysicalType::INT64 + ? FILTERLESS_COUNT_BIN_COUNT + : FILTERLESS_HUGEINT_BIN_COUNT; return BindFilterless(context, arguments, 4, true, bin_count); } -static unique_ptr BindFilterlessCount(ClientContext &context, AggregateFunction &, - vector> &arguments) { +static unique_ptr BindFilterlessCount(BindAggregateFunctionInput &bind_input) { + auto &context = bind_input.GetClientContext(); + auto &arguments = bind_input.GetArguments(); return BindFilterless(context, arguments, 4, false, FILTERLESS_COUNT_BIN_COUNT); } -static unique_ptr BindFilterlessAvg(ClientContext &context, AggregateFunction &, - vector> &arguments) { +static unique_ptr BindFilterlessAvg(BindAggregateFunctionInput &bind_input) { + auto &context = bind_input.GetClientContext(); + auto &arguments = bind_input.GetArguments(); return BindFilterless(context, arguments, 6, true, FILTERLESS_HUGEINT_BIN_COUNT); } @@ -708,28 +713,37 @@ static FilterlessExactResult FinalizeExactComponent(const FilterlessExactCompone SaturatingNoiseScale(static_cast(bound) * bind.max_groups, value_epsilon)}; } -static idx_t FilterlessStateSize(const AggregateFunction &) { +static idx_t FilterlessStateSize(AggregateStateInput &) { return sizeof(FilterlessState); } -static idx_t FilterlessAvgStateSize(const AggregateFunction &) { +static idx_t FilterlessAvgStateSize(AggregateStateInput &) { return sizeof(FilterlessAvgState); } -static idx_t FilterlessExactStateSize(const AggregateFunction &) { +static idx_t FilterlessExactStateSize(AggregateStateInput &) { return sizeof(FilterlessExactState); } -static void FilterlessInitialize(const AggregateFunction &, data_ptr_t state_p) { - memset(state_p, 0, sizeof(FilterlessState)); +static void FilterlessInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_p = states[i]; + memset(state_p, 0, sizeof(FilterlessState)); + } } -static void FilterlessAvgInitialize(const AggregateFunction &, data_ptr_t state_p) { - memset(state_p, 0, sizeof(FilterlessAvgState)); +static void FilterlessAvgInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_p = states[i]; + memset(state_p, 0, sizeof(FilterlessAvgState)); + } } -static void FilterlessExactInitialize(const AggregateFunction &, data_ptr_t state_p) { - memset(state_p, 0, sizeof(FilterlessExactState)); +static void FilterlessExactInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_p = states[i]; + memset(state_p, 0, sizeof(FilterlessExactState)); + } } template @@ -742,12 +756,12 @@ struct FilterlessInputVectors { const INPUT_TYPE *numeric_values[VALUE_COUNT]; FilterlessInputVectors(Vector inputs[], idx_t count) { - inputs[0].ToUnifiedFormat(count, pu); - inputs[1].ToUnifiedFormat(count, active); + inputs[0].ToUnifiedFormat(pu); + inputs[1].ToUnifiedFormat(active); pu_values = UnifiedVectorFormat::GetData(pu); active_values = UnifiedVectorFormat::GetData(active); for (idx_t i = 0; i < VALUE_COUNT; i++) { - inputs[2 + i].ToUnifiedFormat(count, values[i]); + inputs[2 + i].ToUnifiedFormat(values[i]); numeric_values[i] = UnifiedVectorFormat::GetData(values[i]); } } @@ -818,43 +832,52 @@ static void FilterlessUpdateRows(Vector inputs[], AggregateInputData &aggr, idx_ } } -static void FilterlessUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_p, idx_t count) { - auto state = reinterpret_cast(state_p); - FilterlessUpdateRows<2, double>(inputs, aggr, count, [state](idx_t) { return state; }); +static void FilterlessUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { + auto state = reinterpret_cast(state_p); + FilterlessUpdateRows<2, double>(inputs, aggr, count, [state](idx_t) { return state; }); + }); } static void FilterlessScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { UnifiedVectorFormat state_data; - states.ToUnifiedFormat(count, state_data); + states.ToUnifiedFormat(state_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); FilterlessUpdateRows<2, double>(inputs, aggr, count, [&](idx_t row) { return state_ptrs[state_data.sel->get_index(row)]; }); } -static void FilterlessAvgUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_p, idx_t count) { - auto state = reinterpret_cast(state_p); - FilterlessUpdateRows<4, double>(inputs, aggr, count, [state](idx_t) { return state; }); +static void FilterlessAvgUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { + auto state = reinterpret_cast(state_p); + FilterlessUpdateRows<4, double>(inputs, aggr, count, [state](idx_t) { return state; }); + }); } static void FilterlessAvgScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { UnifiedVectorFormat state_data; - states.ToUnifiedFormat(count, state_data); + states.ToUnifiedFormat(state_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); FilterlessUpdateRows<4, double>(inputs, aggr, count, [&](idx_t row) { return state_ptrs[state_data.sel->get_index(row)]; }); } template -static void FilterlessExactUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_p, idx_t count) { - auto state = reinterpret_cast(state_p); - FilterlessUpdateRows<2, INPUT_TYPE>(inputs, aggr, count, [state](idx_t) { return state; }); +static void FilterlessExactUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { + auto state = reinterpret_cast(state_p); + FilterlessUpdateRows<2, INPUT_TYPE>(inputs, aggr, count, [state](idx_t) { return state; }); + }); } template static void FilterlessExactScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { UnifiedVectorFormat state_data; - states.ToUnifiedFormat(count, state_data); + states.ToUnifiedFormat(state_data); auto state_ptrs = UnifiedVectorFormat::GetData(state_data); FilterlessUpdateRows<2, INPUT_TYPE>(inputs, aggr, count, [&](idx_t row) { return state_ptrs[state_data.sel->get_index(row)]; }); @@ -887,41 +910,42 @@ static void FilterlessExactCombine(Vector &source, Vector &target, AggregateInpu static void WriteDebugResult(Vector &result, idx_t row, const FilterlessResult &value) { auto &children = StructVector::GetEntries(result); - FlatVector::GetData(*children[0])[row] = value.lower_bound; - FlatVector::GetData(*children[1])[row] = value.upper_bound; - FlatVector::GetData(*children[2])[row] = value.clipped_value; - FlatVector::GetData(*children[3])[row] = value.noise_scale; - FlatVector::GetData(*children[4])[row] = value.negative_bin; - FlatVector::GetData(*children[5])[row] = value.positive_bin; - FlatVector::GetData(*children[6])[row] = value.negative_support; - FlatVector::GetData(*children[7])[row] = value.positive_support; - FlatVector::GetData(*children[8])[row] = value.active_contributions; - FlatVector::GetData(*children[9])[row] = value.sampled_contributions; + FlatVector::GetDataMutable(children[0])[row] = value.lower_bound; + FlatVector::GetDataMutable(children[1])[row] = value.upper_bound; + FlatVector::GetDataMutable(children[2])[row] = value.clipped_value; + FlatVector::GetDataMutable(children[3])[row] = value.noise_scale; + FlatVector::GetDataMutable(children[4])[row] = value.negative_bin; + FlatVector::GetDataMutable(children[5])[row] = value.positive_bin; + FlatVector::GetDataMutable(children[6])[row] = value.negative_support; + FlatVector::GetDataMutable(children[7])[row] = value.positive_support; + FlatVector::GetDataMutable(children[8])[row] = value.active_contributions; + FlatVector::GetDataMutable(children[9])[row] = value.sampled_contributions; } static void WriteAvgDebugResult(Vector &result, idx_t row, const FilterlessResult &sum, const FilterlessResult &denominator, double released_average) { auto &children = StructVector::GetEntries(result); - FlatVector::GetData(*children[0])[row] = sum.lower_bound; - FlatVector::GetData(*children[1])[row] = sum.upper_bound; - FlatVector::GetData(*children[2])[row] = sum.clipped_value; - FlatVector::GetData(*children[3])[row] = sum.noise_scale; - FlatVector::GetData(*children[4])[row] = denominator.upper_bound; - FlatVector::GetData(*children[5])[row] = denominator.clipped_value; - FlatVector::GetData(*children[6])[row] = denominator.noise_scale; - FlatVector::GetData(*children[7])[row] = released_average; - FlatVector::GetData(*children[8])[row] = sum.negative_bin; - FlatVector::GetData(*children[9])[row] = sum.positive_bin; - FlatVector::GetData(*children[10])[row] = denominator.positive_bin; - FlatVector::GetData(*children[11])[row] = sum.active_contributions; - FlatVector::GetData(*children[12])[row] = sum.sampled_contributions; + FlatVector::GetDataMutable(children[0])[row] = sum.lower_bound; + FlatVector::GetDataMutable(children[1])[row] = sum.upper_bound; + FlatVector::GetDataMutable(children[2])[row] = sum.clipped_value; + FlatVector::GetDataMutable(children[3])[row] = sum.noise_scale; + FlatVector::GetDataMutable(children[4])[row] = denominator.upper_bound; + FlatVector::GetDataMutable(children[5])[row] = denominator.clipped_value; + FlatVector::GetDataMutable(children[6])[row] = denominator.noise_scale; + FlatVector::GetDataMutable(children[7])[row] = released_average; + FlatVector::GetDataMutable(children[8])[row] = sum.negative_bin; + FlatVector::GetDataMutable(children[9])[row] = sum.positive_bin; + FlatVector::GetDataMutable(children[10])[row] = denominator.positive_bin; + FlatVector::GetDataMutable(children[11])[row] = sum.active_contributions; + FlatVector::GetDataMutable(children[12])[row] = sum.sampled_contributions; } template -static void FilterlessFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, idx_t offset) { +static void FilterlessFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, + idx_t offset) { auto state_ptrs = FlatVector::GetData(states); auto &bind = input.bind_data->Cast(); - auto result_data = DEBUG ? nullptr : FlatVector::GetData(result); + auto result_data = DEBUG ? nullptr : FlatVector::GetDataMutable(result); for (idx_t i = 0; i < count; i++) { auto value = FinalizeComponent(state_ptrs[i]->component, bind, bind.epsilon, COUNT); if (DEBUG) { @@ -934,11 +958,11 @@ static void FilterlessFinalize(Vector &states, AggregateInputData &input, Vector } template -static void FilterlessAvgFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void FilterlessAvgFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { auto state_ptrs = FlatVector::GetData(states); auto &bind = input.bind_data->Cast(); - auto result_data = DEBUG ? nullptr : FlatVector::GetData(result); + auto result_data = DEBUG ? nullptr : FlatVector::GetDataMutable(result); for (idx_t i = 0; i < count; i++) { double component_epsilon = bind.epsilon / 2.0; auto sum = FinalizeComponent(state_ptrs[i]->sum_component, bind, component_epsilon, false); @@ -1006,10 +1030,10 @@ static hugeint_t ReleaseExactComponent(const FilterlessExactComponentState &stat } template -static void FilterlessExactFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void FilterlessExactFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { auto state_ptrs = FlatVector::GetData(states); - auto result_data = FlatVector::GetData(result); + auto result_data = FlatVector::GetDataMutable(result); auto &bind = input.bind_data->Cast(); for (idx_t i = 0; i < count; i++) { auto released = @@ -1070,13 +1094,16 @@ struct FilterlessFusedAvgState { }; template -static idx_t FilterlessFusedAvgStateSize(const AggregateFunction &) { +static idx_t FilterlessFusedAvgStateSize(AggregateStateInput &) { return sizeof(FilterlessFusedAvgState); } template -static void FilterlessFusedAvgInitialize(const AggregateFunction &, data_ptr_t state_p) { - memset(state_p, 0, sizeof(FilterlessFusedAvgState)); +static void FilterlessFusedAvgInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_p = states[i]; + memset(state_p, 0, sizeof(FilterlessFusedAvgState)); + } } template @@ -1084,12 +1111,12 @@ static void FilterlessFusedAvgUpdateRows(Vector inputs[], AggregateInputData &ag STATE_GETTER get_state) { UnifiedVectorFormat pu_data, active_data, answer_sum_data, answer_count_data, histogram_sum_data, histogram_count_data; - inputs[0].ToUnifiedFormat(count, pu_data); - inputs[1].ToUnifiedFormat(count, active_data); - inputs[2].ToUnifiedFormat(count, answer_sum_data); - inputs[3].ToUnifiedFormat(count, answer_count_data); - inputs[4].ToUnifiedFormat(count, histogram_sum_data); - inputs[5].ToUnifiedFormat(count, histogram_count_data); + inputs[0].ToUnifiedFormat(pu_data); + inputs[1].ToUnifiedFormat(active_data); + inputs[2].ToUnifiedFormat(answer_sum_data); + inputs[3].ToUnifiedFormat(answer_count_data); + inputs[4].ToUnifiedFormat(histogram_sum_data); + inputs[5].ToUnifiedFormat(histogram_count_data); auto pus = UnifiedVectorFormat::GetData(pu_data); auto active = UnifiedVectorFormat::GetData(active_data); auto answer_sums = UnifiedVectorFormat::GetData(answer_sum_data); @@ -1126,17 +1153,19 @@ static void FilterlessFusedAvgUpdateRows(Vector inputs[], AggregateInputData &ag } template -static void FilterlessFusedAvgUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, data_ptr_t state_p, - idx_t count) { - auto state = reinterpret_cast *>(state_p); - FilterlessFusedAvgUpdateRows(inputs, aggr, count, [state](idx_t) { return state; }); +static void FilterlessFusedAvgUpdate(Vector inputs[], AggregateInputData &aggr, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { + auto state = reinterpret_cast *>(state_p); + FilterlessFusedAvgUpdateRows(inputs, aggr, count, [state](idx_t) { return state; }); + }); } template static void FilterlessFusedAvgScatterUpdate(Vector inputs[], AggregateInputData &aggr, idx_t, Vector &states, idx_t count) { UnifiedVectorFormat state_data; - states.ToUnifiedFormat(count, state_data); + states.ToUnifiedFormat(state_data); auto state_ptrs = UnifiedVectorFormat::GetData *>(state_data); FilterlessFusedAvgUpdateRows(inputs, aggr, count, [&](idx_t row) { return state_ptrs[state_data.sel->get_index(row)]; }); @@ -1153,11 +1182,11 @@ static void FilterlessFusedAvgCombine(Vector &source, Vector &target, AggregateI } template -static void FilterlessFusedAvgFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, +static void FilterlessFusedAvgFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset) { auto state_ptrs = FlatVector::GetData *>(states); - auto result_data = FlatVector::GetData(result); - auto &validity = FlatVector::Validity(result); + auto result_data = FlatVector::GetDataMutable(result); + auto &validity = FlatVector::ValidityMutable(result); auto &bind = input.bind_data->Cast(); double component_epsilon = bind.epsilon / 2.0; for (idx_t i = 0; i < count; i++) { @@ -1348,22 +1377,25 @@ static LogicalType FilterlessLowerPairType(const LogicalType &value_type) { } template -static idx_t FilterlessLowerPairStateSize(const AggregateFunction &) { +static idx_t FilterlessLowerPairStateSize(AggregateStateInput &) { return sizeof(FilterlessLowerPairState); } template -static void FilterlessLowerPairInitialize(const AggregateFunction &, data_ptr_t state_p) { - memset(state_p, 0, sizeof(FilterlessLowerPairState)); +static void FilterlessLowerPairInitialize(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_p = states[i]; + memset(state_p, 0, sizeof(FilterlessLowerPairState)); + } } template static void FilterlessLowerPairUpdateRows(Vector inputs[], idx_t count, ArenaAllocator &allocator, STATE_GETTER get_state) { UnifiedVectorFormat value_data, active_data, sampled_data; - inputs[0].ToUnifiedFormat(count, value_data); - inputs[1].ToUnifiedFormat(count, active_data); - inputs[2].ToUnifiedFormat(count, sampled_data); + inputs[0].ToUnifiedFormat(value_data); + inputs[1].ToUnifiedFormat(active_data); + inputs[2].ToUnifiedFormat(sampled_data); auto values = UnifiedVectorFormat::GetData(value_data); auto active = UnifiedVectorFormat::GetData(active_data); auto sampled = UnifiedVectorFormat::GetData(sampled_data); @@ -1389,17 +1421,19 @@ static void FilterlessLowerPairUpdateRows(Vector inputs[], idx_t count, ArenaAll } template -static void FilterlessLowerPairUpdate(Vector inputs[], AggregateInputData &input, idx_t, data_ptr_t state_p, - idx_t count) { - auto state = reinterpret_cast *>(state_p); - FilterlessLowerPairUpdateRows(inputs, count, input.allocator, [state](idx_t) { return state; }); +static void FilterlessLowerPairUpdate(Vector inputs[], AggregateInputData &input, idx_t input_count, + const ClusteredAggr &clusters, idx_t count) { + ExecutePrivClusters(inputs, input_count, clusters, [&](Vector inputs[], data_ptr_t state_p, idx_t count) { + auto state = reinterpret_cast *>(state_p); + FilterlessLowerPairUpdateRows(inputs, count, input.allocator, [state](idx_t) { return state; }); + }); } template static void FilterlessLowerPairScatterUpdate(Vector inputs[], AggregateInputData &input, idx_t, Vector &states, idx_t count) { UnifiedVectorFormat state_data; - states.ToUnifiedFormat(count, state_data); + states.ToUnifiedFormat(state_data); auto state_ptrs = UnifiedVectorFormat::GetData *>(state_data); FilterlessLowerPairUpdateRows(inputs, count, input.allocator, [&](idx_t row) { return state_ptrs[state_data.sel->get_index(row)]; }); @@ -1417,24 +1451,24 @@ static void FilterlessLowerPairCombine(Vector &source, Vector &target, Aggregate } template -static void FilterlessLowerPairFinalize(Vector &states, AggregateInputData &, Vector &result, idx_t count, +static void FilterlessLowerPairFinalize(Vector &states, AggregateFinalizeInputData &, Vector &result, idx_t count, idx_t offset) { using pair_state_t = FilterlessLowerPairState; auto state_ptrs = FlatVector::GetData(states); auto &children = StructVector::GetEntries(result); - auto answers = FlatVector::GetData(*children[0]); - auto histograms = FlatVector::GetData(*children[1]); + auto answers = FlatVector::GetDataMutable(children[0]); + auto histograms = FlatVector::GetDataMutable(children[1]); for (idx_t i = 0; i < count; i++) { auto row = offset + i; if (OPERATION::IsSet(state_ptrs[i]->answer)) { answers[row] = OPERATION::Finalize(state_ptrs[i]->answer); } else { - FlatVector::Validity(*children[0]).SetInvalid(row); + FlatVector::ValidityMutable(children[0]).SetInvalid(row); } if (OPERATION::IsSet(state_ptrs[i]->histogram)) { histograms[row] = OPERATION::Finalize(state_ptrs[i]->histogram); } else { - FlatVector::Validity(*children[1]).SetInvalid(row); + FlatVector::ValidityMutable(children[1]).SetInvalid(row); } } } @@ -1473,14 +1507,14 @@ static LogicalType FilterlessAvgDebugType() { } template -static AggregateFunction MakeFilterlessExactFunction(const string &name, const LogicalType &input_type, +static AggregateFunction MakeFilterlessExactFunction(const Identifier &name, const LogicalType &input_type, const LogicalType &return_type, bool explicit_config) { vector arguments = {LogicalType::UBIGINT, LogicalType::BOOLEAN, input_type, input_type}; if (explicit_config) { arguments.push_back(LogicalType::DOUBLE); arguments.push_back(LogicalType::DOUBLE); } - return AggregateFunction(name, std::move(arguments), return_type, FilterlessExactStateSize, + return AggregateFunction(Identifier(name), std::move(arguments), return_type, FilterlessExactStateSize, FilterlessExactInitialize, FilterlessExactScatterUpdate, FilterlessExactCombine, FilterlessExactFinalize, FunctionNullHandling::SPECIAL_HANDLING, FilterlessExactUpdate, @@ -1488,8 +1522,8 @@ static AggregateFunction MakeFilterlessExactFunction(const string &name, const L } template -static void AddFilterlessExactOverloads(AggregateFunctionSet &set, const string &name, const LogicalType &input_type, - const LogicalType &return_type) { +static void AddFilterlessExactOverloads(AggregateFunctionSet &set, const Identifier &name, + const LogicalType &input_type, const LogicalType &return_type) { set.AddFunction(MakeFilterlessExactFunction(name, input_type, return_type, false)); set.AddFunction(MakeFilterlessExactFunction(name, input_type, return_type, true)); } @@ -1520,29 +1554,31 @@ static void ConfigureFilterlessDecimalBind(FilterlessBindData &bind, const Logic bind.exact_output_min = Hugeint::Negate(bind.exact_output_max); } -static unique_ptr BindFilterlessDecimalSum(ClientContext &context, AggregateFunction &function, - vector> &arguments) { +static unique_ptr BindFilterlessDecimalSum(BindAggregateFunctionInput &bind_input) { + auto &context = bind_input.GetClientContext(); + auto &function = bind_input.GetBoundFunction(); + auto &arguments = bind_input.GetArguments(); if (arguments.size() != 4 && arguments.size() != 6) { throw InternalException("filterless_sum: unexpected DECIMAL argument count"); } - auto input_type = arguments[2]->return_type; - if (arguments[3]->return_type != input_type) { + auto input_type = arguments[2]->GetReturnType(); + if (arguments[3]->GetReturnType() != input_type) { throw InvalidInputException("filterless_sum: answer and histogram DECIMAL types must match"); } auto return_type = LogicalType::DECIMAL(Decimal::MAX_WIDTH_DECIMAL, DecimalType::GetScale(input_type)); - function = MakeFilterlessDecimalSumFunction(input_type, return_type, arguments.size() == 6); + function.ReplaceImplementation(MakeFilterlessDecimalSumFunction(input_type, return_type, arguments.size() == 6)); auto result = BindFilterless(context, arguments, 4, false, FILTERLESS_HUGEINT_BIN_COUNT); ConfigureFilterlessDecimalBind(result->Cast(), input_type); return result; } -static void AddSumCountOverloads(AggregateFunctionSet &set, const string &name, aggregate_finalize_t finalize, +static void AddSumCountOverloads(AggregateFunctionSet &set, const Identifier &name, aggregate_finalize_t finalize, const LogicalType &return_type, bind_aggregate_function_t bind) { - set.AddFunction( - AggregateFunction(name, {LogicalType::UBIGINT, LogicalType::BOOLEAN, LogicalType::DOUBLE, LogicalType::DOUBLE}, - return_type, FilterlessStateSize, FilterlessInitialize, FilterlessScatterUpdate, - FilterlessCombine, finalize, FunctionNullHandling::SPECIAL_HANDLING, FilterlessUpdate, bind)); - set.AddFunction(AggregateFunction(name, + set.AddFunction(AggregateFunction( + Identifier(name), {LogicalType::UBIGINT, LogicalType::BOOLEAN, LogicalType::DOUBLE, LogicalType::DOUBLE}, + return_type, FilterlessStateSize, FilterlessInitialize, FilterlessScatterUpdate, FilterlessCombine, finalize, + FunctionNullHandling::SPECIAL_HANDLING, FilterlessUpdate, bind)); + set.AddFunction(AggregateFunction(Identifier(name), {LogicalType::UBIGINT, LogicalType::BOOLEAN, LogicalType::DOUBLE, LogicalType::DOUBLE, LogicalType::DOUBLE, LogicalType::DOUBLE}, return_type, FilterlessStateSize, FilterlessInitialize, FilterlessScatterUpdate, @@ -1550,16 +1586,16 @@ static void AddSumCountOverloads(AggregateFunctionSet &set, const string &name, FilterlessUpdate, bind)); } -static void AddAvgOverloads(AggregateFunctionSet &set, const string &name, aggregate_finalize_t finalize, +static void AddAvgOverloads(AggregateFunctionSet &set, const Identifier &name, aggregate_finalize_t finalize, const LogicalType &return_type) { - set.AddFunction(AggregateFunction(name, + set.AddFunction(AggregateFunction(Identifier(name), {LogicalType::UBIGINT, LogicalType::BOOLEAN, LogicalType::DOUBLE, LogicalType::DOUBLE, LogicalType::DOUBLE, LogicalType::DOUBLE}, return_type, FilterlessAvgStateSize, FilterlessAvgInitialize, FilterlessAvgScatterUpdate, FilterlessAvgCombine, finalize, FunctionNullHandling::SPECIAL_HANDLING, FilterlessAvgUpdate, BindFilterlessAvg)); set.AddFunction(AggregateFunction( - name, + Identifier(name), {LogicalType::UBIGINT, LogicalType::BOOLEAN, LogicalType::DOUBLE, LogicalType::DOUBLE, LogicalType::DOUBLE, LogicalType::DOUBLE, LogicalType::DOUBLE, LogicalType::DOUBLE}, return_type, FilterlessAvgStateSize, FilterlessAvgInitialize, FilterlessAvgScatterUpdate, FilterlessAvgCombine, @@ -1580,38 +1616,42 @@ static AggregateFunction MakeFilterlessFusedAvgFunction(const LogicalType &sum_t return function; } -static unique_ptr BindFilterlessFusedApproxAvg(ClientContext &context, AggregateFunction &, - vector> &arguments) { +static unique_ptr BindFilterlessFusedApproxAvg(BindAggregateFunctionInput &bind_input) { + auto &context = bind_input.GetClientContext(); + auto &arguments = bind_input.GetArguments(); return BindFilterless(context, arguments, 6, true, FILTERLESS_HUGEINT_BIN_COUNT); } -static unique_ptr BindFilterlessFusedExactAvg(ClientContext &context, AggregateFunction &, - vector> &arguments) { +static unique_ptr BindFilterlessFusedExactAvg(BindAggregateFunctionInput &bind_input) { + auto &context = bind_input.GetClientContext(); + auto &arguments = bind_input.GetArguments(); return BindFilterless(context, arguments, 6, false, FILTERLESS_HUGEINT_BIN_COUNT); } -static unique_ptr BindFilterlessFusedDecimalAvg(ClientContext &context, AggregateFunction &function, - vector> &arguments) { - auto sum_type = arguments[2]->return_type; - if (arguments[4]->return_type != sum_type) { +static unique_ptr BindFilterlessFusedDecimalAvg(BindAggregateFunctionInput &bind_input) { + auto &context = bind_input.GetClientContext(); + auto &function = bind_input.GetBoundFunction(); + auto &arguments = bind_input.GetArguments(); + auto sum_type = arguments[2]->GetReturnType(); + if (arguments[4]->GetReturnType() != sum_type) { throw InvalidInputException("priv_filterless_avg: answer and histogram DECIMAL types must match"); } switch (sum_type.InternalType()) { case PhysicalType::INT16: - function = MakeFilterlessFusedAvgFunction>( - sum_type, BindFilterlessFusedDecimalAvg); + function.ReplaceImplementation(MakeFilterlessFusedAvgFunction>( + sum_type, BindFilterlessFusedDecimalAvg)); break; case PhysicalType::INT32: - function = MakeFilterlessFusedAvgFunction>( - sum_type, BindFilterlessFusedDecimalAvg); + function.ReplaceImplementation(MakeFilterlessFusedAvgFunction>( + sum_type, BindFilterlessFusedDecimalAvg)); break; case PhysicalType::INT64: - function = MakeFilterlessFusedAvgFunction>( - sum_type, BindFilterlessFusedDecimalAvg); + function.ReplaceImplementation(MakeFilterlessFusedAvgFunction>( + sum_type, BindFilterlessFusedDecimalAvg)); break; case PhysicalType::INT128: - function = MakeFilterlessFusedAvgFunction>( - sum_type, BindFilterlessFusedDecimalAvg); + function.ReplaceImplementation(MakeFilterlessFusedAvgFunction>( + sum_type, BindFilterlessFusedDecimalAvg)); break; default: throw InternalException("priv_filterless_avg: unsupported DECIMAL physical type"); @@ -1622,10 +1662,10 @@ static unique_ptr BindFilterlessFusedDecimalAvg(ClientContext &con } template -static AggregateFunction MakeFilterlessLowerPairFunction(const string &name, const LogicalType &input_type, +static AggregateFunction MakeFilterlessLowerPairFunction(const Identifier &name, const LogicalType &input_type, const LogicalType &return_type) { auto function = - AggregateFunction(name, {input_type, LogicalType::BOOLEAN, LogicalType::BOOLEAN}, + AggregateFunction(Identifier(name), {input_type, LogicalType::BOOLEAN, LogicalType::BOOLEAN}, FilterlessLowerPairType(return_type), FilterlessLowerPairStateSize, FilterlessLowerPairInitialize, FilterlessLowerPairScatterUpdate, FilterlessLowerPairCombine, FilterlessLowerPairFinalize, @@ -1641,22 +1681,23 @@ static AggregateFunction MakeFilterlessExactSumPairFunction(const LogicalType &i "priv_filterless_sum_pair", input_type, return_type); } -static unique_ptr BindFilterlessDecimalSumPair(ClientContext &, AggregateFunction &function, - vector> &arguments) { - auto input_type = arguments[0]->return_type; +static unique_ptr BindFilterlessDecimalSumPair(BindAggregateFunctionInput &bind_input) { + auto &function = bind_input.GetBoundFunction(); + auto &arguments = bind_input.GetArguments(); + auto input_type = arguments[0]->GetReturnType(); auto return_type = LogicalType::DECIMAL(Decimal::MAX_WIDTH_DECIMAL, DecimalType::GetScale(input_type)); switch (input_type.InternalType()) { case PhysicalType::INT16: - function = MakeFilterlessExactSumPairFunction(input_type, return_type); + function.ReplaceImplementation(MakeFilterlessExactSumPairFunction(input_type, return_type)); break; case PhysicalType::INT32: - function = MakeFilterlessExactSumPairFunction(input_type, return_type); + function.ReplaceImplementation(MakeFilterlessExactSumPairFunction(input_type, return_type)); break; case PhysicalType::INT64: - function = MakeFilterlessExactSumPairFunction(input_type, return_type); + function.ReplaceImplementation(MakeFilterlessExactSumPairFunction(input_type, return_type)); break; case PhysicalType::INT128: - function = MakeFilterlessExactSumPairFunction(input_type, return_type); + function.ReplaceImplementation(MakeFilterlessExactSumPairFunction(input_type, return_type)); break; default: throw InternalException("priv_filterless_sum_pair: unsupported DECIMAL physical type"); diff --git a/src/categorical/pac_categorical.cpp b/src/categorical/pac_categorical.cpp index 3211cfc0..245b9aa7 100644 --- a/src/categorical/pac_categorical.cpp +++ b/src/categorical/pac_categorical.cpp @@ -1,3 +1,4 @@ +#include "aggregates/aggregate_state_export.hpp" // // PAC Categorical Query Support - Implementation // @@ -6,6 +7,7 @@ // Created by ila on 1/22/26. // +#include "duckdb/common/vector/list_vector.hpp" #include "categorical/pac_categorical.hpp" #include "duckdb/parser/parsed_data/create_scalar_function_info.hpp" #include "privacy_debug.hpp" @@ -54,31 +56,25 @@ struct PacCategoricalBindData : public FunctionData { uint64_t query_hash; // derived from seed: used as counter selector for NoisySample std::shared_ptr pstate; // p-tracking state shared across query (may be nullptr) double utility_threshold; // z-score threshold for utility NULLing (NaN = disabled) + shared_ptr random_state; // Primary constructor - reads seed from privacy_seed setting, or uses default 42 if not set. - // When mi > 0, seed is randomized per query via the query number. + // Without an explicit seed, mi > 0 mixes in the query number. explicit PacCategoricalBindData(ClientContext &ctx, double mi_val = 0.0, double correction_val = 1.0) - : mi(mi_val), correction(correction_val), utility_threshold(std::numeric_limits::quiet_NaN()) { + : mi(mi_val), correction(correction_val), utility_threshold(std::numeric_limits::quiet_NaN()), + random_state(GetPacRandomState(ctx)) { // Read utility threshold setting Value ut_val; - if (ctx.TryGetCurrentSetting("privacy_min_group_count", ut_val) && !ut_val.IsNull()) { + if (ctx.TryGetCurrentSetting(Identifier("privacy_min_group_count"), ut_val) && !ut_val.IsNull()) { utility_threshold = ut_val.GetValue(); } - Value pac_seed_val; - if (ctx.TryGetCurrentSetting("privacy_seed", pac_seed_val) && !pac_seed_val.IsNull()) { - seed = uint64_t(pac_seed_val.GetValue()); - } else { - seed = 42; - } - if (mi != 0.0) { // randomize seed per query (not in deterministic mode aka mi==0) - seed ^= PAC_MAGIC_HASH * static_cast(ctx.ActiveTransaction().GetActiveQuery()); - } + seed = GetPacSeed(ctx, mi); query_hash = (seed * PAC_MAGIC_HASH) ^ PAC_MAGIC_HASH; // Initialize p-tracking if mi > 0 and pac_ptracking is enabled if (mi > 0.0) { Value pt_val; bool ptracking_enabled = true; - if (ctx.TryGetCurrentSetting("pac_ptracking", pt_val) && !pt_val.IsNull()) { + if (ctx.TryGetCurrentSetting(Identifier("pac_ptracking"), pt_val) && !pt_val.IsNull()) { ptracking_enabled = pt_val.GetValue(); } if (ptracking_enabled) { @@ -102,8 +98,8 @@ struct PacCategoricalBindData : public FunctionData { // Local state for PAC categorical functions (RNG) // ============================================================================ struct PacCategoricalLocalState : public FunctionLocalState { - std::mt19937_64 gen; - explicit PacCategoricalLocalState(uint64_t seed) : gen(seed) { + PacRandomEngine gen; + explicit PacCategoricalLocalState(PacRandomEngine gen_p) : gen(gen_p) { } }; @@ -113,7 +109,7 @@ PacCategoricalInitLocal(ExpressionState &state, const BoundFunctionExpression &e if (bind_data) { seed = bind_data->Cast().seed; } - return make_uniq(seed); + return make_uniq(PacRandomStream(seed, *GetPacRandomState(state.GetContext()))); } // ============================================================================ @@ -141,7 +137,7 @@ static inline uint64_t BoolListToMask(UnifiedVectorFormat &list_data, UnifiedVec // Helper: Common filter logic for mask-based filtering // ============================================================================ static inline bool FilterFromCount(idx_t true_count, idx_t lane_count, double mi, double correction, - std::mt19937_64 &gen) { + PacRandomEngine &gen) { if (lane_count == 0) { return false; } @@ -168,7 +164,7 @@ static inline bool FilterFromCount(idx_t true_count, idx_t lane_count, double mi // correction: considers correction times more non-nulls (increases true probability) // When mi <= 0: deterministic, true when popcount(mask) * correction >= 32 // When mi > 0: probabilistic, P(true) = popcount(mask) * correction / 64 -static inline bool FilterFromMask(uint64_t mask, double mi, double correction, std::mt19937_64 &gen) { +static inline bool FilterFromMask(uint64_t mask, double mi, double correction, PacRandomEngine &gen) { int popcount = Popcount64(mask); return FilterFromCount(popcount, 64, mi, correction, gen); } @@ -184,20 +180,20 @@ static void PacSelectFunction(DataChunk &args, ExpressionState &state, Vector &r idx_t count = args.size(); UnifiedVectorFormat hash_data; - args.data[0].ToUnifiedFormat(count, hash_data); + args.data[0].ToUnifiedFormat(hash_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto &list_vec = args.data[1]; UnifiedVectorFormat list_data; - list_vec.ToUnifiedFormat(count, list_data); + list_vec.ToUnifiedFormat(list_data); - auto &child_vec = ListVector::GetEntry(list_vec); + auto &child_vec = ListVector::GetChild(list_vec); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(list_vec), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); - auto result_data = FlatVector::GetData(result); - auto &result_validity = FlatVector::Validity(result); + auto result_data = FlatVector::GetDataMutable(result); + auto &result_validity = FlatVector::ValidityMutable(result); for (idx_t i = 0; i < count; i++) { auto h_idx = hash_data.sel->get_index(i); @@ -223,7 +219,7 @@ static void PacSelectFunction(DataChunk &args, ExpressionState &state, Vector &r struct PacFilterParams { double mi; double correction; - std::mt19937_64 &gen; + PacRandomEngine &gen; }; static PacFilterParams GetPacFilterParams(ExpressionState &state) { @@ -231,8 +227,8 @@ static PacFilterParams GetPacFilterParams(ExpressionState &state) { double mi = 0.0; double correction = 1.0; auto &function = state.expr.Cast(); - if (function.bind_info) { - auto &bind_data = function.bind_info->Cast(); + if (function.BindInfo()) { + auto &bind_data = function.BindInfo()->Cast(); mi = bind_data.mi; correction = bind_data.correction; } @@ -247,13 +243,13 @@ static void PacFilterFromBoolListFunction(DataChunk &args, ExpressionState &stat auto params = GetPacFilterParams(state); UnifiedVectorFormat list_data; - list_vec.ToUnifiedFormat(count, list_data); + list_vec.ToUnifiedFormat(list_data); - auto result_data = FlatVector::GetData(result); + auto result_data = FlatVector::GetDataMutable(result); - auto &child_vec = ListVector::GetEntry(list_vec); + auto &child_vec = ListVector::GetChild(list_vec); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(list_vec), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); for (idx_t i = 0; i < count; i++) { @@ -286,8 +282,7 @@ static void PacFilterFromBoolListFunction(DataChunk &args, ExpressionState &stat } // Forward declarations for functions used by priv_filter_ registration -static unique_ptr PacCategoricalBind(ClientContext &ctx, ScalarFunction &func, - vector> &args); +static unique_ptr PacCategoricalBind(BindScalarFunctionInput &bind_input); // ============================================================================ // PAC_FILTER_: Compare scalar against counter list and filter in one pass @@ -325,26 +320,26 @@ static void PacFilterCmpFunction(DataChunk &args, ExpressionState &state, Vector // arg0: PAC_FLOAT scalar value, arg1: LIST counters, arg2 (optional): PAC_FLOAT correction UnifiedVectorFormat val_data; - args.data[0].ToUnifiedFormat(count, val_data); + args.data[0].ToUnifiedFormat(val_data); auto vals = UnifiedVectorFormat::GetData(val_data); auto &list_vec = args.data[1]; UnifiedVectorFormat list_data; - list_vec.ToUnifiedFormat(count, list_data); + list_vec.ToUnifiedFormat(list_data); - auto &child_vec = ListVector::GetEntry(list_vec); + auto &child_vec = ListVector::GetChild(list_vec); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(list_vec), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); UnifiedVectorFormat correction_data; const PAC_FLOAT *corrections = nullptr; if (HAS_CORRECTION_ARG) { - args.data[2].ToUnifiedFormat(count, correction_data); + args.data[2].ToUnifiedFormat(correction_data); corrections = UnifiedVectorFormat::GetData(correction_data); } - auto result_data = FlatVector::GetData(result); + auto result_data = FlatVector::GetDataMutable(result); for (idx_t i = 0; i < count; i++) { auto val_idx = val_data.sel->get_index(i); @@ -385,14 +380,15 @@ static void PacFilterCmpFunction(DataChunk &args, ExpressionState &state, Vector } template -static void RegisterPacFilterCmp(ExtensionLoader &loader, const string &name) { +static void RegisterPacFilterCmp(ExtensionLoader &loader, const Identifier &name) { auto pf = PacFloatLogicalType(); auto ldt = LogicalType::LIST(pf); - ScalarFunctionSet fset(name); - fset.AddFunction(ScalarFunction(name, {pf, ldt}, LogicalType::BOOLEAN, PacFilterCmpFunction, - PacCategoricalBind, nullptr, nullptr, PacCategoricalInitLocal)); - fset.AddFunction(ScalarFunction(name, {pf, ldt, pf}, LogicalType::BOOLEAN, PacFilterCmpFunction, - PacCategoricalBind, nullptr, nullptr, PacCategoricalInitLocal)); + ScalarFunctionSet fset {Identifier(name)}; + fset.AddFunction(ScalarFunction(Identifier(name), {pf, ldt}, LogicalType::BOOLEAN, PacFilterCmpFunction, + PacCategoricalBind, nullptr, PacCategoricalInitLocal)); + fset.AddFunction(ScalarFunction(Identifier(name), {pf, ldt, pf}, LogicalType::BOOLEAN, + PacFilterCmpFunction, PacCategoricalBind, nullptr, + PacCategoricalInitLocal)); CreateScalarFunctionInfo info(fset); FunctionDescription desc; desc.description = "[INTERNAL] Categorical comparison + probabilistic filter for PAC queries."; @@ -412,22 +408,22 @@ static void PacSelectCmpFunction(DataChunk &args, ExpressionState &state, Vector // arg0: UBIGINT hash, arg1: PAC_FLOAT scalar value, arg2: LIST counters UnifiedVectorFormat hash_data, val_data; - args.data[0].ToUnifiedFormat(count, hash_data); - args.data[1].ToUnifiedFormat(count, val_data); + args.data[0].ToUnifiedFormat(hash_data); + args.data[1].ToUnifiedFormat(val_data); auto hashes = UnifiedVectorFormat::GetData(hash_data); auto vals = UnifiedVectorFormat::GetData(val_data); auto &list_vec = args.data[2]; UnifiedVectorFormat list_data; - list_vec.ToUnifiedFormat(count, list_data); + list_vec.ToUnifiedFormat(list_data); - auto &child_vec = ListVector::GetEntry(list_vec); + auto &child_vec = ListVector::GetChild(list_vec); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(list_vec), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); - auto result_data = FlatVector::GetData(result); - auto &result_validity = FlatVector::Validity(result); + auto result_data = FlatVector::GetDataMutable(result); + auto &result_validity = FlatVector::ValidityMutable(result); for (idx_t i = 0; i < count; i++) { auto h_idx = hash_data.sel->get_index(i); @@ -458,7 +454,7 @@ static void PacSelectCmpFunction(DataChunk &args, ExpressionState &state, Vector } template -static void RegisterPacSelectCmp(ExtensionLoader &loader, const string &name) { +static void RegisterPacSelectCmp(ExtensionLoader &loader, const Identifier &name) { auto pf = PacFloatLogicalType(); auto ldt = LogicalType::LIST(pf); ScalarFunction f(name, {LogicalType::UBIGINT, pf, ldt}, LogicalType::UBIGINT, PacSelectCmpFunction, @@ -480,10 +476,10 @@ static void PacCoalesceFunction(DataChunk &args, ExpressionState &state, Vector idx_t count = args.size(); UnifiedVectorFormat input_data; - input.ToUnifiedFormat(count, input_data); + input.ToUnifiedFormat(input_data); // Fast path: if no NULL lists, zero-copy pass-through (common case) - if (input_data.validity.AllValid()) { + if (input_data.validity.CannotHaveNull()) { result.Reference(input); return; } @@ -492,8 +488,8 @@ static void PacCoalesceFunction(DataChunk &args, ExpressionState &state, Vector auto input_entries = UnifiedVectorFormat::GetData(input_data); result.SetVectorType(VectorType::FLAT_VECTOR); - auto list_entries = FlatVector::GetData(result); - auto &child_vec = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &child_vec = ListVector::GetChildMutable(result); // First pass: count total elements needed idx_t total_elements = 0; @@ -508,12 +504,12 @@ static void PacCoalesceFunction(DataChunk &args, ExpressionState &state, Vector ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto child_data = FlatVector::GetData(child_vec); - auto &child_validity = FlatVector::Validity(child_vec); + auto child_data = FlatVector::GetDataMutable(child_vec); + auto &child_validity = FlatVector::ValidityMutable(child_vec); child_validity.EnsureWritable(); // Get the source child vector for non-NULL lists - auto &input_child = ListVector::GetEntry(input); + auto &input_child = ListVector::GetChild(input); bool child_is_flat = (input_child.GetVectorType() == VectorType::FLAT_VECTOR); auto flat_src_data = child_is_flat ? FlatVector::GetData(input_child) : nullptr; auto *flat_src_validity = child_is_flat ? &FlatVector::Validity(input_child) : nullptr; @@ -522,7 +518,7 @@ static void PacCoalesceFunction(DataChunk &args, ExpressionState &state, Vector UnifiedVectorFormat input_child_data; const PAC_FLOAT *src_data = nullptr; if (!child_is_flat) { - input_child.ToUnifiedFormat(ListVector::GetListSize(input), input_child_data); + input_child.ToUnifiedFormat(input_child_data); src_data = UnifiedVectorFormat::GetData(input_child_data); } @@ -585,19 +581,18 @@ static void PacNoisedFunction(DataChunk &args, ExpressionState &state, Vector &r idx_t count = args.size(); auto &function = state.expr.Cast(); - auto &bind_data = function.bind_info->Cast(); - - std::mt19937_64 gen(bind_data.seed); + auto &bind_data = function.BindInfo()->Cast(); + auto gen = PacRandomStream(bind_data.seed, *bind_data.random_state); UnifiedVectorFormat list_data; - list_vec.ToUnifiedFormat(count, list_data); + list_vec.ToUnifiedFormat(list_data); - auto result_data = FlatVector::GetData(result); - auto &result_validity = FlatVector::Validity(result); + auto result_data = FlatVector::GetDataMutable(result); + auto &result_validity = FlatVector::ValidityMutable(result); - auto &child_vec = ListVector::GetEntry(list_vec); + auto &child_vec = ListVector::GetChild(list_vec); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(list_vec), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); for (idx_t i = 0; i < count; i++) { @@ -619,17 +614,17 @@ static void PacNoisedFunction(DataChunk &args, ExpressionState &state, Vector &r // Assume all 64 worlds are populated (all bits set) uint64_t key_hash = ~uint64_t(0); - // Check if we should return NULL based on key_hash (uses mi and correction) - if (PacNoiseInNull(key_hash, bind_data.mi, bind_data.correction, gen)) { - result_validity.SetInvalid(i); - continue; - } // Extract counter values from the list PAC_FLOAT counters[64]; for (idx_t j = 0; j < 64; j++) { auto child_idx = child_data.sel->get_index(entry.offset + j); counters[j] = child_values[child_idx]; } + // Check if we should return NULL based on key_hash (uses mi and correction) + if (PacNoiseInNull(key_hash, bind_data.mi, bind_data.correction, gen)) { + result_validity.SetInvalid(i); + continue; + } double noise_var = 0.0; result_data[i] = PacNoisySampleFrom64Counters(counters, bind_data.mi, bind_data.correction, gen, ~key_hash, bind_data.query_hash, bind_data.pstate, &noise_var); @@ -651,14 +646,14 @@ static void PacDivFunction(DataChunk &args, ExpressionState &state, Vector &resu idx_t count = args.size(); UnifiedVectorFormat num_data, den_data; - num_vec.ToUnifiedFormat(count, num_data); - den_vec.ToUnifiedFormat(count, den_data); + num_vec.ToUnifiedFormat(num_data); + den_vec.ToUnifiedFormat(den_data); - auto &num_child = ListVector::GetEntry(num_vec); - auto &den_child = ListVector::GetEntry(den_vec); + auto &num_child = ListVector::GetChild(num_vec); + auto &den_child = ListVector::GetChild(den_vec); UnifiedVectorFormat num_child_data, den_child_data; - num_child.ToUnifiedFormat(ListVector::GetListSize(num_vec), num_child_data); - den_child.ToUnifiedFormat(ListVector::GetListSize(den_vec), den_child_data); + num_child.ToUnifiedFormat(num_child_data); + den_child.ToUnifiedFormat(den_child_data); auto num_values = UnifiedVectorFormat::GetData(num_child_data); auto den_values = UnifiedVectorFormat::GetData(den_child_data); auto num_entries = UnifiedVectorFormat::GetData(num_data); @@ -677,14 +672,14 @@ static void PacDivFunction(DataChunk &args, ExpressionState &state, Vector &resu } result.SetVectorType(VectorType::FLAT_VECTOR); - auto list_entries = FlatVector::GetData(result); - auto &result_validity = FlatVector::Validity(result); - auto &result_child = ListVector::GetEntry(result); + auto list_entries = FlatVector::GetDataMutable(result); + auto &result_validity = FlatVector::ValidityMutable(result); + auto &result_child = ListVector::GetChildMutable(result); ListVector::Reserve(result, total_elements); ListVector::SetListSize(result, total_elements); - auto result_data = FlatVector::GetData(result_child); - auto &result_child_validity = FlatVector::Validity(result_child); + auto result_data = FlatVector::GetDataMutable(result_child); + auto &result_child_validity = FlatVector::ValidityMutable(result_child); result_child_validity.SetAllValid(total_elements); idx_t base = 0; @@ -732,21 +727,21 @@ static void PacNoisedDivFunction(DataChunk &args, ExpressionState &state, Vector idx_t count = args.size(); auto &function = state.expr.Cast(); - auto &bind_data = function.bind_info->Cast(); - std::mt19937_64 gen(bind_data.seed); + auto &bind_data = function.BindInfo()->Cast(); + auto gen = PacRandomStream(bind_data.seed, *bind_data.random_state); UnifiedVectorFormat sum_data, cnt_data; - sum_vec.ToUnifiedFormat(count, sum_data); - cnt_vec.ToUnifiedFormat(count, cnt_data); + sum_vec.ToUnifiedFormat(sum_data); + cnt_vec.ToUnifiedFormat(cnt_data); - auto result_data = FlatVector::GetData(result); - auto &result_validity = FlatVector::Validity(result); + auto result_data = FlatVector::GetDataMutable(result); + auto &result_validity = FlatVector::ValidityMutable(result); - auto &sum_child = ListVector::GetEntry(sum_vec); - auto &cnt_child = ListVector::GetEntry(cnt_vec); + auto &sum_child = ListVector::GetChild(sum_vec); + auto &cnt_child = ListVector::GetChild(cnt_vec); UnifiedVectorFormat sum_child_data, cnt_child_data; - sum_child.ToUnifiedFormat(ListVector::GetListSize(sum_vec), sum_child_data); - cnt_child.ToUnifiedFormat(ListVector::GetListSize(cnt_vec), cnt_child_data); + sum_child.ToUnifiedFormat(sum_child_data); + cnt_child.ToUnifiedFormat(cnt_child_data); auto sum_values = UnifiedVectorFormat::GetData(sum_child_data); auto cnt_values = UnifiedVectorFormat::GetData(cnt_child_data); @@ -768,10 +763,6 @@ static void PacNoisedDivFunction(DataChunk &args, ExpressionState &state, Vector continue; } uint64_t key_hash = ~uint64_t(0); - if (PacNoiseInNull(key_hash, bind_data.mi, bind_data.correction, gen)) { - result_validity.SetInvalid(i); - continue; - } // Element-wise division: sum[j] / cnt[j] for each of the 64 counters PAC_FLOAT counters[64]; for (idx_t j = 0; j < 64; j++) { @@ -781,6 +772,10 @@ static void PacNoisedDivFunction(DataChunk &args, ExpressionState &state, Vector PAC_FLOAT c = cnt_values[c_idx]; counters[j] = s / c; // 0/0 → NaN, which propagates correctly } + if (PacNoiseInNull(key_hash, bind_data.mi, bind_data.correction, gen)) { + result_validity.SetInvalid(i); + continue; + } double noise_var = 0.0; result_data[i] = static_cast(PacNoisySampleFrom64Counters(counters, bind_data.mi, bind_data.correction, gen, ~key_hash, bind_data.query_hash, @@ -794,8 +789,10 @@ static void PacNoisedDivFunction(DataChunk &args, ExpressionState &state, Vector // ============================================================================ // Bind function for priv_filter/priv_select (reads privacy_seed and pac_mi settings) // ============================================================================ -static unique_ptr PacCategoricalBind(ClientContext &ctx, ScalarFunction &func, - vector> &args) { +static unique_ptr PacCategoricalBind(BindScalarFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &func = bind_input.GetBoundFunction(); + auto &args = bind_input.GetArguments(); // Read mi from pac_mi setting double mi = GetPacMiFromSetting(ctx); // Default correction is 1.0 (can be overridden via explicit parameter in some functions) @@ -828,6 +825,19 @@ struct PacListAggregateState { uint64_t counts[64]; // Count of non-null values (for avg/count) }; +struct ListStateCodec { + static void Write(AggregateStateWriter &out, const PacListAggregateState &state) { + out.Write(state.key_hash); + out.Write(state.values); + out.Write(state.counts); + } + static void Read(AggregateStateReader &in, PacListAggregateState &state, ArenaAllocator &) { + in.Read(state.key_hash); + in.Read(state.values); + in.Read(state.counts); + } +}; + // ============================================================================ // Template-specialized operations for each aggregate type // These are designed to be inlined and auto-vectorized by the compiler @@ -1010,13 +1020,16 @@ struct PacListOps { // ============================================================================ template -static void PacListAggregateInit(const AggregateFunction &, data_ptr_t state_ptr) { - auto &state = *reinterpret_cast(state_ptr); - state.key_hash = 0; - PAC_FLOAT init_val = PacListOps::InitValue(); - for (idx_t i = 0; i < 64; i++) { - state.values[i] = init_val; - state.counts[i] = 0; +static void PacListAggregateInit(AggregateStateInput &, data_ptr_t *states, idx_t count) { + for (idx_t i = 0; i < count; i++) { + auto state_ptr = states[i]; + auto &state = *reinterpret_cast(state_ptr); + state.key_hash = 0; + PAC_FLOAT init_val = PacListOps::InitValue(); + for (idx_t i = 0; i < 64; i++) { + state.values[i] = init_val; + state.counts[i] = 0; + } } } @@ -1060,19 +1073,19 @@ static void PacListAggregateUpdate(Vector inputs[], AggregateInputData &, idx_t auto &list_vec = inputs[0]; UnifiedVectorFormat list_data; - list_vec.ToUnifiedFormat(count, list_data); + list_vec.ToUnifiedFormat(list_data); - auto &child_vec = ListVector::GetEntry(list_vec); + auto &child_vec = ListVector::GetChild(list_vec); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(list_vec), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); UnifiedVectorFormat sdata; - state_vector.ToUnifiedFormat(count, sdata); - auto states = reinterpret_cast(sdata.data); + state_vector.ToUnifiedFormat(sdata); + auto states = UnifiedVectorFormat::GetData(sdata); // Check if child vector is flat with all valid entries (enables dense path) - bool child_is_flat = child_vec.GetVectorType() == VectorType::FLAT_VECTOR && child_data.validity.AllValid(); + bool child_is_flat = child_vec.GetVectorType() == VectorType::FLAT_VECTOR && child_data.validity.CannotHaveNull(); for (idx_t i = 0; i < count; i++) { auto list_idx = list_data.sel->get_index(i); @@ -1098,10 +1111,10 @@ static void PacListAggregateUpdate(Vector inputs[], AggregateInputData &, idx_t template static void PacListAggregateCombine(Vector &source_vec, Vector &target_vec, AggregateInputData &, idx_t count) { UnifiedVectorFormat sdata, tdata; - source_vec.ToUnifiedFormat(count, sdata); - target_vec.ToUnifiedFormat(count, tdata); - auto sources = reinterpret_cast(sdata.data); - auto targets = reinterpret_cast(tdata.data); + source_vec.ToUnifiedFormat(sdata); + target_vec.ToUnifiedFormat(tdata); + auto sources = UnifiedVectorFormat::GetData(sdata); + auto targets = UnifiedVectorFormat::GetData(tdata); for (idx_t i = 0; i < count; i++) { auto &source = *sources[sdata.sel->get_index(i)]; @@ -1116,14 +1129,14 @@ static void PacListAggregateCombine(Vector &source_vec, Vector &target_vec, Aggr } template -static void PacListAggregateFinalize(Vector &state_vector, AggregateInputData &, Vector &result, idx_t count, +static void PacListAggregateFinalize(Vector &state_vector, AggregateFinalizeInputData &, Vector &result, idx_t count, idx_t offset) { UnifiedVectorFormat sdata; - state_vector.ToUnifiedFormat(count, sdata); - auto states = reinterpret_cast(sdata.data); + state_vector.ToUnifiedFormat(sdata); + auto states = UnifiedVectorFormat::GetData(sdata); result.SetVectorType(VectorType::FLAT_VECTOR); - auto &result_validity = FlatVector::Validity(result); + auto &result_validity = FlatVector::ValidityMutable(result); for (idx_t i = 0; i < count; i++) { auto &state = *states[sdata.sel->get_index(i)]; @@ -1160,13 +1173,15 @@ static void PacListAggregateDestructor(Vector &, AggregateInputData &, idx_t) { } template -static AggregateFunction CreatePacListAggregate(const string &name) { +static AggregateFunction CreatePacListAggregate(const Identifier &name) { auto list_double_type = LogicalType::LIST(PacFloatLogicalType()); - return AggregateFunction(name, {list_double_type}, list_double_type, - AggregateFunction::StateSize, PacListAggregateInit, - PacListAggregateUpdate, PacListAggregateCombine, - PacListAggregateFinalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, nullptr, - nullptr, PacListAggregateDestructor); + auto function = AggregateFunction( + Identifier(name), {list_double_type}, list_double_type, AggregateFunction::StateSize, + PacListAggregateInit, PacListAggregateUpdate, PacListAggregateCombine, + PacListAggregateFinalize, FunctionNullHandling::DEFAULT_NULL_HANDLING, nullptr, nullptr, + PacListAggregateDestructor); + AggregateStateExport::Register(function); + return function; } void AddPacListAggregateOverload(AggregateFunctionSet &set, const string &aggr_type) { @@ -1203,19 +1218,20 @@ void RegisterPacCategoricalFunctions(ExtensionLoader &loader) { // priv_filter(list) -> BOOLEAN ScalarFunction pac_filter_list("priv_filter", {list_bool_type}, LogicalType::BOOLEAN, PacFilterFromBoolListFunction, - PacCategoricalBind, nullptr, nullptr, PacCategoricalInitLocal); + PacCategoricalBind, nullptr, PacCategoricalInitLocal); loader.RegisterFunction(make_scalar_info( pac_filter_list, "[INTERNAL] Probabilistic filter from boolean list for categorical queries.")); // priv_coalesce(list) -> list auto list_double_type = LogicalType::LIST(PacFloatLogicalType()); ScalarFunction priv_coalesce("priv_coalesce", {list_double_type}, list_double_type, PacCoalesceFunction); - priv_coalesce.null_handling = FunctionNullHandling::SPECIAL_HANDLING; + priv_coalesce.SetNullHandling(FunctionNullHandling::SPECIAL_HANDLING); loader.RegisterFunction(make_scalar_info(priv_coalesce, "[INTERNAL] Replaces NULL counter list with 64 zeros.")); // as_noised(list) -> PAC_FLOAT ScalarFunction as_noised("as_noised", {list_double_type}, PacFloatLogicalType(), PacNoisedFunction, - PacCategoricalBind, nullptr, nullptr, PacCategoricalInitLocal); + PacCategoricalBind, nullptr, PacCategoricalInitLocal); + as_noised.SetStability(FunctionStability::VOLATILE); loader.RegisterFunction( make_scalar_info(as_noised, "[INTERNAL] Applies PAC noise to 64-element counter list, returns scalar.")); @@ -1225,7 +1241,8 @@ void RegisterPacCategoricalFunctions(ExtensionLoader &loader) { // as_noised_div(list sum, list cnt) -> DOUBLE ScalarFunction as_noised_div("as_noised_div", {list_double_type, list_double_type}, LogicalType::DOUBLE, - PacNoisedDivFunction, PacCategoricalBind, nullptr, nullptr, PacCategoricalInitLocal); + PacNoisedDivFunction, PacCategoricalBind, nullptr, PacCategoricalInitLocal); + as_noised_div.SetStability(FunctionStability::VOLATILE); loader.RegisterFunction(make_scalar_info(as_noised_div, "[INTERNAL] Fused counter-list division + noise for AVG.")); // priv_filter_: optimized comparison + filter in a single pass diff --git a/src/categorical/pac_categorical_detection.cpp b/src/categorical/pac_categorical_detection.cpp index 9d22ddc0..ae2b51c5 100644 --- a/src/categorical/pac_categorical_detection.cpp +++ b/src/categorical/pac_categorical_detection.cpp @@ -5,130 +5,89 @@ // // Created by ila on 1/23/26. // +#include "privacy_debug.hpp" #include "categorical/pac_categorical_detection.hpp" +#include "duckdb/planner/operator/logical_cteref.hpp" +#include "duckdb/planner/operator/logical_materialized_cte.hpp" namespace duckdb { -static string FindPacAggregateInExpression(Expression *expr, LogicalOperator *plan_root, - ColumnBinding *out_source_binding = nullptr); // forward declaration - -// Trace a column binding through the plan to find if it comes from a PAC aggregate -// Returns the PAC aggregate name if found (base name without _counters), empty string otherwise -static string TracePacAggregateFromBinding(const ColumnBinding &binding, LogicalOperator *plan_root, - ColumnBinding *out_source_binding = nullptr) { - auto *source_op = FindOperatorByTableIndex(plan_root, binding.table_index); - if (!source_op) { - return ""; +// Follow every source, including both COUNT copies in DuckDB's nullable-count +// wrapper and aggregates over aggregates. A projection may depend on several +// stochastic outputs; retaining only its first source prematurely finalizes the rest. +static void CollectPacSources(const Expression &expr, LogicalOperator &root, vector &sources, + unordered_set &visited) { + if (IsAlreadyWrappedInPacNoised(&expr)) { + return; } - if (source_op->type == LogicalOperatorType::LOGICAL_PROJECTION) { - auto *unwrapped = RecognizeDuckDBScalarWrapper(source_op); - if (unwrapped) { // Scalar subquery wrapper: trace through inner projection's expression - auto &agg_child = source_op->Cast().children[0]->Cast(); - auto &inner_proj = agg_child.children[0]->Cast(); - if (!inner_proj.expressions.empty()) { - return FindPacAggregateInExpression(inner_proj.expressions[0].get(), plan_root, out_source_binding); - } - return ""; + if (expr.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { + auto binding = expr.Cast().Binding(); + if (!visited.insert(HashBinding(binding)).second) { + return; } - auto &proj = source_op->Cast(); - if (binding.column_index < proj.expressions.size()) { - if (IsAlreadyWrappedInPacNoised(proj.expressions[binding.column_index].get())) { - return ""; - } - return FindPacAggregateInExpression(proj.expressions[binding.column_index].get(), plan_root, - out_source_binding); + auto *source = FindOperatorByTableIndex(&root, binding.table_index.index); + if (!source) { + return; } - } else if (source_op->type == LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY) { - auto &aggr = source_op->Cast(); - if (binding.table_index == aggr.aggregate_index) { - if (binding.column_index < aggr.expressions.size()) { - auto &agg_expr = aggr.expressions[binding.column_index]; - if (agg_expr->type == ExpressionType::BOUND_AGGREGATE) { - auto &bound_agg = agg_expr->Cast(); - if (IsAnyPacAggregate(bound_agg.function.name)) { - if (out_source_binding) { - *out_source_binding = binding; - } - return GetBasePacAggregateName(bound_agg.function.name); - } - } + if (source->type == LogicalOperatorType::LOGICAL_PROJECTION) { + auto &projection = source->Cast(); + if (binding.column_index < projection.expressions.size()) { + CollectPacSources(*projection.expressions[binding.column_index], root, sources, visited); } - } - } - return ""; -} - -// Helper to collect ALL distinct PAC aggregate bindings in an expression -// Returns the bindings in the order they were discovered -static void CollectPacBindingsInExpression(Expression *expr, LogicalOperator *root, vector &bindings, - unordered_map &binding_hash_to_index) { - // Check if this is a column reference that traces back to a PAC aggregate - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { - auto &col_ref = expr->Cast(); - ColumnBinding source_binding; - string pac_name = TracePacAggregateFromBinding(col_ref.binding, root, &source_binding); - if (!pac_name.empty()) { - uint64_t binding_hash = HashBinding(col_ref.binding); - if (binding_hash_to_index.find(binding_hash) == binding_hash_to_index.end()) { - PacBindingInfo info; - info.binding = col_ref.binding; - info.index = bindings.size(); - info.source_binding = source_binding; - binding_hash_to_index[binding_hash] = info.index; - bindings.push_back(info); + } else if (source->type == LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY) { + auto &aggregate = source->Cast(); + if (binding.table_index != aggregate.aggregate_index || + binding.column_index >= aggregate.expressions.size()) { + return; + } + auto &aggregate_expr = *aggregate.expressions[binding.column_index]; + if (aggregate_expr.GetExpressionType() == ExpressionType::BOUND_AGGREGATE && + IsAnyPacAggregate( + aggregate_expr.Cast().Function().GetName().GetIdentifierName())) { + sources.push_back(binding); + } else { + CollectPacSources(aggregate_expr, root, sources, visited); + } + } else if (source->type == LogicalOperatorType::LOGICAL_CTE_REF) { + auto &ref = source->Cast(); + PRIVACY_DEBUG_PRINT("Categorical detection: tracing CTE reader #" + std::to_string(ref.table_index.index)); + auto *cte = FindMaterializedCTE(&root, ref.cte_index.index); + if (cte) { + auto bindings = cte->children[0]->GetColumnBindings(); + if (binding.column_index < bindings.size()) { + BoundColumnRefExpression definition_ref(expr.GetReturnType(), bindings[binding.column_index]); + CollectPacSources(definition_ref, root, sources, visited); + } } } + return; } - ExpressionIterator::EnumerateChildren(*expr, [&](Expression &child) { - CollectPacBindingsInExpression(&child, root, bindings, binding_hash_to_index); - }); + ExpressionIterator::EnumerateChildren( + expr, [&](const Expression &child) { CollectPacSources(child, root, sources, visited); }); } -// Find all PAC aggregate bindings in an expression vector FindAllPacBindingsInExpression(Expression *expr, LogicalOperator *plan_root) { - vector bindings; - unordered_map binding_hash_to_index; - CollectPacBindingsInExpression(expr, plan_root, bindings, binding_hash_to_index); - return bindings; -} - -// Recursively search for PAC aggregate in an expression tree, with plan context for tracing column refs -// Returns the base aggregate name (without _counters suffix) -static string FindPacAggregateInExpression(Expression *expr, LogicalOperator *plan_root, - ColumnBinding *out_source_binding) { - // examine specific expression types where the PAC aggregate could be in - if (expr->type == ExpressionType::BOUND_AGGREGATE) { // Base case: direct PAC aggregate - auto &agg_expr = expr->Cast(); - if (IsAnyPacAggregate(agg_expr.function.name)) { - return GetBasePacAggregateName(agg_expr.function.name); - } - } else if (expr->type == ExpressionType::BOUND_COLUMN_REF) { - auto &col_ref = expr->Cast(); - return TracePacAggregateFromBinding(col_ref.binding, plan_root, out_source_binding); - } else if (expr->type == ExpressionType::SUBQUERY) { - auto &subquery_expr = expr->Cast(); - for (auto &child : subquery_expr.children) { - string result = FindPacAggregateInExpression(child.get(), plan_root); - if (!result.empty()) { - return result; // Found PAC aggregate in the subquery's children (for IN, ANY, ALL operators) - } - } - if (subquery_expr.subquery_type == SubqueryType::SCALAR) { - string result = FindPacAggregateInOperator(plan_root); - if (!result.empty()) { - return result; // found PAC aggregate in scalar subquery + vector result; + unordered_set seen; + std::function collect = [&](const Expression &expression) { + if (expression.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { + auto binding = expression.Cast().Binding(); + if (seen.insert(HashBinding(binding)).second) { + vector sources; + unordered_set visited; + CollectPacSources(expression, *plan_root, sources, visited); + for (auto &source : sources) { + PacBindingInfo info; + info.binding = binding; + info.source_binding = source; + info.index = result.size(); + result.push_back(info); + } } } - return ""; - } - // Generic traversal for all other expression types (comparisons, operators, casts, - // functions, constants, CASE, BETWEEN, conjunctions, window functions, etc.) - string result; - ExpressionIterator::EnumerateChildren(*expr, [&](Expression &child) { - if (result.empty()) { - result = FindPacAggregateInExpression(&child, plan_root, out_source_binding); - } - }); + ExpressionIterator::EnumerateChildren(expression, collect); + }; + collect(*expr); return result; } @@ -146,7 +105,7 @@ static bool IsHavingBinding(LogicalOperator *filter_op, const ColumnBinding &bin LogicalOperator *child = filter_op->children[0].get(); while (child && child->type == LogicalOperatorType::LOGICAL_PROJECTION) { auto &proj = child->Cast(); - if (binding.table_index == proj.table_index) { + if (binding.table_index.index == proj.table_index.index) { return true; } if (child->children.empty()) { @@ -156,7 +115,8 @@ static bool IsHavingBinding(LogicalOperator *filter_op, const ColumnBinding &bin } if (child && child->type == LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY) { auto &agg = child->Cast(); - return binding.table_index == agg.aggregate_index || binding.table_index == agg.group_index; + return binding.table_index.index == agg.aggregate_index.index || + binding.table_index.index == agg.group_index.index; } return false; } @@ -178,16 +138,18 @@ void DetectCategoricalPatterns(LogicalOperator *op, LogicalOperator *plan_root, // Fallback: extract from children[0] of the first PAC aggregate expression. if (op->type == LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY) { auto &agg = op->Cast(); - auto it = pac_agg_info.find(agg.aggregate_index); + auto it = pac_agg_info.find(agg.aggregate_index.index); if (it != pac_agg_info.end()) { hash_for_children = it->second.hash_binding; } else { for (auto &agg_expr : agg.expressions) { - if (agg_expr->type == ExpressionType::BOUND_AGGREGATE) { + if (agg_expr->GetExpressionType() == ExpressionType::BOUND_AGGREGATE) { auto &bound_agg = agg_expr->Cast(); - if (IsPacAggregate(bound_agg.function.name) && !bound_agg.children.empty() && - bound_agg.children[0]->type == ExpressionType::BOUND_COLUMN_REF) { - hash_for_children = bound_agg.children[0]->Cast().binding; + if (IsPacAggregate(bound_agg.FunctionMutable().GetName().GetIdentifierName()) && + !bound_agg.GetChildrenMutable().empty() && + bound_agg.GetChildrenMutable()[0]->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { + hash_for_children = + bound_agg.GetChildrenMutable()[0]->Cast().BindingMutable(); break; } } @@ -196,14 +158,14 @@ void DetectCategoricalPatterns(LogicalOperator *op, LogicalOperator *plan_root, } // Resolve hash binding through projections (follow col_ref chains) - if (hash_for_children.table_index != DConstants::INVALID_INDEX && + if (hash_for_children.table_index.index != DConstants::INVALID_INDEX && op->type == LogicalOperatorType::LOGICAL_PROJECTION) { auto &proj = op->Cast(); - if (hash_for_children.table_index == proj.table_index && + if (hash_for_children.table_index.index == proj.table_index.index && hash_for_children.column_index < proj.expressions.size()) { auto &expr = proj.expressions[hash_for_children.column_index]; - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { - hash_for_children = expr->Cast().binding; + if (expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { + hash_for_children = expr->Cast().BindingMutable(); } } } @@ -238,7 +200,7 @@ void DetectCategoricalPatterns(LogicalOperator *op, LogicalOperator *plan_root, info.expr_index = i; info.pac_bindings = std::move(pac_bindings); info.outer_pac_hash = hash_for_children; - info.has_outer_pac_hash = (hash_for_children.table_index != DConstants::INVALID_INDEX); + info.has_outer_pac_hash = (hash_for_children.table_index.index != DConstants::INVALID_INDEX); patterns.push_back(info); } } @@ -247,20 +209,26 @@ void DetectCategoricalPatterns(LogicalOperator *op, LogicalOperator *plan_root, auto &join = op->Cast(); for (idx_t i = 0; i < join.conditions.size(); i++) { auto &cond = join.conditions[i]; - // Check both sides; any PAC aggregate in a join condition is categorical (not HAVING). - for (auto *side : {cond.left.get(), cond.right.get()}) { - string pac_name = FindPacAggregateInExpression(side, plan_root); - if (pac_name.empty()) { - continue; + vector bindings; + unordered_set seen; + for (auto *side : {cond.LeftReference().get(), cond.RightReference().get()}) { + for (auto &binding : FindAllPacBindingsInExpression(side, plan_root)) { + if (seen.insert(HashBinding(binding.source_binding)).second) { + binding.index = bindings.size(); + bindings.push_back(binding); + } } - CategoricalPatternInfo info; - info.parent_op = op; - info.expr_index = i; - info.outer_pac_hash = hash_for_children; - info.has_outer_pac_hash = (hash_for_children.table_index != DConstants::INVALID_INDEX); - patterns.push_back(info); - break; // one match per condition is enough } + if (bindings.empty()) { + continue; + } + CategoricalPatternInfo info; + info.parent_op = op; + info.expr_index = i; + info.pac_bindings = std::move(bindings); + info.outer_pac_hash = hash_for_children; + info.has_outer_pac_hash = (hash_for_children.table_index.index != DConstants::INVALID_INDEX); + patterns.push_back(std::move(info)); } } else if (!inside_aggregate && patterns.empty() && op->type == LogicalOperatorType::LOGICAL_PROJECTION) { // Check projection expressions for arithmetic involving multiple PAC aggregates @@ -271,7 +239,7 @@ void DetectCategoricalPatterns(LogicalOperator *op, LogicalOperator *plan_root, auto &proj = op->Cast(); for (idx_t i = 0; i < proj.expressions.size(); i++) { auto &expr = proj.expressions[i]; - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { + if (expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { continue; // Skip simple column references - they don't need rewriting } // Find ALL PAC aggregate bindings in this expression @@ -287,18 +255,19 @@ void DetectCategoricalPatterns(LogicalOperator *op, LogicalOperator *plan_root, is_arithmetic_with_pac = true; } else if (pac_bindings.size() == 1) { // Single aggregate - check if it's in an arithmetic expression (not just a column ref or simple cast) - if (expr->type != ExpressionType::BOUND_COLUMN_REF && expr->type != ExpressionType::OPERATOR_CAST) { + if (expr->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF && + expr->GetExpressionType() != ExpressionType::OPERATOR_CAST) { is_arithmetic_with_pac = true; - } else if (expr->type == ExpressionType::OPERATOR_CAST) { + } else if (expr->GetExpressionType() == ExpressionType::OPERATOR_CAST) { // Check if cast contains arithmetic - auto &cast_expr = expr->Cast(); - if (cast_expr.child->type != ExpressionType::BOUND_COLUMN_REF) { + auto &cast_expr = expr->Cast(); + if (cast_expr.GetChildrenMutable()[0]->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { is_arithmetic_with_pac = true; } } } if (is_arithmetic_with_pac) { - if (!IsNumericalType(expr->return_type)) { + if (!IsNumericalType(expr->GetReturnType())) { continue; // Only create pattern if result is numerical (as_noised only works on numbers) } CategoricalPatternInfo info; diff --git a/src/categorical/pac_categorical_lambdas.cpp b/src/categorical/pac_categorical_lambdas.cpp index fc6d8f01..a9ba4090 100644 --- a/src/categorical/pac_categorical_lambdas.cpp +++ b/src/categorical/pac_categorical_lambdas.cpp @@ -6,6 +6,7 @@ // Created by ila on 3/13/26. // #include "categorical/pac_categorical_lambdas.hpp" +#include "privacy_debug.hpp" #include "aggregates/as_aggregate.hpp" #include "utils/privacy_helpers.hpp" #include "duckdb/planner/expression/bound_conjunction_expression.hpp" @@ -19,55 +20,45 @@ namespace duckdb { // pac_binding_map: maps binding hash -> struct field index (single: one entry mapping to 0) // struct_type: nullptr = single binding (elem ref), non-null = multi binding (struct field extract) static unique_ptr -CloneForLambdaBody(Expression *expr, const unordered_map &pac_binding_map, +CloneForLambdaBody(ClientContext &context, Expression *expr, const unordered_map &pac_binding_map, vector> &captures, unordered_map &capture_map, const unordered_map &counter_bindings, const LogicalType *struct_type) { // Helper to build the replacement expression for a matched PAC binding auto make_pac_replacement = [&](const unordered_map::const_iterator &it, - const string &alias) -> unique_ptr { + const Identifier &alias) -> unique_ptr { if (struct_type) { idx_t field_idx = it->second; - string field_name = GetStructFieldName(field_idx); auto elem_ref = make_uniq("elem", *struct_type, idx_t(0)); - auto child_types = StructType::GetChildTypes(*struct_type); - LogicalType extract_return_type = PacFloatLogicalType(); - for (idx_t j = 0; j < child_types.size(); j++) { - if (child_types[j].first == field_name) { - extract_return_type = child_types[j].second; - break; - } - } auto extract_func = StructExtractAtFun::GetFunction(); - auto bind_data = StructExtractAtFun::GetBindData(field_idx); + vector> extract_children; extract_children.push_back(std::move(elem_ref)); extract_children.push_back( make_uniq(Value::BIGINT(static_cast(field_idx + 1)))); - return make_uniq(extract_return_type, extract_func, std::move(extract_children), - std::move(bind_data)); + return FunctionBinder(context).BindScalarFunction(extract_func, std::move(extract_children)); } else { return make_uniq(alias, PacFloatLogicalType(), idx_t(0)); } }; // rewrite all kinds of expressions: - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { + if (expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = expr->Cast(); - uint64_t binding_hash = HashBinding(col_ref.binding); + uint64_t binding_hash = HashBinding(col_ref.BindingMutable()); // Direct match auto it = pac_binding_map.find(binding_hash); if (it != pac_binding_map.end()) { return make_pac_replacement(it, "elem"); } // Trace through counter_bindings to find the aggregate-level PAC binding - auto cb_it = counter_bindings.find(HashBinding(col_ref.binding)); + auto cb_it = counter_bindings.find(HashBinding(col_ref.BindingMutable())); if (cb_it != counter_bindings.end()) { auto traced_it = pac_binding_map.find(HashBinding(cb_it->second)); if (traced_it != pac_binding_map.end()) { - return make_pac_replacement(traced_it, col_ref.alias); + return make_pac_replacement(traced_it, col_ref.GetAlias()); } } // Other column ref — capture for use in lambda - uint64_t hash = HashBinding(col_ref.binding); + uint64_t hash = HashBinding(col_ref.BindingMutable()); idx_t capture_idx; auto cap_it = capture_map.find(hash); if (cap_it != capture_map.end()) { @@ -77,128 +68,129 @@ CloneForLambdaBody(Expression *expr, const unordered_map &pac_b capture_map[hash] = capture_idx; captures.push_back(col_ref.Copy()); } - return make_uniq(col_ref.alias, col_ref.return_type, 1 + capture_idx); - } else if (expr->type == ExpressionType::VALUE_CONSTANT) { + return make_uniq(col_ref.GetAlias(), col_ref.GetReturnType(), 1 + capture_idx); + } else if (expr->GetExpressionType() == ExpressionType::VALUE_CONSTANT) { return expr->Copy(); - } else if (expr->type == ExpressionType::OPERATOR_CAST) { - auto &cast = expr->Cast(); - auto child_clone = - CloneForLambdaBody(cast.child.get(), pac_binding_map, captures, capture_map, counter_bindings, struct_type); - if (child_clone->return_type == cast.return_type) { + } else if (expr->GetExpressionType() == ExpressionType::OPERATOR_CAST) { + auto &cast = expr->Cast(); + auto child_clone = CloneForLambdaBody(context, cast.GetChildrenMutable()[0].get(), pac_binding_map, captures, + capture_map, counter_bindings, struct_type); + if (child_clone->GetReturnType() == cast.GetReturnType()) { return child_clone; // If the child's type already matches the target type, skip the cast } // Otherwise, create a new cast with the correct function for the new child type - return BoundCastExpression::AddDefaultCastToType(std::move(child_clone), cast.return_type); - } else if (expr->GetExpressionClass() == ExpressionClass::BOUND_COMPARISON) { - auto &comp = expr->Cast(); - auto left_clone = - CloneForLambdaBody(comp.left.get(), pac_binding_map, captures, capture_map, counter_bindings, struct_type); - auto right_clone = - CloneForLambdaBody(comp.right.get(), pac_binding_map, captures, capture_map, counter_bindings, struct_type); + return BoundCastExpression::AddDefaultCastToType(std::move(child_clone), cast.GetReturnType()); + } else if (BoundComparisonExpression::IsComparison(*expr)) { + auto &comp = expr->Cast(); + auto left_clone = CloneForLambdaBody(context, comp.GetChildrenMutable()[0].get(), pac_binding_map, captures, + capture_map, counter_bindings, struct_type); + auto right_clone = CloneForLambdaBody(context, comp.GetChildrenMutable()[1].get(), pac_binding_map, captures, + capture_map, counter_bindings, struct_type); // Reconcile types if they differ (needed for multi where struct fields are DOUBLE // but original CASTs may introduce DECIMAL) - if (left_clone->return_type != right_clone->return_type) { - if (left_clone->return_type != PacFloatLogicalType()) { + if (left_clone->GetReturnType() != right_clone->GetReturnType()) { + if (left_clone->GetReturnType() != PacFloatLogicalType()) { left_clone = BoundCastExpression::AddDefaultCastToType(std::move(left_clone), PacFloatLogicalType()); } - if (right_clone->return_type != PacFloatLogicalType()) { + if (right_clone->GetReturnType() != PacFloatLogicalType()) { right_clone = BoundCastExpression::AddDefaultCastToType(std::move(right_clone), PacFloatLogicalType()); } } - return make_uniq(expr->type, std::move(left_clone), std::move(right_clone)); - } else if (expr->type == ExpressionType::BOUND_FUNCTION) { + return BoundComparisonExpression::Create(expr->GetExpressionType(), std::move(left_clone), + std::move(right_clone)); + } else if (expr->GetExpressionType() == ExpressionType::BOUND_FUNCTION) { auto &func = expr->Cast(); vector> new_children; bool any_cast_needed = false; - for (idx_t i = 0; i < func.children.size(); i++) { - auto child_clone = CloneForLambdaBody(func.children[i].get(), pac_binding_map, captures, capture_map, - counter_bindings, struct_type); + for (idx_t i = 0; i < func.GetChildrenMutable().size(); i++) { + auto child_clone = CloneForLambdaBody(context, func.GetChildrenMutable()[i].get(), pac_binding_map, + captures, capture_map, counter_bindings, struct_type); // If a child's type changed (e.g., DECIMAL->DOUBLE from PAC counter conversion), // cast it to the type the bound function expects, so the function binding stays valid. - if (i < func.function.arguments.size() && child_clone->return_type != func.function.arguments[i]) { - child_clone = - BoundCastExpression::AddDefaultCastToType(std::move(child_clone), func.function.arguments[i]); + if (i < func.FunctionMutable().GetArguments().size() && + child_clone->GetReturnType() != func.FunctionMutable().GetArguments()[i]) { + child_clone = BoundCastExpression::AddDefaultCastToType(std::move(child_clone), + func.FunctionMutable().GetArguments()[i]); any_cast_needed = true; } new_children.push_back(std::move(child_clone)); } unique_ptr result = - make_uniq(func.return_type, func.function, std::move(new_children), - func.bind_info ? func.bind_info->Copy() : nullptr); + make_uniq(func.FunctionMutable(), std::move(new_children), + func.BindInfoMutable() ? func.BindInfoMutable()->Copy() : nullptr); // Children were cast from PAC_FLOAT to the function's bound types (e.g., DECIMAL). // Cast the result back to PAC_FLOAT so the list_transform output stays LIST. - if (any_cast_needed && result->return_type != PacFloatLogicalType()) { + if (any_cast_needed && result->GetReturnType() != PacFloatLogicalType()) { result = BoundCastExpression::AddDefaultCastToType(std::move(result), PacFloatLogicalType()); } return result; } else if (expr->GetExpressionClass() == ExpressionClass::BOUND_OPERATOR) { // NOT, arithmetic, COALESCE, IS NULL.. auto &op = expr->Cast(); vector> new_children; - for (auto &child : op.children) { - new_children.push_back( - CloneForLambdaBody(child.get(), pac_binding_map, captures, capture_map, counter_bindings, struct_type)); + for (auto &child : op.GetChildrenMutable()) { + new_children.push_back(CloneForLambdaBody(context, child.get(), pac_binding_map, captures, capture_map, + counter_bindings, struct_type)); } // For COALESCE with mismatched child types, cast all to the first child's type - LogicalType result_type = op.return_type; - if (expr->type == ExpressionType::OPERATOR_COALESCE && new_children.size() > 1) { - LogicalType first_type = new_children[0]->return_type; + LogicalType result_type = op.GetReturnType(); + if (expr->GetExpressionType() == ExpressionType::OPERATOR_COALESCE && new_children.size() > 1) { + LogicalType first_type = new_children[0]->GetReturnType(); bool types_mismatch = false; for (idx_t ci = 1; ci < new_children.size(); ci++) { - if (new_children[ci]->return_type != first_type) { + if (new_children[ci]->GetReturnType() != first_type) { types_mismatch = true; break; } } if (types_mismatch) { for (auto &c : new_children) { - if (c->return_type != first_type) { + if (c->GetReturnType() != first_type) { c = BoundCastExpression::AddDefaultCastToType(std::move(c), first_type); } } result_type = first_type; } } - auto op_result = make_uniq(expr->type, result_type); + auto op_result = make_uniq(expr->GetExpressionType(), result_type); for (auto &c : new_children) { - op_result->children.push_back(std::move(c)); + op_result->GetChildrenMutable().push_back(std::move(c)); } return op_result; } else if (expr->GetExpressionClass() == ExpressionClass::BOUND_CONJUNCTION) { // AND, OR auto &conj = expr->Cast(); - auto result = make_uniq(expr->type); - for (auto &child : conj.children) { - result->children.push_back( - CloneForLambdaBody(child.get(), pac_binding_map, captures, capture_map, counter_bindings, struct_type)); + auto result = make_uniq(expr->GetExpressionType()); + for (auto &child : conj.GetChildrenMutable()) { + result->GetChildrenMutable().push_back(CloneForLambdaBody(context, child.get(), pac_binding_map, captures, + capture_map, counter_bindings, struct_type)); } return result; - } else if (expr->type == ExpressionType::CASE_EXPR) { + } else if (expr->GetExpressionType() == ExpressionType::CASE_EXPR) { auto &case_expr = expr->Cast(); if (IsScalarSubqueryWrapper(case_expr)) { - return CloneForLambdaBody(case_expr.else_expr.get(), pac_binding_map, captures, capture_map, - counter_bindings, struct_type); + PRIVACY_DEBUG_PRINT("Categorical lambda: preserving scalar-subquery cardinality check"); } // Regular CASE - recurse into all branches // Start with original return type, will update based on cloned branches - auto result = make_uniq(case_expr.return_type); - for (auto &check : case_expr.case_checks) { + auto result = make_uniq(case_expr.GetReturnType()); + for (auto &check : case_expr.CaseChecksMutable()) { BoundCaseCheck new_check; - new_check.when_expr = CloneForLambdaBody(check.when_expr.get(), pac_binding_map, captures, capture_map, - counter_bindings, struct_type); - new_check.then_expr = CloneForLambdaBody(check.then_expr.get(), pac_binding_map, captures, capture_map, - counter_bindings, struct_type); - result->case_checks.push_back(std::move(new_check)); + new_check.when_expr = CloneForLambdaBody(context, check.when_expr.get(), pac_binding_map, captures, + capture_map, counter_bindings, struct_type); + new_check.then_expr = CloneForLambdaBody(context, check.then_expr.get(), pac_binding_map, captures, + capture_map, counter_bindings, struct_type); + result->CaseChecksMutable().push_back(std::move(new_check)); } - if (case_expr.else_expr) { - result->else_expr = CloneForLambdaBody(case_expr.else_expr.get(), pac_binding_map, captures, capture_map, - counter_bindings, struct_type); + if (case_expr.ElseMutable()) { + result->ElseMutable() = CloneForLambdaBody(context, case_expr.ElseMutable().get(), pac_binding_map, + captures, capture_map, counter_bindings, struct_type); // Update return type to match ELSE branch (the PAC element type) - result->return_type = result->else_expr->return_type; + result->SetReturnType(result->ElseMutable()->GetReturnType()); } // Cast THEN branches to match the return type if needed - for (auto &check : result->case_checks) { - if (check.then_expr && check.then_expr->return_type != result->return_type) { + for (auto &check : result->CaseChecksMutable()) { + if (check.then_expr && check.then_expr->GetReturnType() != result->GetReturnType()) { check.then_expr = - BoundCastExpression::AddDefaultCastToType(std::move(check.then_expr), result->return_type); + BoundCastExpression::AddDefaultCastToType(std::move(check.then_expr), result->GetReturnType()); } } return result; @@ -211,7 +203,7 @@ static unique_ptr BuildPacLambda(unique_ptr lambda_body, vector> captures) { auto lambda = make_uniq(ExpressionType::LAMBDA, LogicalType::LAMBDA, std::move(lambda_body), idx_t(1)); - lambda->captures = std::move(captures); + lambda->CapturesMutable() = std::move(captures); return lambda; } @@ -221,29 +213,21 @@ static unique_ptr BuildListTransformCall(OptimizerExtensionInput &in unique_ptr counters_list, unique_ptr lambda_expr, const LogicalType &element_return_type = LogicalType::BOOLEAN) { - // Get the lambda body for ListLambdaBindData - auto &bound_lambda = lambda_expr->Cast(); - - // Get list_transform function from catalog - auto &catalog = Catalog::GetSystemCatalog(input.context); - auto &func_entry = catalog.GetEntry(input.context, DEFAULT_SCHEMA, "list_transform"); - - // Get the first function overload (list_transform has only one signature pattern) - auto &scalar_func = func_entry.functions.functions[0]; - - // Create the ListLambdaBindData with the lambda body - auto list_return_type = LogicalType::LIST(element_return_type); - auto bind_data = make_uniq(list_return_type, std::move(bound_lambda.lambda_expr), false, false); - - // Build children: [list, captures...] Note: The lambda itself is NOT a child after binding - only its captures are - vector> children; - children.push_back(std::move(counters_list)); - - for (auto &capture : bound_lambda.captures) { - children.push_back(std::move(capture)); // Add captures as children + vector> arguments; + arguments.push_back(std::move(counters_list)); + arguments.push_back(std::move(lambda_expr)); + auto result = + FunctionBinder(input.context).BindScalarFunction(DEFAULT_SCHEMA, "list_transform", std::move(arguments)); + D_ASSERT(result->GetReturnType() == LogicalType::LIST(element_return_type)); + auto &function = result->Cast(); + auto &lambda = function.GetChildrenMutable()[1]->Cast(); + // Match DuckDB's lambda binder: the placeholder stays at argument 1 and + // captured row values are appended as physical children. + for (auto &capture : lambda.CapturesMutable()) { + function.GetChildrenMutable().push_back(std::move(capture)); } - // Create the bound function expression - return make_uniq(list_return_type, scalar_func, std::move(children), std::move(bind_data)); + lambda.CapturesMutable().clear(); + return result; } // Build a list_zip function call combining multiple counter lists @@ -251,41 +235,19 @@ static unique_ptr BuildListTransformCall(OptimizerExtensionInput &in static unique_ptr BuildListZipCall(OptimizerExtensionInput &input, vector> counter_lists, LogicalType &out_struct_type) { - // Build the struct type for list_zip result -- list_zip returns LIST> - child_list_t struct_children; - for (idx_t i = 0; i < counter_lists.size(); i++) { - string field_name = GetStructFieldName(i); - struct_children.push_back( - make_pair(field_name, PacFloatLogicalType())); // All counter lists use PAC_FLOAT element type - } - out_struct_type = LogicalType::STRUCT(struct_children); - auto list_struct_type = LogicalType::LIST(out_struct_type); - - // Get list_zip function from catalog - auto &catalog = Catalog::GetSystemCatalog(input.context); - auto &func_entry = catalog.GetEntry(input.context, DEFAULT_SCHEMA, "list_zip"); - - // Find the appropriate overload (list_zip is variadic) - vector arg_types; - for (auto &list : counter_lists) { - arg_types.push_back(list->return_type); - } - ErrorData error; - FunctionBinder function_binder(input.context); - auto best_function = function_binder.BindFunction(func_entry.name, func_entry.functions, arg_types, error); - if (best_function.IsValid()) { - auto scalar_func = func_entry.functions.GetFunctionByOffset(best_function.GetIndex()); - return make_uniq(list_struct_type, scalar_func, std::move(counter_lists), nullptr); - } - return nullptr; + auto result = + FunctionBinder(input.context).BindScalarFunction(DEFAULT_SCHEMA, "list_zip", std::move(counter_lists)); + out_struct_type = ListType::GetChildType(result->GetReturnType()); + return result; } // Check whether an expression tree contains any null-handling operators // (COALESCE, IS NULL, IS NOT NULL). When absent, priv_coalesce is unnecessary // because NULL propagates identically through arithmetic. static bool ExpressionContainsNullHandling(Expression *expr) { - if (expr->type == ExpressionType::OPERATOR_COALESCE || expr->type == ExpressionType::OPERATOR_IS_NULL || - expr->type == ExpressionType::OPERATOR_IS_NOT_NULL) { + if (expr->GetExpressionType() == ExpressionType::OPERATOR_COALESCE || + expr->GetExpressionType() == ExpressionType::OPERATOR_IS_NULL || + expr->GetExpressionType() == ExpressionType::OPERATOR_IS_NOT_NULL) { return true; } bool found = false; @@ -300,15 +262,15 @@ static bool ExpressionContainsNullHandling(Expression *expr) { // Check which child of a binary function contains a counter binding // Returns 0 if first child has counter, 1 if second child has counter, -1 if neither or both static int FindPacChildIndex(Expression *expr, const unordered_map &counter_bindings) { - if (expr->type != ExpressionType::BOUND_FUNCTION) { + if (expr->GetExpressionType() != ExpressionType::BOUND_FUNCTION) { return -1; } auto &func = expr->Cast(); - if (func.children.size() != 2) { + if (func.GetChildrenMutable().size() != 2) { return -1; } - bool left_has = ReferencesCounterBinding(func.children[0].get(), counter_bindings); - bool right_has = ReferencesCounterBinding(func.children[1].get(), counter_bindings); + bool left_has = ReferencesCounterBinding(func.GetChildrenMutable()[0].get(), counter_bindings); + bool right_has = ReferencesCounterBinding(func.GetChildrenMutable()[1].get(), counter_bindings); if (left_has && !right_has) { return 0; } @@ -339,10 +301,10 @@ unique_ptr BuildCategoricalLambdas(OptimizerExtensionInput &input, pac_binding_map[HashBinding(pac_binding)] = 0; vector> captures; unordered_map capture_map; - auto lambda_body = - CloneForLambdaBody(expr_to_transform, pac_binding_map, captures, capture_map, counter_bindings, nullptr); + auto lambda_body = CloneForLambdaBody(input.context, expr_to_transform, pac_binding_map, captures, capture_map, + counter_bindings, nullptr); // Ensure body returns result_element_type if needed - if (result_element_type == PacFloatLogicalType() && lambda_body->return_type != PacFloatLogicalType()) { + if (result_element_type == PacFloatLogicalType() && lambda_body->GetReturnType() != PacFloatLogicalType()) { lambda_body = BoundCastExpression::AddDefaultCastToType(std::move(lambda_body), PacFloatLogicalType()); } auto lambda = BuildPacLambda(std::move(lambda_body), std::move(captures)); @@ -366,9 +328,9 @@ unique_ptr BuildCategoricalLambdas(OptimizerExtensionInput &input, if (zipped_list) { vector> captures; unordered_map map; - auto lambda_body = - CloneForLambdaBody(expr_to_transform, binding_to_index, captures, map, counter_bindings, &struct_type); - if (result_element_type == PacFloatLogicalType() && lambda_body->return_type != PacFloatLogicalType()) { + auto lambda_body = CloneForLambdaBody(input.context, expr_to_transform, binding_to_index, captures, map, + counter_bindings, &struct_type); + if (result_element_type == PacFloatLogicalType() && lambda_body->GetReturnType() != PacFloatLogicalType()) { lambda_body = BoundCastExpression::AddDefaultCastToType(std::move(lambda_body), PacFloatLogicalType()); } auto lambda = BuildPacLambda(std::move(lambda_body), std::move(captures)); @@ -382,15 +344,15 @@ unique_ptr TryRewriteFilterComparison(OptimizerExtensionInput &input const vector &pac_bindings, Expression *expr, const unordered_map &counter_bindings, PacWrapKind wrap_kind, ColumnBinding priv_hash) { - if (pac_bindings.size() != 1 || expr->GetExpressionClass() != ExpressionClass::BOUND_COMPARISON) { + if (pac_bindings.size() != 1 || !BoundComparisonExpression::IsComparison(*expr)) { return nullptr; // Only works with a comparison against a single PAC binding } - auto &comp = expr->Cast(); - ExpressionType cmp_type = comp.type; + auto &comp = expr->Cast(); + ExpressionType cmp_type = comp.GetExpressionType(); // Determine which side has the PAC binding - auto left_bindings = FindCounterBindings(comp.left.get(), counter_bindings); - auto right_bindings = FindCounterBindings(comp.right.get(), counter_bindings); + auto left_bindings = FindCounterBindings(comp.GetChildrenMutable()[0].get(), counter_bindings); + auto right_bindings = FindCounterBindings(comp.GetChildrenMutable()[1].get(), counter_bindings); bool left_has_pac = !left_bindings.empty(); bool right_has_pac = !right_bindings.empty(); @@ -401,17 +363,17 @@ unique_ptr TryRewriteFilterComparison(OptimizerExtensionInput &input unique_ptr scalar_side; unique_ptr list_side; if (left_has_pac) { // PAC on left: flip comparison - scalar_side = comp.right->Copy(); - list_side = comp.left->Copy(); + scalar_side = comp.GetChildrenMutable()[1]->Copy(); + list_side = comp.GetChildrenMutable()[0]->Copy(); cmp_type = FlipComparison(cmp_type); } else { - scalar_side = comp.left->Copy(); - list_side = comp.right->Copy(); + scalar_side = comp.GetChildrenMutable()[0]->Copy(); + list_side = comp.GetChildrenMutable()[1]->Copy(); } // Iteratively simplify: move arithmetic from list_side to scalar_side - while (list_side->type == ExpressionType::BOUND_FUNCTION) { + while (list_side->GetExpressionType() == ExpressionType::BOUND_FUNCTION) { auto &func = list_side->Cast(); - if (func.children.size() != 2) { + if (func.GetChildrenMutable().size() != 2) { break; } int pac_child = FindPacChildIndex(list_side.get(), counter_bindings); @@ -419,11 +381,12 @@ unique_ptr TryRewriteFilterComparison(OptimizerExtensionInput &input break; } bool needs_positive_check = false; - const char *inverse_op = GetInverseArithmeticOp(func.function.name, pac_child, needs_positive_check); + const char *inverse_op = GetInverseArithmeticOp(func.FunctionMutable().GetName().GetIdentifierName(), pac_child, + needs_positive_check); if (!inverse_op) { break; } - auto &scalar_operand = func.children[1 - pac_child]; + auto &scalar_operand = func.GetChildrenMutable()[1 - pac_child]; if (needs_positive_check && !IsPositiveConstant(scalar_operand.get())) { break; } @@ -437,32 +400,32 @@ unique_ptr TryRewriteFilterComparison(OptimizerExtensionInput &input scalar_side = input.optimizer.BindScalarFunction(inverse_op, std::move(scalar_side), scalar_operand->Copy()); } - list_side = func.children[pac_child]->Copy(); + list_side = func.GetChildrenMutable()[pac_child]->Copy(); } // After simplification, list_side should be a bare column ref (possibly through casts) Expression *stripped = StripCasts(list_side.get()); - if (stripped->type != ExpressionType::BOUND_COLUMN_REF) { + if (stripped->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { // Can't emit priv_filter_, but return the simplified comparison // so the lambda path benefits from the algebraic rewriting - return make_uniq(cmp_type, std::move(scalar_side), std::move(list_side)); + return BoundComparisonExpression::Create(cmp_type, std::move(scalar_side), std::move(list_side)); } // Verify the column ref is a counter binding auto &col_ref = stripped->Cast(); - if (!counter_bindings.count(HashBinding(col_ref.binding))) { - return make_uniq(cmp_type, std::move(scalar_side), std::move(list_side)); + if (!counter_bindings.count(HashBinding(col_ref.BindingMutable()))) { + return BoundComparisonExpression::Create(cmp_type, std::move(scalar_side), std::move(list_side)); } // Build priv_{filter,select}_(scalar_side, counters) or priv_select_(hash, scalar, counters) string privacy_mode = GetPrivacyMode(input.context); const char *func_name = (wrap_kind == PacWrapKind::PAC_SELECT) ? GetPacSelectCmpName(cmp_type, privacy_mode) : GetPacFilterCmpName(cmp_type, privacy_mode); if (!func_name) { - return make_uniq(cmp_type, std::move(scalar_side), std::move(list_side)); + return BoundComparisonExpression::Create(cmp_type, std::move(scalar_side), std::move(list_side)); } - if (scalar_side->return_type != PacFloatLogicalType()) { + if (scalar_side->GetReturnType() != PacFloatLogicalType()) { scalar_side = BoundCastExpression::AddDefaultCastToType(std::move(scalar_side), PacFloatLogicalType()); } - auto counters_ref = - make_uniq("pac_var", LogicalType::LIST(PacFloatLogicalType()), col_ref.binding); + auto counters_ref = make_uniq("pac_var", LogicalType::LIST(PacFloatLogicalType()), + col_ref.BindingMutable()); if (wrap_kind == PacWrapKind::PAC_SELECT) { // priv_select_(hash, scalar, counters) — 3 args, need manual binding @@ -473,21 +436,22 @@ unique_ptr TryRewriteFilterComparison(OptimizerExtensionInput &input args.push_back(std::move(counters_ref)); vector arg_types; for (auto &a : args) { - arg_types.push_back(a->return_type); + arg_types.push_back(a->GetReturnType()); } auto &catalog = Catalog::GetSystemCatalog(input.context); - auto &entry = catalog.GetEntry(input.context, DEFAULT_SCHEMA, func_name); + auto &entry = catalog.GetEntry( + input.context, QualifiedName(Catalog::GetSystemCatalog(input.context).GetName(), + Identifier::DefaultSchema(), Identifier(func_name))); ErrorData error; FunctionBinder func_binder(input.context); - auto best = func_binder.BindFunction(func_name, entry.functions, arg_types, error); + auto best = func_binder.BindFunction(Identifier(func_name), entry.functions, arg_types, error); if (best.IsValid()) { auto func = entry.functions.GetFunctionByOffset(best.GetIndex()); - auto bind_info = func.bind ? func.bind(input.context, func, args) : nullptr; - return make_uniq(func.return_type, func, std::move(args), std::move(bind_info)); + return func_binder.BindScalarFunction(func, std::move(args)); } return nullptr; } - return input.optimizer.BindScalarFunction(func_name, std::move(scalar_side), std::move(counters_ref)); + return input.optimizer.BindScalarFunction(Identifier(func_name), std::move(scalar_side), std::move(counters_ref)); } } // namespace duckdb diff --git a/src/categorical/pac_categorical_rewriter.cpp b/src/categorical/pac_categorical_rewriter.cpp index d2e7cfa2..9c6d2ab4 100644 --- a/src/categorical/pac_categorical_rewriter.cpp +++ b/src/categorical/pac_categorical_rewriter.cpp @@ -5,6 +5,7 @@ // // Created by ila on 1/23/26. // +#include "privacy_debug.hpp" #include "categorical/pac_categorical_lambdas.hpp" #include "aggregates/as_aggregate.hpp" #include "utils/privacy_helpers.hpp" @@ -13,6 +14,8 @@ #include "duckdb/planner/expression/bound_conjunction_expression.hpp" #include "duckdb/planner/operator/logical_order.hpp" #include "duckdb/planner/operator/logical_top_n.hpp" +#include "duckdb/planner/operator/logical_materialized_cte.hpp" +#include "duckdb/planner/operator/logical_cteref.hpp" namespace duckdb { @@ -23,12 +26,12 @@ vector FindCounterBindings(Expression *expr, vector result; unordered_map seen; std::function Collect = [&](Expression *e) { - if (e->type == ExpressionType::BOUND_COLUMN_REF) { + if (e->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e->Cast(); - uint64_t h = HashBinding(col_ref.binding); + uint64_t h = HashBinding(col_ref.BindingMutable()); if (counter_bindings.count(h) && !seen.count(h)) { PacBindingInfo info; - info.binding = col_ref.binding; + info.binding = col_ref.BindingMutable(); info.index = result.size(); seen[h] = info.index; result.push_back(info); @@ -43,7 +46,7 @@ vector FindCounterBindings(Expression *expr, // Register a binding as a counter binding, optionally tracing to its origin. static void AddCounterBinding(unordered_map &counter_bindings, const ColumnBinding &binding, const ColumnBinding &origin = ColumnBinding()) { - counter_bindings[HashBinding(binding)] = (origin.table_index != DConstants::INVALID_INDEX) ? origin : binding; + counter_bindings[HashBinding(binding)] = (origin.table_index.index != DConstants::INVALID_INDEX) ? origin : binding; } // For a projection expression, trace back through counter_bindings to find the origin binding. @@ -51,9 +54,9 @@ static void AddCounterBinding(unordered_map &counter_bi // itself. static ColumnBinding TraceCounterOrigin(const unique_ptr &expr, const unordered_map &counter_bindings) { - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { + if (expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = expr->Cast(); - auto it = counter_bindings.find(HashBinding(col_ref.binding)); + auto it = counter_bindings.find(HashBinding(col_ref.BindingMutable())); if (it != counter_bindings.end()) { return it->second; } @@ -63,8 +66,8 @@ static ColumnBinding TraceCounterOrigin(const unique_ptr &expr, // Check if an expression references any counter binding. bool ReferencesCounterBinding(Expression *expr, const unordered_map &counter_bindings) { - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { - return counter_bindings.count(HashBinding(expr->Cast().binding)) > 0; + if (expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { + return counter_bindings.count(HashBinding(expr->Cast().BindingMutable())) > 0; } bool found = false; ExpressionIterator::EnumerateChildren(*expr, [&](Expression &child) { @@ -82,30 +85,33 @@ bool ReferencesCounterBinding(Expression *expr, const unordered_map, and the aggregate rebound to the _list variant. static bool TryRebindToListVariant(LogicalAggregate &agg, idx_t i, ClientContext &context) { auto &agg_expr = agg.expressions[i]; - if (agg_expr->type != ExpressionType::BOUND_AGGREGATE) { + if (agg_expr->GetExpressionType() != ExpressionType::BOUND_AGGREGATE) { return false; } auto &bound_agg = agg_expr->Cast(); - if (bound_agg.children.empty()) { + if (bound_agg.GetChildrenMutable().empty()) { return false; } // PAC/PAC_counters aggregates have (hash, value); standard aggregates have (value). // For _list variants we only need the value child. - bool is_pac_style = (IsPacAggregate(bound_agg.function.name) || IsPacCountersAggregate(bound_agg.function.name)) && - bound_agg.children.size() > 1; + bool is_pac_style = (IsPacAggregate(bound_agg.FunctionMutable().GetName().GetIdentifierName()) || + IsPacCountersAggregate(bound_agg.FunctionMutable().GetName().GetIdentifierName())) && + bound_agg.GetChildrenMutable().size() > 1; idx_t value_child_idx = is_pac_style ? 1 : 0; - string base_name = GetBasePacAggregateName(bound_agg.function.name); // strips _counters if present + string base_name = GetBasePacAggregateName( + bound_agg.FunctionMutable().GetName().GetIdentifierName()); // strips _counters if present string list_name = GetListAggregateVariant(base_name, GetPrivacyMode(context)); if (list_name.empty()) { // Generic aggregate (e.g., first()) over counter input — rebind with same name - if (bound_agg.children[value_child_idx]->return_type == LogicalType::LIST(PacFloatLogicalType())) { - list_name = bound_agg.function.name; + if (bound_agg.GetChildrenMutable()[value_child_idx]->GetReturnType() == + LogicalType::LIST(PacFloatLogicalType())) { + list_name = bound_agg.FunctionMutable().GetName().GetIdentifierName(); } else { return false; } } vector> children; - auto child_copy = bound_agg.children[value_child_idx]->Copy(); + auto child_copy = bound_agg.GetChildrenMutable()[value_child_idx]->Copy(); children.push_back(std::move(child_copy)); auto new_aggr = RebindAggregate(context, list_name, std::move(children), bound_agg.IsDistinct()); if (!new_aggr) { @@ -130,19 +136,23 @@ static void ReplaceAggregatesOverCounters(LogicalOperator *op, ClientContext &co auto &agg = op->Cast(); for (idx_t i = 0; i < agg.expressions.size(); i++) { auto &agg_expr = agg.expressions[i]; - if (agg_expr->type != ExpressionType::BOUND_AGGREGATE) { + if (agg_expr->GetExpressionType() != ExpressionType::BOUND_AGGREGATE) { continue; } auto &bound_agg = agg_expr->Cast(); - if (bound_agg.children.empty()) { + if (bound_agg.GetChildrenMutable().empty()) { continue; } - idx_t value_child_idx = IsPacAggregate(bound_agg.function.name) && bound_agg.children.size() > 1 ? 1 : 0; - if (!ReferencesCounterBinding(bound_agg.children[value_child_idx].get(), counter_bindings)) { + idx_t value_child_idx = IsPacAggregate(bound_agg.FunctionMutable().GetName().GetIdentifierName()) && + bound_agg.GetChildrenMutable().size() > 1 + ? 1 + : 0; + if (!ReferencesCounterBinding(bound_agg.GetChildrenMutable()[value_child_idx].get(), counter_bindings)) { continue; } if (TryRebindToListVariant(agg, i, context)) { - AddCounterBinding(counter_bindings, ColumnBinding(agg.aggregate_index, i)); + AddCounterBinding(counter_bindings, + ColumnBinding(TableIndex(agg.aggregate_index.index), ProjectionIndex(i))); } } } @@ -161,19 +171,19 @@ static void PrepareFilterPatternSlot(LogicalProjection &proj, idx_t i, // col_ref types are already fixed by the universal type fix above. // Update projection output type if (i < proj.types.size()) { - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { + if (expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { proj.types[i] = list_type; - } else if (expr->type == ExpressionType::BOUND_FUNCTION) { - proj.types[i] = expr->return_type; - } else if (expr->type == ExpressionType::CASE_EXPR && + } else if (expr->GetExpressionType() == ExpressionType::BOUND_FUNCTION) { + proj.types[i] = expr->GetReturnType(); + } else if (expr->GetExpressionType() == ExpressionType::CASE_EXPR && IsScalarSubqueryWrapper(expr->Cast())) { // first(value) now returns a counter list, but DuckDB's scalar-subquery // cardinality-check CASE still has its original scalar result type. auto &wrapper = expr->Cast(); - if (wrapper.else_expr->return_type == list_type) { - wrapper.case_checks[0].then_expr = - BoundCastExpression::AddDefaultCastToType(std::move(wrapper.case_checks[0].then_expr), list_type); - wrapper.return_type = list_type; + if (wrapper.Else().GetReturnType() == list_type) { + wrapper.CaseChecksMutable()[0].then_expr = BoundCastExpression::AddDefaultCastToType( + std::move(wrapper.CaseChecksMutable()[0].then_expr), list_type); + wrapper.SetReturnType(list_type); proj.types[i] = list_type; PRIVACY_DEBUG_PRINT("Categorical rewrite: preserving counter lists through scalar-subquery wrapper"); } @@ -205,8 +215,8 @@ RewriteExpressionWithCounters(OptimizerExtensionInput &input, const vectorGetExpressionClass() == ExpressionClass::BOUND_FUNCTION) { auto &func = cmp_result->Cast(); - if (func.function.name.rfind("priv_filter_", 0) == 0 || - func.function.name.rfind("priv_select_", 0) == 0) { + if (func.FunctionMutable().GetName().GetIdentifierName().rfind("priv_filter_", 0) == 0 || + func.FunctionMutable().GetName().GetIdentifierName().rfind("priv_select_", 0) == 0) { return cmp_result; // Full optimization to priv_{filter,select}_ } } @@ -236,83 +246,6 @@ RewriteExpressionWithCounters(OptimizerExtensionInput &input, const vector(); - if (col_ref.binding == old_binding) { - col_ref.binding = new_binding; - } - } else { - ExpressionIterator::EnumerateChildren( - expr, [&](Expression &child) { ReplaceBindingInExpression(child, old_binding, new_binding); }); - } -} - -// Expand projection references in a filter/join expression. -// For each col_ref to a projection slot with counter arithmetic: -// - create counter pass-through slots in the projection -// - substitute the col_ref with the arithmetic rebased to pass-through bindings -// This lets the filter fold arithmetic + comparison into one list_transform lambda. -static void ExpandProjectionRefsInExpression(unique_ptr &expr, LogicalOperator *plan_root, - unordered_map &counter_bindings) { - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { - auto &col_ref = expr->Cast(); - if (!counter_bindings.count(HashBinding(col_ref.binding))) { - return; // Not a counter binding — nothing to expand - } - auto *source_op = FindOperatorByTableIndex(plan_root, col_ref.binding.table_index); - if (!source_op || source_op->type != LogicalOperatorType::LOGICAL_PROJECTION) { - return; - } - auto &proj = source_op->Cast(); - idx_t slot = col_ref.binding.column_index; - if (slot >= proj.expressions.size()) { - return; - } - // Only expand slots with multi-binding counter arithmetic that needs decomposition. - // Skip: simple col_refs (already pass-throughs) and already-wrapped terminals. - auto &proj_expr = proj.expressions[slot]; - if (proj_expr->type == ExpressionType::BOUND_COLUMN_REF) { - return; - } - if (IsAlreadyWrappedInPacNoised(proj_expr.get())) { - return; - } - auto pac_bindings = FindCounterBindings(proj_expr.get(), counter_bindings); - if (pac_bindings.empty()) { - return; - } - auto list_type = LogicalType::LIST(PacFloatLogicalType()); - // Clone the projection expression for rebasing - auto expanded = proj_expr->Copy(); - // Create pass-through for first binding (reuse this slot) - proj.expressions[slot] = make_uniq("pac_var", list_type, pac_bindings[0].binding); - if (slot < proj.types.size()) { - proj.types[slot] = list_type; - } - ReplaceBindingInExpression(*expanded, pac_bindings[0].binding, ColumnBinding(proj.table_index, slot)); - // Extra pass-throughs for additional bindings - for (idx_t j = 1; j < pac_bindings.size(); j++) { - idx_t extra_slot = proj.expressions.size(); - proj.expressions.push_back( - make_uniq("pac_var", list_type, pac_bindings[j].binding)); - proj.types.push_back(list_type); - auto extra_cb = ColumnBinding(proj.table_index, extra_slot); - auto src_it = counter_bindings.find(HashBinding(pac_bindings[j].binding)); - AddCounterBinding(counter_bindings, extra_cb, - (src_it != counter_bindings.end()) ? src_it->second : pac_bindings[j].binding); - ReplaceBindingInExpression(*expanded, pac_bindings[j].binding, extra_cb); - } - expr = std::move(expanded); - return; - } - ExpressionIterator::EnumerateChildren(*expr, [&](unique_ptr &child) { - ExpandProjectionRefsInExpression(child, plan_root, counter_bindings); - }); -} - // Wrap counter column references with as_noised. // After converting as_noised_sum → as_sum (LIST output), // expressions may still reference the aggregate with the original type (e.g. DECIMAL). @@ -320,16 +253,16 @@ static void ExpandProjectionRefsInExpression(unique_ptr &expr, Logic static void WrapCounterRefsWithNoised(unique_ptr &expr, const unordered_map &counter_bindings, OptimizerExtensionInput &input) { - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { + if (expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = expr->Cast(); - if (counter_bindings.count(HashBinding(col_ref.binding))) { - auto original_type = col_ref.return_type; + if (counter_bindings.count(HashBinding(col_ref.BindingMutable()))) { + auto original_type = col_ref.GetReturnType(); auto list_type = LogicalType::LIST(PacFloatLogicalType()); // If the universal type fix already changed the col_ref to LIST, // we can't cast back to it — just use as_noised's natural scalar return type. bool needs_cast = (original_type != PacFloatLogicalType() && original_type != list_type); // Ensure col_ref type is LIST for correct as_noised binding. - col_ref.return_type = list_type; + col_ref.SetReturnType(list_type); unique_ptr noised = input.optimizer.BindScalarFunction("as_noised", expr->Copy()); if (needs_cast) { noised = BoundCastExpression::AddDefaultCastToType(std::move(noised), original_type); @@ -359,8 +292,8 @@ static bool RewriteProjectionExpression(OptimizerExtensionInput &input, LogicalP // Simple column reference to a counter list. // After the universal col_ref type fix, the col_ref type is LIST. // proj.types[i] still holds the original scalar type (e.g. HUGEINT). - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { - auto original_type = (i < proj.types.size()) ? proj.types[i] : expr->return_type; + if (expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { + auto original_type = (i < proj.types.size()) ? proj.types[i] : expr->GetReturnType(); if (is_terminal) { unique_ptr result = input.optimizer.BindScalarFunction("as_noised", expr->Copy()); if (original_type != PacFloatLogicalType()) { @@ -374,21 +307,33 @@ static bool RewriteProjectionExpression(OptimizerExtensionInput &input, LogicalP } return true; // non-terminal pass-through — still a counter binding } - if (!IsNumericalType(expr->return_type)) { + if (!IsNumericalType(expr->GetReturnType())) { WrapCounterRefsWithNoised(expr, counter_bindings, input); return false; } // Arithmetic over PAC aggregates: build list_transform + as_noised terminal. Expression *expr_to_clone = expr.get(); - if (expr->type == ExpressionType::OPERATOR_CAST && - expr->Cast().return_type != PacFloatLogicalType()) { - expr_to_clone = expr->Cast().child.get(); + if (expr->GetExpressionType() == ExpressionType::OPERATOR_CAST && + expr->Cast().GetReturnType() != PacFloatLogicalType()) { + expr_to_clone = expr->Cast().GetChildrenMutable()[0].get(); + } + if (!is_terminal) { + auto result = + BuildCategoricalLambdas(input, pac_bindings, expr_to_clone, counter_bindings, PacFloatLogicalType()); + if (result) { + PRIVACY_DEBUG_PRINT("Categorical rewrite: preserving per-world projection #" + + std::to_string(proj.table_index.index) + "." + std::to_string(i)); + proj.expressions[i] = std::move(result); + proj.types[i] = proj.expressions[i]->GetReturnType(); + return true; + } + return false; } auto result = RewriteExpressionWithCounters(input, pac_bindings, expr_to_clone, counter_bindings, - PacWrapKind::PAC_NOISED, expr->return_type); + PacWrapKind::PAC_NOISED, expr->GetReturnType()); if (result) { proj.expressions[i] = std::move(result); - proj.types[i] = expr->return_type; + proj.types[i] = expr->GetReturnType(); } return false; } @@ -404,7 +349,7 @@ static ColumnBinding InsertPacSelectProjection(OptimizerExtensionInput &input, u auto old_bindings = filter_child_slot->GetColumnBindings(); idx_t num_cols = old_bindings.size(); - idx_t proj_table_index = input.optimizer.binder.GenerateTableIndex(); + idx_t proj_table_index = input.optimizer.binder.GenerateTableIndex().index; vector> proj_expressions; // Find which position the hash binding occupies @@ -430,7 +375,7 @@ static ColumnBinding InsertPacSelectProjection(OptimizerExtensionInput &input, u proj_expressions.push_back(std::move(pac_select_expr)); } - auto projection = make_uniq(proj_table_index, std::move(proj_expressions)); + auto projection = make_uniq(TableIndex(proj_table_index), std::move(proj_expressions)); projection->children.push_back(std::move(filter_child_slot)); projection->ResolveOperatorTypes(); @@ -440,12 +385,37 @@ static ColumnBinding InsertPacSelectProjection(OptimizerExtensionInput &input, u // Remap old bindings → new projection bindings ColumnBindingReplacer replacer; for (idx_t i = 0; i < num_cols; i++) { - replacer.replacement_bindings.emplace_back(old_bindings[i], ColumnBinding(proj_table_index, i)); + replacer.replacement_bindings.emplace_back(old_bindings[i], + ColumnBinding(TableIndex(proj_table_index), ProjectionIndex(i))); } replacer.stop_operator = proj_ptr; replacer.VisitOperator(*plan); - return ColumnBinding(proj_table_index, hash_position); + return ColumnBinding(TableIndex(proj_table_index), ProjectionIndex(hash_position)); +} + +// CTE readers have independent bindings and cache their input types. Preserve +// counter lineage across the materialization boundary before rewriting consumers. +static void SyncCTECounterBindings(LogicalOperator &op, LogicalMaterializedCTE &cte, + unordered_map &counter_bindings) { + if (op.type == LogicalOperatorType::LOGICAL_CTE_REF) { + auto &ref = op.Cast(); + if (ref.cte_index == cte.table_index) { + auto bindings = cte.children[0]->GetColumnBindings(); + for (idx_t i = 0; i < bindings.size() && i < ref.chunk_types.size(); i++) { + auto origin = counter_bindings.find(HashBinding(bindings[i])); + if (origin != counter_bindings.end()) { + ref.chunk_types[i] = cte.children[0]->types[i]; + AddCounterBinding(counter_bindings, ColumnBinding(ref.table_index, ProjectionIndex(i)), + origin->second); + } + } + ref.ResolveOperatorTypes(); + } + } + for (auto &child : op.children) { + SyncCTECounterBindings(*child, cte, counter_bindings); + } } // Single bottom-up rewrite pass. @@ -464,8 +434,15 @@ static void RewriteBottomUp(unique_ptr &op_ptr, OptimizerExtens for (idx_t ci = 0; ci < op->children.size(); ci++) { // Recurse into children first (bottom-up) // Child 0 of MATERIALIZED_CTE is the CTE definition — its output types must remain // stable (numeric, not LIST) because CTE_SCAN consumers may re-aggregate the results. - bool child_in_cte = - inside_cte_definition || (op->type == LogicalOperatorType::LOGICAL_MATERIALIZED_CTE && ci == 0); + bool is_cte = op->type == LogicalOperatorType::LOGICAL_MATERIALIZED_CTE; + bool generated_correlation_cte = + is_cte && + StringUtil::StartsWith(op->Cast().ctename.GetIdentifierName(), "__duckdb_delim_"); + bool child_in_cte = inside_cte_definition || (is_cte && ci == 0 && !generated_correlation_cte); + if (is_cte && ci == 1) { + op->children[0]->ResolveOperatorTypes(); + SyncCTECounterBindings(*op->children[1], op->Cast(), counter_bindings); + } RewriteBottomUp(op->children[ci], input, plan, pattern_lookup, patterns, replaced_mark_bindings, counter_bindings, child_in_cte); } @@ -477,9 +454,9 @@ static void RewriteBottomUp(unique_ptr &op_ptr, OptimizerExtens auto &filter = op->Cast(); bool has_dangling_mark = false; for (auto &fexpr : filter.expressions) { - if (fexpr->type == ExpressionType::BOUND_COLUMN_REF) { + if (fexpr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &ref = fexpr->Cast(); - if (replaced_mark_bindings.count(HashBinding(ref.binding))) { + if (replaced_mark_bindings.count(HashBinding(ref.BindingMutable()))) { has_dangling_mark = true; break; } @@ -517,10 +494,10 @@ static void RewriteBottomUp(unique_ptr &op_ptr, OptimizerExtens auto list_type = LogicalType::LIST(PacFloatLogicalType()); LogicalOperatorVisitor::EnumerateExpressions(*op, [&](unique_ptr *expr_ptr) { std::function &)> Fix = [&](unique_ptr &e) { - if (e->type == ExpressionType::BOUND_COLUMN_REF) { + if (e->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e->Cast(); - if (counter_bindings.count(HashBinding(col_ref.binding))) { - col_ref.return_type = list_type; + if (counter_bindings.count(HashBinding(col_ref.BindingMutable()))) { + col_ref.SetReturnType(list_type); } } ExpressionIterator::EnumerateChildren(*e, Fix); @@ -541,15 +518,15 @@ static void RewriteBottomUp(unique_ptr &op_ptr, OptimizerExtens // expressions that are sources of categorical patterns. for (idx_t i = 0; i < agg.expressions.size(); i++) { auto &agg_expr = agg.expressions[i]; - if (agg_expr->type != ExpressionType::BOUND_AGGREGATE) { + if (agg_expr->GetExpressionType() != ExpressionType::BOUND_AGGREGATE) { continue; } auto &bound_agg = agg_expr->Cast(); - if (!IsPacAggregate(bound_agg.function.name)) { + if (!IsPacAggregate(bound_agg.FunctionMutable().GetName().GetIdentifierName())) { continue; } // Check if this specific aggregate expression is a source of any pattern - ColumnBinding this_binding(agg.aggregate_index, i); + ColumnBinding this_binding {TableIndex(agg.aggregate_index.index), ProjectionIndex(i)}; bool convert = false; for (auto &p : patterns) { for (auto &bi : p.pac_bindings) { @@ -565,20 +542,21 @@ static void RewriteBottomUp(unique_ptr &op_ptr, OptimizerExtens if (!convert) { continue; } - string counters_name = GetCountersVariant(bound_agg.function.name); + string counters_name = GetCountersVariant(bound_agg.FunctionMutable().GetName().GetIdentifierName()); vector> children; - for (auto &child_expr : bound_agg.children) { + for (auto &child_expr : bound_agg.GetChildrenMutable()) { children.push_back(child_expr->Copy()); } auto new_aggr = RebindAggregate(input.context, counters_name, std::move(children), bound_agg.IsDistinct()); if (!new_aggr) { // Fallback: just rename the function in place - bound_agg.function.name = counters_name; - bound_agg.function.return_type = LogicalType::LIST(PacFloatLogicalType()); - agg_expr->return_type = LogicalType::LIST(PacFloatLogicalType()); + bound_agg.FunctionMutable().SetName(Identifier(counters_name)); + bound_agg.FunctionMutable().SetReturnType(LogicalType::LIST(PacFloatLogicalType())); + agg_expr->SetReturnType(LogicalType::LIST(PacFloatLogicalType())); } else { agg.expressions[i] = std::move(new_aggr); } - AddCounterBinding(counter_bindings, ColumnBinding(agg.aggregate_index, i)); + AddCounterBinding(counter_bindings, + ColumnBinding(TableIndex(agg.aggregate_index.index), ProjectionIndex(i))); idx_t types_index = agg.groups.size() + i; if (types_index < agg.types.size()) { agg.types[types_index] = LogicalType::LIST(PacFloatLogicalType()); @@ -587,55 +565,24 @@ static void RewriteBottomUp(unique_ptr &op_ptr, OptimizerExtens } else if (op->type == LogicalOperatorType::LOGICAL_PROJECTION && !inside_cte_definition) { // === PROJECTION: rewrite PAC expressions === auto &proj = op->Cast(); - // Check if this projection's output is consumed by a filter/join pattern. - // If so, skip terminal wrapping — the parent filter/join handles the rewrite. - bool is_filter_pattern = false; - for (auto &p : patterns) { - if (p.parent_op && pattern_lookup.count(p.parent_op)) { - for (auto &bi : p.pac_bindings) { - if (bi.binding.table_index == proj.table_index) { - is_filter_pattern = true; - break; - } - } - } - if (is_filter_pattern) { + // Intermediate projections compute valid per-world values; only the + // user-visible output projection finalizes them into scalar answers. + auto *output = plan_root; + while (output && !output->children.empty()) { + if (output->type == LogicalOperatorType::LOGICAL_MATERIALIZED_CTE) { + output = output->children[1].get(); + } else if (output->type == LogicalOperatorType::LOGICAL_ORDER_BY || + output->type == LogicalOperatorType::LOGICAL_TOP_N || + output->type == LogicalOperatorType::LOGICAL_LIMIT) { + output = output->children[0].get(); + } else { break; } } - if (is_filter_pattern) { - // Filter-pattern projection: minimal processing only (type updates, no terminal wrapping). - // The parent filter/join will expand these expressions and generate list_transform + priv_filter. - for (idx_t i = 0; i < proj.expressions.size(); i++) { - PrepareFilterPatternSlot(proj, i, counter_bindings); - // All filter-pattern slots with counter content remain counter bindings - auto slot_bindings = FindCounterBindings(proj.expressions[i].get(), counter_bindings); - if (!slot_bindings.empty()) { - auto proj_cb = ColumnBinding(proj.table_index, i); - AddCounterBinding(counter_bindings, proj_cb, - TraceCounterOrigin(proj.expressions[i], counter_bindings)); - } - } - } else { - // Normal projection: full PAC rewriting with as_noised terminal. - bool is_terminal = (op == plan_root); - if (!is_terminal) { - auto *root = plan_root; - while (root && - (root->type == LogicalOperatorType::LOGICAL_ORDER_BY || - root->type == LogicalOperatorType::LOGICAL_TOP_N || - root->type == LogicalOperatorType::LOGICAL_LIMIT) && - !root->children.empty()) { - root = root->children[0].get(); - } - is_terminal = (root == op); - } - for (idx_t i = 0; i < proj.expressions.size(); i++) { - if (RewriteProjectionExpression(input, proj, i, plan_root, is_terminal, counter_bindings)) { - auto proj_cb = ColumnBinding(proj.table_index, i); - AddCounterBinding(counter_bindings, proj_cb, - TraceCounterOrigin(proj.expressions[i], counter_bindings)); - } + for (idx_t i = 0; i < proj.expressions.size(); i++) { + if (RewriteProjectionExpression(input, proj, i, plan_root, output == op, counter_bindings)) { + AddCounterBinding(counter_bindings, ColumnBinding(proj.table_index, ProjectionIndex(i)), + TraceCounterOrigin(proj.expressions[i], counter_bindings)); } } } else if (op->type == LogicalOperatorType::LOGICAL_FILTER && @@ -648,9 +595,7 @@ static void RewriteBottomUp(unique_ptr &op_ptr, OptimizerExtens continue; } auto &filter_expr = filter.expressions[expr_idx]; - // Expand projection arithmetic into the filter expression so arithmetic + comparison - // are folded into one list_transform lambda. Handles both single and multi-binding. - ExpandProjectionRefsInExpression(filter_expr, plan_root, counter_bindings); + // Projection arithmetic has already been lifted to per-world lists. auto pac_bindings = FindCounterBindings(filter_expr.get(), counter_bindings); if (pac_bindings.empty()) { continue; @@ -677,7 +622,7 @@ static void RewriteBottomUp(unique_ptr &op_ptr, OptimizerExtens auto new_hash = InsertPacSelectProjection(input, plan, op->children[0], pattern_hash, std::move(result)); // Replace filter condition with priv_hash <> 0 for majority-vote decision - filter.expressions[expr_idx] = make_uniq( + filter.expressions[expr_idx] = BoundComparisonExpression::Create( ExpressionType::COMPARE_NOTEQUAL, make_uniq("pac_pu", LogicalType::UBIGINT, new_hash), make_uniq(Value::UBIGINT(0))); @@ -723,75 +668,57 @@ static void RewriteBottomUp(unique_ptr &op_ptr, OptimizerExtens } } auto &join = op->Cast(); - for (auto expr_idx : it->second) { - if (expr_idx >= join.conditions.size()) { - continue; - } - auto &cond = join.conditions[expr_idx]; - // Expand projection arithmetic into join conditions - ExpandProjectionRefsInExpression(cond.left, plan_root, counter_bindings); - ExpandProjectionRefsInExpression(cond.right, plan_root, counter_bindings); - auto left_bindings = FindCounterBindings(cond.left.get(), counter_bindings); - auto right_bindings = FindCounterBindings(cond.right.get(), counter_bindings); - bool left_is_list = !left_bindings.empty(); - bool right_is_list = !right_bindings.empty(); - if (!left_is_list && !right_is_list) { - continue; - } - // Collect PAC bindings (deduplicated) - vector all_bindings; - unordered_set seen_hashes; - for (auto &b : left_bindings) { - uint64_t h = HashBinding(b.binding); - if (seen_hashes.insert(h).second) { - all_bindings.push_back(b); - } - } - for (auto &b : right_bindings) { - uint64_t h = HashBinding(b.binding); - if (seen_hashes.insert(h).second) { - all_bindings.push_back(b); - } - } - for (idx_t j = 0; j < all_bindings.size(); j++) { - all_bindings[j].index = j; + unique_ptr comparison; + vector> ordinary_conditions; + for (auto &condition : join.conditions) { + auto expression = BoundComparisonExpression::Create( + condition.GetComparisonType(), condition.LeftReference()->Copy(), condition.RightReference()->Copy()); + if (!ReferencesCounterBinding(expression.get(), counter_bindings)) { + ordinary_conditions.push_back(std::move(expression)); + } else if (comparison) { + comparison = make_uniq(ExpressionType::CONJUNCTION_AND, + std::move(comparison), std::move(expression)); + } else { + comparison = std::move(expression); } - auto comparison = - make_uniq(cond.comparison, cond.left->Copy(), cond.right->Copy()); - auto pac_expr = RewriteExpressionWithCounters(input, all_bindings, comparison.get(), counter_bindings, - join_wrap_kind, PacFloatLogicalType(), join_pattern_hash); - if (pac_expr) { - auto cross_product = - LogicalCrossProduct::Create(std::move(join.children[0]), std::move(join.children[1])); - auto filter_op = make_uniq(); - - if (join_wrap_kind == PacWrapKind::PAC_SELECT) { - // Insert projection below filter replacing hash with priv_select result - auto new_hash = - InsertPacSelectProjection(input, plan, cross_product, join_pattern_hash, std::move(pac_expr)); - filter_op->expressions.push_back(make_uniq( - ExpressionType::COMPARE_NOTEQUAL, - make_uniq("pac_pu", LogicalType::UBIGINT, new_hash), - make_uniq(Value::UBIGINT(0)))); - } else { - filter_op->expressions.push_back(std::move(pac_expr)); - } - filter_op->children.push_back(std::move(cross_product)); - - for (auto &p : patterns) { - if (p.parent_op == op) { - p.parent_op = nullptr; - } - } - // Record the mark column binding so the parent FILTER - // (which references it) can be cleaned up. - if (join.join_type == JoinType::MARK) { - replaced_mark_bindings.insert(HashBinding(ColumnBinding(join.mark_index, 0))); - } - op_ptr = std::move(filter_op); - break; + } + if (!comparison) { + return; + } + // Evaluate all world-dependent conditions together, then make one + // categorical decision. Keep ordinary join predicates outside the lambda + // so DuckDB can still recover its usual hash joins during optimization. + auto bindings = FindCounterBindings(comparison.get(), counter_bindings); + auto pac_expr = RewriteExpressionWithCounters(input, bindings, comparison.get(), counter_bindings, + join_wrap_kind, PacFloatLogicalType(), join_pattern_hash); + if (!pac_expr) { + return; + } + auto cross_product = LogicalCrossProduct::Create(std::move(join.children[0]), std::move(join.children[1])); + auto filter = make_uniq(); + if (join_wrap_kind == PacWrapKind::PAC_SELECT) { + auto new_hash = + InsertPacSelectProjection(input, plan, cross_product, join_pattern_hash, std::move(pac_expr)); + filter->expressions.push_back(BoundComparisonExpression::Create( + ExpressionType::COMPARE_NOTEQUAL, + make_uniq("pac_pu", LogicalType::UBIGINT, new_hash), + make_uniq(Value::UBIGINT(0)))); + } else { + filter->expressions.push_back(std::move(pac_expr)); + } + for (auto &condition : ordinary_conditions) { + filter->expressions.push_back(std::move(condition)); + } + filter->children.push_back(std::move(cross_product)); + for (auto &pattern : patterns) { + if (pattern.parent_op == op) { + pattern.parent_op = nullptr; } } + if (join.join_type == JoinType::MARK) { + replaced_mark_bindings.insert(HashBinding(ColumnBinding(join.mark_index, ProjectionIndex(0)))); + } + op_ptr = std::move(filter); } } @@ -808,26 +735,29 @@ static void ConvertPacAggregatesToListVariants(LogicalOperator &plan, ClientCont } auto &agg = plan.Cast(); for (idx_t i = 0; i < agg.expressions.size(); i++) { - if (agg.expressions[i]->type != ExpressionType::BOUND_AGGREGATE) { + if (agg.expressions[i]->GetExpressionType() != ExpressionType::BOUND_AGGREGATE) { continue; } - if (counter_bindings.count(HashBinding(ColumnBinding(agg.aggregate_index, i)))) { + if (counter_bindings.count( + HashBinding(ColumnBinding(TableIndex(agg.aggregate_index.index), ProjectionIndex(i))))) { continue; } auto &bound_agg = agg.expressions[i]->Cast(); - if (bound_agg.children.empty() || IsPacListAggregate(bound_agg.function.name)) { + if (bound_agg.GetChildrenMutable().empty() || + IsPacListAggregate(bound_agg.FunctionMutable().GetName().GetIdentifierName())) { continue; } - idx_t value_idx = - (IsPacAggregate(bound_agg.function.name) || IsPacCountersAggregate(bound_agg.function.name)) && - bound_agg.children.size() > 1 - ? 1 - : 0; - if (!ReferencesCounterBinding(bound_agg.children[value_idx].get(), counter_bindings)) { + idx_t value_idx = (IsPacAggregate(bound_agg.FunctionMutable().GetName().GetIdentifierName()) || + IsPacCountersAggregate(bound_agg.FunctionMutable().GetName().GetIdentifierName())) && + bound_agg.GetChildrenMutable().size() > 1 + ? 1 + : 0; + if (!ReferencesCounterBinding(bound_agg.GetChildrenMutable()[value_idx].get(), counter_bindings)) { continue; } if (TryRebindToListVariant(agg, i, context)) { - AddCounterBinding(counter_bindings, ColumnBinding(agg.aggregate_index, i)); + AddCounterBinding(counter_bindings, + ColumnBinding(TableIndex(agg.aggregate_index.index), ProjectionIndex(i))); } } } diff --git a/src/compiler/privacy_bitslice_add_fkjoins.cpp b/src/compiler/privacy_bitslice_add_fkjoins.cpp index 4ee68352..6b2dfe35 100644 --- a/src/compiler/privacy_bitslice_add_fkjoins.cpp +++ b/src/compiler/privacy_bitslice_add_fkjoins.cpp @@ -179,7 +179,7 @@ CreateGetsAndChainJoins(const PrivacyCompatibilityResult &check, OptimizerExtens throw InvalidInputException("Privacy compiler: missing table metadata for table: " + table); } auto required_cols = GetRequiredColumnsForTable(check, table, fk_path, privacy_units); - idx_t local_idx = binder.GenerateTableIndex(); + idx_t local_idx = binder.GenerateTableIndex().index; auto get = CreateLogicalGet(input.context, plan, table, local_idx, required_cols); local_get_map[table] = std::move(get); } @@ -209,15 +209,15 @@ static bool FindAccessibleAlternativeAndJoin(const PrivacyCompatibilityResult &c for (auto *table_node : table_nodes) { auto &table_get = table_node->get()->Cast(); - if (table_get.table_index == skip_table_index) { + if (table_get.table_index.index == skip_table_index) { continue; } - if (!AreTableColumnsAccessible(plan.get(), table_get.table_index)) { + if (!AreTableColumnsAccessible(plan.get(), table_get.table_index.index)) { continue; } #if PRIVACY_DEBUG PRIVACY_DEBUG_PRINT("FindAccessibleAlternativeAndJoin: Found accessible table " + candidate + " #" + - std::to_string(table_get.table_index) + " - joining " + + std::to_string(table_get.table_index.index) + " - joining " + std::to_string(tables_to_join.size()) + " tables"); #endif unique_ptr current_node = (*table_node)->Copy(input.context); @@ -272,7 +272,7 @@ void AddMissingFKJoins(const PrivacyCompatibilityResult &check, OptimizerExtensi throw InvalidInputException("Privacy compiler: missing table metadata for missing GET: " + table); } auto required_cols = GetRequiredColumnsForTable(check, table, fk_path, privacy_units); - idx_t idx = binder.GenerateTableIndex(); + idx_t idx = binder.GenerateTableIndex().index; auto get = CreateLogicalGet(input.context, plan, table, idx, required_cols); get_map[table] = std::move(get); ordered_table_names.push_back(table); @@ -351,7 +351,7 @@ void AddMissingFKJoins(const PrivacyCompatibilityResult &check, OptimizerExtensi // Iterate over each instance of the connecting table and add joins for (auto *target_ref : all_connecting_nodes) { auto &target_op = (*target_ref)->Cast(); - idx_t connecting_table_idx = target_op.table_index; + idx_t connecting_table_idx = target_op.table_index.index; bool is_in_subquery = IsInDelimJoinSubqueryBranch(&plan, target_ref->get()); vector tables_to_join_for_instance = ordered_table_names; @@ -428,7 +428,7 @@ void AddMissingFKJoins(const PrivacyCompatibilityResult &check, OptimizerExtensi } for (auto *fk_node : fk_table_nodes) { auto &fk_table_get = fk_node->get()->Cast(); - idx_t fk_table_idx = fk_table_get.table_index; + idx_t fk_table_idx = fk_table_get.table_index.index; if (fk_table_idx == connecting_table_idx) { if (AreTableColumnsAccessible(plan.get(), fk_table_idx)) { @@ -468,7 +468,7 @@ void AddMissingFKJoins(const PrivacyCompatibilityResult &check, OptimizerExtensi string connecting_table_name; auto table_ptr = target_op.GetTable(); if (table_ptr) { - connecting_table_name = table_ptr->name; + connecting_table_name = table_ptr->name.GetIdentifierName(); } // Case 1: connecting table IS the FK table but blocked (SEMI/ANTI join) @@ -579,10 +579,10 @@ void AddMissingFKJoins(const PrivacyCompatibilityResult &check, OptimizerExtensi for (auto *table_node : table_nodes) { auto &table_get = table_node->get()->Cast(); - if (table_get.table_index == connecting_table_idx) { + if (table_get.table_index.index == connecting_table_idx) { continue; } - if (!AreTableColumnsAccessible(plan.get(), table_get.table_index)) { + if (!AreTableColumnsAccessible(plan.get(), table_get.table_index.index)) { continue; } @@ -609,7 +609,7 @@ void AddMissingFKJoins(const PrivacyCompatibilityResult &check, OptimizerExtensi } #if PRIVACY_DEBUG PRIVACY_DEBUG_PRINT("AddMissingFKJoins: Found accessible table " + candidate + " #" + - std::to_string(table_get.table_index) + " - adding full join chain:"); + std::to_string(table_get.table_index.index) + " - adding full join chain:"); for (auto &t : full_tables_to_add) { PRIVACY_DEBUG_PRINT(" - " + t); } diff --git a/src/compiler/privacy_compiler.cpp b/src/compiler/privacy_compiler.cpp index 7b827eb0..8796b196 100644 --- a/src/compiler/privacy_compiler.cpp +++ b/src/compiler/privacy_compiler.cpp @@ -91,7 +91,7 @@ static bool IsDirectlyReachable(LogicalOperator *from, idx_t table_index, bool i } if (from->type == LogicalOperatorType::LOGICAL_GET) { auto &get = from->Cast(); - if (get.table_index == table_index) { + if (get.table_index.index == table_index) { return true; } } @@ -115,9 +115,9 @@ static unique_ptr HandlePassthroughDeepHash(OptimizerExtensionInput std::unordered_map &hash_cache) { auto inner_hash_binding = GetOrInsertHashProjection(input, plan, inner_get, key_cols, false, hash_cache); - auto propagated_to_inner = PropagateSingleBinding(*plan, inner_hash_binding.table_index, inner_hash_binding, + auto propagated_to_inner = PropagateSingleBinding(*plan, inner_hash_binding.table_index.index, inner_hash_binding, LogicalType::UBIGINT, inner_agg); - if (propagated_to_inner.table_index == DConstants::INVALID_INDEX) { + if (propagated_to_inner.table_index.index == DConstants::INVALID_INDEX) { return nullptr; } // Add hash as GROUP BY column in the inner aggregate @@ -125,16 +125,17 @@ static unique_ptr HandlePassthroughDeepHash(OptimizerExtensionInput inner_agg->groups.push_back(make_uniq(LogicalType::UBIGINT, propagated_to_inner)); // Also add to grouping_sets so the physical HASH_GROUP_BY includes this column for (auto &gs : inner_agg->grouping_sets) { - gs.insert(new_group_idx); + gs.insert(ProjectionIndex(new_group_idx)); } inner_agg->ResolveOperatorTypes(); - ColumnBinding inner_hash_output(inner_agg->group_index, inner_agg->groups.size() - 1); + ColumnBinding inner_hash_output {TableIndex(inner_agg->group_index.index), + ProjectionIndex(inner_agg->groups.size() - 1)}; // Propagate from inner aggregate output to outer aggregate - auto propagated_to_outer = PropagateSingleBinding(*plan, inner_hash_output.table_index, inner_hash_output, + auto propagated_to_outer = PropagateSingleBinding(*plan, inner_hash_output.table_index.index, inner_hash_output, LogicalType::UBIGINT, target_agg); - if (propagated_to_outer.table_index == DConstants::INVALID_INDEX) { + if (propagated_to_outer.table_index.index == DConstants::INVALID_INDEX) { return nullptr; } return make_uniq(LogicalType::UBIGINT, propagated_to_outer); @@ -230,7 +231,7 @@ static void CollectCTERefs(LogicalOperator *op, idx_t target_cte_index, vectortype == LogicalOperatorType::LOGICAL_CTE_REF) { auto &ref = op->Cast(); - if (ref.cte_index == target_cte_index) { + if (ref.cte_index.index == target_cte_index) { refs.push_back(&ref); } } @@ -278,7 +279,7 @@ static ColumnBinding PropagateColumnThroughSubtree(LogicalOperator *root, Logica auto &proj = op->Cast(); idx_t new_idx = proj.expressions.size(); proj.expressions.push_back(make_uniq(col_type, current)); - current = ColumnBinding(proj.table_index, new_idx); + current = ColumnBinding(TableIndex(proj.table_index.index), ProjectionIndex(new_idx)); proj.ResolveOperatorTypes(); } else if (op != &get) { // For non-projection operators, check that the binding passes through @@ -292,7 +293,7 @@ static ColumnBinding PropagateColumnThroughSubtree(LogicalOperator *root, Logica } if (!found) { // Binding is not in this operator's output — propagation is impossible - return ColumnBinding(DConstants::INVALID_INDEX, 0); + return ColumnBinding(TableIndex(DConstants::INVALID_INDEX), ProjectionIndex(0)); } } } @@ -305,7 +306,7 @@ static ColumnBinding PropagateColumnThroughSubtree(LogicalOperator *root, Logica static bool ExpandCTEWithColumns(LogicalOperator *plan_root, const LogicalCTERef &ref, const string &table_name, const vector &columns) { // Find the MATERIALIZED_CTE that defines this CTE - auto *mat_cte = FindMaterializedCTE(plan_root, ref.cte_index); + auto *mat_cte = FindMaterializedCTE(plan_root, ref.cte_index.index); if (!mat_cte || mat_cte->children.empty()) { return false; } @@ -347,7 +348,7 @@ static bool ExpandCTEWithColumns(LogicalOperator *plan_root, const LogicalCTERef LogicalType col_type = LogicalType::INVALID; auto &catalog_columns = entry->GetColumns(); for (auto &col : catalog_columns.Logical()) { - if (col.Name() == col_name) { + if (col.Name().GetIdentifierName() == col_name) { col_oid = col.Oid(); col_type = col.Type(); break; @@ -365,22 +366,22 @@ static bool ExpandCTEWithColumns(LogicalOperator *plan_root, const LogicalCTERef // Add to projection_ids if the GET uses them if (!get->projection_ids.empty()) { - get->projection_ids.push_back(new_col_idx); + get->projection_ids.push_back(ProjectionIndex(new_col_idx)); } get->ResolveOperatorTypes(); - // The GET now outputs this column at binding (get->table_index, new_output_idx) + // The GET now outputs this column at binding (get->table_index.index, new_output_idx) // If projection_ids is non-empty, the output index is projection_ids.size()-1 // Otherwise it's column_ids.size()-1 idx_t get_output_idx = get->projection_ids.empty() ? (get->GetColumnIds().size() - 1) : (get->projection_ids.size() - 1); - ColumnBinding get_binding(get->table_index, get_output_idx); + ColumnBinding get_binding {TableIndex(get->table_index.index), ProjectionIndex(get_output_idx)}; // Propagate through intermediate operators (projections, filters) in the CTE definition. // If propagation fails (binding blocked by aggregate, semi-join, etc.), bail out. auto top_binding = PropagateColumnThroughSubtree(mat_cte->children[0].get(), *get, get_binding, col_type); - if (top_binding.table_index == DConstants::INVALID_INDEX) { + if (top_binding.table_index.index == DConstants::INVALID_INDEX) { return false; } @@ -389,15 +390,15 @@ static bool ExpandCTEWithColumns(LogicalOperator *plan_root, const LogicalCTERef // Update ALL CTE_REF nodes that reference this CTE vector all_refs; - CollectCTERefs(plan_root, ref.cte_index, all_refs); + CollectCTERefs(plan_root, ref.cte_index.index, all_refs); for (auto *cte_ref : all_refs) { - cte_ref->bound_columns.push_back(col_name); + cte_ref->bound_columns.emplace_back(col_name); cte_ref->chunk_types.push_back(col_type); cte_ref->ResolveOperatorTypes(); } PRIVACY_DEBUG_PRINT("ExpandCTEWithColumns: added column '" + col_name + "' (oid=" + std::to_string(col_oid) + - ") to CTE index " + std::to_string(ref.cte_index)); + ") to CTE index " + std::to_string(ref.cte_index.index)); } return true; @@ -418,7 +419,7 @@ static CTEHashMatch FindCTEHashSource(LogicalOperator *op, const string &pu_tabl if (!inside_delim_join && op->type == LogicalOperatorType::LOGICAL_CTE_REF) { auto &ref = op->Cast(); - auto cte_it = cte_map.find(ref.cte_index); + auto cte_it = cte_map.find(ref.cte_index.index); if (cte_it != cte_map.end()) { auto &cte_tables = cte_it->second; @@ -509,17 +510,18 @@ class PrivateUnionRewriter { AppendMembership(cte.children[0], mask_column); cte.column_count = value_columns + 1; vector refs; - CollectCTERefs(plan.get(), cte.table_index, refs); + CollectCTERefs(plan.get(), cte.table_index.index, refs); for (auto *ref : refs) { if (ref->chunk_types.size() != value_columns) { throw InvalidInputException("Private UNION CTE has unresolved ownership output columns"); } ref->chunk_types.push_back(LogicalType::UBIGINT); - ref->bound_columns.push_back("private_union_membership"); + ref->bound_columns.emplace_back("private_union_membership"); ref->ResolveOperatorTypes(); - memberships.emplace(ref, ColumnBinding(ref->table_index, mask_column)); + memberships.emplace(ref, ColumnBinding(ref->table_index, ProjectionIndex(mask_column))); } - PRIVACY_DEBUG_PRINT("[PRIVATE UNION] Exported membership through CTE " + cte.ctename); + PRIVACY_DEBUG_PRINT("[PRIVATE UNION] Exported membership through CTE " + + cte.ctename.GetIdentifierName()); } Rewrite(cte.children[1]); cte.ResolveOperatorTypes(); @@ -563,15 +565,15 @@ class PrivateUnionRewriter { vector> found; CollectMemberships(*agg.children[0], found); if (found.empty()) { - return ColumnBinding(DConstants::INVALID_INDEX, DConstants::INVALID_INDEX); + return ColumnBinding(); } if (found.size() != 1 || ContainsOutsideEntity(*agg.children[0])) { throw InvalidInputException("Private UNION cannot combine independent ownership streams above a UNION"); } CheckBranch(*agg.children[0]); - auto mask = - PropagateSingleBinding(*plan, found[0].second.table_index, found[0].second, LogicalType::UBIGINT, &agg); - if (mask.table_index == DConstants::INVALID_INDEX) { + auto mask = PropagateSingleBinding(*plan, found[0].second.table_index.index, found[0].second, + LogicalType::UBIGINT, &agg); + if (mask.table_index.index == DConstants::INVALID_INDEX) { throw InvalidInputException("Private UNION could not propagate membership to the consuming aggregate"); } return mask; @@ -584,7 +586,7 @@ class PrivateUnionRewriter { const vector &pus; std::unordered_map memberships; - void ProjectMembership(unique_ptr &op, idx_t count, idx_t output_index, bool distinct) { + void ProjectMembership(unique_ptr &op, idx_t count, TableIndex output_index, bool distinct) { auto types = op->types; auto bindings = op->GetColumnBindings(); vector> outputs; @@ -600,10 +602,11 @@ class PrivateUnionRewriter { input.optimizer.binder.GenerateTableIndex(), std::move(aggregates)); for (idx_t i = 0; i < count; i++) { dedup->groups.push_back(make_uniq(types[i], bindings[i])); - outputs.push_back(make_uniq(types[i], ColumnBinding(dedup->group_index, i))); + outputs.push_back(make_uniq( + types[i], ColumnBinding(dedup->group_index, ProjectionIndex(i)))); } - outputs.push_back( - make_uniq(LogicalType::UBIGINT, ColumnBinding(dedup->aggregate_index, 0))); + outputs.push_back(make_uniq( + LogicalType::UBIGINT, ColumnBinding(dedup->aggregate_index, ProjectionIndex(0)))); dedup->children.push_back(std::move(op)); op = std::move(dedup); } else { @@ -616,7 +619,7 @@ class PrivateUnionRewriter { projection->children.push_back(std::move(op)); op = std::move(projection); op->ResolveOperatorTypes(); - memberships.emplace(op.get(), ColumnBinding(output_index, count)); + memberships.emplace(op.get(), ColumnBinding(output_index, ProjectionIndex(count))); } void RewriteDistinct(unique_ptr &op) { @@ -636,8 +639,8 @@ class PrivateUnionRewriter { auto old_bindings = op->children[0]->GetColumnBindings(); for (idx_t i = 0; i < count; i++) { auto &expr = distinct.distinct_targets[i]; - if (i >= old_bindings.size() || expr->type != ExpressionType::BOUND_COLUMN_REF || - expr->Cast().binding != old_bindings[i]) { + if (i >= old_bindings.size() || expr->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF || + expr->Cast().Binding() != old_bindings[i]) { throw InvalidInputException("Private UNION DISTINCT requires the visible output columns"); } } @@ -646,7 +649,7 @@ class PrivateUnionRewriter { // DuckDB's own DISTINCT modifier above UNION DISTINCT is already satisfied. if (annotation != memberships.end() && child->type == LogicalOperatorType::LOGICAL_PROJECTION && child->children[0]->type == LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY && - annotation->second.column_index == count) { + annotation->second.column_index == ProjectionIndex(count)) { op = std::move(op->children[0]); return; } @@ -656,7 +659,8 @@ class PrivateUnionRewriter { ProjectMembership(op, count, output_index, true); ColumnBindingReplacer replacer; for (idx_t i = 0; i < count; i++) { - replacer.replacement_bindings.emplace_back(old_bindings[i], ColumnBinding(output_index, i)); + replacer.replacement_bindings.emplace_back(old_bindings[i], + ColumnBinding(output_index, ProjectionIndex(i))); } replacer.stop_operator = op.get(); replacer.VisitOperator(*plan); @@ -794,18 +798,18 @@ class PrivateUnionRewriter { if (!table) { continue; } - if (!AreTableColumnsAccessible(&target, get->table_index)) { + if (!AreTableColumnsAccessible(&target, get->table_index.index)) { if (ContainsUnaggregatedEntityInput(*get, check)) { throw InvalidInputException("Private UNION branch has inaccessible privacy-unit ownership"); } continue; } - auto meta = check.table_metadata.find(table->name); + auto meta = check.table_metadata.find(table->name.GetIdentifierName()); if (meta == check.table_metadata.end()) { continue; } vector candidate; - if (table->name == pus[0]) { + if (table->name.GetIdentifierName() == pus[0]) { candidate = meta->second.pks; } else { for (auto &fk : meta->second.fks) { @@ -838,8 +842,8 @@ class PrivateUnionRewriter { return; } auto propagated = - PropagateSingleBindingToOperator(*plan, source.table_index, source, LogicalType::UBIGINT, &target); - if (propagated.table_index == DConstants::INVALID_INDEX) { + PropagateSingleBindingToOperator(*plan, source.table_index.index, source, LogicalType::UBIGINT, &target); + if (propagated.table_index.index == DConstants::INVALID_INDEX) { throw InvalidInputException("Private UNION could not propagate branch membership"); } target.expressions.push_back(make_uniq(LogicalType::UBIGINT, propagated)); @@ -909,7 +913,7 @@ vector BuildPUHashExpressionsForAggregates(OptimizerExten // Cache for hash projections: get.table_index → hash column binding std::unordered_map hash_cache; - // Cache for CTE hash projections: cte_ref.table_index → hash column binding + // Cache for CTE hash projections: cte_ref.table_index.index → hash column binding // Prevents inserting duplicate projections above the same CTE_SCAN when // multiple aggregates try the CTE path for the same CTE_SCAN node. std::unordered_map cte_hash_cache; @@ -918,7 +922,7 @@ vector BuildPUHashExpressionsForAggregates(OptimizerExten vector result; for (auto *target_agg : target_aggregates) { auto union_mask = unions.FindForAggregate(*target_agg); - if (union_mask.table_index != DConstants::INVALID_INDEX) { + if (union_mask.table_index.index != DConstants::INVALID_INDEX) { result.push_back( {target_agg, make_uniq(LogicalType::UBIGINT, union_mask), union_mask, 1.0}); continue; @@ -937,7 +941,7 @@ vector BuildPUHashExpressionsForAggregates(OptimizerExten if (pu_in_subtree) { get_ptr = FindTableScanInSubtree(target_agg, pu_table_name); if (get_ptr) { - pu_columns_accessible = AreTableColumnsAccessible(target_agg, get_ptr->table_index); + pu_columns_accessible = AreTableColumnsAccessible(target_agg, get_ptr->table_index.index); } } @@ -961,7 +965,7 @@ vector BuildPUHashExpressionsForAggregates(OptimizerExten // Check if the PU table is directly reachable (not behind a nested aggregate). // PU-key passthrough: the table is behind an inner aggregate that groups // by PU key — insert hash deep and propagate through the inner's GROUP BY. - if (!IsDirectlyReachable(target_agg, get.table_index)) { + if (!IsDirectlyReachable(target_agg, get.table_index.index)) { ColumnBinding pk_binding; auto *inner_agg = FindInnerAggregateWithPUKeyGroup(target_agg, check, pu_table_names, pk_binding); if (inner_agg) { @@ -979,17 +983,17 @@ vector BuildPUHashExpressionsForAggregates(OptimizerExten // Direct path: insert hash projection above get and propagate auto hash_binding = GetOrInsertHashProjection(input, plan, get, pks, false, hash_cache); - auto propagated = PropagateSingleBinding(*plan, hash_binding.table_index, hash_binding, + auto propagated = PropagateSingleBinding(*plan, hash_binding.table_index.index, hash_binding, LogicalType::UBIGINT, target_agg); - if (propagated.table_index == DConstants::INVALID_INDEX) { + if (propagated.table_index.index == DConstants::INVALID_INDEX) { // Try DELIM_JOIN fallback before giving up - auto delim_result = AddBindingToDelimJoin(plan, hash_binding.table_index, hash_binding, + auto delim_result = AddBindingToDelimJoin(plan, hash_binding.table_index.index, hash_binding, LogicalType::UBIGINT, target_agg); if (delim_result.IsValid()) { propagated = delim_result.binding; } } - if (propagated.table_index == DConstants::INVALID_INDEX) { + if (propagated.table_index.index == DConstants::INVALID_INDEX) { continue; } auto hash_ref = make_uniq(LogicalType::UBIGINT, propagated); @@ -1042,7 +1046,7 @@ vector BuildPUHashExpressionsForAggregates(OptimizerExten } // PU-key passthrough: if the FK table is behind an inner aggregate, // insert hash deep and propagate through the inner's GROUP BY. - if (!IsDirectlyReachable(target_agg, fk_get_ptr->table_index)) { + if (!IsDirectlyReachable(target_agg, fk_get_ptr->table_index.index)) { ColumnBinding pk_binding; auto *inner_agg = FindInnerAggregateWithPUKeyGroup(target_agg, check, pu_table_names, pk_binding); @@ -1063,18 +1067,19 @@ vector BuildPUHashExpressionsForAggregates(OptimizerExten // Direct path: insert hash projection and propagate auto fk_hash_binding = GetOrInsertHashProjection(input, plan, *fk_get_ptr, fk_cols, false, hash_cache); - auto fk_propagated = PropagateSingleBinding(*plan, fk_hash_binding.table_index, fk_hash_binding, - LogicalType::UBIGINT, target_agg); + auto fk_propagated = PropagateSingleBinding(*plan, fk_hash_binding.table_index.index, + fk_hash_binding, LogicalType::UBIGINT, target_agg); // If propagation failed, try DELIM_JOIN fallback - if (fk_propagated.table_index == DConstants::INVALID_INDEX) { - auto delim_result = AddBindingToDelimJoin( - plan, fk_hash_binding.table_index, fk_hash_binding, LogicalType::UBIGINT, target_agg); + if (fk_propagated.table_index.index == DConstants::INVALID_INDEX) { + auto delim_result = + AddBindingToDelimJoin(plan, fk_hash_binding.table_index.index, fk_hash_binding, + LogicalType::UBIGINT, target_agg); if (delim_result.IsValid()) { fk_propagated = delim_result.binding; } } - if (fk_propagated.table_index == DConstants::INVALID_INDEX) { + if (fk_propagated.table_index.index == DConstants::INVALID_INDEX) { continue; } auto fk_hash_ref = make_uniq(LogicalType::UBIGINT, fk_propagated); @@ -1108,29 +1113,29 @@ vector BuildPUHashExpressionsForAggregates(OptimizerExten continue; } ColumnBinding hash_binding; - auto cte_cache_it = cte_hash_cache.find(match.cte_ref->table_index); + auto cte_cache_it = cte_hash_cache.find(match.cte_ref->table_index.index); if (cte_cache_it != cte_hash_cache.end()) { hash_binding = cte_cache_it->second; } else { hash_binding = InsertHashProjectionAboveCTERef(input, plan, *match.cte_ref, match.hash_columns); - if (hash_binding.table_index == DConstants::INVALID_INDEX) { + if (hash_binding.table_index.index == DConstants::INVALID_INDEX) { continue; } - cte_hash_cache[match.cte_ref->table_index] = hash_binding; + cte_hash_cache[match.cte_ref->table_index.index] = hash_binding; } - auto propagated = PropagateSingleBinding(*plan, hash_binding.table_index, hash_binding, + auto propagated = PropagateSingleBinding(*plan, hash_binding.table_index.index, hash_binding, LogicalType::UBIGINT, target_agg); // If propagation failed (e.g., blocked by DELIM_JOIN), try AddBindingToDelimJoin - if (propagated.table_index == DConstants::INVALID_INDEX) { - auto delim_result = AddBindingToDelimJoin(plan, hash_binding.table_index, hash_binding, + if (propagated.table_index.index == DConstants::INVALID_INDEX) { + auto delim_result = AddBindingToDelimJoin(plan, hash_binding.table_index.index, hash_binding, LogicalType::UBIGINT, target_agg); if (delim_result.IsValid()) { propagated = delim_result.binding; } } - if (propagated.table_index == DConstants::INVALID_INDEX) { + if (propagated.table_index.index == DConstants::INVALID_INDEX) { continue; } auto hash_ref = make_uniq(LogicalType::UBIGINT, propagated); @@ -1151,8 +1156,8 @@ vector BuildPUHashExpressionsForAggregates(OptimizerExten // Record hash binding for categorical rewriter before modifying ColumnBinding hash_binding; - if (combined_hash_expr->type == ExpressionType::BOUND_COLUMN_REF) { - hash_binding = combined_hash_expr->Cast().binding; + if (combined_hash_expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { + hash_binding = combined_hash_expr->Cast().BindingMutable(); } result.push_back({target_agg, std::move(combined_hash_expr), hash_binding, correction}); @@ -1169,12 +1174,12 @@ void ModifyPlanWithPU(OptimizerExtensionInput &input, unique_ptraggregate_index; - info.group_index = hash_info.aggregate->group_index; + info.aggregate_index = hash_info.aggregate->aggregate_index.index; + info.group_index = hash_info.aggregate->group_index.index; info.hash_binding = hash_info.hash_binding; - pac_agg_info[hash_info.aggregate->aggregate_index] = info; + pac_agg_info[hash_info.aggregate->aggregate_index.index] = info; } } } @@ -1309,7 +1314,7 @@ static void CompilePacMechanism(const PrivacyCompatibilityResult &check, Optimiz // clipping variants with per-PU pre-aggregation below. { Value clip_val; - if (input.context.TryGetCurrentSetting("priv_clip_support", clip_val) && !clip_val.IsNull()) { + if (input.context.TryGetCurrentSetting(Identifier("priv_clip_support"), clip_val) && !clip_val.IsNull()) { RewriteClipAggregates(input, plan, check, pu_setup.pu_names); } } diff --git a/src/compiler/privacy_compiler_helpers.cpp b/src/compiler/privacy_compiler_helpers.cpp index 591a0253..c57f2528 100644 --- a/src/compiler/privacy_compiler_helpers.cpp +++ b/src/compiler/privacy_compiler_helpers.cpp @@ -2,6 +2,7 @@ // Created by ila on 12/23/25. // +#include "duckdb/main/database_manager.hpp" #include "compiler/privacy_compiler_helpers.hpp" #include "privacy_debug.hpp" #include "core/privacy_optimizer.hpp" @@ -44,14 +45,16 @@ void BuildJoinConditions(LogicalGet *left_get, LogicalGet *right_get, const vect if (lproj == DConstants::INVALID_INDEX || rproj == DConstants::INVALID_INDEX) { throw InternalException("Privacy compiler: failed to project FK/PK columns for join"); } - JoinCondition cond; - cond.comparison = ExpressionType::COMPARE_EQUAL; + JoinCondition cond(nullptr, nullptr, ExpressionType::COMPARE_EQUAL); + auto left_col_index = left_get->GetColumnIds()[lproj]; auto right_col_index = right_get->GetColumnIds()[rproj]; auto &left_type = left_get->GetColumnType(left_col_index); auto &right_type = right_get->GetColumnType(right_col_index); - cond.left = make_uniq(left_type, ColumnBinding(left_get->table_index, lproj)); - cond.right = make_uniq(right_type, ColumnBinding(right_get->table_index, rproj)); + cond.LeftReference() = make_uniq( + left_type, ColumnBinding(TableIndex(left_get->table_index.index), ProjectionIndex(lproj))); + cond.RightReference() = make_uniq( + right_type, ColumnBinding(TableIndex(right_get->table_index.index), ProjectionIndex(rproj))); conditions.push_back(std::move(cond)); } } @@ -83,8 +86,8 @@ unique_ptr CreateLogicalJoin(const PrivacyCompatibilityResult & throw InternalException( "Privacy compiler: expected both LogicalGet nodes to be bound to tables for join creation"); } - string left_table_name = left_table_ptr->name; - string right_table_name = right_table_ptr->name; + string left_table_name = left_table_ptr->name.GetIdentifierName(); + string right_table_name = right_table_ptr->name.GetIdentifierName(); // Require metadata to be present; compatibility check is responsible for populating it. auto lit = check.table_metadata.find(left_table_name); @@ -160,9 +163,8 @@ unique_ptr CreateLogicalJoin(const PrivacyCompatibilityResult & right_table_name); } - vector> extra; return LogicalComparisonJoin::CreateJoin(JoinType::INNER, JoinRefType::REGULAR, std::move(left_operator), - std::move(right), std::move(conditions), std::move(extra)); + std::move(right), std::move(conditions)); } // Create a LogicalGet operator for a table by name, projecting only the specified columns @@ -173,8 +175,8 @@ unique_ptr CreateLogicalGet(ClientContext &context, unique_ptr CreateLogicalGet(ClientContext &context, unique_ptr bind_data; auto scan_function = table_entry.GetScanFunction(context, bind_data); vector return_types = {}; - vector return_names = {}; + vector return_names = {}; vector column_ids = {}; - vector projection_ids = {}; + vector projection_ids = {}; // Build a set of required column names for fast lookup (case-insensitive) std::unordered_set required_set; @@ -198,15 +200,15 @@ unique_ptr CreateLogicalGet(ClientContext &context, unique_ptr found_columns; for (auto &col : table_entry.GetColumns().Logical()) { - string col_name_lower = StringUtil::Lower(col.Name()); + string col_name_lower = StringUtil::Lower(col.Name().GetIdentifierName()); // If required_columns is empty, include all columns (backward compatible) // Otherwise, only include columns that are in the required set if (required_columns.empty() || required_set.count(col_name_lower) > 0) { return_types.push_back(col.Type()); - return_names.push_back(col.Name()); + return_names.emplace_back(col.Name().GetIdentifierName()); column_ids.push_back(ColumnIndex(col.Oid())); - projection_ids.push_back(column_ids.size() - 1); + projection_ids.push_back(ProjectionIndex(column_ids.size() - 1)); if (!required_columns.empty()) { found_columns.insert(col_name_lower); @@ -238,20 +240,21 @@ unique_ptr CreateLogicalGet(ClientContext &context, unique_ptr get = make_uniq(idx, scan_function, std::move(bind_data), - std::move(return_types), std::move(return_names)); + unique_ptr get = + make_uniq(TableIndex(idx), BoundTableFunction(scan_function), std::move(bind_data), + std::move(return_types), std::move(return_names)); get->SetColumnIds(std::move(column_ids)); get->projection_ids = projection_ids; get->ResolveOperatorTypes(); get->Verify(context); - return get; + return std::move(get); } throw ParserException("PAC: missing internal sample table " + table); @@ -354,7 +357,7 @@ LogicalGet *FindAccessibleGetInSubtree(unique_ptr &plan, Logica FindAllNodesByTable(&plan, table_name, nodes); for (auto *node : nodes) { auto &node_get = node->get()->Cast(); - if (HasTableIndexInSubtree(subtree_root, node_get.table_index)) { + if (HasTableIndexInSubtree(subtree_root, node_get.table_index.index)) { return &node_get; } } diff --git a/src/compiler/privacy_mechanisms.cpp b/src/compiler/privacy_mechanisms.cpp index 9bfa2948..5df90e32 100644 --- a/src/compiler/privacy_mechanisms.cpp +++ b/src/compiler/privacy_mechanisms.cpp @@ -89,7 +89,7 @@ static unique_ptr BindScalarLocal(OptimizerExtensionInput &input, co vector> children) { FunctionBinder function_binder(input.context); ErrorData error; - auto expr = function_binder.BindScalarFunction(DEFAULT_SCHEMA, func_name, std::move(children), error); + auto expr = function_binder.BindScalarFunction(DEFAULT_SCHEMA, Identifier(func_name), std::move(children), error); if (error.HasError()) { throw InternalException("dp: failed to bind scalar function '" + func_name + "': " + error.Message()); } @@ -103,12 +103,14 @@ static unique_ptr BindAggregateLocal(OptimizerExtensionInput &input, vector arg_types; arg_types.reserve(children.size()); for (auto &child : children) { - arg_types.push_back(child->return_type); + arg_types.push_back(child->GetReturnType()); } auto &entry = Catalog::GetSystemCatalog(input.context) - .GetEntry(input.context, DEFAULT_SCHEMA, func_name); - auto best = function_binder.BindFunction(entry.name, entry.functions, arg_types, error); + .GetEntry( + input.context, QualifiedName(Catalog::GetSystemCatalog(input.context).GetName(), + Identifier::DefaultSchema(), Identifier(func_name))); + auto best = function_binder.BindFunction(Identifier(entry.name), entry.functions, arg_types, error); if (!best.IsValid()) { throw InternalException("dp: failed to bind aggregate function '" + func_name + "': " + error.Message()); } @@ -130,15 +132,15 @@ static vector RewriteAvgAggregates(OptimizerExtensionInput &input, Logi idx_t avg_idx = 0; for (idx_t i = 0; i < n_original; i++) { auto &aggr = agg->expressions[i]->Cast(); - if (aggr.function.name != "avg") { + if (aggr.FunctionMutable().GetName().GetIdentifierName() != "avg") { continue; } - auto output_type = aggr.return_type; + auto output_type = aggr.GetReturnType(); // Preserve `FILTER (WHERE ...)` from the original AVG. For fixed-bound means, // both the rewritten SUM and appended COUNT(*) receive the predicate. - auto count_arg = count_non_null_value ? aggr.children[0]->Copy() : nullptr; - auto arg = std::move(aggr.children[0]); - auto filter_for_sum = std::move(aggr.filter); + auto count_arg = count_non_null_value ? aggr.GetChildrenMutable()[0]->Copy() : nullptr; + auto arg = std::move(aggr.GetChildrenMutable()[0]); + auto filter_for_sum = std::move(aggr.GetFilterMutable()); auto filter_for_count = !automatic_bounds_mean && filter_for_sum ? filter_for_sum->Copy() : nullptr; AvgDomain domain {0.0, 0.0}; if (avg_bounds) { @@ -159,14 +161,14 @@ static vector RewriteAvgAggregates(OptimizerExtensionInput &input, Logi } agg->expressions[i] = BindPlainAggregate(input, "sum", std::move(arg)); - agg->expressions[i]->Cast().filter = std::move(filter_for_sum); + agg->expressions[i]->Cast().GetFilterMutable() = std::move(filter_for_sum); idx_t count_pos = DConstants::INVALID_INDEX; if (!automatic_bounds_mean) { count_pos = agg->expressions.size(); agg->expressions.push_back(count_non_null_value ? BindPlainAggregate(input, "count", std::move(count_arg)) : BindPlainAggregate(input, "count_star", nullptr)); - agg->expressions.back()->Cast().filter = std::move(filter_for_count); + agg->expressions.back()->Cast().GetFilterMutable() = std::move(filter_for_count); } AvgInfo info; @@ -202,7 +204,7 @@ static DPFKChain ExtractFKChain(const PrivacyCompatibilityResult &check, const v if (!table) { continue; } - string name = StringUtil::Lower(table->name); + string name = StringUtil::Lower(table->name.GetIdentifierName()); name_count[name]++; if (name_count[name] > 1) { throw InvalidInputException("dp: self-joins are not supported (table '" + table->name + @@ -363,11 +365,11 @@ static void CheckDPAggregateNode(LogicalAggregate *agg, const string &mechanism_ throw InvalidInputException(mechanism_name + ": unsupported non-aggregate expression in aggregate node"); } auto &aggr = expr->Cast(); - string name = StringUtil::Lower(aggr.function.name); + string name = StringUtil::Lower(aggr.FunctionMutable().GetName().GetIdentifierName()); bool is_count = name == "count" || name == "count_star"; if (aggr.IsDistinct()) { if ((mechanism_name != "dp_sass" && mechanism_name != "dp_elastic" && mechanism_name != "dp_standard") || - !is_count || aggr.children.size() != 1 || aggr.filter) { + !is_count || aggr.GetChildrenMutable().size() != 1 || aggr.GetFilterMutable()) { throw InvalidInputException(mechanism_name + ": only COUNT(DISTINCT x) is supported for DISTINCT"); } continue; @@ -383,14 +385,14 @@ static void CheckDPAggregateNode(LogicalAggregate *agg, const string &mechanism_ if (!is_count && !is_sum_avg && !(allow_min_max && is_min_max)) { string supported = allow_min_max ? "COUNT, SUM, AVG, MIN, and MAX" : "COUNT, SUM, and AVG"; throw InvalidInputException(mechanism_name + ": only " + supported + " are supported (got '" + - aggr.function.name + "')"); + aggr.FunctionMutable().GetName().GetIdentifierName() + "')"); } - if ((is_sum_avg || is_min_max) && aggr.children.size() != 1) { + if ((is_sum_avg || is_min_max) && aggr.GetChildrenMutable().size() != 1) { throw InvalidInputException(mechanism_name + ": " + name + " must have exactly one argument"); } - if ((is_sum_avg || is_min_max) && !aggr.children[0]->return_type.IsNumeric()) { + if ((is_sum_avg || is_min_max) && !aggr.GetChildrenMutable()[0]->GetReturnType().IsNumeric()) { throw InvalidInputException(mechanism_name + ": " + name + " argument must be a numeric type (got '" + - aggr.children[0]->return_type.ToString() + "')"); + aggr.GetChildrenMutable()[0]->GetReturnType().ToString() + "')"); } } } @@ -419,7 +421,7 @@ static void CountAccessibleTableScans(LogicalOperator *op, std::unordered_mapCast(); auto table = get.GetTable(); if (table) { - counts[StringUtil::Lower(table->name)]++; + counts[StringUtil::Lower(table->name.GetIdentifierName())]++; } return; } @@ -453,7 +455,7 @@ static double ValidateDPSelfJoins(unique_ptr &plan, const strin for (auto *get : gets) { auto table = get->GetTable(); if (table) { - all_counts[StringUtil::Lower(table->name)]++; + all_counts[StringUtil::Lower(table->name.GetIdentifierName())]++; } } @@ -493,7 +495,8 @@ static double ComputeMfK(ClientContext &context, const string &table_name, const // Table is empty: no rows contribute → max frequency is 0 return 0.0; } - auto val = result->GetValue(0, 0); + auto chunk = result->Fetch(); + auto val = chunk->GetValue(0, 0); return val.IsNull() ? 0.0 : val.GetValue(); } @@ -608,7 +611,8 @@ enum class DPSensitivityKind { GLOBAL, ELASTIC_SMOOTH }; static bool AggregateContainsSum(const LogicalAggregate *agg) { for (auto &expr : agg->expressions) { auto &a = expr->Cast(); - if (a.function.name == "sum" || a.function.name == "avg") { + if (a.FunctionMutable().GetName().GetIdentifierName() == "sum" || + a.FunctionMutable().GetName().GetIdentifierName() == "avg") { return true; } } @@ -618,7 +622,7 @@ static bool AggregateContainsSum(const LogicalAggregate *agg) { static bool AggregateContainsAvg(const LogicalAggregate *agg) { for (auto &expr : agg->expressions) { auto &a = expr->Cast(); - if (a.function.name == "avg") { + if (a.FunctionMutable().GetName().GetIdentifierName() == "avg") { return true; } } @@ -629,7 +633,7 @@ static idx_t CountAvgAggregates(const LogicalAggregate *agg) { idx_t count = 0; for (auto &expr : agg->expressions) { auto &a = expr->Cast(); - if (a.function.name == "avg") { + if (a.FunctionMutable().GetName().GetIdentifierName() == "avg") { count++; } } @@ -640,7 +644,7 @@ static idx_t CountExplicitSumAggregates(const LogicalAggregate *agg) { idx_t count = 0; for (auto &expr : agg->expressions) { auto &a = expr->Cast(); - if (a.function.name == "sum") { + if (a.FunctionMutable().GetName().GetIdentifierName() == "sum") { count++; } } @@ -683,7 +687,7 @@ static vector GetDpSumContributionBounds(ClientContext &context, const s idx_t explicit_sum_count = 0; for (idx_t i = 0; i < agg->expressions.size(); i++) { auto &aggr = agg->expressions[i]->Cast(); - if (aggr.function.name == "sum" && !FindAvgInfoForSumPos(avg_infos, i)) { + if (aggr.FunctionMutable().GetName().GetIdentifierName() == "sum" && !FindAvgInfoForSumPos(avg_infos, i)) { explicit_sum_count++; } } @@ -709,7 +713,7 @@ static vector GetDpSumContributionBounds(ClientContext &context, const s bool needs_scalar_sum_bound = false; for (idx_t i = 0; i < agg->expressions.size(); i++) { auto &aggr = agg->expressions[i]->Cast(); - if (aggr.function.name != "sum") { + if (aggr.FunctionMutable().GetName().GetIdentifierName() != "sum") { continue; } if (FindAvgInfoForSumPos(avg_infos, i)) { @@ -727,7 +731,7 @@ static vector GetDpSumContributionBounds(ClientContext &context, const s idx_t sum_list_idx = 0; for (idx_t i = 0; i < agg->expressions.size(); i++) { auto &aggr = agg->expressions[i]->Cast(); - if (aggr.function.name != "sum") { + if (aggr.FunctionMutable().GetName().GetIdentifierName() != "sum") { continue; } if (FindAvgInfoForSumPos(avg_infos, i)) { @@ -745,7 +749,7 @@ static unique_ptr ClipToBounds(OptimizerExtensionInput &input, uniqu double hi, const LogicalType &restore_type) { auto as_double = BoundCastExpression::AddCastToType(input.context, std::move(expr), LogicalType::DOUBLE); auto is_null = make_uniq(ExpressionType::OPERATOR_IS_NULL, LogicalType::BOOLEAN); - is_null->children.push_back(as_double->Copy()); + is_null->GetChildrenMutable().push_back(as_double->Copy()); auto hi_const = make_uniq(Value::DOUBLE(hi)); auto lo_const = make_uniq(Value::DOUBLE(lo)); auto upper = input.optimizer.BindScalarFunction("least", std::move(as_double), std::move(hi_const)); @@ -754,8 +758,8 @@ static unique_ptr ClipToBounds(OptimizerExtensionInput &input, uniqu BoundCaseCheck null_check; null_check.when_expr = std::move(is_null); null_check.then_expr = make_uniq(Value(LogicalType::DOUBLE)); - preserve_null->case_checks.push_back(std::move(null_check)); - preserve_null->else_expr = std::move(clipped); + preserve_null->CaseChecksMutable().push_back(std::move(null_check)); + preserve_null->ElseMutable() = std::move(clipped); return BoundCastExpression::AddCastToType(input.context, std::move(preserve_null), restore_type); } @@ -763,13 +767,14 @@ static void ClipSumInputs(OptimizerExtensionInput &input, LogicalAggregate *agg, for (idx_t i = 0; i < agg->expressions.size(); i++) { auto &expr = agg->expressions[i]; auto &aggr = expr->Cast(); - if (aggr.function.name != "sum" || aggr.children.empty()) { + if (aggr.FunctionMutable().GetName().GetIdentifierName() != "sum" || aggr.GetChildrenMutable().empty()) { continue; } - auto original_type = aggr.children[0]->return_type; + auto original_type = aggr.GetChildrenMutable()[0]->GetReturnType(); D_ASSERT(i < sum_bounds.size()); double sum_bound = sum_bounds[i]; - aggr.children[0] = ClipToBounds(input, std::move(aggr.children[0]), -sum_bound, sum_bound, original_type); + aggr.GetChildrenMutable()[0] = + ClipToBounds(input, std::move(aggr.GetChildrenMutable()[0]), -sum_bound, sum_bound, original_type); } agg->ResolveOperatorTypes(); } @@ -786,7 +791,7 @@ static double GetRequiredDpBound(ClientContext &context, const string &setting_n } else if (setting_name == "dp_sass_count_output_bound" || setting_name == "dp_sass_sum_output_bound" || setting_name == "dp_sass_avg_lower_bound" || setting_name == "dp_sass_avg_upper_bound") { Value val; - found = context.TryGetCurrentSetting(setting_name, val) && !val.IsNull(); + found = context.TryGetCurrentSetting(Identifier(setting_name), val) && !val.IsNull(); if (found) { bound = val.GetValue(); } @@ -818,7 +823,7 @@ static int64_t GetRequiredDpIntegerBound(ClientContext &context, const string &s static double GetRequiredFiniteSetting(ClientContext &context, const string &setting_name, const string &mechanism_name) { Value val; - if (!context.TryGetCurrentSetting(setting_name, val) || val.IsNull()) { + if (!context.TryGetCurrentSetting(Identifier(setting_name), val) || val.IsNull()) { throw InvalidInputException(mechanism_name + ": " + setting_name + " must be set"); } double bound = val.GetValue(); @@ -831,7 +836,7 @@ static double GetRequiredFiniteSetting(ClientContext &context, const string &set static bool TryGetFiniteSetting(ClientContext &context, const string &setting_name, double &bound) { Value val; - if (!context.TryGetCurrentSetting(setting_name, val) || val.IsNull()) { + if (!context.TryGetCurrentSetting(Identifier(setting_name), val) || val.IsNull()) { return false; } bound = val.GetValue(); @@ -844,7 +849,7 @@ static bool TryGetFiniteSetting(ClientContext &context, const string &setting_na static vector ParseFiniteSettingList(ClientContext &context, const string &setting_name, bool &found) { found = false; Value val; - if (!context.TryGetCurrentSetting(setting_name, val) || val.IsNull()) { + if (!context.TryGetCurrentSetting(Identifier(setting_name), val) || val.IsNull()) { return {}; } string text = val.ToString(); @@ -1029,7 +1034,7 @@ static std::pair GetDpMinMaxBounds(ClientContext &context, const // Elastic Laplace scale = sensitivity / epsilon; sensitivity = es for COUNT, es * value_bound for SUM. static double ElasticSensitivity(idx_t aggregate_pos, const BoundAggregateExpression &aggr, const vector &sum_bounds, double es) { - const string &name = aggr.function.name; + const string &name = aggr.Function().GetName().GetIdentifierName(); if (name == "count" || name == "count_star") { return es; } @@ -1041,7 +1046,8 @@ static double ElasticSensitivity(idx_t aggregate_pos, const BoundAggregateExpres } static bool IsCountAggregate(const BoundAggregateExpression &aggr) { - return aggr.function.name == "count" || aggr.function.name == "count_star"; + return aggr.Function().GetName().GetIdentifierName() == "count" || + aggr.Function().GetName().GetIdentifierName() == "count_star"; } // Does the aggregate contain any COUNT/COUNT(*) (after AVG→SUM+COUNT rewrite)? @@ -1055,7 +1061,8 @@ static bool AggregateContainsCount(const LogicalAggregate *agg) { } static bool IsMinMaxAggregate(const BoundAggregateExpression &aggr) { - return aggr.function.name == "min" || aggr.function.name == "max"; + return aggr.Function().GetName().GetIdentifierName() == "min" || + aggr.Function().GetName().GetIdentifierName() == "max"; } static bool AggregateContainsMinMax(const LogicalAggregate *agg) { @@ -1108,7 +1115,7 @@ static void EnsurePuAdjacentTableAvailable(OptimizerExtensionInput &input, uniqu auto &binder = input.optimizer.binder; unique_ptr current = (*start_ref)->Copy(input.context); for (auto &table : tables_to_join) { - idx_t table_index = binder.GenerateTableIndex(); + idx_t table_index = binder.GenerateTableIndex().index; auto get = CreateLogicalGet(input.context, plan, table, table_index); current = CreateLogicalJoin(check, input.context, std::move(current), std::move(get)); } @@ -1155,10 +1162,10 @@ static unique_ptr BuildPerPuExpressionAtOperator(OptimizerExtensionI } auto col_index = key_get->GetColumnIds()[proj_idx]; auto col_type = key_get->GetColumnType(col_index); - ColumnBinding source_binding(key_get->table_index, proj_idx); + ColumnBinding source_binding {TableIndex(key_get->table_index.index), ProjectionIndex(proj_idx)}; auto propagated = - PropagateSingleBindingToOperator(*plan, key_get->table_index, source_binding, col_type, target); - if (propagated.table_index == DConstants::INVALID_INDEX) { + PropagateSingleBindingToOperator(*plan, key_get->table_index.index, source_binding, col_type, target); + if (propagated.table_index.index == DConstants::INVALID_INDEX) { throw InternalException(mech + ": could not propagate PU key column '" + key_column + "' to aggregate input"); } @@ -1167,18 +1174,19 @@ static unique_ptr BuildPerPuExpressionAtOperator(OptimizerExtensionI auto pu_group_expr = BuildXorHash(input, std::move(key_exprs)); pu_group_expr = BoundCastExpression::AddCastToType(input.context, std::move(pu_group_expr), LogicalType::UBIGINT); - pu_group_expr->alias = mech + "_pu_key"; + pu_group_expr->SetAlias(Identifier(mech + "_pu_key")); return pu_group_expr; } static unique_ptr BuildRankSourceRef(const vector &source_types, idx_t source_table_index, idx_t column) { - return make_uniq(source_types[column], ColumnBinding(source_table_index, column)); + return make_uniq(source_types[column], + ColumnBinding(TableIndex(source_table_index), ProjectionIndex(column))); } static unique_ptr BuildIsNullPredicate(unique_ptr expr) { auto is_null = make_uniq(ExpressionType::OPERATOR_IS_NULL, LogicalType::BOOLEAN); - is_null->children.push_back(std::move(expr)); + is_null->GetChildrenMutable().push_back(std::move(expr)); return is_null; } @@ -1205,22 +1213,23 @@ static unique_ptr BuildRankCapFilter(OptimizerExtensionInput &i return child; } auto &binder = input.optimizer.binder; - idx_t window_index = binder.GenerateTableIndex(); - auto window = make_uniq(window_index); + idx_t window_index = binder.GenerateTableIndex().index; + auto window = make_uniq(TableIndex(window_index)); for (auto &spec : specs) { - auto rank = make_uniq(spec.rank_type, LogicalType::BIGINT, nullptr, nullptr); + auto rank = BindRankingWindow(input.context, spec.rank_type); for (auto col : spec.partition_cols) { - rank->partitions.push_back(BuildRankSourceRef(source_types, source_table_index, col)); + rank->PartitionsMutable().push_back(BuildRankSourceRef(source_types, source_table_index, col)); } for (auto col : spec.keep_null_cols) { - rank->orders.emplace_back(OrderType::ASCENDING, OrderByNullType::NULLS_LAST, - BuildIsNullPredicate(BuildRankSourceRef(source_types, source_table_index, col))); + rank->OrderByMutable().emplace_back( + OrderType::ASCENDING, OrderByNullType::NULLS_LAST, + BuildIsNullPredicate(BuildRankSourceRef(source_types, source_table_index, col))); } if (spec.rank_type == ExpressionType::WINDOW_RANK_DENSE) { for (auto col : spec.order_cols) { - rank->orders.emplace_back(OrderType::ASCENDING, OrderByNullType::NULLS_LAST, - BuildRankSourceRef(source_types, source_table_index, col)); + rank->OrderByMutable().emplace_back(OrderType::ASCENDING, OrderByNullType::NULLS_LAST, + BuildRankSourceRef(source_types, source_table_index, col)); } } else { vector> hash_children; @@ -1229,19 +1238,21 @@ static unique_ptr BuildRankCapFilter(OptimizerExtensionInput &i hash_children.push_back(BuildRankSourceRef(source_types, source_table_index, col)); } auto hash_expr = BindScalarLocal(input, "hash", std::move(hash_children)); - rank->orders.emplace_back(OrderType::ASCENDING, OrderByNullType::NULLS_LAST, std::move(hash_expr)); + rank->OrderByMutable().emplace_back(OrderType::ASCENDING, OrderByNullType::NULLS_LAST, + std::move(hash_expr)); } - rank->start = WindowBoundary::UNBOUNDED_PRECEDING; - rank->end = WindowBoundary::CURRENT_ROW_ROWS; + rank->WindowStartMutable() = WindowBoundary::UNBOUNDED_PRECEDING; + rank->WindowEndMutable() = WindowBoundary::CURRENT_ROW_ROWS; window->expressions.push_back(std::move(rank)); } window->children.push_back(std::move(child)); vector> keep_conditions; for (idx_t i = 0; i < specs.size(); i++) { - auto rank_ref = make_uniq(LogicalType::BIGINT, ColumnBinding(window_index, i)); + auto rank_ref = make_uniq( + LogicalType::BIGINT, ColumnBinding(TableIndex(window_index), ProjectionIndex(i))); auto cap_const = make_uniq(Value::BIGINT(specs[i].cap)); - unique_ptr keep_condition = make_uniq( + unique_ptr keep_condition = BoundComparisonExpression::Create( ExpressionType::COMPARE_LESSTHANOREQUALTO, std::move(rank_ref), std::move(cap_const)); for (auto col : specs[i].keep_null_cols) { keep_condition = make_uniq( @@ -1255,7 +1266,7 @@ static unique_ptr BuildRankCapFilter(OptimizerExtensionInput &i keep_expr = std::move(keep_conditions[0]); } else { auto conj = make_uniq(ExpressionType::CONJUNCTION_AND); - conj->children = std::move(keep_conditions); + conj->GetChildrenMutable() = std::move(keep_conditions); keep_expr = std::move(conj); } auto filter = make_uniq(std::move(keep_expr)); @@ -1265,7 +1276,7 @@ static unique_ptr BuildRankCapFilter(OptimizerExtensionInput &i static unique_ptr BuildIsNotNullPredicate(unique_ptr expr) { auto not_null = make_uniq(ExpressionType::OPERATOR_IS_NOT_NULL, LogicalType::BOOLEAN); - not_null->children.push_back(std::move(expr)); + not_null->GetChildrenMutable().push_back(std::move(expr)); return not_null; } @@ -1283,7 +1294,7 @@ static unique_ptr ConjoinPredicates(vector> p return std::move(present[0]); } auto conjunction = make_uniq(ExpressionType::CONJUNCTION_AND); - conjunction->children = std::move(present); + conjunction->GetChildrenMutable() = std::move(present); return std::move(conjunction); } @@ -1297,35 +1308,36 @@ static vector> PrepareVariableMDistinctListFilters(Optimi } auto &binder = input.optimizer.binder; - idx_t window_index = binder.GenerateTableIndex(); - auto window = make_uniq(window_index); + idx_t window_index = binder.GenerateTableIndex().index; + auto window = make_uniq(TableIndex(window_index)); vector distinct_filter_indices; for (idx_t ai = 0; ai < agg->expressions.size(); ai++) { auto &aggr = agg->expressions[ai]->Cast(); - string name = StringUtil::Lower(aggr.function.name); - if (name != "count" || !aggr.IsDistinct() || aggr.children.empty()) { + string name = StringUtil::Lower(aggr.FunctionMutable().GetName().GetIdentifierName()); + if (name != "count" || !aggr.IsDistinct() || aggr.GetChildrenMutable().empty()) { continue; } - auto rank = - make_uniq(ExpressionType::WINDOW_RANK_DENSE, LogicalType::BIGINT, nullptr, nullptr); + auto rank = BindRankingWindow(input.context, ExpressionType::WINDOW_RANK_DENSE); for (auto &group : agg->groups) { - rank->partitions.push_back(group->Copy()); + rank->PartitionsMutable().push_back(group->Copy()); } - rank->partitions.push_back(pu_hash_expr.Copy()); - rank->orders.emplace_back(OrderType::ASCENDING, OrderByNullType::NULLS_LAST, aggr.children[0]->Copy()); - rank->start = WindowBoundary::UNBOUNDED_PRECEDING; - rank->end = WindowBoundary::CURRENT_ROW_ROWS; + rank->PartitionsMutable().push_back(pu_hash_expr.Copy()); + rank->OrderByMutable().emplace_back(OrderType::ASCENDING, OrderByNullType::NULLS_LAST, + aggr.GetChildrenMutable()[0]->Copy()); + rank->WindowStartMutable() = WindowBoundary::UNBOUNDED_PRECEDING; + rank->WindowEndMutable() = WindowBoundary::CURRENT_ROW_ROWS; idx_t rank_pos = window->expressions.size(); window->expressions.push_back(std::move(rank)); - auto rank_ref = make_uniq(LogicalType::BIGINT, ColumnBinding(window_index, rank_pos)); + auto rank_ref = make_uniq( + LogicalType::BIGINT, ColumnBinding(TableIndex(window_index), ProjectionIndex(rank_pos))); auto cap_const = make_uniq(Value::BIGINT(count_bound)); - auto in_cap = make_uniq(ExpressionType::COMPARE_LESSTHANOREQUALTO, - std::move(rank_ref), std::move(cap_const)); + auto in_cap = BoundComparisonExpression::Create(ExpressionType::COMPARE_LESSTHANOREQUALTO, std::move(rank_ref), + std::move(cap_const)); vector> predicates; - predicates.push_back(BuildIsNotNullPredicate(aggr.children[0]->Copy())); + predicates.push_back(BuildIsNotNullPredicate(aggr.GetChildrenMutable()[0]->Copy())); predicates.push_back(std::move(in_cap)); filters[ai] = ConjoinPredicates(std::move(predicates)); distinct_filter_indices.push_back(ai); @@ -1358,8 +1370,8 @@ static void ApplyMaxGroupsContributed(OptimizerExtensionInput &input, LogicalAgg spec.cap = max_groups; vector specs; specs.push_back(std::move(spec)); - agg->children[0] = - BuildRankCapFilter(input, std::move(agg->children[0]), pre.lower_agg->types, pre.lower_agg->group_index, specs); + agg->children[0] = BuildRankCapFilter(input, std::move(agg->children[0]), pre.lower_agg->types, + pre.lower_agg->group_index.index, specs); } // Standard (global-sensitivity) clipping. Bounds each privacy unit's total contribution to a @@ -1386,17 +1398,17 @@ static void ApplyMaxGroupsContributed(OptimizerExtensionInput &input, LogicalAgg // Returns the value unchanged when there is no filter. static unique_ptr FoldFilterIntoValue(const BoundAggregateExpression &aggr, unique_ptr value, bool else_zero) { - if (!aggr.filter) { + if (!aggr.GetFilter()) { return value; } - auto value_type = value->return_type; + auto value_type = value->GetReturnType(); auto case_expr = make_uniq(value_type); BoundCaseCheck check; - check.when_expr = aggr.filter->Copy(); + check.when_expr = aggr.GetFilter()->Copy(); check.then_expr = std::move(value); - case_expr->case_checks.push_back(std::move(check)); - case_expr->else_expr = make_uniq(else_zero ? Value::Numeric(value_type, 0) - : Value(value_type)); // 0, or NULL of the type + case_expr->CaseChecksMutable().push_back(std::move(check)); + case_expr->ElseMutable() = make_uniq( + else_zero ? Value::Numeric(value_type, 0) : Value(value_type)); // 0, or NULL of the type return std::move(case_expr); } @@ -1453,33 +1465,36 @@ static vector ApplyPerPuClipping(OptimizerExtensionInput &input, unique_ if (aggr.IsDistinct()) { lower_exprs.push_back(agg->expressions[i]->Copy()); // filtered DISTINCT is rejected upstream } else if (IsCountAggregate(aggr)) { - if (aggr.function.name == "count" && !aggr.children.empty()) { + if (aggr.FunctionMutable().GetName().GetIdentifierName() == "count" && !aggr.GetChildrenMutable().empty()) { lower_exprs.push_back(BindPlainAggregate( - input, "count", FoldFilterIntoValue(aggr, aggr.children[0]->Copy(), /*else_zero=*/false))); + input, "count", + FoldFilterIntoValue(aggr, aggr.GetChildrenMutable()[0]->Copy(), /*else_zero=*/false))); } else { // COUNT(*): no value child, so a FILTER clause binds correctly — attach it directly. lower_exprs.push_back(BindPlainAggregate(input, "count_star", nullptr)); - if (aggr.filter) { - lower_exprs.back()->Cast().filter = aggr.filter->Copy(); + if (aggr.GetFilterMutable()) { + lower_exprs.back()->Cast().GetFilterMutable() = + aggr.GetFilterMutable()->Copy(); } } } else if (IsMinMaxAggregate(aggr)) { // Per-PU MIN/MAX partial. else_zero=false → filtered-out rows become NULL, which MIN/MAX skip. lower_exprs.push_back(BindPlainAggregate( - input, aggr.function.name, FoldFilterIntoValue(aggr, aggr.children[0]->Copy(), /*else_zero=*/false))); + input, aggr.FunctionMutable().GetName().GetIdentifierName(), + FoldFilterIntoValue(aggr, aggr.GetChildrenMutable()[0]->Copy(), /*else_zero=*/false))); } else if (avg_sum_info && avg_sum_info->IsAutomaticBoundsMean()) { - auto value = FoldFilterIntoValue(aggr, aggr.children[0]->Copy(), /*else_zero=*/false); - bounds_parameters[i] = GetDpApproxBoundsParameters(value->return_type); + auto value = FoldFilterIntoValue(aggr, aggr.GetChildrenMutable()[0]->Copy(), /*else_zero=*/false); + bounds_parameters[i] = GetDpApproxBoundsParameters(value->GetReturnType()); vector> children; children.push_back( BoundCastExpression::AddCastToType(input.context, std::move(value), LogicalType::DOUBLE)); children.push_back(make_uniq(Value::BIGINT(automatic_mean_count_bound))); lower_exprs.push_back(BindAggregateLocal(input, "dp_bounded_value_list", std::move(children))); } else { // sum - auto value = FoldFilterIntoValue(aggr, aggr.children[0]->Copy(), /*else_zero=*/true); - bounds_parameters[i] = GetDpApproxBoundsParameters(value->return_type); - if (auto_bounds && (value->return_type.InternalType() == PhysicalType::FLOAT || - value->return_type.InternalType() == PhysicalType::DOUBLE)) { + auto value = FoldFilterIntoValue(aggr, aggr.GetChildrenMutable()[0]->Copy(), /*else_zero=*/true); + bounds_parameters[i] = GetDpApproxBoundsParameters(value->GetReturnType()); + if (auto_bounds && (value->GetReturnType().InternalType() == PhysicalType::FLOAT || + value->GetReturnType().InternalType() == PhysicalType::DOUBLE)) { lower_exprs.push_back(BindPlainAggregate(input, "priv_approx_sum", std::move(value))); } else { lower_exprs.push_back(BindPlainAggregate(input, "sum", std::move(value))); @@ -1496,10 +1511,10 @@ static vector ApplyPerPuClipping(OptimizerExtensionInput &input, unique_ auto &orig = agg->expressions[i]->Cast(); bool pos_min_max = IsMinMaxAggregate(orig); auto *avg_sum_info = FindAvgInfoForSumPos(avg_infos, i); - string min_max_name = pos_min_max ? orig.function.name : string(); + string min_max_name = pos_min_max ? orig.FunctionMutable().GetName().GetIdentifierName() : string(); auto lower_type = pre.lower_agg->types[pre.num_original_groups + 1 + i]; - unique_ptr lower_ref = - make_uniq(lower_type, ColumnBinding(pre.lower_agg_index, i)); + unique_ptr lower_ref = make_uniq( + lower_type, ColumnBinding(TableIndex(pre.lower_agg_index), ProjectionIndex(i))); if (pos_min_max) { // Clip the per-PU partial to [mm_lower, mm_upper], then take the overall MIN/MAX. Removing one // PU shifts the extremum by at most mm_range per group → L1 sensitivity mm_range · group_bound. @@ -1572,7 +1587,8 @@ static unique_ptr BuildSupportKeyExpression(unique_ptrGetColumnIds()[proj_idx]; auto col_type = get->GetColumnType(col_index); - return make_uniq(col_type, ColumnBinding(get->table_index, proj_idx)); + return make_uniq( + col_type, ColumnBinding(TableIndex(get->table_index.index), ProjectionIndex(proj_idx))); } static optional_idx AddSupportAggregate(OptimizerExtensionInput &input, unique_ptr &plan, @@ -1602,8 +1618,9 @@ static optional_idx AddSupportAggregate(OptimizerExtensionInput &input, unique_p static LogicalFilter *ApplySupportFilter(OptimizerExtensionInput &input, unique_ptr &plan, LogicalAggregate *agg, idx_t support_pos, double threshold, double noise_scale = 0.0) { - auto support_ref = make_uniq(agg->types[agg->groups.size() + support_pos], - ColumnBinding(agg->aggregate_index, support_pos)); + auto support_ref = make_uniq( + agg->types[agg->groups.size() + support_pos], + ColumnBinding(TableIndex(agg->aggregate_index.index), ProjectionIndex(support_pos))); unique_ptr support_expr = BoundCastExpression::AddCastToType(input.context, std::move(support_ref), LogicalType::DOUBLE); if (noise_scale > 0.0 && std::isfinite(noise_scale)) { @@ -1613,8 +1630,8 @@ static LogicalFilter *ApplySupportFilter(OptimizerExtensionInput &input, unique_ support_expr = BindScalarLocal(input, "dp_noise", std::move(noise_children)); } auto threshold_const = make_uniq(Value::DOUBLE(threshold)); - auto condition = make_uniq(ExpressionType::COMPARE_GREATERTHANOREQUALTO, - std::move(support_expr), std::move(threshold_const)); + auto condition = BoundComparisonExpression::Create(ExpressionType::COMPARE_GREATERTHANOREQUALTO, + std::move(support_expr), std::move(threshold_const)); auto filter = make_uniq(); filter->expressions.push_back(std::move(condition)); @@ -1675,12 +1692,14 @@ static NoiseProjection InsertProjectionAndRemap(unique_ptr &pla ColumnBindingReplacer replacer; for (idx_t gi = 0; gi < n_groups; gi++) { - replacer.replacement_bindings.emplace_back(ColumnBinding(source_group_idx, gi), ColumnBinding(proj_idx, gi), + replacer.replacement_bindings.emplace_back(ColumnBinding(TableIndex(source_group_idx), ProjectionIndex(gi)), + ColumnBinding(TableIndex(proj_idx), ProjectionIndex(gi)), proj_types[gi]); } for (idx_t ai = 0; ai < n_aggs; ai++) { - replacer.replacement_bindings.emplace_back(ColumnBinding(source_agg_idx, source_agg_offset + ai), - ColumnBinding(proj_idx, n_groups + ai), proj_types[n_groups + ai]); + replacer.replacement_bindings.emplace_back( + ColumnBinding(TableIndex(source_agg_idx), ProjectionIndex(source_agg_offset + ai)), + ColumnBinding(TableIndex(proj_idx), ProjectionIndex(n_groups + ai)), proj_types[n_groups + ai]); } replacer.stop_operator = proj_ptr; replacer.VisitOperator(*plan); @@ -1698,8 +1717,8 @@ static NoiseProjection WrapAggregateWithLaplace(OptimizerExtensionInput &input, const vector &output_types) { idx_t n_groups = agg->groups.size(); idx_t n_aggs = agg_scales.size(); - idx_t group_idx = agg->group_index; - idx_t agg_idx = agg->aggregate_index; + idx_t group_idx = agg->group_index.index; + idx_t agg_idx = agg->aggregate_index.index; auto agg_types = agg->types; D_ASSERT(agg_types.size() >= n_groups + n_aggs); @@ -1721,13 +1740,15 @@ static NoiseProjection WrapAggregateWithLaplace(OptimizerExtensionInput &input, } for (idx_t gi = 0; gi < n_groups; gi++) { - proj_exprs.push_back(make_uniq(agg_types[gi], ColumnBinding(group_idx, gi))); + proj_exprs.push_back(make_uniq( + agg_types[gi], ColumnBinding(TableIndex(group_idx), ProjectionIndex(gi)))); } for (idx_t ai = 0; ai < n_aggs; ai++) { auto agg_col_type = agg_types[n_groups + ai]; const auto &out_type = output_types[ai]; - auto col_ref = make_uniq(agg_col_type, ColumnBinding(agg_idx, ai)); + auto col_ref = + make_uniq(agg_col_type, ColumnBinding(TableIndex(agg_idx), ProjectionIndex(ai))); double scale = agg_scales[ai]; if (!std::isfinite(scale)) { // A non-finite scale means the sensitivity or ε overflowed (e.g. an elastic-sensitivity @@ -1765,14 +1786,15 @@ static NoiseProjection WrapAggregateWithLaplace(OptimizerExtensionInput &input, proj_exprs.push_back(std::move(noised)); } - idx_t proj_idx = input.optimizer.binder.GenerateTableIndex(); - auto projection = make_uniq(proj_idx, std::move(proj_exprs)); + idx_t proj_idx = input.optimizer.binder.GenerateTableIndex().index; + auto projection = make_uniq(TableIndex(proj_idx), std::move(proj_exprs)); return InsertProjectionAndRemap(plan, std::move(projection), insert_above, proj_idx, group_idx, agg_idx, n_groups, n_aggs, 0, "dp: could not locate noise insertion slot in plan"); } static bool IsPlainCountStarAggregate(const BoundAggregateExpression &aggr) { - return aggr.function.name == "count_star" && !aggr.IsDistinct() && aggr.children.empty() && !aggr.filter; + return aggr.Function().GetName().GetIdentifierName() == "count_star" && !aggr.IsDistinct() && + aggr.GetChildren().empty() && !aggr.GetFilter(); } static bool IsDPSassFastCountStarEligible(const LogicalAggregate *agg, const DPFKChain &chain) { @@ -1847,8 +1869,8 @@ WrapAvgRatioProjection(OptimizerExtensionInput &input, unique_ptr(noise_proj.proj_types[gi], ColumnBinding(noise_proj.proj_idx, gi))); + proj_exprs.push_back(make_uniq( + noise_proj.proj_types[gi], ColumnBinding(TableIndex(noise_proj.proj_idx), ProjectionIndex(gi)))); } // For each original agg position: pass-through or compute AVG ratio @@ -1857,15 +1879,15 @@ WrapAvgRatioProjection(OptimizerExtensionInput &input, unique_ptr(col_type, ColumnBinding(noise_proj.proj_idx, n_groups + ai))); + proj_exprs.push_back(make_uniq( + col_type, ColumnBinding(TableIndex(noise_proj.proj_idx), ProjectionIndex(n_groups + ai)))); } else { // AVG: compute CAST(noised_sum AS DOUBLE) / CAST(noised_count AS DOUBLE) const AvgInfo &info = *avg_info; const auto &sum_type = noise_proj.proj_types[n_groups + info.sum_pos]; if (info.IsAutomaticBoundsMean()) { auto mean_ref = make_uniq( - sum_type, ColumnBinding(noise_proj.proj_idx, n_groups + info.sum_pos)); + sum_type, ColumnBinding(TableIndex(noise_proj.proj_idx), ProjectionIndex(n_groups + info.sum_pos))); proj_exprs.push_back( BoundCastExpression::AddCastToType(input.context, std::move(mean_ref), LogicalType::DOUBLE)); continue; @@ -1874,12 +1896,13 @@ WrapAvgRatioProjection(OptimizerExtensionInput &input, unique_ptr( - sum_type, ColumnBinding(noise_proj.proj_idx, n_groups + info.sum_pos)); + sum_type, ColumnBinding(TableIndex(noise_proj.proj_idx), ProjectionIndex(n_groups + info.sum_pos))); return BoundCastExpression::AddCastToType(input.context, std::move(sum_ref), LogicalType::DOUBLE); }; auto make_cnt_dbl = [&]() { auto cnt_ref = make_uniq( - cnt_type, ColumnBinding(noise_proj.proj_idx, n_groups + info.count_pos)); + cnt_type, + ColumnBinding(TableIndex(noise_proj.proj_idx), ProjectionIndex(n_groups + info.count_pos))); return BoundCastExpression::AddCastToType(input.context, std::move(cnt_ref), LogicalType::DOUBLE); }; auto add_midpoint = [&](unique_ptr value) -> unique_ptr { @@ -1902,18 +1925,18 @@ WrapAvgRatioProjection(OptimizerExtensionInput &input, unique_ptr(LogicalType::DOUBLE); BoundCaseCheck check; check.when_expr = - make_uniq(ExpressionType::COMPARE_LESSTHANOREQUALTO, make_cnt_dbl(), - make_uniq(Value::DOUBLE(0.0))); + BoundComparisonExpression::Create(ExpressionType::COMPARE_LESSTHANOREQUALTO, make_cnt_dbl(), + make_uniq(Value::DOUBLE(0.0))); check.then_expr = make_uniq(Value(LogicalType::DOUBLE)); - case_expr->case_checks.push_back(std::move(check)); - case_expr->else_expr = maybe_clamp(add_midpoint(std::move(ratio)), info); + case_expr->CaseChecksMutable().push_back(std::move(check)); + case_expr->ElseMutable() = maybe_clamp(add_midpoint(std::move(ratio)), info); proj_exprs.push_back(std::move(case_expr)); } } } - idx_t ratio_idx = input.optimizer.binder.GenerateTableIndex(); - auto ratio_proj = make_uniq(ratio_idx, std::move(proj_exprs)); + idx_t ratio_idx = input.optimizer.binder.GenerateTableIndex().index; + auto ratio_proj = make_uniq(TableIndex(ratio_idx), std::move(proj_exprs)); auto inserted = InsertProjectionAndRemap(plan, std::move(ratio_proj), insert_above, ratio_idx, noise_proj.proj_idx, noise_proj.proj_idx, n_groups, n_original_aggs, n_groups, "dp: could not locate insert point for AVG ratio projection"); @@ -1922,7 +1945,7 @@ WrapAvgRatioProjection(OptimizerExtensionInput &input, unique_ptr(window_index); - auto row_number = - make_uniq(ExpressionType::WINDOW_ROW_NUMBER, LogicalType::BIGINT, nullptr, nullptr); - row_number->start = WindowBoundary::UNBOUNDED_PRECEDING; - row_number->end = WindowBoundary::CURRENT_ROW_ROWS; + idx_t window_index = input.optimizer.binder.GenerateTableIndex().index; + auto window = make_uniq(TableIndex(window_index)); + auto row_number = BindRankingWindow(input.context, ExpressionType::WINDOW_ROW_NUMBER); + row_number->WindowStartMutable() = WindowBoundary::UNBOUNDED_PRECEDING; + row_number->WindowEndMutable() = WindowBoundary::CURRENT_ROW_ROWS; window->expressions.push_back(std::move(row_number)); window->children.push_back(std::move(*slot)); window->ResolveOperatorTypes(); @@ -2014,15 +2036,17 @@ static NoiseProjection WrapSampleMedianProjection( vector> proj_exprs; proj_exprs.reserve(n_groups + n_aggs); for (idx_t gi = 0; gi < n_groups; gi++) { - proj_exprs.push_back(make_uniq(agg->types[gi], ColumnBinding(agg->group_index, gi))); + proj_exprs.push_back(make_uniq( + agg->types[gi], ColumnBinding(TableIndex(agg->group_index.index), ProjectionIndex(gi)))); } int sample_lanes = GetDpSampleLanes(input.context); for (idx_t ai = 0; ai < n_aggs; ai++) { idx_t agg_pos = agg_positions[ai]; vector> children; - children.push_back(make_uniq(agg->types[n_groups + agg_pos], - ColumnBinding(agg->aggregate_index, agg_pos))); + children.push_back(make_uniq( + agg->types[n_groups + agg_pos], + ColumnBinding(TableIndex(agg->aggregate_index.index), ProjectionIndex(agg_pos)))); if (stability_query_mode) { children.push_back(make_uniq(Value::INTEGER(static_cast(ai)))); unique_ptr recorded = BindScalarLocal(input, "dp_sass_record_stability", std::move(children)); @@ -2059,7 +2083,8 @@ static NoiseProjection WrapSampleMedianProjection( children.push_back(make_uniq(Value::INTEGER(static_cast(ai)))); children.push_back(make_uniq(Value::INTEGER(static_cast(n_aggs)))); children.push_back(make_uniq( - LogicalType::BIGINT, ColumnBinding(noise_identity_index.GetIndex(), 0))); + LogicalType::BIGINT, + ColumnBinding(TableIndex(noise_identity_index.GetIndex()), ProjectionIndex(0)))); terminal = "dp_nonempty_mean_noise"; } } else if (noise_scale_query_mode) { @@ -2076,8 +2101,8 @@ static NoiseProjection WrapSampleMedianProjection( D_ASSERT(noise_identity_index.IsValid()); children.push_back(make_uniq(Value::INTEGER(static_cast(ai)))); children.push_back(make_uniq(Value::INTEGER(static_cast(n_aggs)))); - children.push_back(make_uniq(LogicalType::BIGINT, - ColumnBinding(noise_identity_index.GetIndex(), 0))); + children.push_back(make_uniq( + LogicalType::BIGINT, ColumnBinding(TableIndex(noise_identity_index.GetIndex()), ProjectionIndex(0)))); } unique_ptr noised = BindScalarLocal(input, terminal, std::move(children)); if (output_types[ai] != LogicalType::DOUBLE) { @@ -2086,10 +2111,10 @@ static NoiseProjection WrapSampleMedianProjection( proj_exprs.push_back(std::move(noised)); } - idx_t proj_idx = input.optimizer.binder.GenerateTableIndex(); - auto projection = make_uniq(proj_idx, std::move(proj_exprs)); - return InsertProjectionAndRemap(plan, std::move(projection), insert_above, proj_idx, agg->group_index, - agg->aggregate_index, n_groups, n_aggs, 0, + idx_t proj_idx = input.optimizer.binder.GenerateTableIndex().index; + auto projection = make_uniq(TableIndex(proj_idx), std::move(proj_exprs)); + return InsertProjectionAndRemap(plan, std::move(projection), insert_above, proj_idx, agg->group_index.index, + agg->aggregate_index.index, n_groups, n_aggs, 0, "dp_sample_median: could not locate aggregate slot for median projection"); } @@ -2098,8 +2123,8 @@ static bool IsDpSampleCountDistinct(const LogicalAggregate *agg) { return false; } auto &aggr = agg->expressions[0]->Cast(); - string name = StringUtil::Lower(aggr.function.name); - return aggr.IsDistinct() && (name == "count" || name == "count_star") && aggr.children.size() == 1; + string name = StringUtil::Lower(aggr.FunctionMutable().GetName().GetIdentifierName()); + return aggr.IsDistinct() && (name == "count" || name == "count_star") && aggr.GetChildrenMutable().size() == 1; } static bool IsSinglePuKeyDistinctCount(unique_ptr &plan, LogicalAggregate *agg, @@ -2138,15 +2163,15 @@ static bool IsSinglePuKeyDistinctCount(unique_ptr &plan, Logica } auto col_index = key_get->GetColumnIds()[proj_idx]; auto col_type = key_get->GetColumnType(col_index); - ColumnBinding source_binding(key_get->table_index, proj_idx); - auto propagated = PropagateSingleBinding(*plan, key_get->table_index, source_binding, col_type, agg); - if (propagated.table_index == DConstants::INVALID_INDEX) { + ColumnBinding source_binding {TableIndex(key_get->table_index.index), ProjectionIndex(proj_idx)}; + auto propagated = PropagateSingleBinding(*plan, key_get->table_index.index, source_binding, col_type, agg); + if (propagated.table_index.index == DConstants::INVALID_INDEX) { return false; } auto &aggr = agg->expressions[0]->Cast(); auto key_ref = make_uniq(col_type, propagated); - return Expression::Equals(*aggr.children[0], *key_ref); + return Expression::Equals(*aggr.GetChildrenMutable()[0], *key_ref); } struct DpSassContributionBounds { @@ -2198,7 +2223,7 @@ static DpSassContributionBounds GetDpSassContributionBounds(ClientContext &conte idx_t explicit_sum_count = 0; for (idx_t i = 0; i < agg->expressions.size(); i++) { auto &aggr = agg->expressions[i]->Cast(); - string name = StringUtil::Lower(aggr.function.name); + string name = StringUtil::Lower(aggr.FunctionMutable().GetName().GetIdentifierName()); if (name == "sum" && !FindAvgInfoForSumPos(avg_infos, i)) { explicit_sum_count++; } @@ -2224,7 +2249,7 @@ static DpSassContributionBounds GetDpSassContributionBounds(ClientContext &conte idx_t sum_list_idx = 0; for (idx_t i = 0; i < agg->expressions.size(); i++) { auto &aggr = agg->expressions[i]->Cast(); - string name = StringUtil::Lower(aggr.function.name); + string name = StringUtil::Lower(aggr.FunctionMutable().GetName().GetIdentifierName()); if (FindAvgInfoForSumPos(avg_infos, i)) { RequireDpSassSumBound(context, bounds); bounds.sum_bounds_by_aggregate[i] = bounds.sum_bound; @@ -2252,11 +2277,11 @@ static void RewriteDistinctCountToSampleMedian(OptimizerExtensionInput &input, L int64_t group_bound, bool apply_support_filter, bool variable_m) { auto &binder = input.optimizer.binder; auto &old_aggr = agg->expressions[0]->Cast(); - auto distinct_value = old_aggr.children[0]->Copy(); - auto distinct_type = distinct_value->return_type; + auto distinct_value = old_aggr.GetChildrenMutable()[0]->Copy(); + auto distinct_type = distinct_value->GetReturnType(); - idx_t inner_group_index = binder.GenerateTableIndex(); - idx_t inner_agg_index = binder.GenerateTableIndex(); + idx_t inner_group_index = binder.GenerateTableIndex().index; + idx_t inner_agg_index = binder.GenerateTableIndex().index; vector> inner_expressions; if (!variable_m) { vector> mask_children; @@ -2265,7 +2290,8 @@ static void RewriteDistinctCountToSampleMedian(OptimizerExtensionInput &input, L auto bit_or_expr = BindBitOrAggregate(input, std::move(sample_mask), distinct_value->Copy()); inner_expressions.push_back(std::move(bit_or_expr)); } - auto inner_agg = make_uniq(inner_group_index, inner_agg_index, std::move(inner_expressions)); + auto inner_agg = make_uniq(TableIndex(inner_group_index), TableIndex(inner_agg_index), + std::move(inner_expressions)); idx_t num_original_groups = agg->groups.size(); for (auto &group : agg->groups) { @@ -2277,7 +2303,7 @@ static void RewriteDistinctCountToSampleMedian(OptimizerExtensionInput &input, L inner_agg->ResolveOperatorTypes(); auto *inner_agg_ptr = inner_agg.get(); - idx_t inner_gi = inner_agg_ptr->group_index; + idx_t inner_gi = inner_agg_ptr->group_index.index; unique_ptr distinct_rows = std::move(inner_agg); // Each inner row is one (PU, group, distinct value). Bound a PU's contribution with two @@ -2320,20 +2346,19 @@ static void RewriteDistinctCountToSampleMedian(OptimizerExtensionInput &input, L optional_idx support_window_index; if (apply_support_filter) { - idx_t window_index = binder.GenerateTableIndex(); - auto window = make_uniq(window_index); - auto rn = - make_uniq(ExpressionType::WINDOW_ROW_NUMBER, LogicalType::BIGINT, nullptr, nullptr); + idx_t window_index = binder.GenerateTableIndex().index; + auto window = make_uniq(TableIndex(window_index)); + auto rn = BindRankingWindow(input.context, ExpressionType::WINDOW_ROW_NUMBER); for (idx_t gi = 0; gi < num_original_groups; gi++) { - rn->partitions.push_back(BuildRankSourceRef(inner_agg_ptr->types, inner_gi, gi)); + rn->PartitionsMutable().push_back(BuildRankSourceRef(inner_agg_ptr->types, inner_gi, gi)); } - rn->partitions.push_back(BuildRankSourceRef(inner_agg_ptr->types, inner_gi, pu_group_col)); + rn->PartitionsMutable().push_back(BuildRankSourceRef(inner_agg_ptr->types, inner_gi, pu_group_col)); vector> distinct_hash_children; distinct_hash_children.push_back(BuildRankSourceRef(inner_agg_ptr->types, inner_gi, distinct_col)); auto distinct_hash = BindScalarLocal(input, "hash", std::move(distinct_hash_children)); - rn->orders.emplace_back(OrderType::ASCENDING, OrderByNullType::NULLS_LAST, std::move(distinct_hash)); - rn->start = WindowBoundary::UNBOUNDED_PRECEDING; - rn->end = WindowBoundary::CURRENT_ROW_ROWS; + rn->OrderByMutable().emplace_back(OrderType::ASCENDING, OrderByNullType::NULLS_LAST, std::move(distinct_hash)); + rn->WindowStartMutable() = WindowBoundary::UNBOUNDED_PRECEDING; + rn->WindowEndMutable() = WindowBoundary::CURRENT_ROW_ROWS; window->expressions.push_back(std::move(rn)); window->children.push_back(std::move(distinct_rows)); distinct_rows = std::move(window); @@ -2344,46 +2369,49 @@ static void RewriteDistinctCountToSampleMedian(OptimizerExtensionInput &input, L // A per-PU COUNT(DISTINCT) cannot be summed directly: two PUs in the same lane // may carry the same value. Regroup by (query group, lane, value) to produce one // row per lane-level distinct value, then count those rows into the m-lane list. - idx_t lane_group_index = binder.GenerateTableIndex(); - idx_t lane_agg_index = binder.GenerateTableIndex(); + idx_t lane_group_index = binder.GenerateTableIndex().index; + idx_t lane_agg_index = binder.GenerateTableIndex().index; vector> lane_expressions; if (apply_support_filter) { - auto rn_ref = make_uniq(LogicalType::BIGINT, - ColumnBinding(support_window_index.GetIndex(), 0)); - auto is_first = make_uniq(ExpressionType::COMPARE_EQUAL, std::move(rn_ref), - make_uniq(Value::BIGINT(1))); + auto rn_ref = make_uniq( + LogicalType::BIGINT, ColumnBinding(TableIndex(support_window_index.GetIndex()), ProjectionIndex(0))); + auto is_first = BoundComparisonExpression::Create(ExpressionType::COMPARE_EQUAL, std::move(rn_ref), + make_uniq(Value::BIGINT(1))); auto marker = make_uniq(LogicalType::BIGINT); BoundCaseCheck check; check.when_expr = std::move(is_first); check.then_expr = make_uniq(Value::BIGINT(1)); - marker->case_checks.push_back(std::move(check)); - marker->else_expr = make_uniq(Value::BIGINT(0)); + marker->CaseChecksMutable().push_back(std::move(check)); + marker->ElseMutable() = make_uniq(Value::BIGINT(0)); lane_expressions.push_back(BindPlainAggregate(input, "sum", std::move(marker))); } - auto lane_agg = make_uniq(lane_group_index, lane_agg_index, std::move(lane_expressions)); + auto lane_agg = make_uniq(TableIndex(lane_group_index), TableIndex(lane_agg_index), + std::move(lane_expressions)); for (idx_t gi = 0; gi < num_original_groups; gi++) { - lane_agg->groups.push_back(make_uniq(agg->groups[gi]->return_type, - ColumnBinding(inner_group_index, gi))); + lane_agg->groups.push_back(make_uniq( + agg->groups[gi]->GetReturnType(), ColumnBinding(TableIndex(inner_group_index), ProjectionIndex(gi)))); } auto lane_source = BuildRankSourceRef(inner_agg_ptr->types, inner_gi, pu_group_col); auto lane_mask = make_uniq(Value::UBIGINT(static_cast(GetDpSassM(input.context) - 1))); lane_agg->groups.push_back( input.optimizer.BindScalarFunction("&", std::move(lane_source), std::move(lane_mask))); - lane_agg->groups.push_back( - make_uniq(distinct_type, ColumnBinding(inner_group_index, distinct_col))); + lane_agg->groups.push_back(make_uniq( + distinct_type, ColumnBinding(TableIndex(inner_group_index), ProjectionIndex(distinct_col)))); lane_agg->children.push_back(std::move(distinct_rows)); lane_agg->ResolveOperatorTypes(); for (idx_t gi = 0; gi < num_original_groups; gi++) { - agg->groups[gi] = - make_uniq(lane_agg->types[gi], ColumnBinding(lane_group_index, gi)); + agg->groups[gi] = make_uniq( + lane_agg->types[gi], ColumnBinding(TableIndex(lane_group_index), ProjectionIndex(gi))); } - auto lane_ref = make_uniq(lane_agg->types[num_original_groups], - ColumnBinding(lane_group_index, num_original_groups)); + auto lane_ref = make_uniq( + lane_agg->types[num_original_groups], + ColumnBinding(TableIndex(lane_group_index), ProjectionIndex(num_original_groups))); auto distinct_ref = make_uniq( - lane_agg->types[num_original_groups + 1], ColumnBinding(lane_group_index, num_original_groups + 1)); + lane_agg->types[num_original_groups + 1], + ColumnBinding(TableIndex(lane_group_index), ProjectionIndex(num_original_groups + 1))); LogicalType support_type; if (apply_support_filter) { support_type = lane_agg->types[lane_agg->groups.size()]; @@ -2397,7 +2425,8 @@ static void RewriteDistinctCountToSampleMedian(OptimizerExtensionInput &input, L top_expressions.push_back( BindAggregateLocal(input, "as_sample_m_count_distinct_if", std::move(count_children))); if (apply_support_filter) { - auto support_ref = make_uniq(support_type, ColumnBinding(lane_agg_index, 0)); + auto support_ref = make_uniq( + support_type, ColumnBinding(TableIndex(lane_agg_index), ProjectionIndex(0))); top_expressions.push_back(BindPlainAggregate(input, "sum", std::move(support_ref))); } agg->expressions = std::move(top_expressions); @@ -2405,61 +2434,64 @@ static void RewriteDistinctCountToSampleMedian(OptimizerExtensionInput &input, L return; } - idx_t merge_group_index = binder.GenerateTableIndex(); - idx_t merge_agg_index = binder.GenerateTableIndex(); + idx_t merge_group_index = binder.GenerateTableIndex().index; + idx_t merge_agg_index = binder.GenerateTableIndex().index; vector> merge_expressions; - auto mask_ref = make_uniq(LogicalType::UBIGINT, ColumnBinding(inner_agg_index, 0)); + auto mask_ref = make_uniq(LogicalType::UBIGINT, + ColumnBinding(TableIndex(inner_agg_index), ProjectionIndex(0))); merge_expressions.push_back(BindPlainAggregate(input, "bit_or", std::move(mask_ref))); if (apply_support_filter) { - auto rn_ref = - make_uniq(LogicalType::BIGINT, ColumnBinding(support_window_index.GetIndex(), 0)); - auto is_first = make_uniq(ExpressionType::COMPARE_EQUAL, std::move(rn_ref), - make_uniq(Value::BIGINT(1))); + auto rn_ref = make_uniq( + LogicalType::BIGINT, ColumnBinding(TableIndex(support_window_index.GetIndex()), ProjectionIndex(0))); + auto is_first = BoundComparisonExpression::Create(ExpressionType::COMPARE_EQUAL, std::move(rn_ref), + make_uniq(Value::BIGINT(1))); auto marker = make_uniq(LogicalType::BIGINT); BoundCaseCheck check; check.when_expr = std::move(is_first); check.then_expr = make_uniq(Value::BIGINT(1)); - marker->case_checks.push_back(std::move(check)); - marker->else_expr = make_uniq(Value::BIGINT(0)); + marker->CaseChecksMutable().push_back(std::move(check)); + marker->ElseMutable() = make_uniq(Value::BIGINT(0)); merge_expressions.push_back(BindPlainAggregate(input, "sum", std::move(marker))); } - auto merge_agg = make_uniq(merge_group_index, merge_agg_index, std::move(merge_expressions)); + auto merge_agg = make_uniq(TableIndex(merge_group_index), TableIndex(merge_agg_index), + std::move(merge_expressions)); for (idx_t gi = 0; gi < num_original_groups; gi++) { - merge_agg->groups.push_back( - make_uniq(agg->groups[gi]->return_type, ColumnBinding(inner_group_index, gi))); + merge_agg->groups.push_back(make_uniq( + agg->groups[gi]->GetReturnType(), ColumnBinding(TableIndex(inner_group_index), ProjectionIndex(gi)))); } - merge_agg->groups.push_back( - make_uniq(distinct_type, ColumnBinding(inner_group_index, num_original_groups + 1))); + merge_agg->groups.push_back(make_uniq( + distinct_type, ColumnBinding(TableIndex(inner_group_index), ProjectionIndex(num_original_groups + 1)))); merge_agg->children.push_back(std::move(distinct_rows)); merge_agg->ResolveOperatorTypes(); - idx_t distinct_proj_index = binder.GenerateTableIndex(); + idx_t distinct_proj_index = binder.GenerateTableIndex().index; vector> distinct_proj_exprs; distinct_proj_exprs.reserve(num_original_groups + 1 + (apply_support_filter ? 1 : 0)); for (idx_t gi = 0; gi < num_original_groups; gi++) { - distinct_proj_exprs.push_back( - make_uniq(merge_agg->types[gi], ColumnBinding(merge_group_index, gi))); + distinct_proj_exprs.push_back(make_uniq( + merge_agg->types[gi], ColumnBinding(TableIndex(merge_group_index), ProjectionIndex(gi)))); } if (apply_support_filter) { - distinct_proj_exprs.push_back(make_uniq( - merge_agg->types[merge_agg->groups.size() + 1], ColumnBinding(merge_agg_index, 1))); + distinct_proj_exprs.push_back( + make_uniq(merge_agg->types[merge_agg->groups.size() + 1], + ColumnBinding(TableIndex(merge_agg_index), ProjectionIndex(1)))); } - distinct_proj_exprs.push_back(make_uniq(merge_agg->types[merge_agg->groups.size()], - ColumnBinding(merge_agg_index, 0))); - auto distinct_proj = make_uniq(distinct_proj_index, std::move(distinct_proj_exprs)); + distinct_proj_exprs.push_back(make_uniq( + merge_agg->types[merge_agg->groups.size()], ColumnBinding(TableIndex(merge_agg_index), ProjectionIndex(0)))); + auto distinct_proj = make_uniq(TableIndex(distinct_proj_index), std::move(distinct_proj_exprs)); distinct_proj->children.push_back(std::move(merge_agg)); distinct_proj->ResolveOperatorTypes(); unique_ptr final_input = std::move(distinct_proj); for (idx_t gi = 0; gi < num_original_groups; gi++) { - agg->groups[gi] = - make_uniq(final_input->types[gi], ColumnBinding(distinct_proj_index, gi)); + agg->groups[gi] = make_uniq( + final_input->types[gi], ColumnBinding(TableIndex(distinct_proj_index), ProjectionIndex(gi))); } idx_t support_proj_pos = num_original_groups; idx_t mask_proj_pos = apply_support_filter ? num_original_groups + 1 : num_original_groups; - auto merged_mask_ref = - make_uniq(LogicalType::UBIGINT, ColumnBinding(distinct_proj_index, mask_proj_pos)); + auto merged_mask_ref = make_uniq( + LogicalType::UBIGINT, ColumnBinding(TableIndex(distinct_proj_index), ProjectionIndex(mask_proj_pos))); agg->children[0] = std::move(final_input); vector> top_expressions; // expr[0] = the 64-counter mask, released via the smooth-sensitivity median. When partition @@ -2468,8 +2500,9 @@ static void RewriteDistinctCountToSampleMedian(OptimizerExtensionInput &input, L // the same projection-based path the generic aggregates use, so the noise is keyed on the group. top_expressions.push_back(BindPlainAggregate(input, "as_sample_count_mask", std::move(merged_mask_ref))); if (apply_support_filter) { - auto support_ref = make_uniq(agg->children[0]->types[support_proj_pos], - ColumnBinding(distinct_proj_index, support_proj_pos)); + auto support_ref = make_uniq( + agg->children[0]->types[support_proj_pos], + ColumnBinding(TableIndex(distinct_proj_index), ProjectionIndex(support_proj_pos))); top_expressions.push_back(BindPlainAggregate(input, "sum", std::move(support_ref))); } agg->expressions = std::move(top_expressions); @@ -2482,20 +2515,23 @@ static void RewritePuDistinctCountToSampleMedian(OptimizerExtensionInput &input, // PU rows, every remaining row contributes exactly one count to its own PU lane. auto &binder = input.optimizer.binder; auto &old_aggr = agg->expressions[0]->Cast(); - auto distinct_value = old_aggr.children[0]->Copy(); - auto distinct_type = distinct_value->return_type; + auto distinct_value = old_aggr.GetChildrenMutable()[0]->Copy(); + auto distinct_type = distinct_value->GetReturnType(); - idx_t inner_group_index = binder.GenerateTableIndex(); - idx_t inner_agg_index = binder.GenerateTableIndex(); + idx_t inner_group_index = binder.GenerateTableIndex().index; + idx_t inner_agg_index = binder.GenerateTableIndex().index; vector> inner_expressions; - auto inner_agg = make_uniq(inner_group_index, inner_agg_index, std::move(inner_expressions)); + auto inner_agg = make_uniq(TableIndex(inner_group_index), TableIndex(inner_agg_index), + std::move(inner_expressions)); inner_agg->groups.push_back(std::move(pu_hash_expr)); inner_agg->groups.push_back(std::move(distinct_value)); inner_agg->children.push_back(std::move(agg->children[0])); inner_agg->ResolveOperatorTypes(); - auto hash_ref = make_uniq(LogicalType::UBIGINT, ColumnBinding(inner_group_index, 0)); - auto value_ref = make_uniq(distinct_type, ColumnBinding(inner_group_index, 1)); + auto hash_ref = make_uniq( + LogicalType::UBIGINT, ColumnBinding(TableIndex(inner_group_index), ProjectionIndex(0))); + auto value_ref = make_uniq( + distinct_type, ColumnBinding(TableIndex(inner_group_index), ProjectionIndex(1))); vector> top_expressions; vector> count_children; count_children.push_back(std::move(hash_ref)); @@ -2511,11 +2547,11 @@ static void RewritePuDistinctCountToSampleMedian(OptimizerExtensionInput &input, } static void RemapExpressionBindingsLocal(Expression &e, const std::unordered_map &map) { - if (e.type == ExpressionType::BOUND_COLUMN_REF) { + if (e.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e.Cast(); - auto it = map.find(col_ref.binding.table_index); + auto it = map.find(col_ref.BindingMutable().table_index.index); if (it != map.end()) { - col_ref.binding.table_index = it->second; + col_ref.BindingMutable().table_index.index = it->second; } } ExpressionIterator::EnumerateChildren(e, [&](Expression &child) { RemapExpressionBindingsLocal(child, map); }); @@ -2532,10 +2568,10 @@ static unique_ptr CopyMaybeRemapped(const Expression &expr, static unique_ptr CopyFilterMaybeRemapped(const BoundAggregateExpression &aggr, const std::unordered_map *index_map) { - if (!aggr.filter) { + if (!aggr.GetFilter()) { return nullptr; } - return CopyMaybeRemapped(*aggr.filter, index_map); + return CopyMaybeRemapped(*aggr.GetFilter(), index_map); } struct SampleAggregateRewrite { @@ -2556,7 +2592,7 @@ static void BuildSampleMedianLowerExpression(OptimizerExtensionInput &input, con vector> &lower_expressions, idx_t aggregate_pos, vector &rewrites, bool variable_m, bool nullable_sum, const Expression *extra_filter = nullptr) { - string name = StringUtil::Lower(aggr.function.name); + string name = StringUtil::Lower(aggr.Function().GetName().GetIdentifierName()); string sample_sum_name = nullable_sum ? (variable_m ? "as_sample_m_nullable_sum" : "as_sample_nullable_sum") : (variable_m ? "as_sample_m_sum" : "as_sample_sum"); string sample_avg_name = variable_m ? "as_sample_m_avg" : "as_sample_avg"; @@ -2569,8 +2605,8 @@ static void BuildSampleMedianLowerExpression(OptimizerExtensionInput &input, con rewrites.push_back( {sample_sum_name, lower_value_pos, optional_idx(), true, 0.0, bounds.count_bound, false, 0.0, false}); } else if (name == "count") { - if (!aggr.children.empty()) { - children.push_back(CopyMaybeRemapped(*aggr.children[0], index_map)); + if (!aggr.GetChildren().empty()) { + children.push_back(CopyMaybeRemapped(*aggr.GetChildren()[0], index_map)); } if (variable_m && aggr.IsDistinct()) { lower_expressions.push_back( @@ -2580,7 +2616,7 @@ static void BuildSampleMedianLowerExpression(OptimizerExtensionInput &input, con if (extra_filter) { predicates.push_back(extra_filter->Copy()); } - lower_expressions.back()->Cast().filter = + lower_expressions.back()->Cast().GetFilterMutable() = ConjoinPredicates(std::move(predicates)); rewrites.push_back({"as_sample_m_count_distinct_values", lower_value_pos, optional_idx(), false, 0.0, 0.0, false, 0.0, true}); @@ -2593,23 +2629,23 @@ static void BuildSampleMedianLowerExpression(OptimizerExtensionInput &input, con {sample_sum_name, lower_value_pos, optional_idx(), true, 0.0, bounds.count_bound, false, 0.0, false}); } } else if (name == "sum") { - children.push_back(CopyMaybeRemapped(*aggr.children[0], index_map)); + children.push_back(CopyMaybeRemapped(*aggr.GetChildren()[0], index_map)); lower_expressions.push_back(BindPlainAggregate(input, "sum", std::move(children))); double sum_bound = GetDpSassSumContributionBound(bounds, aggregate_pos); rewrites.push_back( {sample_sum_name, lower_value_pos, optional_idx(), true, -sum_bound, sum_bound, false, 0.0, false}); } else if (name == "avg") { - auto sum_arg = CopyMaybeRemapped(*aggr.children[0], index_map); - auto count_arg = CopyMaybeRemapped(*aggr.children[0], index_map); + auto sum_arg = CopyMaybeRemapped(*aggr.GetChildren()[0], index_map); + auto count_arg = CopyMaybeRemapped(*aggr.GetChildren()[0], index_map); auto filter_for_sum = CopyFilterMaybeRemapped(aggr, index_map); auto filter_for_count = CopyFilterMaybeRemapped(aggr, index_map); lower_expressions.push_back(BindPlainAggregate(input, "sum", std::move(sum_arg))); - lower_expressions.back()->Cast().filter = std::move(filter_for_sum); + lower_expressions.back()->Cast().GetFilterMutable() = std::move(filter_for_sum); idx_t lower_count_pos = lower_expressions.size(); lower_expressions.push_back(BindPlainAggregate(input, "count", std::move(count_arg))); - lower_expressions.back()->Cast().filter = std::move(filter_for_count); + lower_expressions.back()->Cast().GetFilterMutable() = std::move(filter_for_count); rewrites.push_back({sample_avg_name, lower_value_pos, optional_idx(lower_count_pos), true, -bounds.sum_bound, bounds.sum_bound, true, bounds.count_bound, false}); return; @@ -2617,7 +2653,7 @@ static void BuildSampleMedianLowerExpression(OptimizerExtensionInput &input, con // Per-PU MIN/MAX of the value; per-lane sample MIN/MAX then median/mean release. No per-lane // magnitude clip — the release terminal already clamps each lane answer to the public // [dp_sass_minmax_lower_bound, dp_sass_minmax_upper_bound] output domain. - children.push_back(CopyMaybeRemapped(*aggr.children[0], index_map)); + children.push_back(CopyMaybeRemapped(*aggr.GetChildren()[0], index_map)); lower_expressions.push_back(BindPlainAggregate(input, name, std::move(children))); rewrites.push_back({name == "min" ? sample_min_name : sample_max_name, lower_value_pos, optional_idx(), false, 0.0, 0.0, false, 0.0, false}); @@ -2625,7 +2661,8 @@ static void BuildSampleMedianLowerExpression(OptimizerExtensionInput &input, con throw InvalidInputException("dp_sample_median: only COUNT, COUNT(DISTINCT), SUM, AVG, MIN, and MAX aggregates " "are supported"); } - lower_expressions.back()->Cast().filter = CopyFilterMaybeRemapped(aggr, index_map); + lower_expressions.back()->Cast().GetFilterMutable() = + CopyFilterMaybeRemapped(aggr, index_map); } static unique_ptr BuildSampleMedianAggregateRef(OptimizerExtensionInput &input, @@ -2633,7 +2670,7 @@ static unique_ptr BuildSampleMedianAggregateRef(OptimizerExtensionIn const SampleAggregateRewrite &rewrite) { auto lower_type = pre_agg.lower_agg->types[pre_agg.num_original_groups + 1 + rewrite.lower_value_pos]; unique_ptr lower_ref = make_uniq( - lower_type, ColumnBinding(pre_agg.lower_agg_index, rewrite.lower_value_pos)); + lower_type, ColumnBinding(TableIndex(pre_agg.lower_agg_index), ProjectionIndex(rewrite.lower_value_pos))); if (!rewrite.raw_lower_value) { lower_ref = BoundCastExpression::AddCastToType(input.context, std::move(lower_ref), LogicalType::DOUBLE); if (rewrite.clip_value) { @@ -2649,7 +2686,8 @@ static unique_ptr BuildSampleMedianAggregateRef(OptimizerExtensionIn auto lower_count_type = pre_agg.lower_agg->types[pre_agg.num_original_groups + 1 + rewrite.lower_count_pos.GetIndex()]; unique_ptr lower_count_ref = make_uniq( - lower_count_type, ColumnBinding(pre_agg.lower_agg_index, rewrite.lower_count_pos.GetIndex())); + lower_count_type, + ColumnBinding(TableIndex(pre_agg.lower_agg_index), ProjectionIndex(rewrite.lower_count_pos.GetIndex()))); auto count_ref = BoundCastExpression::AddCastToType(input.context, std::move(lower_count_ref), LogicalType::DOUBLE); if (rewrite.clip_count) { @@ -2692,7 +2730,7 @@ static void RewriteAggregateToSampleMedian(OptimizerExtensionInput &input, Logic bool has_variable_m_distinct = false; for (auto &expr : agg->expressions) { auto &aggr = expr->Cast(); - string name = StringUtil::Lower(aggr.function.name); + string name = StringUtil::Lower(aggr.FunctionMutable().GetName().GetIdentifierName()); if (name == "count" && aggr.IsDistinct()) { has_variable_m_distinct = true; break; @@ -2723,7 +2761,8 @@ static void RewriteAggregateToSampleMedian(OptimizerExtensionInput &input, Logic for (idx_t ai = 0; ai < agg->expressions.size(); ai++) { auto &rewrite = rewrites[ai]; auto &original = agg->expressions[ai]->Cast(); - PRIVACY_DEBUG_PRINT("[DP_SAMPLE_MEDIAN] rewrite " + original.function.name + " -> " + rewrite.sample_name); + PRIVACY_DEBUG_PRINT("[DP_SAMPLE_MEDIAN] rewrite " + original.Function().GetName().GetIdentifierName() + " -> " + + rewrite.sample_name); agg->expressions[ai] = BuildSampleMedianAggregateRef(input, pre_agg, rewrite); } agg->ResolveOperatorTypes(); @@ -2770,16 +2809,16 @@ BuildStableDpSassPuHashExpression(OptimizerExtensionInput &input, unique_ptrGetColumnIds()[proj_idx]; auto col_type = key_get->GetColumnType(col_index); - ColumnBinding source_binding(key_get->table_index, proj_idx); - auto propagated = PropagateSingleBinding(*plan, key_get->table_index, source_binding, col_type, agg); - if (propagated.table_index == DConstants::INVALID_INDEX) { + ColumnBinding source_binding {TableIndex(key_get->table_index.index), ProjectionIndex(proj_idx)}; + auto propagated = PropagateSingleBinding(*plan, key_get->table_index.index, source_binding, col_type, agg); + if (propagated.table_index.index == DConstants::INVALID_INDEX) { return nullptr; } key_exprs.push_back(make_uniq(col_type, propagated)); } auto pu_hash_expr = BuildXorHash(input, std::move(key_exprs)); pu_hash_expr = BoundCastExpression::AddCastToType(input.context, std::move(pu_hash_expr), LogicalType::UBIGINT); - pu_hash_expr->alias = "dp_sass_lane_key"; + pu_hash_expr->SetAlias(Identifier("dp_sass_lane_key")); PRIVACY_DEBUG_PRINT("[DP_SAMPLE_MEDIAN] using stable hash-based sample lanes"); return pu_hash_expr; } @@ -3037,8 +3076,8 @@ static unique_ptr BuildFilterlessSamplePredicate(OptimizerExtensionI } auto shift = make_uniq(Value::INTEGER(64 - sample_bits)); auto prefix = input.optimizer.BindScalarFunction(">>", pu_hash.Copy(), std::move(shift)); - return make_uniq(ExpressionType::COMPARE_EQUAL, std::move(prefix), - make_uniq(Value::UBIGINT(0))); + return BoundComparisonExpression::Create(ExpressionType::COMPARE_EQUAL, std::move(prefix), + make_uniq(Value::UBIGINT(0))); } static unique_ptr BuildFilterlessEncodedPu(OptimizerExtensionInput &input, const Expression &pu_hash, @@ -3051,11 +3090,11 @@ static unique_ptr BuildFilterlessEncodedPu(OptimizerExtensionInput & BoundCaseCheck active_check; active_check.when_expr = std::move(active); active_check.then_expr = make_uniq(Value::BOOLEAN(true)); - active_value->case_checks.push_back(std::move(active_check)); - active_value->else_expr = make_uniq(Value::BOOLEAN(false)); + active_value->CaseChecksMutable().push_back(std::move(active_check)); + active_value->ElseMutable() = make_uniq(Value::BOOLEAN(false)); auto active_bit = BoundCastExpression::AddCastToType(input.context, std::move(active_value), LogicalType::UBIGINT); auto encoded = input.optimizer.BindScalarFunction("|", std::move(logical_pu), std::move(active_bit)); - encoded->alias = "dp_filterless_encoded_pu"; + encoded->SetAlias(Identifier("dp_filterless_encoded_pu")); return encoded; } @@ -3078,8 +3117,8 @@ static vector InsertFilterlessProjection(OptimizerExtensionInput expressions.push_back(std::move(expression)); } - idx_t projection_index = input.optimizer.binder.GenerateTableIndex(); - auto projection = make_uniq(projection_index, std::move(expressions)); + idx_t projection_index = input.optimizer.binder.GenerateTableIndex().index; + auto projection = make_uniq(TableIndex(projection_index), std::move(expressions)); auto *slot = FindOperatorSlotByPointer(plan, insert_above); if (!slot) { throw InternalException("dp_filterless: could not locate encoded-PU projection point"); @@ -3091,14 +3130,15 @@ static vector InsertFilterlessProjection(OptimizerExtensionInput ColumnBindingReplacer replacer; for (idx_t i = 0; i < old_bindings.size(); i++) { - replacer.replacement_bindings.emplace_back(old_bindings[i], ColumnBinding(projection_index, i), old_types[i]); + replacer.replacement_bindings.emplace_back( + old_bindings[i], ColumnBinding(TableIndex(projection_index), ProjectionIndex(i)), old_types[i]); } replacer.stop_operator = projection_ptr; replacer.VisitOperator(*plan); vector result; result.reserve(extra_expressions.size()); for (idx_t i = 0; i < extra_expressions.size(); i++) { - result.emplace_back(projection_index, old_bindings.size() + i); + result.emplace_back(TableIndex(projection_index), ProjectionIndex(old_bindings.size() + i)); } return result; } @@ -3159,8 +3199,8 @@ static unique_ptr BuildFilterlessEncodedPuInput(OptimizerExtensionIn auto encoded_bindings = InsertFilterlessProjection(input, plan, projection_child, std::move(encoded_expression)); auto encoded_binding = encoded_bindings[0]; auto propagated = - PropagateSingleBinding(*plan, encoded_binding.table_index, encoded_binding, LogicalType::UBIGINT, agg); - if (propagated.table_index == DConstants::INVALID_INDEX) { + PropagateSingleBinding(*plan, encoded_binding.table_index.index, encoded_binding, LogicalType::UBIGINT, agg); + if (propagated.table_index.index == DConstants::INVALID_INDEX) { throw InternalException("dp_filterless: could not propagate encoded PU to aggregate input"); } return make_uniq(LogicalType::UBIGINT, propagated); @@ -3176,16 +3216,19 @@ struct FilterlessPreAggregationInput { vector aggregate_types; unique_ptr GroupRef(idx_t index) const { - return make_uniq(group_types[index], ColumnBinding(group_table_index, index)); + return make_uniq( + group_types[index], ColumnBinding(TableIndex(group_table_index), ProjectionIndex(index))); } unique_ptr PuRef() const { - return make_uniq(LogicalType::UBIGINT, ColumnBinding(group_table_index, group_count)); + return make_uniq( + LogicalType::UBIGINT, ColumnBinding(TableIndex(group_table_index), ProjectionIndex(group_count))); } unique_ptr AggregateRef(idx_t index) const { return make_uniq( - aggregate_types[index], ColumnBinding(aggregate_table_index, component_column_offset + index)); + aggregate_types[index], + ColumnBinding(TableIndex(aggregate_table_index), ProjectionIndex(component_column_offset + index))); } unique_ptr ComponentRef(OptimizerExtensionInput &input, idx_t index, bool answer) const { @@ -3211,7 +3254,7 @@ struct FilterlessPreAggregationInput { static FilterlessPreAggregationInput BuildFilterlessPreAggregationInput(const PuPreAggregationInfo &pre, idx_t component_count) { FilterlessPreAggregationInput result; - result.group_table_index = pre.lower_agg->group_index; + result.group_table_index = pre.lower_agg->group_index.index; result.aggregate_table_index = pre.lower_agg_index; result.group_count = pre.num_original_groups; result.component_column_offset = 0; @@ -3237,18 +3280,19 @@ static FilterlessPreAggregationInput ApplyFilterlessMaxGroups(OptimizerExtension // caps the group set without needing duplicate-aware DENSE_RANK. DuckDB can optimize this // single hashed ordering into a bounded per-PU top-k aggregate. idx_t lower_aggregate_count = component_count + 1; - idx_t projection_index = input.optimizer.binder.GenerateTableIndex(); + idx_t projection_index = input.optimizer.binder.GenerateTableIndex().index; vector> expressions; expressions.reserve(pre.num_original_groups + 1 + lower_aggregate_count); for (idx_t i = 0; i < pre.num_original_groups + 1; i++) { - expressions.push_back( - make_uniq(pre.lower_agg->types[i], ColumnBinding(pre.lower_agg->group_index, i))); + expressions.push_back(make_uniq( + pre.lower_agg->types[i], ColumnBinding(TableIndex(pre.lower_agg->group_index.index), ProjectionIndex(i)))); } for (idx_t i = 0; i < lower_aggregate_count; i++) { - expressions.push_back(make_uniq(pre.lower_agg->types[pre.num_original_groups + 1 + i], - ColumnBinding(pre.lower_agg_index, i))); + expressions.push_back( + make_uniq(pre.lower_agg->types[pre.num_original_groups + 1 + i], + ColumnBinding(TableIndex(pre.lower_agg_index), ProjectionIndex(i)))); } - auto projection = make_uniq(projection_index, std::move(expressions)); + auto projection = make_uniq(TableIndex(projection_index), std::move(expressions)); projection->children.push_back(std::move(agg->children[0])); projection->ResolveOperatorTypes(); auto projection_types = projection->types; @@ -3271,8 +3315,8 @@ static FilterlessPreAggregationInput ApplyFilterlessMaxGroups(OptimizerExtension agg->children[0]->ResolveOperatorTypes(); for (idx_t i = 0; i < pre.num_original_groups; i++) { - agg->groups[i] = - make_uniq(pre.lower_agg->types[i], ColumnBinding(projection_index, i)); + agg->groups[i] = make_uniq( + pre.lower_agg->types[i], ColumnBinding(TableIndex(projection_index), ProjectionIndex(i))); } FilterlessPreAggregationInput result; result.group_table_index = projection_index; @@ -3294,22 +3338,22 @@ static FilterlessPreAggregationInput ApplyFilterlessMaxGroups(OptimizerExtension static unique_ptr BuildFilterlessActiveSupportKey(OptimizerExtensionInput &input, const FilterlessPreAggregationInput &pre_input) { auto active_count = pre_input.ActiveCountRef(); - auto is_active = make_uniq(ExpressionType::COMPARE_GREATERTHAN, std::move(active_count), - make_uniq(Value::BIGINT(0))); + auto is_active = BoundComparisonExpression::Create(ExpressionType::COMPARE_GREATERTHAN, std::move(active_count), + make_uniq(Value::BIGINT(0))); auto active_pu = make_uniq(LogicalType::UBIGINT); BoundCaseCheck check; check.when_expr = std::move(is_active); check.then_expr = pre_input.PuRef(); - active_pu->case_checks.push_back(std::move(check)); - active_pu->else_expr = make_uniq(Value(LogicalType::UBIGINT)); + active_pu->CaseChecksMutable().push_back(std::move(check)); + active_pu->ElseMutable() = make_uniq(Value(LogicalType::UBIGINT)); return active_pu; } static unique_ptr BuildFilterlessIsActive(OptimizerExtensionInput &input, const Expression &encoded_pu) { auto active_bit = input.optimizer.BindScalarFunction("&", encoded_pu.Copy(), make_uniq(Value::UBIGINT(1))); - return make_uniq(ExpressionType::COMPARE_NOTEQUAL, std::move(active_bit), - make_uniq(Value::UBIGINT(0))); + return BoundComparisonExpression::Create(ExpressionType::COMPARE_NOTEQUAL, std::move(active_bit), + make_uniq(Value::UBIGINT(0))); } static unique_ptr BuildFilterlessLogicalPu(OptimizerExtensionInput &input, @@ -3323,13 +3367,13 @@ static unique_ptr BuildFilterlessLowerPair(OptimizerExtensionInput & const Expression &encoded_pu, int sample_bits) { vector> children; if (is_count) { - if (aggregate.function.name == "count" && !aggregate.children.empty()) { - children.push_back(BuildIsNotNullPredicate(aggregate.children[0]->Copy())); + if (aggregate.Function().GetName().GetIdentifierName() == "count" && !aggregate.GetChildren().empty()) { + children.push_back(BuildIsNotNullPredicate(aggregate.GetChildren()[0]->Copy())); } else { children.push_back(make_uniq(Value::BOOLEAN(true))); } } else { - children.push_back(aggregate.children[0]->Copy()); + children.push_back(aggregate.GetChildren()[0]->Copy()); } children.push_back(BuildFilterlessIsActive(input, encoded_pu)); children.push_back(BuildFilterlessSamplePredicate(input, encoded_pu, sample_bits)); @@ -3348,7 +3392,7 @@ void CompileDPFilterlessQuery(const PrivacyCompatibilityResult &check, Optimizer idx_t n_groups = agg->groups.size(); for (auto &expression : agg->expressions) { auto &aggregate = expression->Cast(); - if (aggregate.filter) { + if (aggregate.GetFilterMutable()) { throw NotImplementedException("dp_filterless: aggregate FILTER clauses are not implemented yet"); } } @@ -3408,7 +3452,7 @@ void CompileDPFilterlessQuery(const PrivacyCompatibilityResult &check, Optimizer BuildFilterlessLowerPair(input, aggregate, is_count, *encoded_pu, filterless_settings.sample_bits)); } auto active_count = BindPlainAggregate(input, "count_star", nullptr); - active_count->Cast().filter = BuildFilterlessIsActive(input, *encoded_pu); + active_count->Cast().GetFilterMutable() = BuildFilterlessIsActive(input, *encoded_pu); lower_expressions.push_back(std::move(active_count)); auto logical_pu = BuildFilterlessLogicalPu(input, std::move(encoded_pu)); @@ -3421,8 +3465,8 @@ void CompileDPFilterlessQuery(const PrivacyCompatibilityResult &check, Optimizer upper_expressions.reserve(original_aggregate_count + (n_groups > 0 ? 1 : 0)); for (idx_t i = 0; i < original_aggregate_count; i++) { auto is_active = - make_uniq(ExpressionType::COMPARE_GREATERTHAN, pre_input.ActiveCountRef(), - make_uniq(Value::BIGINT(0))); + BoundComparisonExpression::Create(ExpressionType::COMPARE_GREATERTHAN, pre_input.ActiveCountRef(), + make_uniq(Value::BIGINT(0))); vector> children; children.push_back(pre_input.PuRef()); children.push_back(std::move(is_active)); @@ -3649,7 +3693,7 @@ void CompileDPSampleMedianQuery(const PrivacyCompatibilityResult &check, Optimiz idx_t sum_count = 0; for (idx_t ai = 0; ai < agg->expressions.size(); ai++) { auto &aggr = agg->expressions[ai]->Cast(); - string name = StringUtil::Lower(aggr.function.name); + string name = StringUtil::Lower(aggr.FunctionMutable().GetName().GetIdentifierName()); if (avg_components.count(ai)) { continue; } @@ -3665,9 +3709,9 @@ void CompileDPSampleMedianQuery(const PrivacyCompatibilityResult &check, Optimiz idx_t sum_output_idx = 0; for (idx_t ai = 0; ai < agg->expressions.size(); ai++) { auto &aggr = agg->expressions[ai]->Cast(); - string name = StringUtil::Lower(aggr.function.name); + string name = StringUtil::Lower(aggr.FunctionMutable().GetName().GetIdentifierName()); bool avg_sum_component = FindAvgInfoForSumPos(avg_infos, ai) != nullptr; - output_types.push_back(name == "avg" ? LogicalType::DOUBLE : aggr.return_type); + output_types.push_back(name == "avg" ? LogicalType::DOUBLE : aggr.GetReturnType()); if (avg_sum_component) { double output_bound = GetRequiredDpBound(input.context, "dp_sass_sum_output_bound", "dp_sass"); lower_bounds.push_back(-output_bound); diff --git a/src/core/privacy_extension.cpp b/src/core/privacy_extension.cpp index b4258c7e..7c54ef48 100644 --- a/src/core/privacy_extension.cpp +++ b/src/core/privacy_extension.cpp @@ -2,6 +2,7 @@ #include "privacy_extension.hpp" #include "duckdb.hpp" +#include "duckdb/main/database_manager.hpp" #include "duckdb/common/exception.hpp" #include "duckdb/function/scalar_function.hpp" #include @@ -31,6 +32,7 @@ #include "table_functions/dp_sass_stability_query.hpp" #include "query_processing/pac_topk_rewriter.hpp" #include "query_processing/pac_avg_rewriter.hpp" +#include "query_processing/as_correlation_optimizer.hpp" #include "privacy_debug.hpp" namespace duckdb { @@ -92,7 +94,7 @@ static string QuoteIdentifier(const string &identifier) { static void SetExtensionSetting(ClientContext &context, const string &name, const Value &value) { auto &config = DBConfig::GetConfig(context); ExtensionOption extension_option; - if (!config.TryGetExtensionOption(name, extension_option)) { + if (!config.TryGetExtensionOption(Identifier(name), extension_option)) { throw InternalException("refresh_dp_stats: extension setting '" + name + "' was not registered"); } auto target_value = value.CastAs(context, extension_option.type); @@ -286,6 +288,8 @@ static void LoadInternal(ExtensionLoader &loader) { #endif } + RegisterASCorrelationOptimizer(db); + // Register PAC optimizer rule auto priv_rewrite_rule = PACRewriteRule(); // attach PAC-specific optimizer info so the extension can coordinate replan state diff --git a/src/core/privacy_optimizer.cpp b/src/core/privacy_optimizer.cpp index 65df410c..b01ce4b0 100644 --- a/src/core/privacy_optimizer.cpp +++ b/src/core/privacy_optimizer.cpp @@ -22,7 +22,7 @@ #include "duckdb/optimizer/optimizer.hpp" // Include DuckDB headers for DROP operation handling #include "duckdb/parser/parsed_data/drop_info.hpp" -#include "duckdb/planner/operator/logical_simple.hpp" +#include "duckdb/planner/operator/logical_drop.hpp" // Include DuckDB headers for CTAS metadata propagation #include "duckdb/planner/operator/logical_create_table.hpp" #include "duckdb/planner/operator/logical_get.hpp" @@ -47,7 +47,7 @@ void PACDropTableRule::PACDropTableRuleFunction(OptimizerExtensionInput &input, return; } // Cast to LogicalSimple to access the parse info (DROP operations use LogicalSimple) - auto &simple = plan->Cast(); + auto &simple = plan->Cast(); if (simple.info->info_type != ParseInfoType::DROP_INFO) { return; } @@ -56,7 +56,7 @@ void PACDropTableRule::PACDropTableRuleFunction(OptimizerExtensionInput &input, if (drop_info.type != CatalogType::TABLE_ENTRY) { return; // Only handle DROP TABLE operations } - string table_name = drop_info.name; + string table_name = drop_info.GetQualifiedName().Name().GetIdentifierName(); // Check if this table has PAC metadata auto &metadata_mgr = PrivacyMetadataManager::Get(); @@ -140,7 +140,7 @@ void PACDropTableRule::PACDropTableRuleFunction(OptimizerExtensionInput &input, static void PropagateCTASMetadata(unique_ptr &outer_plan, unique_ptr &select_plan, OptimizerExtensionInput &input) { auto &create_table = outer_plan->Cast(); - string new_table = create_table.info->Base().table; + string new_table = create_table.info->Base().GetQualifiedName().Name().GetIdentifierName(); auto &mgr = PrivacyMetadataManager::Get(); if (mgr.HasMetadata(new_table)) { return; @@ -167,7 +167,7 @@ static void PropagateCTASMetadata(unique_ptr &outer_plan, uniqu if (!entry) { continue; } - auto *meta = mgr.GetTableMetadata(entry->name); + auto *meta = mgr.GetTableMetadata(entry->name.GetIdentifierName()); if (meta && (meta->is_privacy_unit || meta->derived_pu)) { source_meta = meta; source_get = &get; @@ -187,7 +187,7 @@ static void PropagateCTASMetadata(unique_ptr &outer_plan, uniqu std::unordered_map oid_to_name; if (entry) { for (auto &col : entry->GetColumns().Logical()) { - oid_to_name[col.Oid()] = col.Name(); + oid_to_name[col.Oid()] = col.Name().GetIdentifierName(); } } @@ -215,7 +215,7 @@ static void PropagateCTASMetadata(unique_ptr &outer_plan, uniqu for (auto &new_col : new_columns.Logical()) { auto name_it = get_col_idx_to_name.find(dest_idx); if (name_it != get_col_idx_to_name.end()) { - src_to_dest_name[name_it->second] = new_col.Name(); + src_to_dest_name[name_it->second] = new_col.Name().GetIdentifierName(); } dest_idx++; } @@ -247,7 +247,7 @@ static void PropagateCTASMetadata(unique_ptr &outer_plan, uniqu // fall back to a simple name check against the output column names. std::unordered_set new_col_names; for (auto &new_col : new_columns.Logical()) { - new_col_names.insert(new_col.Name()); + new_col_names.insert(new_col.Name().GetIdentifierName()); } for (auto &prot : source_meta->protected_columns) { if (new_col_names.count(prot) > 0) { @@ -495,7 +495,7 @@ void PACRewriteRule::PACPreOptimizeFunction(OptimizerExtensionInput &input, uniq bool do_diff = false; { Value val; - if (input.context.TryGetCurrentSetting("privacy_diffcols", val) && !val.IsNull()) { + if (input.context.TryGetCurrentSetting(Identifier("privacy_diffcols"), val) && !val.IsNull()) { string val_str = val.ToString(); if (!val_str.empty()) { if (std::isdigit(val_str[0])) { @@ -541,10 +541,10 @@ void PACRewriteRule::PACPreOptimizeFunction(OptimizerExtensionInput &input, uniq string target_table; if (outer_plan->type == LogicalOperatorType::LOGICAL_CREATE_TABLE) { auto &create = outer_plan->Cast(); - target_table = create.info->Base().table; + target_table = create.info->Base().GetQualifiedName().Name().GetIdentifierName(); } else if (outer_plan->type == LogicalOperatorType::LOGICAL_INSERT) { auto &insert = outer_plan->Cast(); - target_table = insert.table.name; + target_table = insert.table.name.GetIdentifierName(); } if (!target_table.empty()) { auto &mgr = PrivacyMetadataManager::Get(); @@ -570,7 +570,7 @@ void PACRewriteRule::PACPreOptimizeFunction(OptimizerExtensionInput &input, uniq col.SetType(plan_types[i]); } if (plan_types[i] == counter_list_type && mut_meta) { - mut_meta->counter_columns.push_back(col.Name()); + mut_meta->counter_columns.push_back(col.Name().GetIdentifierName()); } } } diff --git a/src/diff/pac_utility_diff.cpp b/src/diff/pac_utility_diff.cpp index fa3f0ae4..524deb60 100644 --- a/src/diff/pac_utility_diff.cpp +++ b/src/diff/pac_utility_diff.cpp @@ -60,14 +60,13 @@ static bool IsNumericType(const LogicalType &type) { } // Create a ROW_NUMBER() OVER () window expression for positional matching. -static unique_ptr AddRowNumber(unique_ptr child, idx_t table_index, - ColumnBinding &rn_binding_out) { - auto window = make_uniq(table_index); +static unique_ptr AddRowNumber(ClientContext &context, unique_ptr child, + idx_t table_index, ColumnBinding &rn_binding_out) { + auto window = make_uniq(TableIndex(table_index)); - auto row_number = - make_uniq(ExpressionType::WINDOW_ROW_NUMBER, LogicalType::BIGINT, nullptr, nullptr); + auto row_number = BindRankingWindow(context, ExpressionType::WINDOW_ROW_NUMBER); - rn_binding_out = ColumnBinding(table_index, 0); + rn_binding_out = ColumnBinding(TableIndex(table_index), ProjectionIndex(0)); window->expressions.push_back(std::move(row_number)); window->children.push_back(std::move(child)); @@ -77,7 +76,7 @@ static unique_ptr AddRowNumber(unique_ptr chil // Create an IS NULL check expression. static unique_ptr MakeIsNull(unique_ptr child) { auto op = make_uniq(ExpressionType::OPERATOR_IS_NULL, LogicalType::BOOLEAN); - op->children.push_back(std::move(child)); + op->GetChildrenMutable().push_back(std::move(child)); return std::move(op); } @@ -95,7 +94,8 @@ static vector ExtractColumnNames(LogicalOperator *node) { if (node && node->type == LogicalOperatorType::LOGICAL_PROJECTION) { auto &proj = node->Cast(); for (auto &expr : proj.expressions) { - names.push_back(expr->alias.empty() ? expr->GetName() : expr->alias); + names.push_back(expr->GetAlias().empty() ? expr->GetName().GetIdentifierName() + : expr->GetAlias().GetIdentifierName()); } } return names; @@ -153,13 +153,14 @@ void ApplyUtilityDiff(OptimizerExtensionInput &input, unique_ptrCast(); auto child_bindings = node->children[0]->GetColumnBindings(); for (auto &o : order_op.orders) { - if (o.expression->type == ExpressionType::BOUND_COLUMN_REF) { + if (o.expression->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &ref = o.expression->Cast(); for (idx_t i = 0; i < child_bindings.size(); i++) { - if (child_bindings[i] == ref.binding) { + if (child_bindings[i] == ref.BindingMutable()) { saved_orders.emplace_back( o.type, o.null_order, - make_uniq(ref.return_type, ColumnBinding(0, i))); + make_uniq(ref.GetReturnType(), + ColumnBinding(TableIndex(0), ProjectionIndex(i)))); break; } } @@ -217,7 +218,7 @@ void ApplyUtilityDiff(OptimizerExtensionInput &input, unique_ptr(col_types[i], bindings[i])); } - auto proj = make_uniq(binder.GenerateTableIndex(), std::move(exprs)); + auto proj = make_uniq(TableIndex(binder.GenerateTableIndex().index), std::move(exprs)); proj->children.push_back(std::move(child)); return std::move(proj); }; @@ -231,25 +232,25 @@ void ApplyUtilityDiff(OptimizerExtensionInput &input, unique_ptrGetColumnBindings(); if (num_key_cols == 0) { - pac_side = AddRowNumber(std::move(pac_side), binder.GenerateTableIndex(), pac_rn_binding); - ref_side = AddRowNumber(std::move(ref_side), binder.GenerateTableIndex(), ref_rn_binding); + pac_side = AddRowNumber(input.context, std::move(pac_side), binder.GenerateTableIndex().index, pac_rn_binding); + ref_side = AddRowNumber(input.context, std::move(ref_side), binder.GenerateTableIndex().index, ref_rn_binding); } // ---- Step 2: Create the FULL OUTER JOIN ---- auto join = make_uniq(JoinType::OUTER); if (num_key_cols == 0) { - JoinCondition cond; - cond.comparison = ExpressionType::COMPARE_EQUAL; - cond.left = make_uniq(LogicalType::BIGINT, pac_rn_binding); - cond.right = make_uniq(LogicalType::BIGINT, ref_rn_binding); + JoinCondition cond(nullptr, nullptr, ExpressionType::COMPARE_EQUAL); + + cond.LeftReference() = make_uniq(LogicalType::BIGINT, pac_rn_binding); + cond.RightReference() = make_uniq(LogicalType::BIGINT, ref_rn_binding); join->conditions.push_back(std::move(cond)); } else { for (idx_t i = 0; i < num_key_cols; i++) { - JoinCondition cond; - cond.comparison = ExpressionType::COMPARE_EQUAL; - cond.left = make_uniq(col_types[i], pac_bindings[i]); - cond.right = make_uniq(col_types[i], ref_bindings[i]); + JoinCondition cond(nullptr, nullptr, ExpressionType::COMPARE_EQUAL); + + cond.LeftReference() = make_uniq(col_types[i], pac_bindings[i]); + cond.RightReference() = make_uniq(col_types[i], ref_bindings[i]); join->conditions.push_back(std::move(cond)); } } @@ -283,16 +284,16 @@ void ApplyUtilityDiff(OptimizerExtensionInput &input, unique_ptr(probe_type, pac_probe_binding)); check.then_expr = make_uniq(Value(col_types[i])); - case_expr->case_checks.push_back(std::move(check)); + case_expr->CaseChecksMutable().push_back(std::move(check)); } { BoundCaseCheck check; check.when_expr = MakeIsNull(make_uniq(probe_type, ref_probe_binding)); check.then_expr = make_uniq(col_types[i], pac_bindings[i]); - case_expr->case_checks.push_back(std::move(check)); + case_expr->CaseChecksMutable().push_back(std::move(check)); } - case_expr->else_expr = make_uniq(col_types[i], ref_bindings[i]); - case_expr->alias = col_names[i]; + case_expr->ElseMutable() = make_uniq(col_types[i], ref_bindings[i]); + case_expr->SetAlias(Identifier(col_names[i])); new_expressions.push_back(std::move(case_expr)); } else if (is_numeric) { // Numeric measure columns: @@ -318,7 +319,7 @@ void ApplyUtilityDiff(OptimizerExtensionInput &input, unique_ptr(probe_type, ref_probe_binding)); check.then_expr = make_uniq(Value(col_types[i])); - case_expr->case_checks.push_back(std::move(check)); + case_expr->CaseChecksMutable().push_back(std::move(check)); } // "-" row → 0 (missing from PAC result = zero utility) @@ -327,7 +328,7 @@ void ApplyUtilityDiff(OptimizerExtensionInput &input, unique_ptr(probe_type, pac_probe_binding)); check.then_expr = BoundCastExpression::AddDefaultCastToType( make_uniq(Value::DOUBLE(0.0)), col_types[i]); - case_expr->case_checks.push_back(std::move(check)); + case_expr->CaseChecksMutable().push_back(std::move(check)); } // Both NULL → 0 (perfect match) @@ -339,7 +340,7 @@ void ApplyUtilityDiff(OptimizerExtensionInput &input, unique_ptr(Value::DOUBLE(0.0)), col_types[i]); - case_expr->case_checks.push_back(std::move(check)); + case_expr->CaseChecksMutable().push_back(std::move(check)); } // One NULL → 50 @@ -351,7 +352,7 @@ void ApplyUtilityDiff(OptimizerExtensionInput &input, unique_ptr(Value::DOUBLE(50.0)), col_types[i]); - case_expr->case_checks.push_back(std::move(check)); + case_expr->CaseChecksMutable().push_back(std::move(check)); } // ELSE: CAST(100 * |ref - pac| / max(0.00001, |ref|) AS original_type) @@ -360,12 +361,12 @@ void ApplyUtilityDiff(OptimizerExtensionInput &input, unique_ptr(LogicalType::DOUBLE); { BoundCaseCheck check; - check.when_expr = make_uniq(ExpressionType::COMPARE_GREATERTHANOREQUALTO, - make_ref_dbl(), make_pac_dbl()); + check.when_expr = BoundComparisonExpression::Create(ExpressionType::COMPARE_GREATERTHANOREQUALTO, + make_ref_dbl(), make_pac_dbl()); check.then_expr = optimizer.BindScalarFunction("-", make_ref_dbl(), make_pac_dbl()); - abs_diff->case_checks.push_back(std::move(check)); + abs_diff->CaseChecksMutable().push_back(std::move(check)); } - abs_diff->else_expr = optimizer.BindScalarFunction("-", make_pac_dbl(), make_ref_dbl()); + abs_diff->ElseMutable() = optimizer.BindScalarFunction("-", make_pac_dbl(), make_ref_dbl()); // max(0.00001, |ref|): // CASE WHEN ref >= 0.00001 THEN ref @@ -374,32 +375,32 @@ void ApplyUtilityDiff(OptimizerExtensionInput &input, unique_ptr(LogicalType::DOUBLE); { BoundCaseCheck check; - check.when_expr = make_uniq( - ExpressionType::COMPARE_GREATERTHANOREQUALTO, make_ref_dbl(), - make_uniq(Value::DOUBLE(0.00001))); + check.when_expr = + BoundComparisonExpression::Create(ExpressionType::COMPARE_GREATERTHANOREQUALTO, make_ref_dbl(), + make_uniq(Value::DOUBLE(0.00001))); check.then_expr = make_ref_dbl(); - denom->case_checks.push_back(std::move(check)); + denom->CaseChecksMutable().push_back(std::move(check)); } { BoundCaseCheck check; - check.when_expr = make_uniq( - ExpressionType::COMPARE_LESSTHANOREQUALTO, make_ref_dbl(), - make_uniq(Value::DOUBLE(-0.00001))); + check.when_expr = + BoundComparisonExpression::Create(ExpressionType::COMPARE_LESSTHANOREQUALTO, make_ref_dbl(), + make_uniq(Value::DOUBLE(-0.00001))); check.then_expr = optimizer.BindScalarFunction( "-", make_uniq(Value::DOUBLE(0.0)), make_ref_dbl()); - denom->case_checks.push_back(std::move(check)); + denom->CaseChecksMutable().push_back(std::move(check)); } - denom->else_expr = make_uniq(Value::DOUBLE(0.00001)); + denom->ElseMutable() = make_uniq(Value::DOUBLE(0.00001)); // 100 * abs_diff / denom auto hundred = make_uniq(Value::DOUBLE(100.0)); auto numerator = optimizer.BindScalarFunction("*", std::move(hundred), std::move(abs_diff)); auto result = optimizer.BindScalarFunction("/", std::move(numerator), std::move(denom)); - case_expr->else_expr = BoundCastExpression::AddDefaultCastToType(std::move(result), col_types[i]); + case_expr->ElseMutable() = BoundCastExpression::AddDefaultCastToType(std::move(result), col_types[i]); } - case_expr->alias = col_names[i]; + case_expr->SetAlias(Identifier(col_names[i])); new_expressions.push_back(std::move(case_expr)); } else { // Non-numeric, non-key columns: @@ -411,35 +412,37 @@ void ApplyUtilityDiff(OptimizerExtensionInput &input, unique_ptr(probe_type, ref_probe_binding)); check.then_expr = make_uniq(Value(col_types[i])); - case_expr->case_checks.push_back(std::move(check)); + case_expr->CaseChecksMutable().push_back(std::move(check)); } { BoundCaseCheck check; check.when_expr = MakeIsNull(make_uniq(probe_type, pac_probe_binding)); check.then_expr = make_uniq(Value(col_types[i])); - case_expr->case_checks.push_back(std::move(check)); + case_expr->CaseChecksMutable().push_back(std::move(check)); } auto coalesce = make_uniq(ExpressionType::OPERATOR_COALESCE, col_types[i]); - coalesce->children.push_back(make_uniq(col_types[i], pac_bindings[i])); - coalesce->children.push_back(make_uniq(col_types[i], ref_bindings[i])); - case_expr->else_expr = std::move(coalesce); - case_expr->alias = col_names[i]; + coalesce->GetChildrenMutable().push_back( + make_uniq(col_types[i], pac_bindings[i])); + coalesce->GetChildrenMutable().push_back( + make_uniq(col_types[i], ref_bindings[i])); + case_expr->ElseMutable() = std::move(coalesce); + case_expr->SetAlias(Identifier(col_names[i])); new_expressions.push_back(std::move(case_expr)); } } // ---- Step 4: Create a new top-level diff projection ---- - idx_t diff_proj_idx = binder.GenerateTableIndex(); - plan = make_uniq(diff_proj_idx, std::move(new_expressions)); + idx_t diff_proj_idx = binder.GenerateTableIndex().index; + plan = make_uniq(TableIndex(diff_proj_idx), std::move(new_expressions)); plan->children.push_back(std::move(join)); // ---- Step 4b: Re-apply ORDER BY if the original query had one ---- if (!saved_orders.empty()) { for (auto &o : saved_orders) { auto &ref = o.expression->Cast(); - idx_t col_idx = ref.binding.column_index; - ref.binding = ColumnBinding(diff_proj_idx, col_idx); - ref.return_type = col_types[col_idx]; + idx_t col_idx = ref.BindingMutable().column_index; + ref.BindingMutable() = ColumnBinding(TableIndex(diff_proj_idx), ProjectionIndex(col_idx)); + ref.SetReturnType(col_types[col_idx]); } auto order = make_uniq(std::move(saved_orders)); order->children.push_back(std::move(plan)); diff --git a/src/include/aggregates/aggregate_state_export.hpp b/src/include/aggregates/aggregate_state_export.hpp new file mode 100644 index 00000000..ffd33298 --- /dev/null +++ b/src/include/aggregates/aggregate_state_export.hpp @@ -0,0 +1,111 @@ +#pragma once + +#include "aggregates/as_aggregate.hpp" +#include "duckdb/common/vector/string_vector.hpp" + +namespace duckdb { + +// Cold-path state transport shared by the stochastic aggregate implementations. +class AggregateStateWriter { +public: + string data; + + template + void Write(const T &value) { + WriteData(&value, sizeof(value)); + } + + void WriteData(const void *ptr, idx_t size) { + data.append(reinterpret_cast(ptr), size); + } +}; + +class AggregateStateReader { +public: + explicit AggregateStateReader(const string_t &blob) : ptr(blob.GetData()), remaining(blob.GetSize()) { + } + + template + void Read(T &value) { + ReadData(&value, sizeof(value)); + } + + void ReadData(void *value, idx_t size) { + if (remaining < size) { + throw InvalidInputException("Truncated stochastic aggregate state"); + } + memcpy(value, ptr, size); + ptr += size; + remaining -= size; + } + + bool ReadFlag() { + uint8_t flag; + Read(flag); + if (flag > 1) { + throw InvalidInputException("Invalid stochastic aggregate state flag"); + } + return flag != 0; + } + + void Finish() { + if (remaining) { + throw InvalidInputException("Unexpected trailing stochastic aggregate state data"); + } + } + +private: + const char *ptr; + idx_t remaining; +}; + +template +struct AggregateStateExport { + static AggregateStateLayout Layout(AggregateLayoutInput &input) { + auto layout = AggregateStateLayout(LogicalType::BLOB, sizeof(STATE)); + auto &arguments = input.function.GetLogicalArguments(); + if (arguments.size() == 3 && input.bind_data) { + layout.constant_parameters.emplace(2, Value::DOUBLE(input.bind_data->Cast().correction)); + } + return layout; + } + + static void Export(Vector &states, AggregateFinalizeInputData &, Vector &result, idx_t count, idx_t offset) { + UnifiedVectorFormat format; + states.ToUnifiedFormat(format); + auto addresses = UnifiedVectorFormat::GetData(format); + auto writer = FlatVector::Writer(result, count, offset); + for (idx_t i = 0; i < count; i++) { + AggregateStateWriter out; + CODEC::Write(out, *addresses[format.sel->get_index(i)]); + writer.WriteValue(StringVector::AddStringOrBlob(result, out.data)); + } + } + + static void Import(AggregateImportInputData &input) { + auto values = input.input_vec.Values(); + for (idx_t i = 0; i < input.input_vec.size(); i++) { + auto value = values[i]; + if (!value.IsValid()) { + // Keep the caller-initialized empty state. finalize propagates NULL, + // while combine treats one NULL input as the identity. + continue; + } + AggregateStateReader in(value.GetValue()); + auto &state = *reinterpret_cast(input.dest_buffer + i * input.layout.total_state_size); + // Other aggregates can leave this shared arena at an odd byte offset. + // Buffered stochastic states store their count in the pointer's low bits. + input.allocator.AlignNext(); + CODEC::Read(in, state, input.allocator); + in.Finish(); + } + } + + static void Register(BaseAggregateFunction &function) { + function.GetCallbacks().get_state_type = Layout; + function.SetExportAggregateStateCallback(Export); + function.SetImportAggregateStateCallback(Import); + } +}; + +} // namespace duckdb diff --git a/src/include/aggregates/as_aggregate.hpp b/src/include/aggregates/as_aggregate.hpp index 1dc10ee9..e062f81a 100644 --- a/src/include/aggregates/as_aggregate.hpp +++ b/src/include/aggregates/as_aggregate.hpp @@ -3,8 +3,10 @@ #include "duckdb.hpp" #include "duckdb/common/random_engine.hpp" +#include "duckdb/common/clustered_aggregate.hpp" #include "duckdb/execution/expression_executor.hpp" #include "duckdb/transaction/meta_transaction.hpp" +#include "pcg_random.hpp" #include #include #include @@ -48,7 +50,8 @@ static inline int ValidateDpSampleLanes(int64_t sample_lanes) { static inline int GetDpSampleLanes(duckdb::ClientContext &ctx) { duckdb::Value sample_lanes_val; - if (ctx.TryGetCurrentSetting("dp_sample_lanes", sample_lanes_val) && !sample_lanes_val.IsNull()) { + if (ctx.TryGetCurrentSetting(duckdb::Identifier("dp_sample_lanes"), sample_lanes_val) && + !sample_lanes_val.IsNull()) { return ValidateDpSampleLanes(sample_lanes_val.GetValue()); } return DP_SAMPLE_DEFAULT_LANES; @@ -63,7 +66,7 @@ static inline int ValidateDpSassM(int64_t m) { static inline int GetDpSassM(duckdb::ClientContext &ctx) { duckdb::Value m_val; - if (ctx.TryGetCurrentSetting("dp_sass_m", m_val) && !m_val.IsNull()) { + if (ctx.TryGetCurrentSetting(duckdb::Identifier("dp_sass_m"), m_val) && !m_val.IsNull()) { return ValidateDpSassM(m_val.GetValue()); } return DP_SASS_DEFAULT_M; @@ -87,7 +90,7 @@ static inline double DpSampleRescale(int sample_lanes) { static inline bool GetDpSassRescale(duckdb::ClientContext &ctx) { duckdb::Value rescale_val; - if (ctx.TryGetCurrentSetting("dp_sass_rescale", rescale_val) && !rescale_val.IsNull()) { + if (ctx.TryGetCurrentSetting(duckdb::Identifier("dp_sass_rescale"), rescale_val) && !rescale_val.IsNull()) { return rescale_val.GetValue(); } return true; @@ -102,7 +105,7 @@ static inline int ValidatePacM(int64_t m) { static inline int GetPacSampleM(duckdb::ClientContext &ctx) { duckdb::Value m_val; - if (ctx.TryGetCurrentSetting("pac_m", m_val) && !m_val.IsNull()) { + if (ctx.TryGetCurrentSetting(duckdb::Identifier("pac_m"), m_val) && !m_val.IsNull()) { return ValidatePacM(m_val.GetValue()); } return DP_SASS_DEFAULT_M; @@ -110,7 +113,7 @@ static inline int GetPacSampleM(duckdb::ClientContext &ctx) { static inline bool GetPacHashRepair(duckdb::ClientContext &ctx) { duckdb::Value hash_repair_val; - if (ctx.TryGetCurrentSetting("pac_hash_repair", hash_repair_val) && !hash_repair_val.IsNull()) { + if (ctx.TryGetCurrentSetting(duckdb::Identifier("pac_hash_repair"), hash_repair_val) && !hash_repair_val.IsNull()) { return hash_repair_val.GetValue(); } return true; @@ -184,6 +187,54 @@ static inline int pac_ctzll(uint64_t x) { namespace duckdb { +// Preserve the single-state kernels for ungrouped input and selected group runs. +template +void ExecutePrivClusters(Vector inputs[], idx_t input_count, const ClusteredAggr &clusters, UPDATE update) { + vector selected; + sel_t selection_data[STANDARD_VECTOR_SIZE]; + SelectionVector selection(selection_data, STANDARD_VECTOR_SIZE); + for (auto &run : clusters.runs()) { + if (!run.sel) { + update(inputs, run.state, run.count); + continue; + } + if (selected.empty()) { + selected.reserve(input_count); + for (idx_t col = 0; col < input_count; col++) { + selected.emplace_back(Vector::Ref(inputs[col])); + } + } + for (idx_t row = 0; row < run.count; row++) { + selection.set_index(row, run.sel[row]); + } + for (idx_t col = 0; col < input_count; col++) { + selected[col].Slice(inputs[col], selection, run.count); + } + update(selected.data(), run.state, run.count); + } +} + +// DuckDB's PCG32 core, exposed as a 64-bit URBG for the standard distributions. +class PacRandomEngine { +public: + using result_type = uint64_t; + explicit PacRandomEngine(uint64_t seed, uint64_t stream) : engine(seed, stream) { + } + static constexpr result_type min() { + return 0; + } + static constexpr result_type max() { + return ~result_type(0); + } + result_type operator()() { + const auto high = engine(); + return (result_type(high) << 32) | engine(); + } + +private: + pcg32 engine; +}; + // Returns the DuckDB LogicalType corresponding to PAC_FLOAT (FLOAT or DOUBLE). static inline LogicalType PacFloatLogicalType() { return std::is_same::value ? LogicalType::FLOAT : LogicalType::DOUBLE; @@ -237,7 +288,7 @@ void RegisterPacFinalizeFunction(ExtensionLoader &loader); // correction: factor to multiply values by after compacting NULLs but before adding noise // pstate: optional shared p-tracking state for persistent secret composition (may be nullptr) // out_noise_variance: if non-null, receives the noise variance used (for utility NULLing) -PAC_FLOAT PacNoisySampleFrom64Counters(const PAC_FLOAT counters[64], double mi, double correction, std::mt19937_64 &gen, +PAC_FLOAT PacNoisySampleFrom64Counters(const PAC_FLOAT counters[64], double mi, double correction, PacRandomEngine &gen, uint64_t is_null = 0, uint64_t counter_selector = 0, const std::shared_ptr &pstate = nullptr, double *out_noise_variance = nullptr); @@ -253,7 +304,7 @@ static inline bool PacNoisedSelect(uint64_t key_hash, uint64_t rnd) { // correction: reduces NULL probability by this factor (considers correction times more non-nulls) // Probabilistic: P(NULL) = popcount(~key_hash) / (64 * correction) // Deterministic: NULL when popcount(key_hash) * correction < 1 -bool PacNoiseInNull(uint64_t key_hash, double mi, double correction, std::mt19937_64 &gen); +bool PacNoiseInNull(uint64_t key_hash, double mi, double correction, PacRandomEngine &gen); // PacUtilityNull: post-processing NULLing of low-SNR cells. // Returns true (NULL the cell) probabilistically based on z-score = |noised_value| / noise_std_dev. @@ -261,7 +312,14 @@ bool PacNoiseInNull(uint64_t key_hash, double mi, double correction, std::mt1993 // Safe by data processing inequality: this is a function of the already-noised output + independent coins. // threshold: z-score center (e.g., 4.0 means ~20% expected relative error cutoff). NaN = disabled. // noise_variance: the PAC noise variance used for this cell (delta = sigma^2 / (2*mi)) -bool PacUtilityNull(double noised_value, double noise_variance, double threshold, std::mt19937_64 &gen); +bool PacUtilityNull(double noised_value, double noise_variance, double threshold, PacRandomEngine &gen); + +// Streams advance across cells and reset for each query execution. Fixed seeds +// reproduce draws for the same execution order, not after row/worker reordering. +struct PacRandomState; +shared_ptr GetPacRandomState(ClientContext &context); +uint64_t GetPacSeed(ClientContext &context, double mi); +PacRandomEngine PacRandomStream(uint64_t seed, PacRandomState &state); // Minimum total rows across all groups before the diversity check applies // With hash32_32 (hash_repair), single-PU groups always have exactly 32 zero bits @@ -284,20 +342,20 @@ uint64_t PacGenerateRandomSeed(); // Returns 1.0/128 if not found or null. inline double GetPacMiFromSetting(ClientContext &ctx) { Value pac_mi_val; - if (ctx.TryGetCurrentSetting("pac_mi", pac_mi_val) && !pac_mi_val.IsNull()) { + if (ctx.TryGetCurrentSetting(Identifier("pac_mi"), pac_mi_val) && !pac_mi_val.IsNull()) { return pac_mi_val.GetValue(); } return 1.0 / 128; // default } // Bind data used by privacy aggregates to carry `mi`, `correction`, and sampling parameters. -// Reads seed from privacy_seed setting (or uses query-id if not set) and computes query_hash. +// Reads an explicit privacy_seed, otherwise mixes in query-id when noise is enabled. // query_hash is XOR'd with per-row key_hash inside priv_hash() (centralized) and used as // the counter selector for PacNoisySampleFrom64Counters in finalize functions. struct PrivBindData : public FunctionData { double mi; // mutual information parameter from pac_mi setting (controls noise/NULL probability) double correction; // correction factor: multiplies sum/avg/count results, reduces NULL prob for min/max - uint64_t seed; // RNG seed: privacy_seed setting value, or query-id if not set + uint64_t seed; // RNG seed: explicit privacy_seed, otherwise default seed mixed with query-id uint64_t query_hash; // derived from seed: used inside priv_hash() for XOR and as counter selector double scale_divisor; // for DECIMAL as_avg: divide result by 10^scale (default 1.0) bool hash_repair; // if true, priv_hash() repairs hash to exactly 32 bits set @@ -306,6 +364,7 @@ struct PrivBindData : public FunctionData { int sample_count; // SASS mode: total number of sample lanes (m); 64 uses the SIMD bitmask path bool sample_rescale; // SASS mode: rescale SUM/COUNT samples to full-dataset-estimator scale double utility_threshold; // z-score threshold for utility NULLing (NaN = disabled, any value = enabled) + shared_ptr random_state; // query-local stream allocation, shared across aggregate bindings // Persistent secret p-tracking: shared across all aggregates in the same query (same query_hash). // When active (mi > 0 and pac_ptracking enabled), noise calibration uses p-weighted variance @@ -317,7 +376,7 @@ struct PrivBindData : public FunctionData { mutable uint64_t suspicious_count; // groups with [29..35] zero bits AND all-same-value mutable uint64_t nonsuspicious_count; // groups that are not suspicious - // Primary constructor - reads seed from privacy_seed setting, or uses query-id if not set. + // Primary constructor - an explicit privacy_seed is stable across queries. // All aggregates in the same query get the same seed and query_hash. explicit PrivBindData(ClientContext &ctx, double mi_val, double correction_val = 1.0, double scale_div = 1.0, bool hash_repair_val = false, int sample_lanes_val = 0, @@ -325,32 +384,24 @@ struct PrivBindData : public FunctionData { : mi(mi_val), correction(correction_val), scale_divisor(scale_div), hash_repair(hash_repair_val), sample_diversity_check(true), sample_lanes(sample_lanes_val), sample_count(sample_count_val), sample_rescale(sample_rescale_val), utility_threshold(std::numeric_limits::quiet_NaN()), - total_update_count(0), suspicious_count(0), nonsuspicious_count(0) { + random_state(GetPacRandomState(ctx)), total_update_count(0), suspicious_count(0), nonsuspicious_count(0) { Value sd_val; - if (ctx.TryGetCurrentSetting("pac_sample_diversity_check", sd_val) && !sd_val.IsNull()) { + if (ctx.TryGetCurrentSetting(Identifier("pac_sample_diversity_check"), sd_val) && !sd_val.IsNull()) { sample_diversity_check = sd_val.GetValue(); } // Read utility threshold: if set (non-null), enables probabilistic NULLing of low-SNR cells Value ut_val; - if (ctx.TryGetCurrentSetting("privacy_min_group_count", ut_val) && !ut_val.IsNull()) { + if (ctx.TryGetCurrentSetting(Identifier("privacy_min_group_count"), ut_val) && !ut_val.IsNull()) { utility_threshold = ut_val.GetValue(); } - Value pac_seed_val; - if (ctx.TryGetCurrentSetting("privacy_seed", pac_seed_val) && !pac_seed_val.IsNull()) { - seed = uint64_t(pac_seed_val.GetValue()); - } else { - seed = 42; - } - if (mi != 0.0) { // randomize seed per query (not in deterministic mode aka mi==0) - seed ^= PAC_MAGIC_HASH * static_cast(ctx.ActiveTransaction().GetActiveQuery()); - } + seed = GetPacSeed(ctx, mi); query_hash = (seed * PAC_MAGIC_HASH) ^ PAC_MAGIC_HASH; // Initialize p-tracking if mi > 0 and pac_ptracking is enabled if (mi > 0.0) { Value pt_val; bool ptracking_enabled = true; // default: enabled - if (ctx.TryGetCurrentSetting("pac_ptracking", pt_val) && !pt_val.IsNull()) { + if (ctx.TryGetCurrentSetting(Identifier("pac_ptracking"), pt_val) && !pt_val.IsNull()) { ptracking_enabled = pt_val.GetValue(); } if (ptracking_enabled) { diff --git a/src/include/aggregates/as_clip_aggr.hpp b/src/include/aggregates/as_clip_aggr.hpp index 323b11bd..e96d5c2c 100644 --- a/src/include/aggregates/as_clip_aggr.hpp +++ b/src/include/aggregates/as_clip_aggr.hpp @@ -158,29 +158,32 @@ static unique_ptr PacClipBindWithScale(ClientContext &ctx, vector< } int clip_support = 0; Value dc_val; - if (ctx.TryGetCurrentSetting("priv_clip_support", dc_val) && !dc_val.IsNull()) { + if (ctx.TryGetCurrentSetting(Identifier("priv_clip_support"), dc_val) && !dc_val.IsNull()) { clip_support = static_cast(dc_val.GetValue()); } bool clip_scale_val = false; Value cs_val; - if (ctx.TryGetCurrentSetting("priv_clip_scale", cs_val) && !cs_val.IsNull()) { + if (ctx.TryGetCurrentSetting(Identifier("priv_clip_scale"), cs_val) && !cs_val.IsNull()) { clip_scale_val = cs_val.GetValue(); } return make_uniq(ctx, mi, correction, clip_support, float_scale, clip_scale_val); } -static unique_ptr PacClipBind(ClientContext &ctx, AggregateFunction &, - vector> &args) { +static unique_ptr PacClipBind(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &args = bind_input.GetArguments(); return PacClipBindWithScale(ctx, args, "as_clip"); } -static unique_ptr PacClipBindFloat(ClientContext &ctx, AggregateFunction &, - vector> &args) { +static unique_ptr PacClipBindFloat(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &args = bind_input.GetArguments(); return PacClipBindWithScale(ctx, args, "as_clip", CLIP_FLOAT_SCALE); } -static unique_ptr PacClipBindDouble(ClientContext &ctx, AggregateFunction &, - vector> &args) { +static unique_ptr PacClipBindDouble(BindAggregateFunctionInput &bind_input) { + auto &ctx = bind_input.GetClientContext(); + auto &args = bind_input.GetArguments(); return PacClipBindWithScale(ctx, args, "as_clip", CLIP_DOUBLE_SCALE); } diff --git a/src/include/aggregates/as_sum.hpp b/src/include/aggregates/as_sum.hpp index 6d893bcd..44b1ca20 100644 --- a/src/include/aggregates/as_sum.hpp +++ b/src/include/aggregates/as_sum.hpp @@ -804,15 +804,15 @@ SubtractNegStateTotals(State *, PAC_FLOAT *, uint64_t &, ArenaAllocator &) { // Generic finalize template (used by as_sum) template -void PacSumFinalize(Vector &states, AggregateInputData &input, Vector &result, idx_t count, idx_t offset); +void PacSumFinalize(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, idx_t offset); // Bind for DECIMAL as_sum -unique_ptr BindDecimalPacSum(ClientContext &ctx, AggregateFunction &function, - vector> &args); +unique_ptr BindDecimalPacSum(BindAggregateFunctionInput &input); // FinalizeCounters for as_sum template -void PacSumAvgFinalizeCounters(Vector &states, AggregateInputData &input, Vector &result, idx_t count, idx_t offset); +void PacSumAvgFinalizeCounters(Vector &states, AggregateFinalizeInputData &input, Vector &result, idx_t count, + idx_t offset); } // namespace duckdb #endif // PAC_GODBOLT diff --git a/src/include/aggregates/as_sum_state_export.hpp b/src/include/aggregates/as_sum_state_export.hpp new file mode 100644 index 00000000..beb3b2bf --- /dev/null +++ b/src/include/aggregates/as_sum_state_export.hpp @@ -0,0 +1,9 @@ +#pragma once + +#include "duckdb/function/aggregate_function.hpp" + +namespace duckdb { + +void RegisterSumStateExport(AggregateFunction &function); + +} // namespace duckdb diff --git a/src/include/categorical/pac_categorical_detection.hpp b/src/include/categorical/pac_categorical_detection.hpp index 003b7b8a..401a57cc 100644 --- a/src/include/categorical/pac_categorical_detection.hpp +++ b/src/include/categorical/pac_categorical_detection.hpp @@ -46,8 +46,8 @@ namespace duckdb { // Metadata about a PAC aggregate produced by the bitslice compiler. // Keyed by aggregate_index in PacAggregateInfoMap. struct PacAggregateInfo { - idx_t aggregate_index; // agg->aggregate_index - idx_t group_index; // agg->group_index + idx_t aggregate_index; // agg->aggregate_index.index + idx_t group_index; // agg->group_index.index ColumnBinding hash_binding; // the hash column binding used for this aggregate's children[0] }; // Map from aggregate_index → PacAggregateInfo @@ -179,7 +179,7 @@ inline bool IsNumericalType(const LogicalType &type) { // Hash a column binding to a unique 64-bit value for use in maps/sets static inline uint64_t HashBinding(const ColumnBinding &binding) { - return (uint64_t(binding.table_index) << 32) | binding.column_index; + return (uint64_t(binding.table_index.index) << 32) | binding.column_index.GetIndexUnsafe(); } // Rebind an aggregate function by name (catalog lookup + FunctionBinder). @@ -187,27 +187,29 @@ static inline uint64_t HashBinding(const ColumnBinding &binding) { static inline unique_ptr RebindAggregate(ClientContext &context, const string &func_name, vector> children, bool is_distinct) { auto &catalog = Catalog::GetSystemCatalog(context); - auto &func_entry = catalog.GetEntry(context, DEFAULT_SCHEMA, func_name); + auto &func_entry = catalog.GetEntry( + context, QualifiedName(Catalog::GetSystemCatalog(context).GetName(), Identifier::DefaultSchema(), + Identifier(func_name))); vector arg_types; for (auto &child : children) { - arg_types.push_back(child->return_type); + arg_types.push_back(child->GetReturnType()); } ErrorData error; FunctionBinder function_binder(context); - auto best_function = function_binder.BindFunction(func_name, func_entry.functions, arg_types, error); + auto best_function = function_binder.BindFunction(Identifier(func_name), func_entry.functions, arg_types, error); if (!best_function.IsValid()) { return nullptr; } - AggregateFunction func = func_entry.functions.GetFunctionByOffset(best_function.GetIndex()); - return function_binder.BindAggregateFunction(func, std::move(children), nullptr, + const auto &func = func_entry.functions.GetFunctionByOffset(best_function.GetIndex()); + return function_binder.BindAggregateFunction(*func, std::move(children), nullptr, is_distinct ? AggregateType::DISTINCT : AggregateType::NON_DISTINCT); } // Walk through cast expressions to find the underlying expression. // Returns the innermost non-cast expression. static inline Expression *StripCasts(Expression *expr) { - while (expr->type == ExpressionType::OPERATOR_CAST) { - expr = expr->Cast().child.get(); + while (expr->GetExpressionType() == ExpressionType::OPERATOR_CAST) { + expr = expr->Cast().GetChildrenMutable()[0].get(); } return expr; } @@ -215,24 +217,27 @@ static inline Expression *StripCasts(Expression *expr) { // Check if an expression is already wrapped in a categorical rewrite terminal function // This includes as_noised, priv_filter, list_transform, and list_zip // These functions indicate that the expression has already been processed by categorical rewriting -static inline bool IsAlreadyWrappedInPacNoised(Expression *expr) { +static inline bool IsAlreadyWrappedInPacNoised(const Expression *expr) { if (!expr) { return false; } - if (expr->type == ExpressionType::BOUND_FUNCTION) { + if (expr->GetExpressionType() == ExpressionType::BOUND_FUNCTION) { auto &func = expr->Cast(); // Check for all categorical rewrite terminal functions - if (func.function.name == "as_noised" || func.function.name == "as_noised_div" || - StringUtil::StartsWith(func.function.name, "priv_filter") || - StringUtil::StartsWith(func.function.name, "priv_select") || func.function.name == "priv_coalesce" || - func.function.name == "list_transform" || func.function.name == "list_zip") { + if (func.Function().GetName().GetIdentifierName() == "as_noised" || + func.Function().GetName().GetIdentifierName() == "as_noised_div" || + StringUtil::StartsWith(func.Function().GetName().GetIdentifierName(), "priv_filter") || + StringUtil::StartsWith(func.Function().GetName().GetIdentifierName(), "priv_select") || + func.Function().GetName().GetIdentifierName() == "priv_coalesce" || + func.Function().GetName().GetIdentifierName() == "list_transform" || + func.Function().GetName().GetIdentifierName() == "list_zip") { return true; } } // Check for CAST(terminal_function(...)) - if (expr->type == ExpressionType::OPERATOR_CAST) { - auto &cast = expr->Cast(); - return IsAlreadyWrappedInPacNoised(cast.child.get()); + if (expr->GetExpressionType() == ExpressionType::OPERATOR_CAST) { + auto &cast = expr->Cast(); + return IsAlreadyWrappedInPacNoised(cast.GetChildren()[0].get()); } return false; } @@ -240,15 +245,15 @@ static inline bool IsAlreadyWrappedInPacNoised(Expression *expr) { // Check if a CASE expression is DuckDB's scalar subquery wrapper pattern: // CASE WHEN count>1 THEN error(...) ELSE value END static inline bool IsScalarSubqueryWrapper(const BoundCaseExpression &case_expr) { - if (case_expr.case_checks.size() != 1 || !case_expr.else_expr) { + if (case_expr.CaseChecks().size() != 1 || !case_expr.ElseExpression()) { return false; } - auto &check = case_expr.case_checks[0]; + auto &check = case_expr.CaseChecks()[0]; // Counter-list wrappers cast the error branch to the new result type. auto *then_expr = check.then_expr ? StripCasts(check.then_expr.get()) : nullptr; - if (then_expr && then_expr->type == ExpressionType::BOUND_FUNCTION) { + if (then_expr && then_expr->GetExpressionType() == ExpressionType::BOUND_FUNCTION) { auto &func = then_expr->Cast(); - return func.function.name == "error"; + return func.FunctionMutable().GetName().GetIdentifierName() == "error"; } return false; } @@ -264,10 +269,10 @@ static inline string FindPacAggregateInOperator(LogicalOperator *op) { if (op->type == LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY) { auto &agg = op->Cast(); for (auto &agg_expr : agg.expressions) { - if (agg_expr->type == ExpressionType::BOUND_AGGREGATE) { + if (agg_expr->GetExpressionType() == ExpressionType::BOUND_AGGREGATE) { auto &bound_agg = agg_expr->Cast(); - if (IsAnyPacAggregate(bound_agg.function.name)) { - return GetBasePacAggregateName(bound_agg.function.name); + if (IsAnyPacAggregate(bound_agg.FunctionMutable().GetName().GetIdentifierName())) { + return GetBasePacAggregateName(bound_agg.FunctionMutable().GetName().GetIdentifierName()); } } } @@ -294,7 +299,7 @@ static inline LogicalOperator *RecognizeDuckDBScalarWrapper(LogicalOperator *op) return nullptr; } auto &expr = outer_proj.expressions[0]; - if (expr->type != ExpressionType::CASE_EXPR) { + if (expr->GetExpressionType() != ExpressionType::CASE_EXPR) { return nullptr; } if (!IsScalarSubqueryWrapper(expr->Cast())) { @@ -312,11 +317,11 @@ static inline LogicalOperator *RecognizeDuckDBScalarWrapper(LogicalOperator *op) return nullptr; } auto &agg_expr = agg.expressions[0]; - if (agg_expr->type != ExpressionType::BOUND_AGGREGATE) { + if (agg_expr->GetExpressionType() != ExpressionType::BOUND_AGGREGATE) { return nullptr; } auto &bound_agg = agg_expr->Cast(); - if (StringUtil::Lower(bound_agg.function.name) != "first") { + if (StringUtil::Lower(bound_agg.FunctionMutable().GetName().GetIdentifierName()) != "first") { return nullptr; } if (agg.children.empty()) { diff --git a/src/include/categorical/pac_categorical_rewriter.hpp b/src/include/categorical/pac_categorical_rewriter.hpp index 070444d7..2e795425 100644 --- a/src/include/categorical/pac_categorical_rewriter.hpp +++ b/src/include/categorical/pac_categorical_rewriter.hpp @@ -142,20 +142,20 @@ static inline ExpressionType FlipComparison(ExpressionType cmp_type) { // Try to extract a double value from an expression, looking through casts. // Returns true and sets out_val if the expression is a constant that can be cast to DOUBLE. static inline bool TryGetConstantDouble(Expression *expr, double &out_val) { - while (expr->type == ExpressionType::OPERATOR_CAST) { - expr = expr->Cast().child.get(); + while (expr->GetExpressionType() == ExpressionType::OPERATOR_CAST) { + expr = expr->Cast().GetChildrenMutable()[0].get(); } - if (expr->type != ExpressionType::VALUE_CONSTANT) { + if (expr->GetExpressionType() != ExpressionType::VALUE_CONSTANT) { return false; } auto &const_expr = expr->Cast(); - if (const_expr.value.IsNull()) { + if (const_expr.GetValue().IsNull()) { return false; } - Value double_val; string error_msg; - if (const_expr.value.DefaultTryCastAs(LogicalType::DOUBLE, double_val, &error_msg)) { - out_val = double_val.GetValue(); + auto double_val = const_expr.GetValue().DefaultTryCastAs(LogicalType::DOUBLE, &error_msg); + if (double_val) { + out_val = double_val->GetValue(); return true; } return false; diff --git a/src/include/parser/privacy_parser.hpp b/src/include/parser/privacy_parser.hpp index 7d66d4fb..0eb30c7c 100644 --- a/src/include/parser/privacy_parser.hpp +++ b/src/include/parser/privacy_parser.hpp @@ -74,11 +74,12 @@ struct PACParseData : public ParserExtensionParseData { class PrivacyParserExtension : public ParserExtension { public: PrivacyParserExtension() { - parse_function = PACParseFunction; + parse_function = PACParseTokens; plan_function = PACPlanFunction; } static ParserExtensionParseResult PACParseFunction(ParserExtensionInfo *info, const string &query); + static ParserExtensionParseResult PACParseTokens(ParserExtensionInfo *info, const vector &tokens); static ParserExtensionPlanResult PACPlanFunction(ParserExtensionInfo *info, ClientContext &context, unique_ptr parse_data); diff --git a/src/include/query_processing/as_correlation_optimizer.hpp b/src/include/query_processing/as_correlation_optimizer.hpp new file mode 100644 index 00000000..b31f759c --- /dev/null +++ b/src/include/query_processing/as_correlation_optimizer.hpp @@ -0,0 +1,8 @@ +#pragma once + +namespace duckdb { +class DatabaseInstance; + +void RegisterASCorrelationOptimizer(DatabaseInstance &db); + +} // namespace duckdb diff --git a/src/include/query_processing/pac_subquery_handler.hpp b/src/include/query_processing/pac_subquery_handler.hpp index 7999a6cd..6360fcdf 100644 --- a/src/include/query_processing/pac_subquery_handler.hpp +++ b/src/include/query_processing/pac_subquery_handler.hpp @@ -24,7 +24,7 @@ struct DelimColumnResult { LogicalType type; bool IsValid() const { - return binding.table_index != DConstants::INVALID_INDEX; + return binding.table_index.index != DConstants::INVALID_INDEX; } }; diff --git a/src/include/query_processing/privacy_expression_builder.hpp b/src/include/query_processing/privacy_expression_builder.hpp index 75bc2d84..49888aab 100644 --- a/src/include/query_processing/privacy_expression_builder.hpp +++ b/src/include/query_processing/privacy_expression_builder.hpp @@ -38,7 +38,7 @@ unique_ptr BuildXorHashFromPKs(OptimizerExtensionInput &input, Logic ColumnBinding InsertHashProjectionAboveGet(OptimizerExtensionInput &input, unique_ptr &plan, LogicalGet &get, const vector &key_columns, bool use_rowid); -// Caching wrapper around InsertHashProjectionAboveGet. Keyed by get.table_index. +// Caching wrapper around InsertHashProjectionAboveGet. Keyed by get.table_index.index. // Returns cached binding if hash projection was already inserted for this get. ColumnBinding GetOrInsertHashProjection(OptimizerExtensionInput &input, unique_ptr &plan, LogicalGet &get, const vector &key_columns, bool use_rowid, diff --git a/src/include/query_processing/privacy_fk_resolver.hpp b/src/include/query_processing/privacy_fk_resolver.hpp index a23586c1..c4733e7e 100644 --- a/src/include/query_processing/privacy_fk_resolver.hpp +++ b/src/include/query_processing/privacy_fk_resolver.hpp @@ -38,7 +38,7 @@ FKToPUResult GetFKColumnsToPU(const string &table_name, const vector &pr // This is used when the primary FK table is blocked by SEMI/ANTI joins struct AccessibleFKTableResult { LogicalGet *get_node; - idx_t table_index; + TableIndex table_index; string table_name; vector fk_columns; bool found; diff --git a/src/include/utils/privacy_helpers.hpp b/src/include/utils/privacy_helpers.hpp index bf6482c8..65d1d229 100644 --- a/src/include/utils/privacy_helpers.hpp +++ b/src/include/utils/privacy_helpers.hpp @@ -8,6 +8,9 @@ namespace duckdb { class Binder; +class BoundWindowExpression; + +unique_ptr BindRankingWindow(ClientContext &context, ExpressionType type); // Sanitize a string to be used as a PAC privacy unit or table name (alphanumeric + underscores only) string Sanitize(const string &in); diff --git a/src/metadata/privacy_compatibility_check.cpp b/src/metadata/privacy_compatibility_check.cpp index bcac178b..28828249 100644 --- a/src/metadata/privacy_compatibility_check.cpp +++ b/src/metadata/privacy_compatibility_check.cpp @@ -103,7 +103,7 @@ bool ContainsUnaggregatedEntityInput(const LogicalOperator &op, const PrivacyCom } if (op.type == LogicalOperatorType::LOGICAL_GET) { auto table = op.Cast().GetTable(); - if (table && (check.fk_paths.count(table->name) || + if (table && (check.fk_paths.count(table->name.GetIdentifierName()) || std::find(check.scanned_pu_tables.begin(), check.scanned_pu_tables.end(), table->name) != check.scanned_pu_tables.end())) { return true; @@ -135,6 +135,11 @@ static bool IsDPPrivacyMode(const string &privacy_mode) { } static bool ContainsUnsupportedDPShape(const LogicalOperator &op, string &shape) { + if (op.type == LogicalOperatorType::LOGICAL_MATERIALIZED_CTE && + StringUtil::StartsWith(op.Cast().ctename.GetIdentifierName(), "__duckdb_delim_")) { + shape = "lateral or correlated subqueries"; + return true; + } if (op.type == LogicalOperatorType::LOGICAL_MATERIALIZED_CTE || op.type == LogicalOperatorType::LOGICAL_CTE_REF || op.type == LogicalOperatorType::LOGICAL_RECURSIVE_CTE) { shape = "CTEs"; @@ -181,7 +186,7 @@ static bool ContainsAggregation(const LogicalOperator &op) { for (auto &expr : aggr.expressions) { if (expr && expr->IsAggregate()) { auto &ag = expr->Cast(); - if (IsAllowedAggregate(ag.function.name)) { + if (IsAllowedAggregate(ag.FunctionMutable().GetName().GetIdentifierName())) { return true; } } @@ -201,7 +206,7 @@ static void GetScannedTablesInScope(const LogicalOperator &op, std::unordered_se auto &get = op.Cast(); auto table_entry = get.GetTable(); if (table_entry) { - tables.insert(table_entry->name); + tables.insert(table_entry->name.GetIdentifierName()); } } @@ -250,10 +255,10 @@ static bool CheckPacAggregatesHaveProperJoins(const LogicalOperator &op, for (auto &expr : aggr.expressions) { if (expr && expr->IsAggregate()) { auto &ag = expr->Cast(); - if (IsPacAggregate(ag.function.name)) { + if (IsPacAggregate(ag.FunctionMutable().GetName().GetIdentifierName())) { has_pac_aggregate = true; found_pac_aggregate = true; - } else if (IsAllowedAggregate(ag.function.name)) { + } else if (IsAllowedAggregate(ag.FunctionMutable().GetName().GetIdentifierName())) { has_regular_aggregate = true; } } @@ -359,7 +364,7 @@ static bool CheckPacAggregatesHaveProperJoins(const LogicalOperator &op, static void TraceBindingToSources(LogicalOperator &root, const ColumnBinding &binding, const std::function &visit_source_binding) { // Find the operator that produces this binding's table_index - auto *source_op = FindOperatorByTableIndex(&root, binding.table_index); + auto *source_op = FindOperatorByTableIndex(&root, binding.table_index.index); if (!source_op) { return; // Can't find source, assume safe } @@ -370,35 +375,36 @@ static void TraceBindingToSources(LogicalOperator &root, const ColumnBinding &bi auto &aggr = source_op->Cast(); // Check if this binding is from the aggregate's group output or aggregate output - if (binding.table_index == aggr.group_index) { + if (binding.table_index.index == aggr.group_index.index) { // This is a grouped column - trace it further down idx_t group_idx = binding.column_index; if (group_idx < aggr.groups.size() && aggr.groups[group_idx]) { // Find column refs in this group expression and trace them ExpressionIterator::EnumerateExpression( const_cast &>(aggr.groups[group_idx]), [&](Expression &expr) { - if (expr.type == ExpressionType::BOUND_COLUMN_REF) { + if (expr.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = expr.Cast(); - TraceBindingToSources(root, col_ref.binding, visit_source_binding); + TraceBindingToSources(root, col_ref.BindingMutable(), visit_source_binding); } }); } } - // If binding.table_index == aggr.aggregate_index, it's an aggregate result - safe, don't trace further + // If binding.table_index.index == aggr.aggregate_index.index, it's an aggregate result - safe, don't trace + // further } else if (source_op->type == LogicalOperatorType::LOGICAL_PROJECTION) { auto &proj = source_op->Cast(); // Trace the expression that produces this column if (binding.column_index < proj.expressions.size() && proj.expressions[binding.column_index]) { ExpressionIterator::EnumerateExpression(proj.expressions[binding.column_index], [&](Expression &expr) { - if (expr.type == ExpressionType::BOUND_COLUMN_REF) { + if (expr.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = expr.Cast(); - TraceBindingToSources(root, col_ref.binding, visit_source_binding); + TraceBindingToSources(root, col_ref.BindingMutable(), visit_source_binding); } }); } } else if (source_op->type == LogicalOperatorType::LOGICAL_CTE_REF) { auto &cte_ref = source_op->Cast(); - auto *cte_def = FindMaterializedCTE(&root, cte_ref.cte_index); + auto *cte_def = FindMaterializedCTE(&root, cte_ref.cte_index.index); if (cte_def && !cte_def->children.empty() && cte_def->children[0]) { auto &cte_body = *cte_def->children[0]; auto body_bindings = cte_body.GetColumnBindings(); @@ -440,7 +446,7 @@ static std::pair GetProtectedColumnInfo(LogicalOperator &root, const ColumnBinding &binding, const std::unordered_map> &protected_columns) { // Find the LogicalGet that produces this binding - auto *source_op = FindOperatorByTableIndex(&root, binding.table_index); + auto *source_op = FindOperatorByTableIndex(&root, binding.table_index.index); if (!source_op || source_op->type != LogicalOperatorType::LOGICAL_GET) { return {"", ""}; } @@ -457,18 +463,18 @@ GetProtectedColumnInfo(LogicalOperator &root, const ColumnBinding &binding, string col_name; const auto &column_ids = get.GetColumnIds(); if (binding.column_index < column_ids.size()) { - col_name = get.GetColumnName(column_ids[binding.column_index]); + col_name = get.GetColumnName(column_ids[binding.column_index]).GetIdentifierName(); } if (col_name.empty()) { return {"", ""}; } // Check against the unified protected_columns set - string table_lower = StringUtil::Lower(table_entry->name); + string table_lower = StringUtil::Lower(table_entry->name.GetIdentifierName()); string col_lower = StringUtil::Lower(col_name); auto it = protected_columns.find(table_lower); if (it != protected_columns.end() && it->second.count(col_lower) > 0) { - return {table_entry->name, col_name}; + return {table_entry->name.GetIdentifierName(), col_name}; } return {"", ""}; @@ -508,9 +514,9 @@ static void CheckOutputColumnsNotFromPU(LogicalOperator ¤t_op, LogicalOper continue; } ExpressionIterator::EnumerateExpression(expr, [&](Expression &e) { - if (e.type == ExpressionType::BOUND_COLUMN_REF) { + if (e.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e.Cast(); - TraceBindingToPUTable(plan_root, col_ref.binding, actual_pu_tables); + TraceBindingToPUTable(plan_root, col_ref.BindingMutable(), actual_pu_tables); } }); } @@ -529,7 +535,7 @@ static void CheckOutputColumnsNotFromPU(LogicalOperator ¤t_op, LogicalOper auto table_entry = get.GetTable(); if (table_entry) { for (auto &pu_table : actual_pu_tables) { - if (table_entry->name == pu_table) { + if (table_entry->name.GetIdentifierName() == pu_table) { throw InvalidInputException( "Privacy rewrite: columns from privacy unit tables can only be accessed inside aggregate " "functions (e.g., SUM, COUNT, AVG, MIN, MAX)"); @@ -581,9 +587,9 @@ CheckOutputColumnsNotProtected(LogicalOperator ¤t_op, LogicalOperator &pla continue; } ExpressionIterator::EnumerateExpression(expr, [&](Expression &e) { - if (e.type == ExpressionType::BOUND_COLUMN_REF) { + if (e.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e.Cast(); - TraceBindingForProtectedColumns(plan_root, col_ref.binding, protected_columns); + TraceBindingForProtectedColumns(plan_root, col_ref.BindingMutable(), protected_columns); } }); } @@ -600,18 +606,18 @@ CheckOutputColumnsNotProtected(LogicalOperator ¤t_op, LogicalOperator &pla auto &get = current_op.Cast(); auto table_entry = get.GetTable(); if (table_entry) { - string table_lower = StringUtil::Lower(table_entry->name); + string table_lower = StringUtil::Lower(table_entry->name.GetIdentifierName()); auto it = protected_columns.find(table_lower); if (it != protected_columns.end()) { const auto &column_ids = get.GetColumnIds(); for (const auto &col_idx : column_ids) { - const string &col_name = get.GetColumnName(col_idx); + const string &col_name = get.GetColumnName(col_idx).GetIdentifierName(); string col_lower = StringUtil::Lower(col_name); if (it->second.count(col_lower) > 0) { throw InvalidInputException( "Privacy rewrite: protected column '%s.%s' can only be accessed inside aggregate " "functions (e.g., SUM, COUNT, AVG, MIN, MAX)", - table_entry->name.c_str(), col_name.c_str()); + table_entry->name.GetIdentifierName().c_str(), col_name.c_str()); } } } @@ -656,9 +662,10 @@ CheckFiltersNotUsingProtectedColumns(LogicalOperator &op, LogicalOperator &plan_ continue; } ExpressionIterator::EnumerateExpression(expr, [&](Expression &e) { - if (e.type == ExpressionType::BOUND_COLUMN_REF) { + if (e.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e.Cast(); - auto protected_info = GetProtectedColumnInfo(plan_root, col_ref.binding, protected_columns); + auto protected_info = + GetProtectedColumnInfo(plan_root, col_ref.BindingMutable(), protected_columns); if (!protected_info.first.empty()) { throw InvalidInputException( "Privacy rewrite: protected column '%s.%s' cannot be used in WHERE/HAVING filters. " @@ -686,7 +693,7 @@ void CountScans(const LogicalOperator &op, std::unordered_map &co auto &scan = op.Cast(); auto table_entry = scan.GetTable(); if (table_entry) { - counts[table_entry->name]++; + counts[table_entry->name.GetIdentifierName()]++; } } // Handle CTEs: traverse into CTE definitions to find base table scans diff --git a/src/metadata/privacy_metadata_manager.cpp b/src/metadata/privacy_metadata_manager.cpp index 081ce493..f792af92 100644 --- a/src/metadata/privacy_metadata_manager.cpp +++ b/src/metadata/privacy_metadata_manager.cpp @@ -166,7 +166,7 @@ void PrivacyMetadataManager::Clear() { */ string PrivacyMetadataManager::GetMetadataFilePath(ClientContext &context) { // Get the database path from the default catalog - auto &db_name = DatabaseManager::GetDefaultDatabase(context); + auto db_name = DatabaseManager::GetDefaultDatabase(context); auto &catalog = Catalog::GetCatalog(context, db_name); string db_path = catalog.GetDBPath(); @@ -181,7 +181,7 @@ string PrivacyMetadataManager::GetMetadataFilePath(ClientContext &context) { CatalogSearchPath search_path(context); const auto &entries = search_path.Get(); if (!entries.empty()) { - schema_name = entries[0].schema; + schema_name = entries[0].GetSchema().GetIdentifierName(); } } catch (...) { // If we can't get search path, use default schema diff --git a/src/parser/privacy_parser.cpp b/src/parser/privacy_parser.cpp index 1e1c3324..7c2d5936 100644 --- a/src/parser/privacy_parser.cpp +++ b/src/parser/privacy_parser.cpp @@ -46,7 +46,7 @@ static vector FindMissingColumns(const TableCatalogEntry &table, const v for (const auto &col_name : column_names) { bool found = false; for (auto &col : table.GetColumns().Logical()) { - if (StringUtil::Lower(col.GetName()) == StringUtil::Lower(col_name)) { + if (StringUtil::Lower(col.GetName().GetIdentifierName()) == StringUtil::Lower(col_name)) { found = true; break; } @@ -210,7 +210,7 @@ static thread_local string g_pac_pending_table_name; * it persists across database sessions. For in-memory databases, no file is saved. */ static unique_ptr PACDDLBindFunction(ClientContext &context, TableFunctionBindInput &input, - vector &return_types, vector &names) { + vector &return_types, vector &names) { // Get the pending SQL from thread-local storage string sql_to_execute = g_pac_pending_sql; string table_name = g_pac_pending_table_name; @@ -255,7 +255,7 @@ static unique_ptr PACDDLBindFunction(ClientContext &context, Table */ static void PACDDLExecuteFunction(ClientContext &context, TableFunctionInput &data, DataChunk &output) { // Nothing to output - DDL was already executed in bind - output.SetCardinality(0); + output.SetChildCardinality(0); } // ============================================================================ @@ -276,6 +276,30 @@ static void PACDDLExecuteFunction(ClientContext &context, TableFunctionInput &da * The order matters: DROP must be checked before ADD because DROP statements * also contain keywords like "protected" that might match ADD patterns. */ +ParserExtensionParseResult PrivacyParserExtension::PACParseTokens(ParserExtensionInfo *info, + const vector &tokens) { + string query; + idx_t consumed = 0; + for (auto &token : tokens) { + if (token.type == TokenType::END_OF_INPUT) { + break; + } + consumed++; + if (token.type == TokenType::TERMINATOR) { + break; + } + if (!query.empty()) { + query += ' '; + } + query += token.text; + } + auto result = PACParseFunction(info, query); + if (result.type == ParserExtensionResultType::PARSE_SUCCESSFUL) { + result.consumed_tokens = NumericCast(consumed); + } + return result; +} + ParserExtensionParseResult PrivacyParserExtension::PACParseFunction(ParserExtensionInfo *info, const string &query) { // Clean up query - preserve spaces but remove semicolons, comments, and newlines string clean_query = query; @@ -362,8 +386,9 @@ ParserExtensionPlanResult PrivacyParserExtension::PACPlanFunction(ParserExtensio // Check if table exists in the catalog try { auto &catalog = Catalog::GetCatalog(context, DatabaseManager::GetDefaultDatabase(context)); - auto table_entry = catalog.GetEntry(context, CatalogType::TABLE_ENTRY, DEFAULT_SCHEMA, - pac_data.metadata.table_name, OnEntryNotFound::RETURN_NULL); + auto table_entry = + catalog.GetEntry(context, CatalogType::TABLE_ENTRY, DEFAULT_SCHEMA, + Identifier(pac_data.metadata.table_name), OnEntryNotFound::RETURN_NULL); if (!table_entry) { throw CatalogException("Table '" + pac_data.metadata.table_name + "' does not exist"); } @@ -411,8 +436,9 @@ ParserExtensionPlanResult PrivacyParserExtension::PACPlanFunction(ParserExtensio "table '" + pac_data.metadata.table_name + "'. PAC LINK was not added."); // Validate that referenced table exists - auto ref_table_entry = catalog.GetEntry(context, CatalogType::TABLE_ENTRY, DEFAULT_SCHEMA, - link.referenced_table, OnEntryNotFound::RETURN_NULL); + auto ref_table_entry = + catalog.GetEntry(context, CatalogType::TABLE_ENTRY, DEFAULT_SCHEMA, + Identifier(link.referenced_table), OnEntryNotFound::RETURN_NULL); if (!ref_table_entry) { throw CatalogException("Referenced table '" + link.referenced_table + "' does not exist"); } @@ -444,8 +470,9 @@ ParserExtensionPlanResult PrivacyParserExtension::PACPlanFunction(ParserExtensio auto &catalog = Catalog::GetCatalog(context, DatabaseManager::GetDefaultDatabase(context)); for (const auto &link : pac_data.metadata.links) { // Check if referenced table exists and is a PAC table - auto ref_table_entry = catalog.GetEntry(context, CatalogType::TABLE_ENTRY, DEFAULT_SCHEMA, - link.referenced_table, OnEntryNotFound::RETURN_NULL); + auto ref_table_entry = + catalog.GetEntry(context, CatalogType::TABLE_ENTRY, DEFAULT_SCHEMA, + Identifier(link.referenced_table), OnEntryNotFound::RETURN_NULL); if (ref_table_entry) { auto ref_metadata = PrivacyMetadataManager::Get().GetTableMetadata(link.referenced_table); if (ref_metadata && ref_metadata->is_privacy_unit) { diff --git a/src/parser/privacy_parser_helpers.cpp b/src/parser/privacy_parser_helpers.cpp index c46c3ca3..ce6348d6 100644 --- a/src/parser/privacy_parser_helpers.cpp +++ b/src/parser/privacy_parser_helpers.cpp @@ -12,6 +12,7 @@ // IMPORTANT: must be included BEFORE duckdb headers on Windows MSVC // because DuckDB defines its own std namespace that conflicts with #include +#include #include "parser/privacy_parser.hpp" #include "parser/privacy_parser_helpers.hpp" diff --git a/src/query_processing/as_correlation_optimizer.cpp b/src/query_processing/as_correlation_optimizer.cpp new file mode 100644 index 00000000..43be87ac --- /dev/null +++ b/src/query_processing/as_correlation_optimizer.cpp @@ -0,0 +1,469 @@ +#include "query_processing/as_correlation_optimizer.hpp" + +#include "duckdb/common/enums/optimizer_type.hpp" +#include "duckdb/execution/operator/join/physical_comparison_join.hpp" +#include "duckdb/execution/expression_executor.hpp" +#include "duckdb/function/builtin_function_lookup.hpp" +#include "duckdb/function/function_binder.hpp" +#include "duckdb/main/config.hpp" +#include "duckdb/main/settings.hpp" +#include "duckdb/optimizer/column_binding_replacer.hpp" +#include "duckdb/optimizer/join_filter_pushdown_optimizer.hpp" +#include "duckdb/optimizer/optimizer.hpp" +#include "duckdb/optimizer/optimizer_extension.hpp" +#include "duckdb/planner/expression/bound_aggregate_expression.hpp" +#include "duckdb/planner/expression/bound_columnref_expression.hpp" +#include "duckdb/planner/expression/bound_constant_expression.hpp" +#include "duckdb/planner/expression/bound_cast_expression.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/planner/expression/expression_barrier.hpp" +#include "duckdb/planner/expression_iterator.hpp" +#include "duckdb/planner/logical_operator_visitor.hpp" +#include "duckdb/planner/operator/logical_comparison_join.hpp" +#include "duckdb/planner/operator/logical_aggregate.hpp" +#include "duckdb/planner/operator/logical_dependent_join.hpp" +#include "duckdb/planner/operator/logical_filter.hpp" +#include "duckdb/planner/operator/logical_get.hpp" +#include "duckdb/planner/operator/logical_projection.hpp" +#include "duckdb/planner/planner_extension.hpp" +#include "privacy_debug.hpp" + +namespace duckdb { + +static bool ContainsASAggregate(LogicalOperator &op) { + bool found = false; + LogicalOperatorVisitor::EnumerateExpressions(op, [&](unique_ptr *expr) { + ExpressionIterator::VisitExpression(**expr, [&](const BoundAggregateExpression &agg) { + found |= StringUtil::StartsWith(agg.Function().GetName().GetIdentifierName(), "as_"); + }); + }); + for (auto &child : op.children) { + found |= ContainsASAggregate(*child); + } + return found; +} + +static bool ContainsReorderingBarrier(LogicalOperator &op) { + bool found = false; + LogicalOperatorVisitor::EnumerateExpressions(op, [&](unique_ptr *expr) { + found |= (*expr)->IsVolatile() || ExpressionBarrier::Contains(**expr); + }); + for (auto &child : op.children) { + found |= ContainsReorderingBarrier(*child); + } + return found; +} + +// The outer filter may still be one AND expression containing the noisy +// predicate. Separate safe predicates on passthrough columns before the input +// becomes a materialized correlation domain, so noise is drawn only after them. +static void PushCorrelatedInputFilters(ClientContext &context, LogicalFilter &filter) { + if (filter.children[0]->type != LogicalOperatorType::LOGICAL_PROJECTION) { + return; + } + auto &projection = filter.children[0]->Cast(); + if (projection.children[0]->type != LogicalOperatorType::LOGICAL_DEPENDENT_JOIN || projection.expressions.empty() || + projection.expressions[0]->GetExpressionClass() != ExpressionClass::BOUND_FUNCTION || + !projection.expressions[0]->IsVolatile()) { + return; + } + auto release_name = + projection.expressions[0]->Cast().Function().GetName().GetIdentifierName(); + if (release_name != "as_noised" && release_name != "as_noised_div") { + return; + } + auto &join = projection.children[0]->Cast(); + if (join.join_type != JoinType::SINGLE || !join.perform_delim || join.HasProjectionMap()) { + return; + } + auto input_bindings = join.children[0]->GetColumnBindings(); + filter.SplitPredicates(); + vector> remaining; + vector> selected; + for (auto &expression : filter.expressions) { + // Literal casts have not passed through constant folding at post-bind. + // Evaluate only these casts before checking whether a predicate can throw. + ExpressionIterator::EnumerateExpression(expression, [&](unique_ptr &node) { + if (BoundCastExpression::IsCast(*node) && node->IsFoldable() && + BoundCastExpression::Child(node->Cast()).GetExpressionClass() == + ExpressionClass::BOUND_CONSTANT) { + Value value; + if (ExpressionExecutor::TryEvaluateScalar(context, *node, value)) { + node = make_uniq(std::move(value)); + } + } + }); + bool safe = !expression->IsVolatile() && !expression->CanThrow() && !ExpressionBarrier::Contains(*expression); + ExpressionIterator::VisitExpression( + *expression, [&](const BoundColumnRefExpression &ref) { + if (ref.Depth() != 0 || ref.Binding().table_index != projection.table_index || + projection.GetExpression(ref.Binding()).GetExpressionClass() != ExpressionClass::BOUND_COLUMN_REF) { + safe = false; + return; + } + auto binding = projection.GetExpression(ref.Binding()).Cast().Binding(); + safe &= std::find(input_bindings.begin(), input_bindings.end(), binding) != input_bindings.end(); + }); + if (!safe) { + remaining.push_back(std::move(expression)); + continue; + } + ExpressionIterator::VisitExpressionMutable( + expression, [&](BoundColumnRefExpression &ref, unique_ptr &expr) { + expr = projection.GetExpression(ref.Binding()).Copy(); + }); + selected.push_back(std::move(expression)); + } + filter.expressions = std::move(remaining); + if (!selected.empty()) { + auto input_filter = make_uniq(); + input_filter->expressions = std::move(selected); + input_filter->children.push_back(std::move(join.children[0])); + join.children[0] = std::move(input_filter); + PRIVACY_DEBUG_PRINT("[AS] Applied deterministic outer filters before shared counter aggregation"); + } +} + +// A volatile scalar release belongs to each outer row, but the counter +// aggregates underneath it are deterministic and can be shared by correlation +// key. Move only the release above SINGLE before DuckDB chooses its domain. +static void LiftCorrelatedNoise(PlannerExtensionInput &input, unique_ptr &op, + BindingReplacementGraph &replacements) { + for (idx_t i = 0; i < op->children.size(); i++) { + BindingReplacementGraph child_replacements; + ColumnBindingRewrite::RewriteChild( + op, i, [&](unique_ptr &child, BindingReplacementGraph &child_changes) { + LiftCorrelatedNoise(input, child, child_changes); + child_replacements.Merge(child_changes); + }); + replacements.Merge(child_replacements); + } + if (op->type == LogicalOperatorType::LOGICAL_FILTER) { + PushCorrelatedInputFilters(input.context, op->Cast()); + } + if (op->type != LogicalOperatorType::LOGICAL_DEPENDENT_JOIN || op->HasProjectionMap()) { + return; + } + auto &join = op->Cast(); + if (join.join_type != JoinType::SINGLE || join.condition || join.correlated_columns.empty() || + join.children[1]->type != LogicalOperatorType::LOGICAL_PROJECTION) { + return; + } + auto &projection = join.children[1]->Cast(); + if (projection.expressions.size() != 1 || + projection.expressions[0]->GetExpressionClass() != ExpressionClass::BOUND_FUNCTION || + projection.children[0]->type != LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY || + !projection.children[0]->Cast().groups.empty() || + ContainsReorderingBarrier(*projection.children[0])) { + return; + } + auto &release = projection.expressions[0]->Cast(); + auto name = release.Function().GetName().GetIdentifierName(); + if ((name != "as_noised" && name != "as_noised_div") || + release.Function().GetStability() != FunctionStability::VOLATILE) { + return; + } + for (auto &argument : release.GetChildren()) { + if (argument->IsVolatile() || ExpressionBarrier::Contains(*argument)) { + return; + } + bool correlated = false; + ExpressionIterator::VisitExpression( + *argument, [&](const BoundColumnRefExpression &ref) { correlated |= ref.Depth() != 0; }); + if (correlated) { + return; + } + } + CorrelatedColumns shared_columns; + for (idx_t i = 0; i < join.correlated_columns.size(); i++) { + if (!join.perform_delim && i == join.correlated_columns.GetDelimIndex()) { + continue; + } + auto &column = join.correlated_columns[i]; + // Keep DuckDB's row-identity strategy for nested, non-deduplicable keys. + if (column.type.IsNested()) { + return; + } + shared_columns.AddColumnToBack(column); + } + if (shared_columns.empty()) { + return; + } + join.ResolveOperatorTypes(); + auto left_bindings = join.children[0]->GetColumnBindings(); + auto release_index = projection.table_index; + projection.table_index = input.binder.GenerateTableIndex(); + auto expression = std::move(projection.expressions[0]); + projection.expressions = std::move(expression->Cast().GetChildrenMutable()); + for (idx_t i = 0; i < projection.expressions.size(); i++) { + expression->Cast().GetChildrenMutable().push_back(make_uniq( + projection.expressions[i]->GetReturnType(), ColumnBinding(projection.table_index, ProjectionIndex(i)))); + } + vector> expressions; + expressions.push_back(std::move(expression)); + for (idx_t i = 0; i < left_bindings.size(); i++) { + expressions.push_back(make_uniq(join.children[0]->types[i], left_bindings[i])); + replacements.Add(left_bindings[i], ColumnBinding(release_index, ProjectionIndex(i + 1))); + } + join.correlated_columns = std::move(shared_columns); + join.perform_delim = true; + auto lifted = make_uniq(release_index, std::move(expressions)); + lifted->children.push_back(std::move(op)); + op = std::move(lifted); + op->ResolveOperatorTypes(); + PRIVACY_DEBUG_PRINT( + "[AS] Shared deterministic correlated counters while retaining one volatile release per outer row"); +} + +static void ShareCorrelatedNoiseInputs(PlannerExtensionInput &input, BoundStatement &statement) { + if (!statement.plan || !Settings::Get(input.context) || + DBConfig::GetConfig(input.context).IsOptimizerDisabled(OptimizerType::EXTENSION) || + !ContainsASAggregate(*statement.plan)) { + return; + } + BindingReplacementGraph replacements; + LiftCorrelatedNoise(input, statement.plan, replacements); +} + +// Unlike an ordinary anti join, RIGHT_ANTI preserves the build side. Probe rows +// with no possible build-side match may therefore be filtered without changing +// any output row. Carry that restriction to both sides of an inner ownership join. +static void GetOwnershipFilterTargets(LogicalOperator &op, vector columns, + vector &targets) { + if (columns.empty()) { + return; + } + if (op.type == LogicalOperatorType::LOGICAL_PROJECTION) { + auto &projection = op.Cast(); + vector projected; + for (auto &column : columns) { + if (column.probe_column_index.table_index == projection.table_index && + JoinFilterPushdownUtil::PushdownJoinFilterExpression( + projection.GetExpression(column.probe_column_index), column)) { + projected.push_back(std::move(column)); + } + } + if (!projected.empty()) { + GetOwnershipFilterTargets(*op.children[0], std::move(projected), targets); + } + return; + } + if (op.type == LogicalOperatorType::LOGICAL_FILTER) { + GetOwnershipFilterTargets(*op.children[0], std::move(columns), targets); + return; + } + if (op.type == LogicalOperatorType::LOGICAL_COMPARISON_JOIN) { + auto &join = op.Cast(); + if (join.join_type != JoinType::INNER) { + return; + } + for (auto &condition : join.conditions) { + if (!condition.IsComparison() || condition.GetComparisonType() != ExpressionType::COMPARE_EQUAL || + condition.GetLHS().GetExpressionClass() != ExpressionClass::BOUND_COLUMN_REF || + condition.GetRHS().GetExpressionClass() != ExpressionClass::BOUND_COLUMN_REF || + condition.GetLHS().GetReturnType() != condition.GetRHS().GetReturnType()) { + continue; + } + auto left = condition.GetLHS().Cast().Binding(); + auto right = condition.GetRHS().Cast().Binding(); + auto count = columns.size(); + for (idx_t i = 0; i < count; i++) { + if (columns[i].probe_column_index != left && columns[i].probe_column_index != right) { + continue; + } + auto other = columns[i]; + other.probe_column_index = other.probe_column_index == left ? right : left; + bool present = false; + for (auto &column : columns) { + present |= column.probe_column_index == other.probe_column_index && + column.join_filter_idx == other.join_filter_idx; + } + if (!present) { + columns.push_back(std::move(other)); + } + } + } + for (auto &child : op.children) { + auto bindings = child->GetColumnBindings(); + vector child_columns; + for (auto &column : columns) { + for (auto &binding : bindings) { + if (binding == column.probe_column_index) { + child_columns.push_back(column); + break; + } + } + } + if (!child_columns.empty()) { + GetOwnershipFilterTargets(*child, std::move(child_columns), targets); + } + } + return; + } + if (op.type == LogicalOperatorType::LOGICAL_GET) { + JoinFilterPushdownOptimizer::GetPushdownFilterTargets(op, std::move(columns), targets); + } +} + +// A scan on the build side can finish before its parent anti join publishes +// runtime filters. Keep a wide, weakly filtered scan on the probe side when a +// smaller key-only source can build the ownership lookup instead. +static idx_t ScanCardinality(LogicalOperator &op) { + if (op.type == LogicalOperatorType::LOGICAL_GET) { + return op.has_estimated_cardinality ? op.estimated_cardinality : 0; + } + if ((op.type == LogicalOperatorType::LOGICAL_PROJECTION || op.type == LogicalOperatorType::LOGICAL_FILTER) && + op.children.size() == 1) { + return ScanCardinality(*op.children[0]); + } + return 0; +} + +static void PrepareOwnershipProbe(OptimizerExtensionInput &input, unique_ptr &op) { + if (op->type == LogicalOperatorType::LOGICAL_PROJECTION || op->type == LogicalOperatorType::LOGICAL_FILTER) { + PrepareOwnershipProbe(input, op->children[0]); + return; + } + if (op->type != LogicalOperatorType::LOGICAL_COMPARISON_JOIN) { + return; + } + auto &join = op->Cast(); + if (join.join_type != JoinType::INNER || join.children[0]->types.size() != 1 || + join.children[1]->types.size() <= 1 || !JoinFilterPushdownOptimizer::IsFiltering(join.children[1]) || + ContainsReorderingBarrier(*join.children[0])) { + return; + } + auto left_rows = ScanCardinality(*join.children[0]); + auto right_rows = ScanCardinality(*join.children[1]); + if (!left_rows || left_rows >= right_rows / 2 || !join.children[1]->has_estimated_cardinality || + join.children[1]->estimated_cardinality < left_rows / 2) { + return; + } + auto old_bindings = join.GetColumnBindings(); + std::swap(join.children[0], join.children[1]); + std::swap(join.left_projection_map, join.right_projection_map); + for (auto &condition : join.conditions) { + if (condition.IsComparison()) { + condition.Swap(); + } + } + join.filter_pushdown.reset(); + JoinFilterPushdownOptimizer filters(input.optimizer); + filters.VisitOperator(join); + auto bindings = join.GetColumnBindings(); + auto layout = make_uniq(); + for (auto &old_binding : old_bindings) { + for (idx_t i = 0; i < bindings.size(); i++) { + if (bindings[i] == old_binding) { + layout->projection_map.emplace_back(i); + break; + } + } + } + D_ASSERT(layout->projection_map.size() == old_bindings.size()); + layout->children.push_back(std::move(op)); + op = std::move(layout); + op->ResolveOperatorTypes(); + PRIVACY_DEBUG_PRINT("[AS] Kept the filtered ownership scan on the anti-join probe pipeline"); +} + +static void AddAntiJoinFilters(OptimizerExtensionInput &input, LogicalOperator &op) { + for (auto &child : op.children) { + AddAntiJoinFilters(input, *child); + } + if (op.type != LogicalOperatorType::LOGICAL_COMPARISON_JOIN) { + return; + } + auto &join = op.Cast(); + if (join.join_type != JoinType::RIGHT_ANTI || join.filter_pushdown || + !JoinFilterPushdownOptimizer::IsFiltering(join.children[1]) || ContainsReorderingBarrier(*join.children[0])) { + return; + } + // Null-safe equality can be normalized only with a non-null proof. Without + // that proof, filtering NULL probe keys could remove a real NULL=NULL match. + for (auto &condition : join.conditions) { + if (!condition.IsComparison() || condition.GetComparisonType() != ExpressionType::COMPARE_NOT_DISTINCT_FROM) { + continue; + } + if ((condition.GetLeftStats() && !condition.GetLeftStats()->CanHaveNull()) || + (condition.GetRightStats() && !condition.GetRightStats()->CanHaveNull())) { + condition = + JoinCondition(condition.GetLHS().Copy(), condition.GetRHS().Copy(), ExpressionType::COMPARE_EQUAL); + } + } + PhysicalComparisonJoin::ReorderConditions(join.conditions); + auto info = make_uniq(); + vector columns; + for (idx_t i = 0; i < join.conditions.size(); i++) { + auto &condition = join.conditions[i]; + if (!condition.IsComparison() || condition.GetComparisonType() != ExpressionType::COMPARE_EQUAL) { + continue; + } + JoinFilterPushdownColumn column; + if (!JoinFilterPushdownUtil::PushdownJoinFilterExpression(condition.GetLHS(), column)) { + continue; + } + column.join_filter_idx = info->join_condition.size(); + columns.push_back(std::move(column)); + info->join_condition.push_back(i); + } + if (columns.empty()) { + return; + } + if (!DBConfig::GetConfig(input.context).IsOptimizerDisabled(OptimizerType::BUILD_SIDE_PROBE_SIDE)) { + PrepareOwnershipProbe(input, join.children[0]); + } + vector targets; + GetOwnershipFilterTargets(*join.children[0], std::move(columns), targets); + if (targets.empty()) { + return; + } + for (auto &index : info->join_condition) { + for (auto name : {"min", "max"}) { + vector> arguments; + arguments.push_back(join.conditions[index].GetRHS().Copy()); + auto function = GetBuiltinAggregateFunction(input.context, name, {arguments[0]->GetReturnType()}); + FunctionBinder binder(input.context); + auto aggregate = binder.BindAggregateFunction(std::move(function), std::move(arguments), nullptr, + AggregateType::NON_DISTINCT); + if (aggregate->GetChildren().size() != 1) { + return; + } + info->min_max_aggregates.push_back(std::move(aggregate)); + } + } + for (auto &target : targets) { + auto &get = target.get; + if (!get.dynamic_filters) { + get.dynamic_filters = make_shared_ptr(); + } + JoinFilterPushdownFilter filter; + filter.dynamic_filters = get.dynamic_filters; + filter.columns = std::move(target.columns); + if (get.has_estimated_cardinality) { + filter.estimated_scan_cardinality = get.estimated_cardinality; + } + info->probe_info.push_back(std::move(filter)); + } + info->build_side_has_filter = true; + join.filter_pushdown = std::move(info); + PRIVACY_DEBUG_PRINT("[AS] Pushed right-anti membership keys into ownership join inputs"); +} + +static void OptimizeASCorrelations(OptimizerExtensionInput &input, unique_ptr &plan) { + if (!DBConfig::GetConfig(input.context).IsOptimizerDisabled(OptimizerType::JOIN_FILTER_PUSHDOWN) && + ContainsASAggregate(*plan)) { + AddAntiJoinFilters(input, *plan); + } +} + +void RegisterASCorrelationOptimizer(DatabaseInstance &db) { + PlannerExtension planner; + planner.post_bind_function = ShareCorrelatedNoiseInputs; + PlannerExtension::Register(DBConfig::GetConfig(db), std::move(planner)); + OptimizerExtension extension; + extension.optimize_function = OptimizeASCorrelations; + OptimizerExtension::Register(DBConfig::GetConfig(db), std::move(extension)); +} + +} // namespace duckdb diff --git a/src/query_processing/pac_avg_rewriter.cpp b/src/query_processing/pac_avg_rewriter.cpp index f5f5cb75..a090df3a 100644 --- a/src/query_processing/pac_avg_rewriter.cpp +++ b/src/query_processing/pac_avg_rewriter.cpp @@ -4,6 +4,7 @@ // works for both compiler-generated and user-written as_avg() SQL. // +#include "duckdb/planner/expression_iterator.hpp" #include "query_processing/pac_avg_rewriter.hpp" #include "query_processing/privacy_expression_builder.hpp" #include "privacy_debug.hpp" @@ -21,10 +22,10 @@ namespace duckdb { // with a copy of the corresponding expression value. Recurses into children. static void ReplaceAvgRefs(unique_ptr &expr, const unordered_map> &replacements, idx_t table_index) { - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { + if (expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = expr->Cast(); - if (col_ref.binding.table_index == table_index) { - auto it = replacements.find(col_ref.binding.column_index); + if (col_ref.BindingMutable().table_index.index == table_index) { + auto it = replacements.find(col_ref.BindingMutable().column_index); if (it != replacements.end()) { expr = it->second->Copy(); return; @@ -61,7 +62,8 @@ static void RewritePacAvgAggregate(OptimizerExtensionInput &input, unique_ptrCast(); - if (aggr.function.name == "as_noised_avg" || aggr.function.name == "as_avg") { + if (aggr.FunctionMutable().GetName().GetIdentifierName() == "as_noised_avg" || + aggr.FunctionMutable().GetName().GetIdentifierName() == "as_avg") { avg_positions.push_back(i); } } @@ -69,7 +71,7 @@ static void RewritePacAvgAggregate(OptimizerExtensionInput &input, unique_ptrCast(); - bool is_noised = (old_aggr.function.name == "as_noised_avg"); + bool is_noised = (old_aggr.FunctionMutable().GetName().GetIdentifierName() == "as_noised_avg"); // Extract hash and value children: pac_[noised_]avg(hash, col[, correction]) - auto hash_child = old_aggr.children[0]->Copy(); - auto value_child = old_aggr.children[1]->Copy(); + auto hash_child = old_aggr.GetChildrenMutable()[0]->Copy(); + auto value_child = old_aggr.GetChildrenMutable()[1]->Copy(); // Extract optional correction factor (3rd arg), default to constant 1.0 unique_ptr corr_expr; - if (old_aggr.children.size() >= 3) { - corr_expr = old_aggr.children[2]->Copy(); + if (old_aggr.GetChildrenMutable().size() >= 3) { + corr_expr = old_aggr.GetChildrenMutable()[2]->Copy(); } string sum_name = is_noised ? "as_noised_sum" : "as_sum"; @@ -115,11 +117,13 @@ static void RewritePacAvgAggregate(OptimizerExtensionInput &input, unique_ptrreturn_type; - auto count_type = agg.expressions[info.count_col]->return_type; + auto sum_type = agg.expressions[info.sum_col]->GetReturnType(); + auto count_type = agg.expressions[info.count_col]->GetReturnType(); - auto sum_ref = make_uniq(sum_type, ColumnBinding(agg_index, info.sum_col)); - auto cnt_ref = make_uniq(count_type, ColumnBinding(agg_index, info.count_col)); + auto sum_ref = make_uniq( + sum_type, ColumnBinding(TableIndex(agg_index), ProjectionIndex(info.sum_col))); + auto cnt_ref = make_uniq( + count_type, ColumnBinding(TableIndex(agg_index), ProjectionIndex(info.count_col))); unique_ptr div_expr; if (info.is_noised) { diff --git a/src/query_processing/pac_derived_rewriter.cpp b/src/query_processing/pac_derived_rewriter.cpp index 1a94d8ab..16a287a1 100644 --- a/src/query_processing/pac_derived_rewriter.cpp +++ b/src/query_processing/pac_derived_rewriter.cpp @@ -62,8 +62,8 @@ static LogicalType CounterListType() { // Returns true if the expression tree contains a priv_finalize call. static bool ExpressionContainsPacFinalize(Expression &e) { - if (e.type == ExpressionType::BOUND_FUNCTION && - e.Cast().function.name == "priv_finalize") { + if (e.GetExpressionType() == ExpressionType::BOUND_FUNCTION && + e.Cast().FunctionMutable().GetName().GetIdentifierName() == "priv_finalize") { return true; } bool found = false; @@ -76,9 +76,9 @@ static bool ExpressionContainsPacFinalize(Expression &e) { } static bool ExpressionReferencesFinalizedBinding(Expression &e, const unordered_set &finalized_bindings) { - if (e.type == ExpressionType::BOUND_COLUMN_REF) { + if (e.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e.Cast(); - return finalized_bindings.count(HashBinding(col_ref.binding)) > 0; + return finalized_bindings.count(HashBinding(col_ref.BindingMutable())) > 0; } bool found = false; ExpressionIterator::EnumerateChildren(e, [&](Expression &child) { @@ -110,34 +110,35 @@ static void ConvertAggregatesRecursive(OptimizerExtensionInput &input, LogicalOp for (idx_t i = 0; i < agg.expressions.size(); i++) { auto &expr = agg.expressions[i]; - if (expr->type != ExpressionType::BOUND_AGGREGATE) { + if (expr->GetExpressionType() != ExpressionType::BOUND_AGGREGATE) { continue; } auto &bound_agg = expr->Cast(); - if (!IsPacAggregate(bound_agg.function.name)) { + if (!IsPacAggregate(bound_agg.FunctionMutable().GetName().GetIdentifierName())) { continue; } - string counters_name = GetCountersVariant(bound_agg.function.name); + string counters_name = GetCountersVariant(bound_agg.FunctionMutable().GetName().GetIdentifierName()); if (counters_name.empty()) { continue; } vector> children; - for (auto &child_expr : bound_agg.children) { + for (auto &child_expr : bound_agg.GetChildrenMutable()) { children.push_back(child_expr->Copy()); } auto new_aggr = RebindAggregate(input.context, counters_name, std::move(children), bound_agg.IsDistinct()); if (new_aggr) { - PRIVACY_DEBUG_PRINT("[PAC DERIVED WRITE] Rebound " + bound_agg.function.name + " → " + - new_aggr->Cast().function.name); + PRIVACY_DEBUG_PRINT("[PAC DERIVED WRITE] Rebound " + bound_agg.Function().GetName().GetIdentifierName() + + " → " + + new_aggr->Cast().Function().GetName().GetIdentifierName()); agg.expressions[i] = std::move(new_aggr); } else { - bound_agg.function.name = counters_name; - bound_agg.function.return_type = list_type; - expr->return_type = list_type; + bound_agg.FunctionMutable().SetName(Identifier(counters_name)); + bound_agg.FunctionMutable().SetReturnType(list_type); + expr->SetReturnType(list_type); } idx_t types_index = agg.groups.size() + i; @@ -145,7 +146,7 @@ static void ConvertAggregatesRecursive(OptimizerExtensionInput &input, LogicalOp agg.types[types_index] = list_type; } - converted.insert(HashBinding(ColumnBinding(agg.aggregate_index, i))); + converted.insert(HashBinding(ColumnBinding(TableIndex(agg.aggregate_index.index), ProjectionIndex(i)))); } } @@ -156,10 +157,10 @@ static void FixStaleColumnRefTypes(LogicalOperator *op, const unordered_set *expr_ptr) { std::function &)> Fix = [&](unique_ptr &e) { - if (e->type == ExpressionType::BOUND_COLUMN_REF) { + if (e->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e->Cast(); - if (converted.count(HashBinding(col_ref.binding))) { - col_ref.return_type = list_type; + if (converted.count(HashBinding(col_ref.BindingMutable()))) { + col_ref.SetReturnType(list_type); } } ExpressionIterator::EnumerateChildren(*e, Fix); @@ -180,10 +181,10 @@ static void StripRedundantCasts(LogicalOperator *op) { } std::function &)> Strip = [&](unique_ptr &e) { ExpressionIterator::EnumerateChildren(*e, Strip); - if (e->GetExpressionClass() == ExpressionClass::BOUND_CAST) { - auto &cast = e->Cast(); - if (cast.child->return_type == cast.return_type) { - e = std::move(cast.child); + if (e->GetExpressionType() == ExpressionType::OPERATOR_CAST) { + auto &cast = e->Cast(); + if (cast.GetChildrenMutable()[0]->GetReturnType() == cast.GetReturnType()) { + e = std::move(cast.GetChildrenMutable()[0]); } } }; @@ -220,14 +221,14 @@ static void FindDerivedPuGets(LogicalOperator *op, unordered_set &derived return; } auto &mgr = PrivacyMetadataManager::Get(); - auto *meta = mgr.GetTableMetadata(entry->name); + auto *meta = mgr.GetTableMetadata(entry->name.GetIdentifierName()); if (!meta || !meta->derived_pu) { return; } auto list_type = CounterListType(); for (idx_t i = 0; i < get.returned_types.size(); i++) { if (get.returned_types[i] == list_type) { - derived_table_indices.insert(get.table_index); + derived_table_indices.insert(get.table_index.index); return; } } @@ -235,25 +236,25 @@ static void FindDerivedPuGets(LogicalOperator *op, unordered_set &derived // Returns true if the expression is (or contains at its root) a counter column ref. static bool IsCounterColumnRef(Expression &e, const LogicalType &list_type) { - if (e.type == ExpressionType::BOUND_COLUMN_REF && e.return_type == list_type) { + if (e.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF && e.GetReturnType() == list_type) { return true; } // Also catch CAST(counter_col → scalar) inserted by the binder - if (e.GetExpressionClass() == ExpressionClass::BOUND_CAST) { - auto &cast = e.Cast(); - return IsCounterColumnRef(*cast.child, list_type); + if (e.GetExpressionType() == ExpressionType::OPERATOR_CAST) { + auto &cast = e.Cast(); + return IsCounterColumnRef(*cast.GetChildrenMutable()[0], list_type); } return false; } // Extract the raw counter column ref (stripping binder-inserted casts). static unique_ptr ExtractCounterRef(unique_ptr e, const LogicalType &list_type) { - if (e->type == ExpressionType::BOUND_COLUMN_REF && e->return_type == list_type) { + if (e->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF && e->GetReturnType() == list_type) { return e; } - if (e->GetExpressionClass() == ExpressionClass::BOUND_CAST) { - auto &cast = e->Cast(); - return ExtractCounterRef(std::move(cast.child), list_type); + if (e->GetExpressionType() == ExpressionType::OPERATOR_CAST) { + auto &cast = e->Cast(); + return ExtractCounterRef(std::move(cast.GetChildrenMutable()[0]), list_type); } return e; } @@ -262,28 +263,28 @@ static unique_ptr ExtractCounterRef(unique_ptr e, const // Strips binder-inserted CAST(scalar→FLOAT[]) and FLOAT[] constants back to scalars. static unique_ptr CastScalarToPacFloat(unique_ptr scalar, const LogicalType &pac_float_type, const LogicalType &list_type) { - if (scalar->return_type == pac_float_type) { + if (scalar->GetReturnType() == pac_float_type) { return scalar; } // Strip any binder-inserted CAST to FLOAT[] (scalar→LIST) - if (scalar->GetExpressionClass() == ExpressionClass::BOUND_CAST) { - auto &cast = scalar->Cast(); - if (cast.return_type == list_type && cast.child->return_type != list_type) { - scalar = std::move(cast.child); + if (scalar->GetExpressionType() == ExpressionType::OPERATOR_CAST) { + auto &cast = scalar->Cast(); + if (cast.GetReturnType() == list_type && cast.GetChildrenMutable()[0]->GetReturnType() != list_type) { + scalar = std::move(cast.GetChildrenMutable()[0]); } } // Also unwrap FLOAT[] constants back to scalar - if (scalar->type == ExpressionType::VALUE_CONSTANT) { + if (scalar->GetExpressionType() == ExpressionType::VALUE_CONSTANT) { auto &constant = scalar->Cast(); - if (constant.return_type == list_type && !constant.value.IsNull()) { - auto &list_children = ListValue::GetChildren(constant.value); + if (constant.GetReturnType() == list_type && !constant.GetValue().IsNull()) { + auto &list_children = ListValue::GetChildren(constant.GetValue()); if (!list_children.empty()) { auto scalar_val = list_children[0].DefaultCastAs(pac_float_type); return make_uniq(std::move(scalar_val)); } } } - if (scalar->return_type != pac_float_type) { + if (scalar->GetReturnType() != pac_float_type) { return BoundCastExpression::AddDefaultCastToType(std::move(scalar), pac_float_type); } return scalar; @@ -297,7 +298,7 @@ static void SyncMaterializedCTERefTypes(LogicalOperator *op, idx_t cte_index, co } if (op->type == LogicalOperatorType::LOGICAL_CTE_REF) { auto &cte_ref = op->Cast(); - if (cte_ref.cte_index == cte_index) { + if (cte_ref.cte_index.index == cte_index) { auto bindings = cte_ref.GetColumnBindings(); for (idx_t i = 0; i < cte_types.size() && i < cte_ref.chunk_types.size(); i++) { if (cte_ref.chunk_types[i] == list_type && cte_types[i] == finalized_type && i < bindings.size()) { @@ -317,11 +318,11 @@ static void SyncMaterializedCTERefTypes(LogicalOperator *op, idx_t cte_index, co // Rewrite a comparison expression on a counter column to a priv_filter_ call. // Returns the rewritten expression, or nullptr if the comparison doesn't involve a counter column. -static unique_ptr RewriteComparisonToFilterCmp(OptimizerExtensionInput &input, - BoundComparisonExpression &cmp, const LogicalType &list_type, +static unique_ptr RewriteComparisonToFilterCmp(OptimizerExtensionInput &input, BoundFunctionExpression &cmp, + const LogicalType &list_type, const LogicalType &finalized_type) { - bool left_is_counter = IsCounterColumnRef(*cmp.left, list_type); - bool right_is_counter = IsCounterColumnRef(*cmp.right, list_type); + bool left_is_counter = IsCounterColumnRef(*cmp.GetChildrenMutable()[0], list_type); + bool right_is_counter = IsCounterColumnRef(*cmp.GetChildrenMutable()[1], list_type); if (!left_is_counter && !right_is_counter) { return nullptr; } @@ -330,17 +331,19 @@ static unique_ptr RewriteComparisonToFilterCmp(OptimizerExtensionInp return nullptr; } bool counter_on_left = left_is_counter; - ExpressionType effective = counter_on_left ? FlipComparison(cmp.type) : cmp.type; + ExpressionType effective = counter_on_left ? FlipComparison(cmp.GetExpressionType()) : cmp.GetExpressionType(); const char *resolved = GetPacFilterCmpName(effective, GetPrivacyMode(input.context)); if (!resolved) { return nullptr; } string func_name = resolved; - auto counter_ref = ExtractCounterRef(std::move(counter_on_left ? cmp.left : cmp.right), list_type); + auto counter_ref = ExtractCounterRef( + std::move(counter_on_left ? cmp.GetChildrenMutable()[0] : cmp.GetChildrenMutable()[1]), list_type); auto scalar_expr = - CastScalarToPacFloat(std::move(counter_on_left ? cmp.right : cmp.left), finalized_type, list_type); + CastScalarToPacFloat(std::move(counter_on_left ? cmp.GetChildrenMutable()[1] : cmp.GetChildrenMutable()[0]), + finalized_type, list_type); PRIVACY_DEBUG_PRINT("[PAC DERIVED READ] Rewriting comparison to " + func_name); - return input.optimizer.BindScalarFunction(func_name, std::move(scalar_expr), std::move(counter_ref)); + return input.optimizer.BindScalarFunction(Identifier(func_name), std::move(scalar_expr), std::move(counter_ref)); } // Rewrite a BETWEEN expression on a counter column to priv_filter_ AND conjunction. @@ -360,8 +363,10 @@ RewriteBetweenToFilterCmp(OptimizerExtensionInput &input, unique_ptr string upper_func = GetPacFilterCmpName(upper_inclusive ? ExpressionType::COMPARE_GREATERTHANOREQUALTO : ExpressionType::COMPARE_GREATERTHAN, privacy_mode); - auto lower_call = input.optimizer.BindScalarFunction(lower_func, std::move(lower_scalar), std::move(counter_ref1)); - auto upper_call = input.optimizer.BindScalarFunction(upper_func, std::move(upper_scalar), std::move(counter_ref2)); + auto lower_call = + input.optimizer.BindScalarFunction(Identifier(lower_func), std::move(lower_scalar), std::move(counter_ref1)); + auto upper_call = + input.optimizer.BindScalarFunction(Identifier(upper_func), std::move(upper_scalar), std::move(counter_ref2)); PRIVACY_DEBUG_PRINT("[PAC DERIVED READ] Rewriting BETWEEN to " + lower_func + " AND " + upper_func); return make_uniq(ExpressionType::CONJUNCTION_AND, std::move(lower_call), std::move(upper_call)); @@ -394,8 +399,8 @@ static void WrapCounterRefsWithFinalize(OptimizerExtensionInput &input, LogicalO auto &cte = op->Cast(); WrapCounterRefsWithFinalize(input, op->children[0].get(), derived_table_indices, finalized_bindings, child_suppress); - SyncMaterializedCTERefTypes(op->children[1].get(), cte.table_index, op->children[0]->types, CounterListType(), - PacFloatLogicalType(), finalized_bindings); + SyncMaterializedCTERefTypes(op->children[1].get(), cte.table_index.index, op->children[0]->types, + CounterListType(), PacFloatLogicalType(), finalized_bindings); WrapCounterRefsWithFinalize(input, op->children[1].get(), derived_table_indices, finalized_bindings, child_suppress); } else { @@ -415,21 +420,21 @@ static void WrapCounterRefsWithFinalize(OptimizerExtensionInput &input, LogicalO // Fix stale column refs that were finalized by a child operator std::function &)> FixExpr = [&](unique_ptr &e) { - if (e->type == ExpressionType::BOUND_COLUMN_REF) { + if (e->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e->Cast(); - if (finalized_bindings.count(HashBinding(col_ref.binding))) { - col_ref.return_type = finalized_type; + if (finalized_bindings.count(HashBinding(col_ref.BindingMutable()))) { + col_ref.SetReturnType(finalized_type); } } ExpressionIterator::EnumerateChildren(*e, FixExpr); - if (e->return_type == list_type && ExpressionReferencesFinalizedBinding(*e, finalized_bindings)) { - e->return_type = finalized_type; + if (e->GetReturnType() == list_type && ExpressionReferencesFinalizedBinding(*e, finalized_bindings)) { + e->SetReturnType(finalized_type); } - if (e->GetExpressionClass() == ExpressionClass::BOUND_BETWEEN) { - auto &between = e->Cast(); - FixExpr(between.input); - FixExpr(between.lower); - FixExpr(between.upper); + if (e->GetExpressionType() == ExpressionType::COMPARE_BETWEEN) { + auto &between = e->Cast(); + FixExpr(between.GetChildrenMutable()[0]); + FixExpr(between.GetChildrenMutable()[1]); + FixExpr(between.GetChildrenMutable()[2]); } }; LogicalOperatorVisitor::EnumerateExpressions(*op, [&](unique_ptr *expr_ptr) { FixExpr(*expr_ptr); }); @@ -438,30 +443,30 @@ static void WrapCounterRefsWithFinalize(OptimizerExtensionInput &input, LogicalO // and wrap remaining counter column refs with priv_finalize. std::function &)> WrapExpr = [&](unique_ptr &e) { // Skip already-wrapped expressions - if (e->type == ExpressionType::BOUND_FUNCTION) { - auto &func_name = e->Cast().function.name; + if (e->GetExpressionType() == ExpressionType::BOUND_FUNCTION) { + auto &func_name = e->Cast().FunctionMutable().GetName().GetIdentifierName(); if (func_name == "priv_finalize" || StringUtil::StartsWith(func_name, "priv_filter_")) { return; } } // Skip binder-inserted CAST(FLOAT[] → scalar) — the registered cast handles finalization - if (e->GetExpressionClass() == ExpressionClass::BOUND_CAST) { - auto &cast = e->Cast(); - if (cast.child->return_type == list_type && cast.return_type != list_type) { + if (e->GetExpressionType() == ExpressionType::OPERATOR_CAST) { + auto &cast = e->Cast(); + if (cast.GetChildrenMutable()[0]->GetReturnType() == list_type && cast.GetReturnType() != list_type) { return; } } // Rewrite BoundComparisonExpression where one side is a counter column - if (e->GetExpressionClass() == ExpressionClass::BOUND_COMPARISON) { - auto &cmp = e->Cast(); + if (BoundComparisonExpression::IsComparison(*e)) { + auto &cmp = e->Cast(); // Recurse into the scalar side first (it may contain nested expressions) - bool left_is_counter = IsCounterColumnRef(*cmp.left, list_type); - bool right_is_counter = IsCounterColumnRef(*cmp.right, list_type); + bool left_is_counter = IsCounterColumnRef(*cmp.GetChildrenMutable()[0], list_type); + bool right_is_counter = IsCounterColumnRef(*cmp.GetChildrenMutable()[1], list_type); if (left_is_counter && !right_is_counter) { - WrapExpr(cmp.right); + WrapExpr(cmp.GetChildrenMutable()[1]); } else if (right_is_counter && !left_is_counter) { - WrapExpr(cmp.left); + WrapExpr(cmp.GetChildrenMutable()[0]); } auto rewritten = RewriteComparisonToFilterCmp(input, cmp, list_type, finalized_type); if (rewritten) { @@ -472,50 +477,50 @@ static void WrapCounterRefsWithFinalize(OptimizerExtensionInput &input, LogicalO } // Rewrite BoundBetweenExpression where input is a counter column - if (e->GetExpressionClass() == ExpressionClass::BOUND_BETWEEN) { - auto &between = e->Cast(); - if (IsCounterColumnRef(*between.input, list_type)) { - WrapExpr(between.lower); - WrapExpr(between.upper); - auto counter_ref1 = ExtractCounterRef(between.input->Copy(), list_type); - auto counter_ref2 = ExtractCounterRef(std::move(between.input), list_type); - e = RewriteBetweenToFilterCmp(input, std::move(counter_ref1), std::move(counter_ref2), - std::move(between.lower), std::move(between.upper), - between.lower_inclusive, between.upper_inclusive, finalized_type, - list_type); + if (e->GetExpressionType() == ExpressionType::COMPARE_BETWEEN) { + auto &between = e->Cast(); + if (IsCounterColumnRef(*between.GetChildrenMutable()[0], list_type)) { + WrapExpr(between.GetChildrenMutable()[1]); + WrapExpr(between.GetChildrenMutable()[2]); + auto counter_ref1 = ExtractCounterRef(between.GetChildrenMutable()[0]->Copy(), list_type); + auto counter_ref2 = ExtractCounterRef(std::move(between.GetChildrenMutable()[0]), list_type); + e = RewriteBetweenToFilterCmp( + input, std::move(counter_ref1), std::move(counter_ref2), std::move(between.GetChildrenMutable()[1]), + std::move(between.GetChildrenMutable()[2]), BoundBetweenExpression::LowerInclusive(between), + BoundBetweenExpression::UpperInclusive(between), finalized_type, list_type); return; } // Non-counter BETWEEN: recurse into fields manually (not visited by EnumerateChildren) - WrapExpr(between.input); - WrapExpr(between.lower); - WrapExpr(between.upper); + WrapExpr(between.GetChildrenMutable()[0]); + WrapExpr(between.GetChildrenMutable()[1]); + WrapExpr(between.GetChildrenMutable()[2]); } ExpressionIterator::EnumerateChildren(*e, WrapExpr); - if (e->type == ExpressionType::BOUND_COLUMN_REF) { + if (e->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e->Cast(); - if (col_ref.return_type == list_type) { + if (col_ref.GetReturnType() == list_type) { PRIVACY_DEBUG_PRINT("[PAC DERIVED READ] Wrapping binding (" + - std::to_string(col_ref.binding.table_index) + "," + - std::to_string(col_ref.binding.column_index) + ") with priv_finalize"); + std::to_string(col_ref.Binding().table_index.index) + "," + + std::to_string(col_ref.Binding().column_index) + ") with priv_finalize"); e = input.optimizer.BindScalarFunction("priv_finalize", std::move(e)); } } // Retarget CAST(scalar→FLOAT[]) to CAST(scalar→FLOAT), strip no-op casts - if (e->GetExpressionClass() == ExpressionClass::BOUND_CAST) { - auto &cast = e->Cast(); - if (cast.return_type == list_type && cast.child->return_type != list_type) { - e = BoundCastExpression::AddDefaultCastToType(std::move(cast.child), finalized_type); - } else if (cast.child->return_type == cast.return_type) { - e = std::move(cast.child); + if (e->GetExpressionType() == ExpressionType::OPERATOR_CAST) { + auto &cast = e->Cast(); + if (cast.GetReturnType() == list_type && cast.GetChildrenMutable()[0]->GetReturnType() != list_type) { + e = BoundCastExpression::AddDefaultCastToType(std::move(cast.GetChildrenMutable()[0]), finalized_type); + } else if (cast.GetChildrenMutable()[0]->GetReturnType() == cast.GetReturnType()) { + e = std::move(cast.GetChildrenMutable()[0]); } } // Convert FLOAT[] constants back to scalar - if (e->type == ExpressionType::VALUE_CONSTANT) { + if (e->GetExpressionType() == ExpressionType::VALUE_CONSTANT) { auto &constant = e->Cast(); - if (constant.return_type == list_type && !constant.value.IsNull()) { - auto &list_children = ListValue::GetChildren(constant.value); + if (constant.GetReturnType() == list_type && !constant.GetValue().IsNull()) { + auto &list_children = ListValue::GetChildren(constant.GetValue()); if (!list_children.empty()) { auto scalar_val = list_children[0].DefaultCastAs(finalized_type); e = make_uniq(std::move(scalar_val)); @@ -530,21 +535,22 @@ static void WrapCounterRefsWithFinalize(OptimizerExtensionInput &input, LogicalO auto &proj = op->Cast(); for (idx_t i = 0; i < proj.expressions.size(); i++) { if (i < proj.types.size()) { - proj.types[i] = proj.expressions[i]->return_type; + proj.types[i] = proj.expressions[i]->GetReturnType(); } // Only mark as finalized if the expression actually produces a counter type. // Expressions that CONTAIN priv_finalize in sub-expressions (e.g. CASE WHEN priv_finalize(x) > 100 ...) // may return a different type (e.g. VARCHAR) and should NOT be marked finalized. - bool is_counter_output = - (proj.expressions[i]->return_type == finalized_type) || - (proj.expressions[i]->return_type == list_type && ExpressionContainsPacFinalize(*proj.expressions[i])); + bool is_counter_output = (proj.expressions[i]->GetReturnType() == finalized_type) || + (proj.expressions[i]->GetReturnType() == list_type && + ExpressionContainsPacFinalize(*proj.expressions[i])); if (is_counter_output) { - finalized_bindings.insert(HashBinding(ColumnBinding(proj.table_index, i))); + finalized_bindings.insert( + HashBinding(ColumnBinding(TableIndex(proj.table_index.index), ProjectionIndex(i)))); if (i < proj.types.size() && proj.types[i] == list_type) { proj.types[i] = finalized_type; } - if (proj.expressions[i]->return_type == list_type) { - proj.expressions[i]->return_type = finalized_type; + if (proj.expressions[i]->GetReturnType() == list_type) { + proj.expressions[i]->SetReturnType(finalized_type); } } } @@ -582,15 +588,17 @@ static void RewriteCounterFiltersFromGets(OptimizerExtensionInput &input, unique return; } auto &get = op->Cast(); - if (!derived_table_indices.count(get.table_index)) { + if (!derived_table_indices.count(get.table_index.index)) { return; } auto list_type = CounterListType(); auto finalized_type = PacFloatLogicalType(); vector to_pull; - for (auto &entry : get.table_filters.filters) { - auto col_idx = entry.first; - if (col_idx < get.returned_types.size() && get.returned_types[col_idx] == list_type) { + const auto &col_ids = get.GetColumnIds(); + for (auto &entry : get.table_filters) { + auto col_idx = entry.GetIndex(); + // DuckDB 2.0 keys filters by scan projection, not physical table column. + if (col_idx < col_ids.size() && get.GetColumnType(col_ids[col_idx]) == list_type) { to_pull.push_back(col_idx); } } @@ -598,43 +606,47 @@ static void RewriteCounterFiltersFromGets(OptimizerExtensionInput &input, unique return; } auto out_bindings = op->GetColumnBindings(); - auto &col_ids = get.GetColumnIds(); vector> filter_exprs; for (auto filter_col_idx : to_pull) { - auto &filter = get.table_filters.filters[filter_col_idx]; - auto col_type = get.returned_types[filter_col_idx]; + auto &filter = get.table_filters.GetFilterByColumnIndexMutable(ProjectionIndex(filter_col_idx)); + auto col_type = get.GetColumnType(col_ids[filter_col_idx]); // Find the output binding for this column ColumnBinding binding; bool found = false; for (idx_t i = 0; i < col_ids.size() && i < out_bindings.size(); i++) { - if (col_ids[i].GetPrimaryIndex() == filter_col_idx) { + if ((get.projection_ids.empty() ? i : idx_t(get.projection_ids[i])) == filter_col_idx) { binding = out_bindings[i]; found = true; break; } } if (!found) { - get.table_filters.filters.erase(filter_col_idx); - continue; + // Filter-only columns may have been pruned from the scan output. Expose + // them again so the pulled-up privacy predicate can evaluate them. + D_ASSERT(!get.projection_ids.empty()); + get.projection_ids.emplace_back(filter_col_idx); + binding = ColumnBinding(get.table_index, ProjectionIndex(filter_col_idx)); + get.ResolveOperatorTypes(); } auto col_ref = make_uniq(col_type, binding); - auto expr = filter->ToExpression(*col_ref); + auto expr = filter.ToExpression(*col_ref); // Rewrite the expression to priv_filter_ calls - if (expr->GetExpressionClass() == ExpressionClass::BOUND_BETWEEN) { - auto &between = expr->Cast(); + if (expr->GetExpressionType() == ExpressionType::COMPARE_BETWEEN) { + auto &between = expr->Cast(); auto counter1 = make_uniq(col_type, binding); auto counter2 = make_uniq(col_type, binding); filter_exprs.push_back(RewriteBetweenToFilterCmp( - input, std::move(counter1), std::move(counter2), std::move(between.lower), std::move(between.upper), - between.lower_inclusive, between.upper_inclusive, finalized_type, list_type)); - } else if (expr->GetExpressionClass() == ExpressionClass::BOUND_COMPARISON) { - auto &cmp = expr->Cast(); + input, std::move(counter1), std::move(counter2), std::move(between.GetChildrenMutable()[1]), + std::move(between.GetChildrenMutable()[2]), BoundBetweenExpression::LowerInclusive(between), + BoundBetweenExpression::UpperInclusive(between), finalized_type, list_type)); + } else if (BoundComparisonExpression::IsComparison(*expr)) { + auto &cmp = expr->Cast(); auto rewritten = RewriteComparisonToFilterCmp(input, cmp, list_type, finalized_type); filter_exprs.push_back(rewritten ? std::move(rewritten) : std::move(expr)); } else { filter_exprs.push_back(std::move(expr)); } - get.table_filters.filters.erase(filter_col_idx); + get.table_filters.RemoveFilterByColumnIndex(ProjectionIndex(filter_col_idx)); PRIVACY_DEBUG_PRINT("[PAC DERIVED READ] Rewrote table_filter to priv_filter_ for counter column " + std::to_string(filter_col_idx)); } diff --git a/src/query_processing/pac_projection_propagation.cpp b/src/query_processing/pac_projection_propagation.cpp index 4148f95c..2ae22543 100644 --- a/src/query_processing/pac_projection_propagation.cpp +++ b/src/query_processing/pac_projection_propagation.cpp @@ -13,6 +13,8 @@ #include "duckdb/planner/operator/logical_filter.hpp" #include "duckdb/planner/operator/logical_comparison_join.hpp" #include "duckdb/planner/operator/logical_delim_get.hpp" +#include "duckdb/planner/operator/logical_materialized_cte.hpp" +#include "duckdb/planner/operator/logical_cteref.hpp" #include "duckdb/planner/expression/bound_columnref_expression.hpp" #include "duckdb/planner/expression_iterator.hpp" @@ -41,26 +43,31 @@ static bool FindDirectPathToSource(LogicalOperator *current, idx_t target_table_ // Check if this operator matches the target table_index if (current->type == LogicalOperatorType::LOGICAL_GET) { auto &get = current->Cast(); - if (get.table_index == target_table_index) { + if (get.table_index.index == target_table_index) { return true; } } else if (current->type == LogicalOperatorType::LOGICAL_CTE_REF) { - if (current->Cast().table_index == target_table_index) { + if (current->Cast().table_index.index == target_table_index) { return true; } } else if (current->type == LogicalOperatorType::LOGICAL_PROJECTION) { auto &proj = current->Cast(); - if (proj.table_index == target_table_index) { + if (proj.table_index.index == target_table_index) { return true; } } + else if (current->type == LogicalOperatorType::LOGICAL_CTE_REF && + current->Cast().table_index.index == target_table_index) { + return true; + } + // If this is a nested aggregate, check if it IS the target or if the target is below it. // Nested aggregates don't pass through bindings, but PropagateSingleBinding handles // them by adding the binding as a group column — so we continue the search. if (!is_start && current->type == LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY) { auto &agg = current->Cast(); - if (agg.group_index == target_table_index || agg.aggregate_index == target_table_index) { + if (agg.group_index.index == target_table_index || agg.aggregate_index.index == target_table_index) { return true; } // Continue searching through the aggregate's children — PropagateSingleBinding @@ -81,7 +88,7 @@ static bool FindDirectPathToSource(LogicalOperator *current, idx_t target_table_ } // Ensure a single column_index is present in a projection map. Returns true if it was already there. -static bool EnsureInProjectionMap(vector &proj_map, idx_t column_index) { +static bool EnsureInProjectionMap(vector &proj_map, ProjectionIndex column_index) { if (proj_map.empty()) { return true; // Empty map means all columns pass through } @@ -98,20 +105,67 @@ static bool EnsureInProjectionMap(vector &proj_map, idx_t column_index) { static bool BindingInOutput(LogicalOperator &op, const ColumnBinding &binding) { auto bindings = op.GetColumnBindings(); for (auto &b : bindings) { - if (b.table_index == binding.table_index && b.column_index == binding.column_index) { + if (b.table_index.index == binding.table_index.index && b.column_index == binding.column_index) { return true; } } return false; } -ColumnBinding PropagateSingleBindingToOperator(LogicalOperator &plan_root, idx_t source_table_index, - ColumnBinding source_binding, const LogicalType &source_type, - LogicalOperator *target) { - ColumnBinding invalid(DConstants::INVALID_INDEX, DConstants::INVALID_INDEX); +// Find a CTE reader on the output tuple stream, without crossing an aggregation +// or the definition side of another CTE. Correlation-key deduplication must not +// acquire the entity mask as an extra grouping key. +static LogicalCTERef *FindOutputCTERef(LogicalOperator &op, TableIndex cte_index) { + if (op.type == LogicalOperatorType::LOGICAL_CTE_REF) { + auto &ref = op.Cast(); + return ref.cte_index == cte_index ? &ref : nullptr; + } + if (op.type == LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY) { + return nullptr; + } + if (op.type == LogicalOperatorType::LOGICAL_MATERIALIZED_CTE) { + return FindOutputCTERef(*op.children[1], cte_index); + } + for (idx_t i = 0; i < op.children.size(); i++) { + if (IsJoinWithProjectionMap(op.type)) { + auto join_type = op.Cast().join_type; + if ((i == 1 && + (join_type == JoinType::SEMI || join_type == JoinType::ANTI || join_type == JoinType::MARK)) || + (i == 0 && (join_type == JoinType::RIGHT_SEMI || join_type == JoinType::RIGHT_ANTI))) { + continue; + } + } + auto *ref = FindOutputCTERef(*op.children[i], cte_index); + if (ref) { + return ref; + } + } + return nullptr; +} + +static void SyncCTERefColumns(LogicalOperator &op, TableIndex cte_index, const vector &types) { + if (op.type == LogicalOperatorType::LOGICAL_CTE_REF) { + auto &ref = op.Cast(); + if (ref.cte_index == cte_index) { + while (ref.chunk_types.size() < types.size()) { + ref.bound_columns.emplace_back("__privacy_mask_" + std::to_string(ref.chunk_types.size())); + ref.chunk_types.push_back(types[ref.chunk_types.size()]); + } + ref.ResolveOperatorTypes(); + } + } + for (auto &child : op.children) { + SyncCTERefColumns(*child, cte_index, types); + } +} + +static ColumnBinding PropagateBinding(LogicalOperator &plan_root, idx_t source_table_index, + ColumnBinding source_binding, const LogicalType &source_type, + LogicalOperator *target, bool include_target) { + ColumnBinding invalid {TableIndex(DConstants::INVALID_INDEX), ProjectionIndex(DConstants::INVALID_INDEX)}; vector path_ops; - if (!FindDirectPathToSource(target, source_table_index, path_ops, true)) { + if (!FindDirectPathToSource(target, source_table_index, path_ops, !include_target)) { #if PRIVACY_DEBUG PRIVACY_DEBUG_PRINT("PropagateSingleBinding: No direct path from aggregate to source #" + std::to_string(source_table_index)); @@ -137,9 +191,9 @@ ColumnBinding PropagateSingleBindingToOperator(LogicalOperator &plan_root, idx_t // Check if this binding is already in the projection's expressions idx_t existing_idx = DConstants::INVALID_INDEX; for (idx_t i = 0; i < proj.expressions.size(); i++) { - if (proj.expressions[i]->type == ExpressionType::BOUND_COLUMN_REF) { + if (proj.expressions[i]->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &expr_ref = proj.expressions[i]->Cast(); - if (expr_ref.binding == current) { + if (expr_ref.BindingMutable() == current) { existing_idx = i; break; } @@ -147,15 +201,41 @@ ColumnBinding PropagateSingleBindingToOperator(LogicalOperator &plan_root, idx_t } if (existing_idx != DConstants::INVALID_INDEX) { - current = ColumnBinding(proj.table_index, existing_idx); + current = ColumnBinding(TableIndex(proj.table_index.index), ProjectionIndex(existing_idx)); } else { // Add a passthrough column ref for this binding auto col_ref = make_uniq(source_type, current); proj.expressions.push_back(std::move(col_ref)); - current = ColumnBinding(proj.table_index, proj.expressions.size() - 1); + current = + ColumnBinding(TableIndex(proj.table_index.index), ProjectionIndex(proj.expressions.size() - 1)); } proj.ResolveOperatorTypes(); + } else if (op->type == LogicalOperatorType::LOGICAL_MATERIALIZED_CTE && child_idx == 0) { + auto &cte = op->Cast(); + auto *ref = FindOutputCTERef(*cte.children[1], cte.table_index); + if (!ref) { + return invalid; + } + cte.children[0]->ResolveOperatorTypes(); + auto bindings = cte.children[0]->GetColumnBindings(); + auto column = std::find(bindings.begin(), bindings.end(), current); + if (column == bindings.end()) { + return invalid; + } + auto position = idx_t(column - bindings.begin()); + cte.column_count = bindings.size(); + SyncCTERefColumns(*cte.children[1], cte.table_index, cte.children[0]->types); + current = PropagateBinding(plan_root, ref->table_index.index, + ColumnBinding(ref->table_index, ProjectionIndex(position)), source_type, + cte.children[1].get(), true); + if (current.table_index.index == DConstants::INVALID_INDEX) { + return invalid; + } + PRIVACY_DEBUG_PRINT("PropagateSingleBinding: propagated mask through CTE #" + + std::to_string(cte.table_index.index)); + cte.ResolveOperatorTypes(); + } else if (IsJoinWithProjectionMap(op->type)) { auto &join = op->Cast(); bool is_delim_join = (op->type == LogicalOperatorType::LOGICAL_DELIM_JOIN); @@ -190,8 +270,8 @@ ColumnBinding PropagateSingleBindingToOperator(LogicalOperator &plan_root, idx_t if (!BindingInOutput(join, current)) { #if PRIVACY_DEBUG PRIVACY_DEBUG_PRINT("PropagateSingleBinding: WARNING - binding [" + - std::to_string(current.table_index) + "." + std::to_string(current.column_index) + - "] not found in join output"); + std::to_string(current.table_index.index) + "." + + std::to_string(current.column_index) + "] not found in join output"); #endif } // Binding identity preserved through joins @@ -207,10 +287,10 @@ ColumnBinding PropagateSingleBindingToOperator(LogicalOperator &plan_root, idx_t // Check if the binding already passes through as a group column bool found_in_groups = false; for (idx_t i = 0; i < agg.groups.size(); i++) { - if (agg.groups[i]->type == ExpressionType::BOUND_COLUMN_REF) { + if (agg.groups[i]->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = agg.groups[i]->Cast(); - if (col_ref.binding == current) { - current = ColumnBinding(agg.group_index, i); + if (col_ref.BindingMutable() == current) { + current = ColumnBinding(TableIndex(agg.group_index.index), ProjectionIndex(i)); found_in_groups = true; break; } @@ -229,6 +309,12 @@ ColumnBinding PropagateSingleBindingToOperator(LogicalOperator &plan_root, idx_t return current; } +ColumnBinding PropagateSingleBindingToOperator(LogicalOperator &plan_root, idx_t source_table_index, + ColumnBinding source_binding, const LogicalType &source_type, + LogicalOperator *target) { + return PropagateBinding(plan_root, source_table_index, source_binding, source_type, target, false); +} + ColumnBinding PropagateSingleBinding(LogicalOperator &plan_root, idx_t source_table_index, ColumnBinding source_binding, const LogicalType &source_type, LogicalAggregate *target_agg) { return PropagateSingleBindingToOperator(plan_root, source_table_index, source_binding, source_type, target_agg); diff --git a/src/query_processing/pac_subquery_handler.cpp b/src/query_processing/pac_subquery_handler.cpp index f8a15743..14c2875a 100644 --- a/src/query_processing/pac_subquery_handler.cpp +++ b/src/query_processing/pac_subquery_handler.cpp @@ -26,15 +26,15 @@ namespace duckdb { // Traces a binding from the source operator up through the tree, modifying operators as needed. static ColumnBinding EnsureBindingFlowsThroughRecursive(LogicalOperator *op, idx_t source_table_index, ColumnBinding source_binding, const LogicalType &source_type) { - ColumnBinding invalid(DConstants::INVALID_INDEX, DConstants::INVALID_INDEX); + ColumnBinding invalid {TableIndex(DConstants::INVALID_INDEX), ProjectionIndex(DConstants::INVALID_INDEX)}; // Base case: match any operator by table_index if (op->type == LogicalOperatorType::LOGICAL_GET) { - if (op->Cast().table_index == source_table_index) { + if (op->Cast().table_index.index == source_table_index) { return source_binding; } } else if (op->type == LogicalOperatorType::LOGICAL_PROJECTION) { - if (op->Cast().table_index == source_table_index) { + if (op->Cast().table_index.index == source_table_index) { return source_binding; } } @@ -43,12 +43,12 @@ static ColumnBinding EnsureBindingFlowsThroughRecursive(LogicalOperator *op, idx ColumnBinding child_result = invalid; for (auto &child : op->children) { child_result = EnsureBindingFlowsThroughRecursive(child.get(), source_table_index, source_binding, source_type); - if (child_result.table_index != DConstants::INVALID_INDEX) { + if (child_result.table_index.index != DConstants::INVALID_INDEX) { break; } } - if (child_result.table_index == DConstants::INVALID_INDEX) { + if (child_result.table_index.index == DConstants::INVALID_INDEX) { return child_result; } @@ -59,10 +59,10 @@ static ColumnBinding EnsureBindingFlowsThroughRecursive(LogicalOperator *op, idx case LogicalOperatorType::LOGICAL_PROJECTION: { auto &proj = op->Cast(); for (idx_t i = 0; i < proj.expressions.size(); i++) { - if (proj.expressions[i]->type == ExpressionType::BOUND_COLUMN_REF) { + if (proj.expressions[i]->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = proj.expressions[i]->Cast(); - if (col_ref.binding == child_result) { - return ColumnBinding(proj.table_index, i); + if (col_ref.BindingMutable() == child_result) { + return ColumnBinding(TableIndex(proj.table_index.index), ProjectionIndex(i)); } } } @@ -72,10 +72,10 @@ static ColumnBinding EnsureBindingFlowsThroughRecursive(LogicalOperator *op, idx auto &agg = op->Cast(); // Check if binding is already in groups for (idx_t i = 0; i < agg.groups.size(); i++) { - if (agg.groups[i]->type == ExpressionType::BOUND_COLUMN_REF) { + if (agg.groups[i]->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = agg.groups[i]->Cast(); - if (col_ref.binding == child_result) { - return ColumnBinding(agg.group_index, i); + if (col_ref.BindingMutable() == child_result) { + return ColumnBinding(TableIndex(agg.group_index.index), ProjectionIndex(i)); } } } @@ -84,7 +84,7 @@ static ColumnBinding EnsureBindingFlowsThroughRecursive(LogicalOperator *op, idx idx_t new_group_idx = agg.groups.size(); agg.groups.push_back(std::move(group_col_ref)); agg.ResolveOperatorTypes(); - return ColumnBinding(agg.group_index, new_group_idx); + return ColumnBinding(TableIndex(agg.group_index.index), ProjectionIndex(new_group_idx)); } default: return child_result; @@ -105,11 +105,11 @@ static bool HasSourceTableInSubtree(LogicalOperator *op, idx_t source_table_inde return false; } if (op->type == LogicalOperatorType::LOGICAL_GET) { - if (op->Cast().table_index == source_table_index) { + if (op->Cast().table_index.index == source_table_index) { return true; } } else if (op->type == LogicalOperatorType::LOGICAL_PROJECTION) { - if (op->Cast().table_index == source_table_index) { + if (op->Cast().table_index.index == source_table_index) { return true; } } @@ -183,13 +183,13 @@ static void UpdateDelimGetsInSubtree(LogicalOperator *op, const LogicalType &sou if (op->type == LogicalOperatorType::LOGICAL_DELIM_GET) { auto &dg = op->Cast(); #if PRIVACY_DEBUG - PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: Updating DELIM_GET #" + std::to_string(dg.table_index) + + PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: Updating DELIM_GET #" + std::to_string(dg.table_index.index) + " chunk_types before=" + std::to_string(dg.chunk_types.size())); #endif dg.chunk_types.push_back(source_type); op->ResolveOperatorTypes(); #if PRIVACY_DEBUG - PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: DELIM_GET #" + std::to_string(dg.table_index) + + PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: DELIM_GET #" + std::to_string(dg.table_index.index) + " chunk_types after=" + std::to_string(dg.chunk_types.size()) + " types after=" + std::to_string(dg.types.size())); #endif @@ -204,7 +204,8 @@ DelimColumnResult AddBindingToDelimJoin(unique_ptr &plan, idx_t ColumnBinding source_binding, const LogicalType &source_type, LogicalAggregate *target_agg) { DelimColumnResult invalid_result; - invalid_result.binding = ColumnBinding(DConstants::INVALID_INDEX, DConstants::INVALID_INDEX); + invalid_result.binding = + ColumnBinding(TableIndex(DConstants::INVALID_INDEX), ProjectionIndex(DConstants::INVALID_INDEX)); invalid_result.type = LogicalType::INVALID; auto *delim_join = FindDelimJoinForSource(plan.get(), source_table_index, target_agg); @@ -217,9 +218,10 @@ DelimColumnResult AddBindingToDelimJoin(unique_ptr &plan, idx_t } #if PRIVACY_DEBUG - PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: Found DELIM_JOIN, source_table_index=" + - std::to_string(source_table_index) + " binding=[" + std::to_string(source_binding.table_index) + - "." + std::to_string(source_binding.column_index) + "]"); + PRIVACY_DEBUG_PRINT( + "AddBindingToDelimJoin: Found DELIM_JOIN, source_table_index=" + std::to_string(source_table_index) + + " binding=[" + std::to_string(source_binding.table_index.index) + "." + + std::to_string(source_binding.column_index) + "]"); PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: Left child type=" + std::to_string(static_cast(delim_join->children[0]->type))); #endif @@ -229,7 +231,7 @@ DelimColumnResult AddBindingToDelimJoin(unique_ptr &plan, idx_t ColumnBinding output_binding = EnsureBindingFlowsThrough(left_child, source_table_index, source_binding, source_type); - if (output_binding.table_index == DConstants::INVALID_INDEX) { + if (output_binding.table_index.index == DConstants::INVALID_INDEX) { #if PRIVACY_DEBUG PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: EnsureBindingFlowsThrough failed"); #endif @@ -237,20 +239,21 @@ DelimColumnResult AddBindingToDelimJoin(unique_ptr &plan, idx_t } #if PRIVACY_DEBUG - PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: output_binding=[" + std::to_string(output_binding.table_index) + "." + - std::to_string(output_binding.column_index) + "]"); + PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: output_binding=[" + std::to_string(output_binding.table_index.index) + + "." + std::to_string(output_binding.column_index) + "]"); PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: duplicate_eliminated_columns.size()=" + std::to_string(delim_join->duplicate_eliminated_columns.size())); for (idx_t i = 0; i < delim_join->duplicate_eliminated_columns.size(); i++) { auto &dec = delim_join->duplicate_eliminated_columns[i]; PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: existing dup_elim[" + std::to_string(i) + "] = " + dec->ToString() + - " type=" + dec->return_type.ToString()); + " type=" + dec->GetReturnType().ToString()); } // Dump DELIM_JOIN conditions auto &dj = delim_join->Cast(); for (idx_t i = 0; i < dj.conditions.size(); i++) { - PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: condition[" + std::to_string(i) + "] left=" + - dj.conditions[i].left->ToString() + " right=" + dj.conditions[i].right->ToString()); + PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: condition[" + std::to_string(i) + + "] left=" + dj.conditions[i].LeftReference()->ToString() + + " right=" + dj.conditions[i].RightReference()->ToString()); } #endif @@ -271,13 +274,13 @@ DelimColumnResult AddBindingToDelimJoin(unique_ptr &plan, idx_t } #if PRIVACY_DEBUG - PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: DELIM_GET #" + std::to_string(delim_get->table_index) + " now has " + - std::to_string(delim_get->chunk_types.size()) + " chunk_types, " + + PRIVACY_DEBUG_PRINT("AddBindingToDelimJoin: DELIM_GET #" + std::to_string(delim_get->table_index.index) + + " now has " + std::to_string(delim_get->chunk_types.size()) + " chunk_types, " + std::to_string(delim_get->types.size()) + " types"); #endif DelimColumnResult result; - result.binding = ColumnBinding(delim_get->table_index, new_col_idx); + result.binding = ColumnBinding(TableIndex(delim_get->table_index.index), ProjectionIndex(new_col_idx)); result.type = source_type; return result; } @@ -286,7 +289,8 @@ DelimColumnResult AddBindingToDelimJoin(unique_ptr &plan, idx_t DelimColumnResult AddColumnToDelimJoin(unique_ptr &plan, LogicalGet &source_get, const string &column_name, LogicalAggregate *target_agg) { DelimColumnResult invalid_result; - invalid_result.binding = ColumnBinding(DConstants::INVALID_INDEX, DConstants::INVALID_INDEX); + invalid_result.binding = + ColumnBinding(TableIndex(DConstants::INVALID_INDEX), ProjectionIndex(DConstants::INVALID_INDEX)); invalid_result.type = LogicalType::INVALID; // Ensure the column is projected in source_get @@ -312,14 +316,14 @@ DelimColumnResult AddColumnToDelimJoin(unique_ptr &plan, Logica } } if (output_col_idx == DConstants::INVALID_INDEX) { - source_get.projection_ids.push_back(col_proj_idx); + source_get.projection_ids.push_back(ProjectionIndex(col_proj_idx)); output_col_idx = col_proj_idx; source_get.ResolveOperatorTypes(); } } - ColumnBinding source_binding(source_get.table_index, output_col_idx); - return AddBindingToDelimJoin(plan, source_get.table_index, source_binding, col_type, target_agg); + ColumnBinding source_binding {TableIndex(source_get.table_index.index), ProjectionIndex(output_col_idx)}; + return AddBindingToDelimJoin(plan, source_get.table_index.index, source_binding, col_type, target_agg); } } // namespace duckdb diff --git a/src/query_processing/pac_topk_rewriter.cpp b/src/query_processing/pac_topk_rewriter.cpp index 13e1085f..bceeb63a 100644 --- a/src/query_processing/pac_topk_rewriter.cpp +++ b/src/query_processing/pac_topk_rewriter.cpp @@ -5,6 +5,7 @@ // See pac_topk_rewriter.hpp for design documentation. // +#include "duckdb/common/vector/list_vector.hpp" #include "query_processing/pac_topk_rewriter.hpp" #include "privacy_debug.hpp" #include "duckdb/parser/parsed_data/create_scalar_function_info.hpp" @@ -39,14 +40,14 @@ static void PacMeanFunction(DataChunk &args, ExpressionState &state, Vector &res idx_t count = args.size(); UnifiedVectorFormat list_data; - list_vec.ToUnifiedFormat(count, list_data); + list_vec.ToUnifiedFormat(list_data); - auto result_data = FlatVector::GetData(result); - auto &result_validity = FlatVector::Validity(result); + auto result_data = FlatVector::GetDataMutable(result); + auto &result_validity = FlatVector::ValidityMutable(result); - auto &child_vec = ListVector::GetEntry(list_vec); + auto &child_vec = ListVector::GetChild(list_vec); UnifiedVectorFormat child_data; - child_vec.ToUnifiedFormat(ListVector::GetListSize(list_vec), child_data); + child_vec.ToUnifiedFormat(child_data); auto child_values = UnifiedVectorFormat::GetData(child_data); auto list_entries = UnifiedVectorFormat::GetData(list_data); @@ -108,11 +109,11 @@ struct PacTopKAggInfo { }; static void RemapExpressionTableIndices(Expression &expr, const std::unordered_map &index_map) { - if (expr.type == ExpressionType::BOUND_COLUMN_REF) { + if (expr.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = expr.Cast(); - auto it = index_map.find(col_ref.binding.table_index); + auto it = index_map.find(col_ref.BindingMutable().table_index.index); if (it != index_map.end()) { - col_ref.binding.table_index = it->second; + col_ref.BindingMutable().table_index.index = it->second; } } ExpressionIterator::EnumerateChildren(expr, @@ -127,49 +128,49 @@ static unique_ptr CopyExpressionWithTableMap(const Expression &expr, } static Expression *StripCastsMutable(Expression *expr) { - while (expr && expr->type == ExpressionType::OPERATOR_CAST) { - expr = expr->Cast().child.get(); + while (expr && expr->GetExpressionType() == ExpressionType::OPERATOR_CAST) { + expr = expr->Cast().GetChildrenMutable()[0].get(); } return expr; } static const Expression *StripCastsConst(const Expression *expr) { - while (expr && expr->type == ExpressionType::OPERATOR_CAST) { - expr = expr->Cast().child.get(); + while (expr && expr->GetExpressionType() == ExpressionType::OPERATOR_CAST) { + expr = expr->Cast().GetChildren()[0].get(); } return expr; } static bool GetDirectColumnBinding(const Expression &expr, ColumnBinding &binding) { auto stripped = StripCastsConst(&expr); - if (!stripped || stripped->type != ExpressionType::BOUND_COLUMN_REF) { + if (!stripped || stripped->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { return false; } - binding = stripped->Cast().binding; + binding = stripped->Cast().Binding(); return true; } static bool RemapDirectOrderExpression(unique_ptr &expr, const std::unordered_map &map) { auto *stripped = StripCastsMutable(expr.get()); - if (!stripped || stripped->type != ExpressionType::BOUND_COLUMN_REF) { + if (!stripped || stripped->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { return false; } auto &ref = stripped->Cast(); - auto it = map.find(HashBinding(ref.binding)); + auto it = map.find(HashBinding(ref.BindingMutable())); if (it == map.end()) { return false; } - ref.binding = it->second; + ref.BindingMutable() = it->second; return true; } static void RemapExpressionBindingsExact(Expression &expr, const std::unordered_map &map) { - if (expr.type == ExpressionType::BOUND_COLUMN_REF) { + if (expr.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = expr.Cast(); - auto it = map.find(HashBinding(col_ref.binding)); + auto it = map.find(HashBinding(col_ref.BindingMutable())); if (it != map.end()) { - col_ref.binding = it->second; + col_ref.BindingMutable() = it->second; } } ExpressionIterator::EnumerateChildren(expr, [&](Expression &child) { RemapExpressionBindingsExact(child, map); }); @@ -262,11 +263,12 @@ static TopKContext FindTopKContext(const unique_ptr &plan) { // tracing through any number of intermediate projections. static bool OrderExprReferencesPacAgg(Expression &expr, const vector &pac_aggs, const vector &projections, idx_t proj_depth) { - if (expr.type == ExpressionType::BOUND_COLUMN_REF) { + if (expr.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = expr.Cast(); // If there are more projections to trace through, resolve the column ref - if (proj_depth < projections.size() && col_ref.binding.table_index == projections[proj_depth]->table_index) { - auto col_idx = col_ref.binding.column_index; + if (proj_depth < projections.size() && + col_ref.BindingMutable().table_index.index == projections[proj_depth]->table_index.index) { + auto col_idx = col_ref.BindingMutable().column_index; if (col_idx < projections[proj_depth]->expressions.size()) { auto &proj_expr = projections[proj_depth]->expressions[col_idx]; return OrderExprReferencesPacAgg(*proj_expr, pac_aggs, projections, proj_depth + 1); @@ -274,15 +276,15 @@ static bool OrderExprReferencesPacAgg(Expression &expr, const vector(); - return OrderExprReferencesPacAgg(*cast.child, pac_aggs, projections, proj_depth); + if (expr.GetExpressionType() == ExpressionType::OPERATOR_CAST) { + auto &cast = expr.Cast(); + return OrderExprReferencesPacAgg(*cast.GetChildrenMutable()[0], pac_aggs, projections, proj_depth); } return false; } @@ -292,16 +294,16 @@ static vector FindPacAggregates(LogicalAggregate &agg) { vector result; for (idx_t i = 0; i < agg.expressions.size(); i++) { auto &expr = agg.expressions[i]; - if (expr->type != ExpressionType::BOUND_AGGREGATE) { + if (expr->GetExpressionType() != ExpressionType::BOUND_AGGREGATE) { continue; } auto &bound_agg = expr->Cast(); - if (IsPacAggregate(bound_agg.function.name)) { + if (IsPacAggregate(bound_agg.FunctionMutable().GetName().GetIdentifierName())) { PacTopKAggInfo info; info.agg_index = i; - info.original_name = bound_agg.function.name; - info.original_type = expr->return_type; - info.agg_binding = ColumnBinding(agg.aggregate_index, i); + info.original_name = bound_agg.FunctionMutable().GetName().GetIdentifierName(); + info.original_type = expr->GetReturnType(); + info.agg_binding = ColumnBinding(TableIndex(agg.aggregate_index.index), ProjectionIndex(i)); result.push_back(std::move(info)); } } @@ -322,7 +324,7 @@ static bool TraceBindingThroughProjections(const ColumnBinding &binding, const v return true; } auto *proj = projections[depth]; - if (binding.table_index != proj->table_index) { + if (binding.table_index.index != proj->table_index.index) { out_binding = binding; return true; } @@ -331,11 +333,11 @@ static bool TraceBindingThroughProjections(const ColumnBinding &binding, const v return false; } auto &proj_expr = proj->expressions[col_idx]; - if (proj_expr->type != ExpressionType::BOUND_COLUMN_REF) { + if (proj_expr->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { return false; } auto &inner_ref = proj_expr->Cast(); - return TraceBindingThroughProjections(inner_ref.binding, projections, depth + 1, out_binding); + return TraceBindingThroughProjections(inner_ref.BindingMutable(), projections, depth + 1, out_binding); } // Remap column bindings in an operator's expressions after a child was rewritten. @@ -359,11 +361,11 @@ static void RemapBindingsInOperator(LogicalOperator &op, const vector &)> Rewrite = [&](unique_ptr &e) { - if (e->type == ExpressionType::BOUND_COLUMN_REF) { + if (e->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e->Cast(); - auto it = remap.find(HashBinding(col_ref.binding)); + auto it = remap.find(HashBinding(col_ref.BindingMutable())); if (it != remap.end()) { - col_ref.binding = it->second; + col_ref.BindingMutable() = it->second; } } ExpressionIterator::EnumerateChildren(*e, Rewrite); @@ -385,8 +387,8 @@ static void RemapBindingsInOperator(LogicalOperator &op, const vector(); for (auto &cond : join.conditions) { - Rewrite(cond.left); - Rewrite(cond.right); + Rewrite(cond.LeftReference()); + Rewrite(cond.RightReference()); } } @@ -415,7 +417,7 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq std::unordered_set pac_agg_indices; for (auto &info : pac_aggs) { if (info.agg_index >= agg.expressions.size() || - agg.expressions[info.agg_index]->type != ExpressionType::BOUND_AGGREGATE) { + agg.expressions[info.agg_index]->GetExpressionType() != ExpressionType::BOUND_AGGREGATE) { return false; } auto &bound_agg = agg.expressions[info.agg_index]->Cast(); @@ -430,17 +432,19 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq for (auto &order : topn.orders) { ColumnBinding proj_binding; if (!GetDirectColumnBinding(*order.expression, proj_binding) || - proj_binding.table_index != proj.table_index || proj_binding.column_index >= proj.expressions.size()) { + proj_binding.table_index.index != proj.table_index.index || + proj_binding.column_index >= proj.expressions.size()) { return false; } ColumnBinding agg_binding; if (!GetDirectColumnBinding(*proj.expressions[proj_binding.column_index], agg_binding)) { return false; } - if (agg_binding.table_index == agg.group_index) { + if (agg_binding.table_index.index == agg.group_index.index) { continue; } - if (agg_binding.table_index != agg.aggregate_index || agg_binding.column_index >= agg.expressions.size() || + if (agg_binding.table_index.index != agg.aggregate_index.index || + agg_binding.column_index >= agg.expressions.size() || pac_agg_indices.find(agg_binding.column_index) == pac_agg_indices.end()) { return false; } @@ -460,7 +464,7 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq for (idx_t pi = 0; pi < proj.expressions.size(); pi++) { ColumnBinding binding; if (GetDirectColumnBinding(*proj.expressions[pi], binding) && - binding == ColumnBinding(agg.group_index, gi)) { + binding == ColumnBinding(TableIndex(agg.group_index.index), ProjectionIndex(gi))) { group_proj_cols.push_back(pi); found = true; break; @@ -495,9 +499,9 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq #endif auto &binder = input.optimizer.binder; - auto original_group_index = agg.group_index; - auto original_agg_index = agg.aggregate_index; - auto original_proj_index = proj.table_index; + auto original_group_index = agg.group_index.index; + auto original_agg_index = agg.aggregate_index.index; + auto original_proj_index = proj.table_index.index; vector original_agg_types = agg.types; vector original_proj_types = proj.types; @@ -533,18 +537,18 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq std::unordered_map agg_to_top_remap; for (idx_t gi = 0; gi < agg.groups.size(); gi++) { - agg_to_top_remap[HashBinding(ColumnBinding(original_group_index, gi))] = - ColumnBinding(original_group_index, gi); + agg_to_top_remap[HashBinding(ColumnBinding(TableIndex(original_group_index), ProjectionIndex(gi)))] = + ColumnBinding(TableIndex(original_group_index), ProjectionIndex(gi)); } for (auto &kv : order_old_to_new) { - agg_to_top_remap[HashBinding(ColumnBinding(original_agg_index, kv.first))] = - ColumnBinding(original_agg_index, kv.second); + agg_to_top_remap[HashBinding(ColumnBinding(TableIndex(original_agg_index), ProjectionIndex(kv.first)))] = + ColumnBinding(TableIndex(original_agg_index), ProjectionIndex(kv.second)); } agg.expressions = std::move(top_expressions); agg.types.clear(); for (auto &group : agg.groups) { - agg.types.push_back(group->return_type); + agg.types.push_back(group->GetReturnType()); } for (auto old_idx : order_agg_indices) { agg.types.push_back(original_agg_types[agg.groups.size() + old_idx]); @@ -566,8 +570,8 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq } std::unordered_map proj_to_top_remap; for (auto &kv : proj_old_to_new) { - proj_to_top_remap[HashBinding(ColumnBinding(original_proj_index, kv.first))] = - ColumnBinding(original_proj_index, kv.second); + proj_to_top_remap[HashBinding(ColumnBinding(TableIndex(original_proj_index), ProjectionIndex(kv.first)))] = + ColumnBinding(TableIndex(original_proj_index), ProjectionIndex(kv.second)); } for (auto &order : topn.orders) { if (!RemapDirectOrderExpression(order.expression, proj_to_top_remap)) { @@ -581,12 +585,13 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq auto left_topn = std::move(plan); auto join = make_uniq(JoinType::INNER); for (idx_t gi = 0; gi < group_proj_cols.size(); gi++) { - JoinCondition cond; + JoinCondition cond(nullptr, nullptr, ExpressionType::COMPARE_NOT_DISTINCT_FROM); auto left_col = proj_old_to_new[group_proj_cols[gi]]; - cond.left = make_uniq(original_proj_types[group_proj_cols[gi]], - ColumnBinding(original_proj_index, left_col)); - cond.right = CopyExpressionWithTableMap(*original_groups[gi], right_index_map); - cond.comparison = ExpressionType::COMPARE_NOT_DISTINCT_FROM; + cond.LeftReference() = make_uniq( + original_proj_types[group_proj_cols[gi]], + ColumnBinding(TableIndex(original_proj_index), ProjectionIndex(left_col))); + cond.RightReference() = CopyExpressionWithTableMap(*original_groups[gi], right_index_map); + join->conditions.push_back(std::move(cond)); } join->children.push_back(std::move(left_topn)); @@ -601,30 +606,32 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq payload_expressions.push_back(CopyExpressionWithTableMap(*original_aggs[old_idx], right_index_map)); } - auto payload_group_index = binder.GenerateTableIndex(); - auto payload_agg_index = binder.GenerateTableIndex(); + auto payload_group_index = binder.GenerateTableIndex().index; + auto payload_agg_index = binder.GenerateTableIndex().index; vector> payload_groups; for (auto old_col : group_proj_cols) { auto new_col = proj_old_to_new[old_col]; - payload_groups.push_back(make_uniq(original_proj_types[old_col], - ColumnBinding(original_proj_index, new_col))); + payload_groups.push_back(make_uniq( + original_proj_types[old_col], + ColumnBinding(TableIndex(original_proj_index), ProjectionIndex(new_col)))); } for (auto old_col : order_proj_cols) { auto new_col = proj_old_to_new[old_col]; - payload_groups.push_back(make_uniq(original_proj_types[old_col], - ColumnBinding(original_proj_index, new_col))); + payload_groups.push_back(make_uniq( + original_proj_types[old_col], + ColumnBinding(TableIndex(original_proj_index), ProjectionIndex(new_col)))); } - auto payload_agg = - make_uniq(payload_group_index, payload_agg_index, std::move(payload_expressions)); + auto payload_agg = make_uniq(TableIndex(payload_group_index), TableIndex(payload_agg_index), + std::move(payload_expressions)); payload_agg->groups = std::move(payload_groups); payload_agg->children.push_back(std::move(join)); payload_agg->ResolveOperatorTypes(); std::unordered_map final_expr_remap; for (idx_t gi = 0; gi < group_proj_cols.size(); gi++) { - final_expr_remap[HashBinding(ColumnBinding(original_group_index, gi))] = - ColumnBinding(payload_group_index, gi); + final_expr_remap[HashBinding(ColumnBinding(TableIndex(original_group_index), ProjectionIndex(gi)))] = + ColumnBinding(TableIndex(payload_group_index), ProjectionIndex(gi)); } for (idx_t oi = 0; oi < order_proj_cols.size(); oi++) { ColumnBinding agg_binding; @@ -632,22 +639,22 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq return false; } final_expr_remap[HashBinding(agg_binding)] = - ColumnBinding(payload_group_index, group_proj_cols.size() + oi); + ColumnBinding(TableIndex(payload_group_index), ProjectionIndex(group_proj_cols.size() + oi)); } for (auto &kv : payload_old_to_new) { - final_expr_remap[HashBinding(ColumnBinding(original_agg_index, kv.first))] = - ColumnBinding(payload_agg_index, kv.second); + final_expr_remap[HashBinding(ColumnBinding(TableIndex(original_agg_index), ProjectionIndex(kv.first)))] = + ColumnBinding(TableIndex(payload_agg_index), ProjectionIndex(kv.second)); } - auto final_proj_index = binder.GenerateTableIndex(); + auto final_proj_index = binder.GenerateTableIndex().index; vector> final_exprs; std::unordered_map final_order_remap; for (idx_t pi = 0; pi < original_proj_exprs.size(); pi++) { - final_order_remap[HashBinding(ColumnBinding(original_proj_index, pi))] = - ColumnBinding(final_proj_index, final_exprs.size()); + final_order_remap[HashBinding(ColumnBinding(TableIndex(original_proj_index), ProjectionIndex(pi)))] = + ColumnBinding(TableIndex(final_proj_index), ProjectionIndex(final_exprs.size())); final_exprs.push_back(CopyExpressionWithBindingMap(*original_proj_exprs[pi], final_expr_remap)); } - auto final_proj = make_uniq(final_proj_index, std::move(final_exprs)); + auto final_proj = make_uniq(TableIndex(final_proj_index), std::move(final_exprs)); final_proj->children.push_back(std::move(payload_agg)); final_proj->ResolveOperatorTypes(); @@ -684,7 +691,7 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq return false; } auto &expr = agg.expressions[info.agg_index]; - if (expr->type != ExpressionType::BOUND_AGGREGATE) { + if (expr->GetExpressionType() != ExpressionType::BOUND_AGGREGATE) { return false; } auto &bound_agg = expr->Cast(); @@ -700,13 +707,13 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq if (!GetDirectColumnBinding(*order.expression, binding)) { return false; } - if (binding.table_index == agg.group_index) { + if (binding.table_index.index == agg.group_index.index) { if (binding.column_index >= agg.groups.size()) { return false; } continue; } - if (binding.table_index != agg.aggregate_index || binding.column_index >= agg.expressions.size()) { + if (binding.table_index.index != agg.aggregate_index.index || binding.column_index >= agg.expressions.size()) { return false; } if (pac_agg_indices.find(binding.column_index) == pac_agg_indices.end()) { @@ -744,8 +751,8 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq #endif auto &binder = input.optimizer.binder; - auto original_group_index = agg.group_index; - auto original_agg_index = agg.aggregate_index; + auto original_group_index = agg.group_index.index; + auto original_agg_index = agg.aggregate_index.index; vector original_types = agg.types; vector> original_groups; @@ -779,11 +786,12 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq std::unordered_map top_order_remap; for (idx_t gi = 0; gi < agg.groups.size(); gi++) { - top_order_remap[HashBinding(ColumnBinding(original_group_index, gi))] = ColumnBinding(original_group_index, gi); + top_order_remap[HashBinding(ColumnBinding(TableIndex(original_group_index), ProjectionIndex(gi)))] = + ColumnBinding(TableIndex(original_group_index), ProjectionIndex(gi)); } for (auto &kv : order_old_to_new) { - top_order_remap[HashBinding(ColumnBinding(original_agg_index, kv.first))] = - ColumnBinding(original_agg_index, kv.second); + top_order_remap[HashBinding(ColumnBinding(TableIndex(original_agg_index), ProjectionIndex(kv.first)))] = + ColumnBinding(TableIndex(original_agg_index), ProjectionIndex(kv.second)); } for (auto &order : topn.orders) { @@ -795,7 +803,7 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq agg.expressions = std::move(top_expressions); agg.types.clear(); for (auto &group : agg.groups) { - agg.types.push_back(group->return_type); + agg.types.push_back(group->GetReturnType()); } for (auto old_idx : order_agg_indices) { agg.types.push_back(original_types[agg.groups.size() + old_idx]); @@ -806,11 +814,11 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq auto left_topn = std::move(plan); auto join = make_uniq(JoinType::INNER); for (idx_t gi = 0; gi < original_groups.size(); gi++) { - JoinCondition cond; - cond.left = make_uniq(original_groups[gi]->return_type, - ColumnBinding(original_group_index, gi)); - cond.right = CopyExpressionWithTableMap(*original_groups[gi], right_index_map); - cond.comparison = ExpressionType::COMPARE_NOT_DISTINCT_FROM; + JoinCondition cond(nullptr, nullptr, ExpressionType::COMPARE_NOT_DISTINCT_FROM); + cond.LeftReference() = make_uniq( + original_groups[gi]->GetReturnType(), ColumnBinding(TableIndex(original_group_index), ProjectionIndex(gi))); + cond.RightReference() = CopyExpressionWithTableMap(*original_groups[gi], right_index_map); + join->conditions.push_back(std::move(cond)); } join->children.push_back(std::move(left_topn)); @@ -826,56 +834,60 @@ static bool TryRewriteLatePayloadAggregates(OptimizerExtensionInput &input, uniq payload_expressions.push_back(CopyExpressionWithTableMap(*original_aggs[old_idx], right_index_map)); } - auto payload_group_index = binder.GenerateTableIndex(); - auto payload_agg_index = binder.GenerateTableIndex(); + auto payload_group_index = binder.GenerateTableIndex().index; + auto payload_agg_index = binder.GenerateTableIndex().index; vector> payload_groups; payload_groups.reserve(original_groups.size() + order_agg_indices.size()); for (idx_t gi = 0; gi < original_groups.size(); gi++) { - payload_groups.push_back(make_uniq(original_groups[gi]->return_type, - ColumnBinding(original_group_index, gi))); + payload_groups.push_back( + make_uniq(original_groups[gi]->GetReturnType(), + ColumnBinding(TableIndex(original_group_index), ProjectionIndex(gi)))); } for (auto old_idx : order_agg_indices) { auto new_idx = order_old_to_new[old_idx]; - payload_groups.push_back(make_uniq(original_types[original_groups.size() + old_idx], - ColumnBinding(original_agg_index, new_idx))); + payload_groups.push_back(make_uniq( + original_types[original_groups.size() + old_idx], + ColumnBinding(TableIndex(original_agg_index), ProjectionIndex(new_idx)))); } - auto payload_agg = - make_uniq(payload_group_index, payload_agg_index, std::move(payload_expressions)); + auto payload_agg = make_uniq(TableIndex(payload_group_index), TableIndex(payload_agg_index), + std::move(payload_expressions)); payload_agg->groups = std::move(payload_groups); payload_agg->children.push_back(std::move(join)); payload_agg->ResolveOperatorTypes(); - auto final_proj_index = binder.GenerateTableIndex(); + auto final_proj_index = binder.GenerateTableIndex().index; vector> final_exprs; final_exprs.reserve(original_types.size()); std::unordered_map final_order_remap; for (idx_t gi = 0; gi < original_groups.size(); gi++) { - auto final_binding = ColumnBinding(final_proj_index, final_exprs.size()); - final_order_remap[HashBinding(ColumnBinding(original_group_index, gi))] = final_binding; - final_exprs.push_back( - make_uniq(original_types[gi], ColumnBinding(payload_group_index, gi))); + auto final_binding = ColumnBinding(TableIndex(final_proj_index), ProjectionIndex(final_exprs.size())); + final_order_remap[HashBinding(ColumnBinding(TableIndex(original_group_index), ProjectionIndex(gi)))] = + final_binding; + final_exprs.push_back(make_uniq( + original_types[gi], ColumnBinding(TableIndex(payload_group_index), ProjectionIndex(gi)))); } for (idx_t ai = 0; ai < original_aggs.size(); ai++) { - auto final_binding = ColumnBinding(final_proj_index, final_exprs.size()); - final_order_remap[HashBinding(ColumnBinding(original_agg_index, ai))] = final_binding; + auto final_binding = ColumnBinding(TableIndex(final_proj_index), ProjectionIndex(final_exprs.size())); + final_order_remap[HashBinding(ColumnBinding(TableIndex(original_agg_index), ProjectionIndex(ai)))] = + final_binding; auto original_type = original_types[original_groups.size() + ai]; auto order_it = order_old_to_new.find(ai); if (order_it != order_old_to_new.end()) { auto group_pos = original_groups.size() + order_it->second; - final_exprs.push_back( - make_uniq(original_type, ColumnBinding(payload_group_index, group_pos))); + final_exprs.push_back(make_uniq( + original_type, ColumnBinding(TableIndex(payload_group_index), ProjectionIndex(group_pos)))); continue; } auto payload_it = payload_old_to_new.find(ai); if (payload_it == payload_old_to_new.end()) { return false; } - final_exprs.push_back( - make_uniq(original_type, ColumnBinding(payload_agg_index, payload_it->second))); + final_exprs.push_back(make_uniq( + original_type, ColumnBinding(TableIndex(payload_agg_index), ProjectionIndex(payload_it->second)))); } - auto final_proj = make_uniq(final_proj_index, std::move(final_exprs)); + auto final_proj = make_uniq(TableIndex(final_proj_index), std::move(final_exprs)); final_proj->children.push_back(std::move(payload_agg)); final_proj->ResolveOperatorTypes(); @@ -976,7 +988,7 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique for (auto &info : pac_aggs) { PRIVACY_DEBUG_PRINT("PACTopKRule: Found PAC aggregate: " + info.original_name + " at index " + std::to_string(info.agg_index) + " binding=[" + - std::to_string(info.agg_binding.table_index) + "." + + std::to_string(info.agg_binding.table_index.index) + "." + std::to_string(info.agg_binding.column_index) + "]"); } #endif @@ -1027,7 +1039,7 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique double expansion = 1.0; { Value exp_val; - if (input.context.TryGetCurrentSetting("pac_topk_expansion", exp_val) && !exp_val.IsNull()) { + if (input.context.TryGetCurrentSetting(Identifier("pac_topk_expansion"), exp_val) && !exp_val.IsNull()) { expansion = exp_val.GetValue(); if (expansion < 1.0) { expansion = 1.0; @@ -1054,11 +1066,12 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique auto &expr = agg.expressions[info.agg_index]; auto &bound_agg = expr->Cast(); vector> children; - for (auto &child_expr : bound_agg.children) { + for (auto &child_expr : bound_agg.GetChildrenMutable()) { children.push_back(child_expr->Copy()); } - auto new_aggr = RebindAggregate(input.context, GetCountersVariant(bound_agg.function.name), std::move(children), - bound_agg.IsDistinct()); + auto new_aggr = RebindAggregate(input.context, + GetCountersVariant(bound_agg.FunctionMutable().GetName().GetIdentifierName()), + std::move(children), bound_agg.IsDistinct()); if (new_aggr) { #if PRIVACY_DEBUG PRIVACY_DEBUG_PRINT("PACTopKRule: Rebound " + info.original_name + " -> " + @@ -1069,11 +1082,12 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique // Fallback: rename in place #if PRIVACY_DEBUG PRIVACY_DEBUG_PRINT("PACTopKRule: Fallback rename " + info.original_name + " -> " + - GetCountersVariant(bound_agg.function.name)); + GetCountersVariant(bound_agg.Function().GetName().GetIdentifierName())); #endif - bound_agg.function.name = GetCountersVariant(bound_agg.function.name); - bound_agg.function.return_type = list_type; - expr->return_type = list_type; + bound_agg.FunctionMutable().SetName( + Identifier(GetCountersVariant(bound_agg.FunctionMutable().GetName().GetIdentifierName()))); + bound_agg.FunctionMutable().SetReturnType(list_type); + expr->SetReturnType(list_type); } // Update types in the aggregate idx_t types_index = agg.groups.size() + info.agg_index; @@ -1090,7 +1104,7 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique #if PRIVACY_DEBUG PRIVACY_DEBUG_PRINT("PACTopKRule: Step 2 - Building mean projection"); #endif - idx_t mean_proj_idx = binder.GenerateTableIndex(); + idx_t mean_proj_idx = binder.GenerateTableIndex().index; vector> mean_proj_exprs; // Get all column bindings from the aggregate @@ -1115,7 +1129,7 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique pac_mean_column_map[HashBinding(info.agg_binding)] = mean_col_idx; } - auto mean_proj = make_uniq(mean_proj_idx, std::move(mean_proj_exprs)); + auto mean_proj = make_uniq(TableIndex(mean_proj_idx), std::move(mean_proj_exprs)); if (ctx.has_intermediate_projection) { // ================================================================== @@ -1144,14 +1158,14 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique unordered_map order_pac_refs; for (idx_t oi = 0; oi < ctx.topn->orders.size(); oi++) { auto *stripped = StripCasts(ctx.topn->orders[oi].expression.get()); - auto *col_ref = stripped->type == ExpressionType::BOUND_COLUMN_REF + auto *col_ref = stripped->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF ? &stripped->Cast() : nullptr; if (!col_ref) { continue; } ColumnBinding resolved; - if (TraceBindingThroughProjections(col_ref->binding, ctx.intermediate_projections, 0, resolved)) { + if (TraceBindingThroughProjections(col_ref->BindingMutable(), ctx.intermediate_projections, 0, resolved)) { for (idx_t j = 0; j < pac_aggs.size(); j++) { if (resolved == pac_aggs[j].agg_binding) { order_pac_refs[oi] = j; @@ -1164,12 +1178,12 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique // Map: outermost projection column index -> pac_aggs index (for PAC aggregate columns) unordered_map outermost_pac_col_map; for (idx_t i = 0; i < outermost->expressions.size(); i++) { - if (outermost->expressions[i]->type != ExpressionType::BOUND_COLUMN_REF) { + if (outermost->expressions[i]->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { continue; // function calls (e.g. decompress) are non-PAC group columns } auto &ref = outermost->expressions[i]->Cast(); ColumnBinding resolved; - if (TraceBindingThroughProjections(ref.binding, ctx.intermediate_projections, 1, resolved)) { + if (TraceBindingThroughProjections(ref.BindingMutable(), ctx.intermediate_projections, 1, resolved)) { for (idx_t j = 0; j < pac_aggs.size(); j++) { if (resolved == pac_aggs[j].agg_binding) { outermost_pac_col_map[i] = j; @@ -1194,12 +1208,12 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique // Rewrite all ColRefs in innermost projection: aggregate bindings -> MeanProj bindings. // Also update return_type since PAC columns changed from e.g. BIGINT to LIST. std::function &)> UpdateRefs = [&](unique_ptr &e) { - if (e->type == ExpressionType::BOUND_COLUMN_REF) { + if (e->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e->Cast(); for (idx_t i = 0; i < agg_bindings.size(); i++) { - if (col_ref.binding == agg_bindings[i]) { - col_ref.binding = ColumnBinding(mean_proj_idx, i); - col_ref.return_type = agg.types[i]; + if (col_ref.BindingMutable() == agg_bindings[i]) { + col_ref.BindingMutable() = ColumnBinding(TableIndex(mean_proj_idx), ProjectionIndex(i)); + col_ref.SetReturnType(agg.types[i]); break; } } @@ -1228,13 +1242,14 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique // Innermost: reference MeanProj's priv_mean column directly idx_t mean_col = pac_mean_column_map[HashBinding(pac_aggs[j].agg_binding)]; proj->expressions.push_back(make_uniq( - PacFloatLogicalType(), ColumnBinding(mean_proj_idx, mean_col))); + PacFloatLogicalType(), ColumnBinding(TableIndex(mean_proj_idx), ProjectionIndex(mean_col)))); } else { // Reference the projection below's priv_mean passthrough column auto *below = ctx.intermediate_projections[pi + 1]; idx_t below_col = pac_mean_col_indices[pi + 1][j]; proj->expressions.push_back(make_uniq( - PacFloatLogicalType(), ColumnBinding(below->table_index, below_col))); + PacFloatLogicalType(), + ColumnBinding(TableIndex(below->table_index.index), ProjectionIndex(below_col)))); } } } @@ -1248,11 +1263,11 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique auto &child_types = proj->children[0]->types; // Update ColRef return_types based on the child operator's resolved types std::function &)> FixTypes = [&](unique_ptr &e) { - if (e->type == ExpressionType::BOUND_COLUMN_REF) { + if (e->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e->Cast(); for (idx_t ci = 0; ci < child_bindings.size(); ci++) { - if (col_ref.binding == child_bindings[ci] && ci < child_types.size()) { - col_ref.return_type = child_types[ci]; + if (col_ref.BindingMutable() == child_bindings[ci] && ci < child_types.size()) { + col_ref.SetReturnType(child_types[ci]); break; } } @@ -1281,15 +1296,16 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique auto &expr = proj->expressions[i]; #if PRIVACY_DEBUG PRIVACY_DEBUG_PRINT("PACTopKRule: A3b proj expr[" + std::to_string(i) + - "] type=" + std::to_string((int)expr->type) + - " return_type=" + expr->return_type.ToString() + " expr=" + expr->ToString()); + "] type=" + std::to_string((int)expr->GetExpressionType()) + + " return_type=" + expr->GetReturnType().ToString() + " expr=" + expr->ToString()); #endif - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { + if (expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { continue; } std::function &)> WrapListRefs = [&](unique_ptr &e) { - if (e->type == ExpressionType::BOUND_COLUMN_REF && e->return_type.id() == LogicalTypeId::LIST) { - auto binding = e->Cast().binding; + if (e->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF && + e->GetReturnType().id() == LogicalTypeId::LIST) { + auto binding = e->Cast().BindingMutable(); #if PRIVACY_DEBUG PRIVACY_DEBUG_PRINT("PACTopKRule: A3b WRAPPING LIST colref: " + e->ToString()); #endif @@ -1301,7 +1317,7 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique bool traced = TraceBindingThroughProjections(binding, ctx.intermediate_projections, pi + 1, resolved); for (auto &info : pac_aggs) { - if (info.original_type != e->return_type) { + if (info.original_type != e->GetReturnType()) { if (traced && resolved == info.agg_binding) { e = BoundCastExpression::AddCastToType(input.context, std::move(e), info.original_type); @@ -1324,16 +1340,18 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique // If this is a CAST whose child is a LIST colref, replace the // entire CAST so the bound cast function matches the new source // type (FLOAT from as_noised instead of original INT64/DECIMAL). - if (e->GetExpressionClass() == ExpressionClass::BOUND_CAST) { - auto &cast_expr = e->Cast(); - if (cast_expr.child->type == ExpressionType::BOUND_COLUMN_REF && - cast_expr.child->return_type.id() == LogicalTypeId::LIST) { - auto target_type = cast_expr.return_type; + if (e->GetExpressionType() == ExpressionType::OPERATOR_CAST) { + auto &cast_expr = e->Cast(); + if (cast_expr.GetChildrenMutable()[0]->GetExpressionType() == + ExpressionType::BOUND_COLUMN_REF && + cast_expr.GetChildrenMutable()[0]->GetReturnType().id() == LogicalTypeId::LIST) { + auto target_type = cast_expr.GetReturnType(); #if PRIVACY_DEBUG PRIVACY_DEBUG_PRINT("PACTopKRule: A3b REPLACING CAST(LIST->" + target_type.ToString() + ") with CAST(as_noised()->" + target_type.ToString() + ")"); #endif - auto noised = input.optimizer.BindScalarFunction("as_noised", std::move(cast_expr.child)); + auto noised = input.optimizer.BindScalarFunction( + "as_noised", std::move(cast_expr.GetChildrenMutable()[0])); e = BoundCastExpression::AddCastToType(input.context, std::move(noised), target_type); return; } @@ -1354,12 +1372,13 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique // ------------------------------------------------------------------ for (auto &kv : order_pac_refs) { auto *stripped = StripCasts(ctx.topn->orders[kv.first].expression.get()); - auto *col_ref = stripped->type == ExpressionType::BOUND_COLUMN_REF + auto *col_ref = stripped->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF ? &stripped->Cast() : nullptr; if (col_ref) { - col_ref->binding = ColumnBinding(outermost->table_index, pac_mean_col_indices[0][kv.second]); - col_ref->return_type = PacFloatLogicalType(); + col_ref->BindingMutable() = ColumnBinding(TableIndex(outermost->table_index.index), + ProjectionIndex(pac_mean_col_indices[0][kv.second])); + col_ref->SetReturnType(PacFloatLogicalType()); } } @@ -1373,34 +1392,35 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique // priv_mean() applied then cast back to the original return type. // Non-PAC columns are passed through unchanged. // ------------------------------------------------------------------ - idx_t noised_proj_idx = binder.GenerateTableIndex(); + idx_t noised_proj_idx = binder.GenerateTableIndex().index; vector> noised_proj_exprs; for (idx_t i = 0; i < original_outermost_col_count; i++) { auto it_pac = outermost_pac_col_map.find(i); if (it_pac != outermost_pac_col_map.end()) { // PAC aggregate column: apply as_noised to counter LIST, cast back to original type - auto counters_ref = - make_uniq(list_type, ColumnBinding(outermost->table_index, i)); + auto counters_ref = make_uniq( + list_type, ColumnBinding(TableIndex(outermost->table_index.index), ProjectionIndex(i))); auto noised = input.optimizer.BindScalarFunction("as_noised", std::move(counters_ref)); auto &orig_type = pac_aggs[it_pac->second].original_type; - if (orig_type != noised->return_type) { + if (orig_type != noised->GetReturnType()) { noised = BoundCastExpression::AddCastToType(input.context, std::move(noised), orig_type); } noised_proj_exprs.push_back(std::move(noised)); } else { // Non-PAC column: passthrough (preserves decompress result, etc.) - auto ref = - make_uniq(outermost->types[i], ColumnBinding(outermost->table_index, i)); + auto ref = make_uniq( + outermost->types[i], ColumnBinding(TableIndex(outermost->table_index.index), ProjectionIndex(i))); noised_proj_exprs.push_back(std::move(ref)); } } - auto noised_proj = make_uniq(noised_proj_idx, std::move(noised_proj_exprs)); + auto noised_proj = make_uniq(TableIndex(noised_proj_idx), std::move(noised_proj_exprs)); // Build order remap: original outermost projection columns -> NoisedProj columns for (idx_t i = 0; i < original_outermost_col_count; i++) { - order_remap[HashBinding(ColumnBinding(outermost->table_index, i))] = ColumnBinding(noised_proj_idx, i); + order_remap[HashBinding(ColumnBinding(TableIndex(outermost->table_index.index), ProjectionIndex(i)))] = + ColumnBinding(TableIndex(noised_proj_idx), ProjectionIndex(i)); } // ------------------------------------------------------------------ @@ -1430,16 +1450,16 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique // Step B1: Rewrite TopN ORDER BY to reference MeanProj std::function &)> RewriteOrderExpr = [&](unique_ptr &expr) { - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { + if (expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = expr->Cast(); - auto it = pac_mean_column_map.find(HashBinding(col_ref.binding)); + auto it = pac_mean_column_map.find(HashBinding(col_ref.BindingMutable())); if (it != pac_mean_column_map.end()) { - col_ref.binding = ColumnBinding(mean_proj_idx, it->second); - col_ref.return_type = PacFloatLogicalType(); + col_ref.BindingMutable() = ColumnBinding(TableIndex(mean_proj_idx), ProjectionIndex(it->second)); + col_ref.SetReturnType(PacFloatLogicalType()); } else { for (idx_t i = 0; i < agg_bindings.size(); i++) { - if (agg_bindings[i] == col_ref.binding) { - col_ref.binding = ColumnBinding(mean_proj_idx, i); + if (agg_bindings[i] == col_ref.BindingMutable()) { + col_ref.BindingMutable() = ColumnBinding(TableIndex(mean_proj_idx), ProjectionIndex(i)); break; } } @@ -1455,18 +1475,19 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique } // Step B2: Build NoisedProj - idx_t noised_proj_idx = binder.GenerateTableIndex(); + idx_t noised_proj_idx = binder.GenerateTableIndex().index; vector> noised_proj_exprs; for (idx_t i = 0; i < agg_bindings.size(); i++) { auto &binding = agg_bindings[i]; auto it = pac_mean_column_map.find(HashBinding(binding)); if (it != pac_mean_column_map.end()) { - auto counters_ref = make_uniq(list_type, ColumnBinding(mean_proj_idx, i)); + auto counters_ref = make_uniq( + list_type, ColumnBinding(TableIndex(mean_proj_idx), ProjectionIndex(i))); auto noised = input.optimizer.BindScalarFunction("as_noised", std::move(counters_ref)); // Cast back to the original aggregate return type (e.g. BIGINT for count) for (auto &info : pac_aggs) { - if (info.agg_binding == binding && info.original_type != noised->return_type) { + if (info.agg_binding == binding && info.original_type != noised->GetReturnType()) { noised = BoundCastExpression::AddCastToType(input.context, std::move(noised), info.original_type); break; @@ -1474,18 +1495,20 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique } noised_proj_exprs.push_back(std::move(noised)); } else { - auto ref = make_uniq(agg.types[i], ColumnBinding(mean_proj_idx, i)); + auto ref = make_uniq( + agg.types[i], ColumnBinding(TableIndex(mean_proj_idx), ProjectionIndex(i))); noised_proj_exprs.push_back(std::move(ref)); } } - auto noised_proj = make_uniq(noised_proj_idx, std::move(noised_proj_exprs)); + auto noised_proj = make_uniq(TableIndex(noised_proj_idx), std::move(noised_proj_exprs)); // Build order remap: original aggregate bindings -> NoisedProj columns { idx_t noised_col = 0; for (idx_t i = 0; i < agg_bindings.size(); i++) { - order_remap[HashBinding(agg_bindings[i])] = ColumnBinding(noised_proj_idx, noised_col); + order_remap[HashBinding(agg_bindings[i])] = + ColumnBinding(TableIndex(noised_proj_idx), ProjectionIndex(noised_col)); noised_col++; } } @@ -1518,11 +1541,11 @@ void PACTopKRule::PACTopKOptimizeFunction(OptimizerExtensionInput &input, unique for (auto &saved : saved_orders) { auto expr = saved.expression->Copy(); std::function &)> RemapToNoised = [&](unique_ptr &e) { - if (e->type == ExpressionType::BOUND_COLUMN_REF) { + if (e->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &ref = e->Cast(); - auto it = order_remap.find(HashBinding(ref.binding)); + auto it = order_remap.find(HashBinding(ref.BindingMutable())); if (it != order_remap.end()) { - ref.binding = it->second; + ref.BindingMutable() = it->second; } } ExpressionIterator::EnumerateChildren(*e, RemapToNoised); diff --git a/src/query_processing/privacy_expression_builder.cpp b/src/query_processing/privacy_expression_builder.cpp index 41d42229..e523d982 100644 --- a/src/query_processing/privacy_expression_builder.cpp +++ b/src/query_processing/privacy_expression_builder.cpp @@ -43,7 +43,7 @@ static void EnsureProjectionIdsPopulated(LogicalGet &get) { } if (get.projection_ids.empty()) { for (idx_t i = 0; i < column_ids.size(); i++) { - get.projection_ids.push_back(i); + get.projection_ids.push_back(ProjectionIndex(i)); } } } @@ -53,7 +53,7 @@ static void EnsureProjectionIdsPopulated(LogicalGet &get) { static LogicalGet *FindLogicalGetByTableIndex(LogicalOperator &op, idx_t table_index) { if (op.type == LogicalOperatorType::LOGICAL_GET) { auto &get = op.Cast(); - if (get.table_index == table_index) { + if (get.table_index.index == table_index) { return &get; } } @@ -68,19 +68,19 @@ static LogicalGet *FindLogicalGetByTableIndex(LogicalOperator &op, idx_t table_i // Helper function to resolve a column name from a binding by finding the corresponding LogicalGet static string ResolveColumnNameFromBinding(LogicalOperator &root, const ColumnBinding &binding) { - auto get = FindLogicalGetByTableIndex(root, binding.table_index); + auto get = FindLogicalGetByTableIndex(root, binding.table_index.index); if (!get) { - return "[" + std::to_string(binding.table_index) + "." + std::to_string(binding.column_index) + "]"; + return "[" + std::to_string(binding.table_index.index) + "." + std::to_string(binding.column_index) + "]"; } const auto &column_ids = get->GetColumnIds(); if (binding.column_index >= column_ids.size()) { - return "[" + std::to_string(binding.table_index) + "." + std::to_string(binding.column_index) + "]"; + return "[" + std::to_string(binding.table_index.index) + "." + std::to_string(binding.column_index) + "]"; } - auto col_name = get->GetColumnName(column_ids[binding.column_index]); + auto col_name = get->GetColumnName(column_ids[binding.column_index]).GetIdentifierName(); if (col_name.empty()) { - return "[" + std::to_string(binding.table_index) + "." + std::to_string(binding.column_index) + "]"; + return "[" + std::to_string(binding.table_index.index) + "." + std::to_string(binding.column_index) + "]"; } return col_name; @@ -93,7 +93,7 @@ idx_t EnsureProjectedColumn(LogicalGet &g, const string &col_name) { for (idx_t cid = 0; cid < g.GetColumnIds().size(); ++cid) { auto col_idx = g.GetColumnIds()[cid]; if (!col_idx.IsVirtualColumn()) { - const auto &existing_name = g.GetColumnName(col_idx); + const auto &existing_name = g.GetColumnName(col_idx).GetIdentifierName(); if (!existing_name.empty() && StringUtil::CIEquals(existing_name, col_name)) { return cid; } @@ -108,7 +108,7 @@ idx_t EnsureProjectedColumn(LogicalGet &g, const string &col_name) { LogicalType logical_type = LogicalType::INVALID; idx_t ti = 0; for (auto &col : table_entry->GetColumns().Logical()) { - if (StringUtil::CIEquals(col.Name(), col_name)) { + if (StringUtil::CIEquals(col.Name().GetIdentifierName(), col_name)) { logical_idx = ti; logical_type = col.Type(); break; @@ -130,7 +130,7 @@ idx_t EnsureProjectedColumn(LogicalGet &g, const string &col_name) { // The projection index is the position in the output, which equals the new size - 1 idx_t new_proj_idx = g.GetColumnIds().size() - 1; - g.projection_ids.push_back(new_proj_idx); + g.projection_ids.push_back(ProjectionIndex(new_proj_idx)); // ResolveOperatorTypes() calls the protected ResolveTypes() which rebuilds // the types vector from column_ids/projection_ids @@ -162,7 +162,7 @@ void AddRowIDColumn(LogicalGet &get) { EnsureProjectionIdsPopulated(get); get.AddColumnId(COLUMN_IDENTIFIER_ROW_ID); - get.projection_ids.push_back(get.GetColumnIds().size() - 1); + get.projection_ids.push_back(ProjectionIndex(get.GetColumnIds().size() - 1)); // ResolveOperatorTypes() calls the protected ResolveTypes() which rebuilds // the types vector from column_ids/projection_ids @@ -191,7 +191,7 @@ unique_ptr BuildXorHashFromPKs(OptimizerExtensionInput &input, Logic if (proj_idx == DConstants::INVALID_INDEX) { throw InternalException("Privacy compiler: failed to find PK column " + pk); } - auto col_binding = ColumnBinding(get.table_index, proj_idx); + auto col_binding = ColumnBinding(TableIndex(get.table_index.index), ProjectionIndex(proj_idx)); auto col_index_obj = get.GetColumnIds()[proj_idx]; auto &col_type = get.GetColumnType(col_index_obj); pk_cols.push_back(make_uniq(col_type, col_binding)); @@ -205,7 +205,8 @@ ColumnBinding InsertHashProjectionAboveGet(OptimizerExtensionInput &input, uniqu unique_ptr hash_expr; if (use_rowid) { AddRowIDColumn(get); - auto rowid_binding = ColumnBinding(get.table_index, get.GetColumnIds().size() - 1); + auto rowid_binding = + ColumnBinding(TableIndex(get.table_index.index), ProjectionIndex(get.GetColumnIds().size() - 1)); auto rowid_col = make_uniq(LogicalType::BIGINT, rowid_binding); hash_expr = input.optimizer.BindScalarFunction("hash", std::move(rowid_col)); } else { @@ -215,14 +216,14 @@ ColumnBinding InsertHashProjectionAboveGet(OptimizerExtensionInput &input, uniqu // 1b. Wrap in priv_hash() to XOR with query_hash and optionally repair to 32 bits set hash_expr = input.optimizer.BindScalarFunction("priv_hash", std::move(hash_expr)); - hash_expr->alias = "pac_pu"; + hash_expr->SetAlias(Identifier("pac_pu")); // 2. Find the unique_ptr slot holding this get in the plan tree vector *> get_nodes; - FindAllNodesByTableIndex(&plan, get.table_index, get_nodes); + FindAllNodesByTableIndex(&plan, get.table_index.index, get_nodes); if (get_nodes.empty()) { throw InternalException("Privacy compiler: InsertHashProjectionAboveGet could not find get #" + - std::to_string(get.table_index)); + std::to_string(get.table_index.index)); } auto &get_slot = *get_nodes[0]; @@ -232,20 +233,20 @@ ColumnBinding InsertHashProjectionAboveGet(OptimizerExtensionInput &input, uniqu // 4. Create projection with passthrough + hash column auto &binder = input.optimizer.binder; - idx_t proj_table_index = binder.GenerateTableIndex(); + idx_t proj_table_index = binder.GenerateTableIndex().index; vector> proj_expressions; // Passthrough expressions for each existing get output column for (idx_t i = 0; i < num_get_cols; i++) { auto col_type = get_slot->types[i]; - auto binding = ColumnBinding(get.table_index, i); + auto binding = ColumnBinding(TableIndex(get.table_index.index), ProjectionIndex(i)); proj_expressions.push_back(make_uniq(col_type, binding)); } // Hash expression as the last column proj_expressions.push_back(std::move(hash_expr)); - auto projection = make_uniq(proj_table_index, std::move(proj_expressions)); + auto projection = make_uniq(TableIndex(proj_table_index), std::move(proj_expressions)); // 5. Move get into projection as child, place projection in slot projection->children.push_back(std::move(get_slot)); @@ -255,29 +256,29 @@ ColumnBinding InsertHashProjectionAboveGet(OptimizerExtensionInput &input, uniqu LogicalOperator *proj_ptr = projection.get(); get_slot = std::move(projection); - // 6. Use ColumnBindingReplacer to remap [get.table_index, i] → [proj.table_index, i] + // 6. Use ColumnBindingReplacer to remap [get.table_index.index, i] → [proj.table_index.index, i] ColumnBindingReplacer replacer; for (idx_t i = 0; i < num_get_cols; i++) { ColumnBinding old_b = old_bindings[i]; - ColumnBinding new_b(proj_table_index, i); + ColumnBinding new_b {TableIndex(proj_table_index), ProjectionIndex(i)}; replacer.replacement_bindings.emplace_back(old_b, new_b); } replacer.stop_operator = proj_ptr; replacer.VisitOperator(*plan); // 7. Return binding for the hash column (last expression in projection) - return ColumnBinding(proj_table_index, num_get_cols); + return ColumnBinding(TableIndex(proj_table_index), ProjectionIndex(num_get_cols)); } ColumnBinding GetOrInsertHashProjection(OptimizerExtensionInput &input, unique_ptr &plan, LogicalGet &get, const vector &key_columns, bool use_rowid, std::unordered_map &cache) { - auto it = cache.find(get.table_index); + auto it = cache.find(get.table_index.index); if (it != cache.end()) { return it->second; } auto binding = InsertHashProjectionAboveGet(input, plan, get, key_columns, use_rowid); - cache[get.table_index] = binding; + cache[get.table_index.index] = binding; return binding; } @@ -291,7 +292,7 @@ static unique_ptr *FindCTERefSlot(unique_ptr & } if (root->type == LogicalOperatorType::LOGICAL_CTE_REF) { auto &ref = root->Cast(); - if (ref.table_index == table_index) { + if (ref.table_index.index == table_index) { return &root; } } @@ -307,8 +308,9 @@ static unique_ptr *FindCTERefSlot(unique_ptr & ColumnBinding InsertHashProjectionAboveCTERef(OptimizerExtensionInput &input, unique_ptr &plan, LogicalCTERef &cte_ref, const vector &key_columns) { #if PRIVACY_DEBUG - PRIVACY_DEBUG_PRINT("InsertHashProjectionAboveCTERef: CTE_SCAN table_index=" + std::to_string(cte_ref.table_index) + - " cte_index=" + std::to_string(cte_ref.cte_index)); + PRIVACY_DEBUG_PRINT( + "InsertHashProjectionAboveCTERef: CTE_SCAN table_index=" + std::to_string(cte_ref.table_index.index) + + " cte_index=" + std::to_string(cte_ref.cte_index.index)); #endif // 1. Resolve key column indices in the CTE_SCAN's bound_columns @@ -317,7 +319,7 @@ ColumnBinding InsertHashProjectionAboveCTERef(OptimizerExtensionInput &input, un bool found = false; for (idx_t i = 0; i < cte_ref.bound_columns.size(); i++) { if (cte_ref.bound_columns[i] == key) { - auto binding = ColumnBinding(cte_ref.table_index, i); + auto binding = ColumnBinding(TableIndex(cte_ref.table_index.index), ProjectionIndex(i)); key_col_exprs.push_back(make_uniq(cte_ref.chunk_types[i], binding)); found = true; break; @@ -327,7 +329,7 @@ ColumnBinding InsertHashProjectionAboveCTERef(OptimizerExtensionInput &input, un #if PRIVACY_DEBUG PRIVACY_DEBUG_PRINT("InsertHashProjectionAboveCTERef: key column '" + key + "' not found in bound_columns"); #endif - return ColumnBinding(DConstants::INVALID_INDEX, 0); + return ColumnBinding(TableIndex(DConstants::INVALID_INDEX), ProjectionIndex(0)); } } @@ -336,13 +338,13 @@ ColumnBinding InsertHashProjectionAboveCTERef(OptimizerExtensionInput &input, un // Wrap in priv_hash() to XOR with query_hash and optionally repair to 32 bits set hash_expr = input.optimizer.BindScalarFunction("priv_hash", std::move(hash_expr)); - hash_expr->alias = "pac_pu"; + hash_expr->SetAlias(Identifier("pac_pu")); // 3. Find the unique_ptr slot holding this CTE_SCAN in the plan tree - auto *cte_slot_ptr = FindCTERefSlot(plan, cte_ref.table_index); + auto *cte_slot_ptr = FindCTERefSlot(plan, cte_ref.table_index.index); if (!cte_slot_ptr) { throw InternalException("Privacy compiler: InsertHashProjectionAboveCTERef could not find CTE_SCAN #" + - std::to_string(cte_ref.table_index)); + std::to_string(cte_ref.table_index.index)); } auto &cte_slot = *cte_slot_ptr; @@ -351,14 +353,14 @@ ColumnBinding InsertHashProjectionAboveCTERef(OptimizerExtensionInput &input, un idx_t num_cte_cols = old_bindings.size(); auto &binder = input.optimizer.binder; - idx_t proj_table_index = binder.GenerateTableIndex(); + idx_t proj_table_index = binder.GenerateTableIndex().index; vector> proj_expressions; for (idx_t i = 0; i < num_cte_cols; i++) { proj_expressions.push_back(make_uniq(cte_slot->types[i], old_bindings[i])); } proj_expressions.push_back(std::move(hash_expr)); - auto projection = make_uniq(proj_table_index, std::move(proj_expressions)); + auto projection = make_uniq(TableIndex(proj_table_index), std::move(proj_expressions)); projection->children.push_back(std::move(cte_slot)); projection->ResolveOperatorTypes(); @@ -368,18 +370,19 @@ ColumnBinding InsertHashProjectionAboveCTERef(OptimizerExtensionInput &input, un ColumnBindingReplacer replacer; for (idx_t i = 0; i < num_cte_cols; i++) { - replacer.replacement_bindings.emplace_back(old_bindings[i], ColumnBinding(proj_table_index, i)); + replacer.replacement_bindings.emplace_back(old_bindings[i], + ColumnBinding(TableIndex(proj_table_index), ProjectionIndex(i))); } replacer.stop_operator = proj_ptr; replacer.VisitOperator(*plan); #if PRIVACY_DEBUG PRIVACY_DEBUG_PRINT("InsertHashProjectionAboveCTERef: inserted hash projection (table_index=" + - std::to_string(proj_table_index) + ") above CTE_SCAN #" + std::to_string(cte_ref.table_index) + - ", hash at column " + std::to_string(num_cte_cols)); + std::to_string(proj_table_index) + ") above CTE_SCAN #" + + std::to_string(cte_ref.table_index.index) + ", hash at column " + std::to_string(num_cte_cols)); #endif - return ColumnBinding(proj_table_index, num_cte_cols); + return ColumnBinding(TableIndex(proj_table_index), ProjectionIndex(num_cte_cols)); } // Build AND expression from multiple hash expressions (for multiple PUs) @@ -409,14 +412,16 @@ unique_ptr BindPacAggregate(OptimizerExtensionInput &input, const st unique_ptr correction_expr) { FunctionBinder function_binder(input.context); ErrorData error; - vector arg_types = {hash_expr->return_type, value_expr->return_type}; + vector arg_types = {hash_expr->GetReturnType(), value_expr->GetReturnType()}; if (correction_expr) { - arg_types.push_back(correction_expr->return_type); + arg_types.push_back(correction_expr->GetReturnType()); } auto &entry = Catalog::GetSystemCatalog(input.context) - .GetEntry(input.context, DEFAULT_SCHEMA, pac_func_name); - auto best = function_binder.BindFunction(entry.name, entry.functions, arg_types, error); + .GetEntry( + input.context, QualifiedName(Catalog::GetSystemCatalog(input.context).GetName(), + Identifier::DefaultSchema(), Identifier(pac_func_name))); + auto best = function_binder.BindFunction(Identifier(entry.name), entry.functions, arg_types, error); if (!best.IsValid()) { throw InternalException("Privacy compiler: failed to bind " + pac_func_name); } @@ -435,11 +440,12 @@ PuPreAggregationInfo InsertPuPreAggregation(OptimizerExtensionInput &input, Logi vector> lower_expressions, unique_ptr pu_hash_expr) { auto &binder = input.optimizer.binder; - idx_t lower_group_index = binder.GenerateTableIndex(); - idx_t lower_agg_index = binder.GenerateTableIndex(); + idx_t lower_group_index = binder.GenerateTableIndex().index; + idx_t lower_agg_index = binder.GenerateTableIndex().index; idx_t num_original_groups = agg->groups.size(); - auto lower_agg = make_uniq(lower_group_index, lower_agg_index, std::move(lower_expressions)); + auto lower_agg = make_uniq(TableIndex(lower_group_index), TableIndex(lower_agg_index), + std::move(lower_expressions)); for (auto &g : agg->groups) { lower_agg->groups.push_back(g->Copy()); } @@ -448,11 +454,13 @@ PuPreAggregationInfo InsertPuPreAggregation(OptimizerExtensionInput &input, Logi lower_agg->ResolveOperatorTypes(); for (idx_t i = 0; i < num_original_groups; i++) { - agg->groups[i] = make_uniq(lower_agg->types[i], ColumnBinding(lower_group_index, i)); + agg->groups[i] = make_uniq( + lower_agg->types[i], ColumnBinding(TableIndex(lower_group_index), ProjectionIndex(i))); } - auto pu_hash_ref = make_uniq(lower_agg->types[num_original_groups], - ColumnBinding(lower_group_index, num_original_groups)); + auto pu_hash_ref = make_uniq( + lower_agg->types[num_original_groups], + ColumnBinding(TableIndex(lower_group_index), ProjectionIndex(num_original_groups))); auto *lower_agg_ptr = lower_agg.get(); agg->children[0] = std::move(lower_agg); @@ -464,10 +472,13 @@ unique_ptr BindBitOrAggregate(OptimizerExtensionInput &input, unique unique_ptr filter_col_expr) { FunctionBinder function_binder(input.context); ErrorData error; - vector bit_or_types = {hash_expr->return_type}; + vector bit_or_types = {hash_expr->GetReturnType()}; auto &bit_or_entry = Catalog::GetSystemCatalog(input.context) - .GetEntry(input.context, DEFAULT_SCHEMA, "bit_or"); - auto bit_or_best = function_binder.BindFunction(bit_or_entry.name, bit_or_entry.functions, bit_or_types, error); + .GetEntry( + input.context, QualifiedName(Catalog::GetSystemCatalog(input.context).GetName(), + Identifier::DefaultSchema(), "bit_or")); + auto bit_or_best = + function_binder.BindFunction(Identifier(bit_or_entry.name), bit_or_entry.functions, bit_or_types, error); if (!bit_or_best.IsValid()) { throw InternalException("Privacy compiler: failed to bind bit_or"); } @@ -481,8 +492,8 @@ unique_ptr BindBitOrAggregate(OptimizerExtensionInput &input, unique if (filter_col_expr) { auto not_null_filter = make_uniq(ExpressionType::OPERATOR_IS_NOT_NULL, LogicalType::BOOLEAN); - not_null_filter->children.push_back(std::move(filter_col_expr)); - bit_or_expr->Cast().filter = std::move(not_null_filter); + not_null_filter->GetChildrenMutable().push_back(std::move(filter_col_expr)); + bit_or_expr->Cast().GetFilterMutable() = std::move(not_null_filter); } return bit_or_expr; @@ -524,8 +535,8 @@ static void InsertDistinctPreAggregation(OptimizerExtensionInput &input, Logical auto &binder = input.optimizer.binder; // 1. Create inner aggregate with new table indices - idx_t inner_group_index = binder.GenerateTableIndex(); - idx_t inner_agg_index = binder.GenerateTableIndex(); + idx_t inner_group_index = binder.GenerateTableIndex().index; + idx_t inner_agg_index = binder.GenerateTableIndex().index; // 2. Bind bit_or aggregate on the hash expression with IS NOT NULL filter auto bit_or_expr = BindBitOrAggregate(input, hash_input_expr->Copy(), distinct_value_expr->Copy()); @@ -533,7 +544,8 @@ static void InsertDistinctPreAggregation(OptimizerExtensionInput &input, Logical // 3. Build the inner aggregate node vector> inner_expressions; inner_expressions.push_back(std::move(bit_or_expr)); - auto inner_agg_node = make_uniq(inner_group_index, inner_agg_index, std::move(inner_expressions)); + auto inner_agg_node = make_uniq(TableIndex(inner_group_index), TableIndex(inner_agg_index), + std::move(inner_expressions)); // Copy outer groups into inner aggregate (passthrough group keys) idx_t num_original_groups = agg->groups.size(); @@ -552,20 +564,21 @@ static void InsertDistinctPreAggregation(OptimizerExtensionInput &input, Logical // 5. Remap outer aggregate's groups to reference inner aggregate's group output for (idx_t i = 0; i < num_original_groups; i++) { - auto gtype = agg->groups[i]->return_type; - agg->groups[i] = make_uniq(gtype, ColumnBinding(inner_group_index, i)); + auto gtype = agg->groups[i]->GetReturnType(); + agg->groups[i] = make_uniq( + gtype, ColumnBinding(TableIndex(inner_group_index), ProjectionIndex(i))); } // Inner aggregate output layout: // groups: [g1, ..., gN, distinct_col] at inner_group_index // aggregates: [bit_or(hash)] at inner_agg_index - ColumnBinding combined_hash_binding(inner_agg_index, 0); - ColumnBinding distinct_val_binding(inner_group_index, num_original_groups); + ColumnBinding combined_hash_binding {TableIndex(inner_agg_index), ProjectionIndex(0)}; + ColumnBinding distinct_val_binding {TableIndex(inner_group_index), ProjectionIndex(num_original_groups)}; // 6. Replace each aggregate expression with the corresponding PAC function for (idx_t i = 0; i < agg->expressions.size(); i++) { auto &old_aggr = agg->expressions[i]->Cast(); - string function_name = old_aggr.function.name; + string function_name = old_aggr.FunctionMutable().GetName().GetIdentifierName(); string pac_name = GetPacAggregateFunctionName(function_name, &input.context); auto hash_ref = make_uniq(LogicalType::UBIGINT, combined_hash_binding); @@ -574,7 +587,7 @@ static void InsertDistinctPreAggregation(OptimizerExtensionInput &input, Logical value_ref = make_uniq_base(Value::BIGINT(1)); } else { // SUM/AVG/MIN/MAX DISTINCT: value comes from the distinct column (now a group key in inner agg) - value_ref = make_uniq(distinct_value_expr->return_type, distinct_val_binding); + value_ref = make_uniq(distinct_value_expr->GetReturnType(), distinct_val_binding); } agg->expressions[i] = BindPacAggregate(input, pac_name, std::move(hash_ref), std::move(value_ref)); @@ -599,14 +612,15 @@ BuildDistinctBranch(OptimizerExtensionInput &input, unique_ptr const vector> &agg_specs, // (original_agg_idx, original_func_name) idx_t &out_group_index, idx_t &out_agg_index) { auto &binder = input.optimizer.binder; - idx_t inner_group_index = binder.GenerateTableIndex(); - idx_t inner_agg_index = binder.GenerateTableIndex(); + idx_t inner_group_index = binder.GenerateTableIndex().index; + idx_t inner_agg_index = binder.GenerateTableIndex().index; // Inner aggregate: bit_or(hash) FILTER (distinct_col IS NOT NULL) GROUP BY [groups..., distinct_col] auto bit_or_expr = BindBitOrAggregate(input, hash_expr->Copy(), distinct_value_expr->Copy()); vector> inner_expressions; inner_expressions.push_back(std::move(bit_or_expr)); - auto inner_agg = make_uniq(inner_group_index, inner_agg_index, std::move(inner_expressions)); + auto inner_agg = make_uniq(TableIndex(inner_group_index), TableIndex(inner_agg_index), + std::move(inner_expressions)); idx_t num_groups = groups.size(); for (auto &g : groups) { @@ -617,11 +631,11 @@ BuildDistinctBranch(OptimizerExtensionInput &input, unique_ptr inner_agg->ResolveOperatorTypes(); // Outer aggregate: pac_*(combined_hash, value) GROUP BY [groups...] - idx_t outer_group_index = binder.GenerateTableIndex(); - idx_t outer_agg_index = binder.GenerateTableIndex(); + idx_t outer_group_index = binder.GenerateTableIndex().index; + idx_t outer_agg_index = binder.GenerateTableIndex().index; - ColumnBinding combined_hash_binding(inner_agg_index, 0); - ColumnBinding distinct_val_binding(inner_group_index, num_groups); + ColumnBinding combined_hash_binding {TableIndex(inner_agg_index), ProjectionIndex(0)}; + ColumnBinding distinct_val_binding {TableIndex(inner_group_index), ProjectionIndex(num_groups)}; vector> outer_expressions; for (auto &spec : agg_specs) { @@ -631,15 +645,17 @@ BuildDistinctBranch(OptimizerExtensionInput &input, unique_ptr if (spec.second == "count" || spec.second == "count_star") { value_ref = make_uniq_base(Value::BIGINT(1)); } else { - value_ref = make_uniq(distinct_value_expr->return_type, distinct_val_binding); + value_ref = make_uniq(distinct_value_expr->GetReturnType(), distinct_val_binding); } outer_expressions.push_back(BindPacAggregate(input, pac_name, std::move(hash_ref), std::move(value_ref))); } - auto outer_agg = make_uniq(outer_group_index, outer_agg_index, std::move(outer_expressions)); + auto outer_agg = make_uniq(TableIndex(outer_group_index), TableIndex(outer_agg_index), + std::move(outer_expressions)); for (idx_t i = 0; i < num_groups; i++) { - auto gtype = groups[i]->return_type; - outer_agg->groups.push_back(make_uniq(gtype, ColumnBinding(inner_group_index, i))); + auto gtype = groups[i]->GetReturnType(); + outer_agg->groups.push_back(make_uniq( + gtype, ColumnBinding(TableIndex(inner_group_index), ProjectionIndex(i)))); } outer_agg->children.push_back(std::move(inner_agg)); outer_agg->ResolveOperatorTypes(); @@ -657,26 +673,26 @@ static unique_ptr BuildNonDistinctBranch( const vector> &agg_specs, // (original_agg_idx, original_expr) idx_t &out_group_index, idx_t &out_agg_index) { auto &binder = input.optimizer.binder; - idx_t group_index = binder.GenerateTableIndex(); - idx_t agg_index = binder.GenerateTableIndex(); + idx_t group_index = binder.GenerateTableIndex().index; + idx_t agg_index = binder.GenerateTableIndex().index; vector> expressions; for (auto &spec : agg_specs) { auto &old_aggr = *spec.second; - string function_name = old_aggr.function.name; + string function_name = old_aggr.Function().GetName().GetIdentifierName(); string pac_name = GetPacAggregateFunctionName(function_name, &input.context); unique_ptr value_child; - if (old_aggr.children.empty()) { + if (old_aggr.GetChildren().empty()) { value_child = make_uniq_base(Value::BIGINT(1)); } else { - value_child = old_aggr.children[0]->Copy(); + value_child = old_aggr.GetChildren()[0]->Copy(); } expressions.push_back(BindPacAggregate(input, pac_name, hash_expr->Copy(), std::move(value_child))); } - auto agg_node = make_uniq(group_index, agg_index, std::move(expressions)); + auto agg_node = make_uniq(TableIndex(group_index), TableIndex(agg_index), std::move(expressions)); for (auto &g : groups) { agg_node->groups.push_back(g->Copy()); } @@ -690,11 +706,11 @@ static unique_ptr BuildNonDistinctBranch( // Helper: remap expression bindings using an index map static void RemapExpressionBindings(Expression &e, const std::unordered_map &map) { - if (e.type == ExpressionType::BOUND_COLUMN_REF) { + if (e.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e.Cast(); - auto it = map.find(col_ref.binding.table_index); + auto it = map.find(col_ref.BindingMutable().table_index.index); if (it != map.end()) { - col_ref.binding.table_index = it->second; + col_ref.BindingMutable().table_index.index = it->second; } } ExpressionIterator::EnumerateChildren(e, [&](Expression &child) { RemapExpressionBindings(child, map); }); @@ -729,20 +745,21 @@ JoinedDistinctAggregateBranches BuildJoinedDistinctAggregateBranches(OptimizerEx std::unordered_map distinct_col_to_branch; for (idx_t i = 0; i < agg->expressions.size(); i++) { auto &aggr = agg->expressions[i]->Cast(); - bool is_distinct = aggr.IsDistinct() && !aggr.children.empty(); + bool is_distinct = aggr.IsDistinct() && !aggr.GetChildrenMutable().empty(); if (is_distinct) { - string col_str = aggr.children[0]->ToString(); + string col_str = aggr.GetChildrenMutable()[0]->ToString(); auto it = distinct_col_to_branch.find(col_str); if (it != distinct_col_to_branch.end()) { - branches[it->second].distinct_agg_specs.emplace_back(i, aggr.function.name); + branches[it->second].distinct_agg_specs.emplace_back( + i, aggr.FunctionMutable().GetName().GetIdentifierName()); branches[it->second].original_agg_indices.push_back(i); } else { BranchInfo info; info.is_distinct = true; info.distinct_col_str = col_str; - info.distinct_col_expr = aggr.children[0]->Copy(); - info.distinct_agg_specs.emplace_back(i, aggr.function.name); + info.distinct_col_expr = aggr.GetChildrenMutable()[0]->Copy(); + info.distinct_agg_specs.emplace_back(i, aggr.FunctionMutable().GetName().GetIdentifierName()); info.original_agg_indices.push_back(i); distinct_col_to_branch[col_str] = branches.size(); branches.push_back(std::move(info)); @@ -810,7 +827,7 @@ JoinedDistinctAggregateBranches BuildJoinedDistinctAggregateBranches(OptimizerEx } auto &branch_agg = built.subtree->Cast(); for (auto &expr : branch_agg.expressions) { - built.aggregate_types.push_back(expr->return_type); + built.aggregate_types.push_back(expr->GetReturnType()); } built_branches.push_back(std::move(built)); } @@ -823,12 +840,13 @@ JoinedDistinctAggregateBranches BuildJoinedDistinctAggregateBranches(OptimizerEx if (has_groups) { auto join = make_uniq(JoinType::INNER); for (idx_t gi = 0; gi < agg->groups.size(); gi++) { - JoinCondition cond; - cond.left = make_uniq(agg->groups[gi]->return_type, - ColumnBinding(left_group_index, gi)); - cond.right = make_uniq(agg->groups[gi]->return_type, - ColumnBinding(right.group_index, gi)); - cond.comparison = ExpressionType::COMPARE_NOT_DISTINCT_FROM; + JoinCondition cond(nullptr, nullptr, ExpressionType::COMPARE_NOT_DISTINCT_FROM); + cond.LeftReference() = make_uniq( + agg->groups[gi]->GetReturnType(), ColumnBinding(TableIndex(left_group_index), ProjectionIndex(gi))); + cond.RightReference() = make_uniq( + agg->groups[gi]->GetReturnType(), + ColumnBinding(TableIndex(right.group_index), ProjectionIndex(gi))); + join->conditions.push_back(std::move(cond)); } join->children.push_back(std::move(joined)); @@ -884,8 +902,8 @@ static void InsertMultiBranchPreAggregation(OptimizerExtensionInput &input, uniq auto &branch_agg = subtree->Cast(); for (auto &expr : branch_agg.expressions) { auto &bound_agg = expr->Cast(); - if (bound_agg.children.size() > 1) { - RemapExpressionBindings(*bound_agg.children[1], *index_map); + if (bound_agg.GetChildrenMutable().size() > 1) { + RemapExpressionBindings(*bound_agg.GetChildrenMutable()[1], *index_map); } } subtree->ResolveOperatorTypes(); @@ -897,13 +915,14 @@ static void InsertMultiBranchPreAggregation(OptimizerExtensionInput &input, uniq PacBranchBuilder builder; auto branch_plan = BuildJoinedDistinctAggregateBranches(input, plan, agg, hash_input_expr->Copy(), builder); - idx_t proj_table_index = input.optimizer.binder.GenerateTableIndex(); + idx_t proj_table_index = input.optimizer.binder.GenerateTableIndex().index; vector> proj_expressions; idx_t num_groups = agg->groups.size(); for (idx_t gi = 0; gi < num_groups; gi++) { proj_expressions.push_back(make_uniq( - agg->groups[gi]->return_type, ColumnBinding(branch_plan.left_group_index, gi))); + agg->groups[gi]->GetReturnType(), + ColumnBinding(TableIndex(branch_plan.left_group_index), ProjectionIndex(gi)))); } for (idx_t ai = 0; ai < agg->expressions.size(); ai++) { @@ -915,22 +934,22 @@ static void InsertMultiBranchPreAggregation(OptimizerExtensionInput &input, uniq idx_t pos_in_branch = it->second.second; auto &branch = branch_plan.branches[branch_idx]; - auto original_type = agg->expressions[ai]->return_type; + auto original_type = agg->expressions[ai]->GetReturnType(); auto branch_type = branch.aggregate_types[pos_in_branch]; PRIVACY_DEBUG_PRINT("MultiBranch output: ai=" + std::to_string(ai) + " branch=" + std::to_string(branch_idx) + " pos=" + std::to_string(pos_in_branch) + " agg_index=" + std::to_string(branch.agg_index) + " original_type=" + original_type.ToString() + " branch_type=" + branch_type.ToString()); - unique_ptr ref = - make_uniq(branch_type, ColumnBinding(branch.agg_index, pos_in_branch)); + unique_ptr ref = make_uniq( + branch_type, ColumnBinding(TableIndex(branch.agg_index), ProjectionIndex(pos_in_branch))); if (branch_type != original_type) { ref = BoundCastExpression::AddCastToType(input.context, std::move(ref), original_type); } proj_expressions.push_back(std::move(ref)); } - auto projection = make_uniq(proj_table_index, std::move(proj_expressions)); + auto projection = make_uniq(TableIndex(proj_table_index), std::move(proj_expressions)); projection->children.push_back(std::move(branch_plan.joined)); projection->ResolveOperatorTypes(); @@ -939,18 +958,19 @@ static void InsertMultiBranchPreAggregation(OptimizerExtensionInput &input, uniq throw InternalException("Privacy compiler: InsertMultiBranchPreAggregation could not find aggregate in plan"); } - idx_t old_group_index = agg->group_index; - idx_t old_agg_index = agg->aggregate_index; + idx_t old_group_index = agg->group_index.index; + idx_t old_agg_index = agg->aggregate_index.index; idx_t num_agg_exprs = agg->expressions.size(); ColumnBindingReplacer replacer; for (idx_t gi = 0; gi < num_groups; gi++) { - replacer.replacement_bindings.emplace_back(ColumnBinding(old_group_index, gi), - ColumnBinding(proj_table_index, gi)); + replacer.replacement_bindings.emplace_back(ColumnBinding(TableIndex(old_group_index), ProjectionIndex(gi)), + ColumnBinding(TableIndex(proj_table_index), ProjectionIndex(gi))); } for (idx_t ai = 0; ai < num_agg_exprs; ai++) { - replacer.replacement_bindings.emplace_back(ColumnBinding(old_agg_index, ai), - ColumnBinding(proj_table_index, num_groups + ai)); + replacer.replacement_bindings.emplace_back( + ColumnBinding(TableIndex(old_agg_index), ProjectionIndex(ai)), + ColumnBinding(TableIndex(proj_table_index), ProjectionIndex(num_groups + ai))); } LogicalOperator *proj_ptr = projection.get(); @@ -1020,16 +1040,16 @@ void ModifyAggregatesWithPacFunctions(OptimizerExtensionInput &input, LogicalAgg // Debug: Print hash input expression details PRIVACY_DEBUG_PRINT("ModifyAggregatesWithPacFunctions: Hash input expression: " + hash_input_expr->ToString()); PRIVACY_DEBUG_PRINT("ModifyAggregatesWithPacFunctions: Hash input type: " + - hash_input_expr->return_type.ToString()); + hash_input_expr->GetReturnType().ToString()); // Debug: Print all column references in the hash expression ExpressionIterator::EnumerateExpression(hash_input_expr, [&](Expression &expr) { - if (expr.type == ExpressionType::BOUND_COLUMN_REF) { + if (expr.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = expr.Cast(); PRIVACY_DEBUG_PRINT("ModifyAggregatesWithPacFunctions: Hash expr references column [" + - std::to_string(col_ref.binding.table_index) + "." + - std::to_string(col_ref.binding.column_index) + - "] type=" + col_ref.return_type.ToString()); + std::to_string(col_ref.Binding().table_index.index) + "." + + std::to_string(col_ref.Binding().column_index) + + "] type=" + col_ref.GetReturnType().ToString()); } }); @@ -1041,8 +1061,8 @@ void ModifyAggregatesWithPacFunctions(OptimizerExtensionInput &input, LogicalAgg } // Debug: Print aggregate's group_index and aggregate_index - PRIVACY_DEBUG_PRINT("ModifyAggregatesWithPacFunctions: group_index=" + std::to_string(agg->group_index) + - ", aggregate_index=" + std::to_string(agg->aggregate_index)); + PRIVACY_DEBUG_PRINT("ModifyAggregatesWithPacFunctions: group_index=" + std::to_string(agg->group_index.index) + + ", aggregate_index=" + std::to_string(agg->aggregate_index.index)); // Debug: Print child operator info if (!agg->children.empty()) { @@ -1051,7 +1071,7 @@ void ModifyAggregatesWithPacFunctions(OptimizerExtensionInput &input, LogicalAgg if (agg->children[0]->type == LogicalOperatorType::LOGICAL_GET) { auto &child_get = agg->children[0]->Cast(); PRIVACY_DEBUG_PRINT("ModifyAggregatesWithPacFunctions: Child is GET with table_index=" + - std::to_string(child_get.table_index) + + std::to_string(child_get.table_index.index) + ", columns=" + std::to_string(child_get.GetColumnIds().size())); } } @@ -1064,7 +1084,8 @@ void ModifyAggregatesWithPacFunctions(OptimizerExtensionInput &input, LogicalAgg throw NotImplementedException("Not found expected aggregate expression in PAC compiler"); } auto &check_aggr = agg->expressions[i]->Cast(); - GetPacAggregateFunctionName(check_aggr.function.name); // throws if unsupported + GetPacAggregateFunctionName( + check_aggr.FunctionMutable().GetName().GetIdentifierName()); // throws if unsupported } // === DISTINCT pre-aggregation === @@ -1079,11 +1100,11 @@ void ModifyAggregatesWithPacFunctions(OptimizerExtensionInput &input, LogicalAgg for (idx_t i = 0; i < agg->expressions.size(); i++) { auto &check = agg->expressions[i]->Cast(); - if (check.IsDistinct() && !check.children.empty()) { + if (check.IsDistinct() && !check.GetChildrenMutable().empty()) { has_any_distinct = true; if (!common_distinct_value) { - common_distinct_value = check.children[0]->Copy(); - } else if (check.children[0]->ToString() != common_distinct_value->ToString()) { + common_distinct_value = check.GetChildrenMutable()[0]->Copy(); + } else if (check.GetChildrenMutable()[0]->ToString() != common_distinct_value->ToString()) { all_same_col = false; } } else { @@ -1110,7 +1131,7 @@ void ModifyAggregatesWithPacFunctions(OptimizerExtensionInput &input, LogicalAgg } auto &old_aggr = agg->expressions[i]->Cast(); - string function_name = old_aggr.function.name; + string function_name = old_aggr.FunctionMutable().GetName().GetIdentifierName(); #if PRIVACY_DEBUG PRIVACY_DEBUG_PRINT("ModifyAggregatesWithPacFunctions: Transforming " + function_name + " to PAC function"); @@ -1118,14 +1139,14 @@ void ModifyAggregatesWithPacFunctions(OptimizerExtensionInput &input, LogicalAgg // Extract the original aggregate's value child expression unique_ptr value_child; - if (old_aggr.children.empty()) { + if (old_aggr.GetChildrenMutable().empty()) { if (function_name == "count_star" || function_name == "count") { value_child = make_uniq_base(Value::BIGINT(1)); } else { throw InternalException("Privacy compiler: expected aggregate to have a child expression"); } } else { - value_child = old_aggr.children[0]->Copy(); + value_child = old_aggr.GetChildrenMutable()[0]->Copy(); } string pac_function_name = GetPacAggregateFunctionName(function_name, &input.context); @@ -1180,11 +1201,13 @@ unique_ptr BindPlainAggregate(OptimizerExtensionInput &input, const ErrorData error; vector arg_types; for (auto &child : children) { - arg_types.push_back(child->return_type); + arg_types.push_back(child->GetReturnType()); } auto &entry = Catalog::GetSystemCatalog(input.context) - .GetEntry(input.context, DEFAULT_SCHEMA, func_name); - auto best = function_binder.BindFunction(entry.name, entry.functions, arg_types, error); + .GetEntry( + input.context, QualifiedName(Catalog::GetSystemCatalog(input.context).GetName(), + Identifier::DefaultSchema(), Identifier(func_name))); + auto best = function_binder.BindFunction(Identifier(entry.name), entry.functions, arg_types, error); if (!best.IsValid()) { throw InternalException("Privacy compiler: failed to bind aggregate '" + func_name + "'"); } @@ -1208,7 +1231,7 @@ static bool IsPacAggregate(LogicalAggregate *agg) { continue; } auto &aggr = expr->Cast(); - if (!GetClipVariant(aggr.function.name).empty()) { + if (!GetClipVariant(aggr.FunctionMutable().GetName().GetIdentifierName()).empty()) { return true; } } @@ -1241,13 +1264,13 @@ void RewriteClipAggregates(OptimizerExtensionInput &input, unique_ptrexpressions[i]->Cast(); - string clip_name = GetClipVariant(aggr.function.name); + string clip_name = GetClipVariant(aggr.FunctionMutable().GetName().GetIdentifierName()); if (clip_name.empty()) { continue; } // Rebind with the clip variant name vector> children; - for (auto &child : aggr.children) { + for (auto &child : aggr.GetChildrenMutable()) { children.push_back(child->Copy()); } agg->expressions[i] = RebindAggregate(input.context, clip_name, std::move(children), false); @@ -1263,8 +1286,9 @@ void RewriteClipAggregates(OptimizerExtensionInput &input, unique_ptrCast(); - if (!GetClipVariant(aggr.function.name).empty() && !aggr.children.empty()) { - pu_hash_expr = aggr.children[0]->Copy(); + if (!GetClipVariant(aggr.FunctionMutable().GetName().GetIdentifierName()).empty() && + !aggr.GetChildrenMutable().empty()) { + pu_hash_expr = aggr.GetChildrenMutable()[0]->Copy(); break; } } @@ -1281,21 +1305,22 @@ void RewriteClipAggregates(OptimizerExtensionInput &input, unique_ptr> lower_expressions; for (idx_t i = 0; i < agg->expressions.size(); i++) { auto &aggr = agg->expressions[i]->Cast(); - string orig_name = GetOriginalAggregate(aggr.function.name); + string orig_name = GetOriginalAggregate(aggr.FunctionMutable().GetName().GetIdentifierName()); if (orig_name.empty()) { - throw InternalException("PAC clip rewrite: unexpected aggregate " + aggr.function.name); + throw InternalException("PAC clip rewrite: unexpected aggregate " + + aggr.FunctionMutable().GetName().GetIdentifierName()); } vector> plain_children; - if (orig_name == "count" && (aggr.children.size() <= 1)) { + if (orig_name == "count" && (aggr.GetChildrenMutable().size() <= 1)) { // as_noised_count(hash) or as_count(hash) → count_star() // as_noised_count(hash, col) → count(col) — but children[1] might be constant 1 - if (aggr.children.size() >= 2) { - auto &val_child = aggr.children[1]; + if (aggr.GetChildrenMutable().size() >= 2) { + auto &val_child = aggr.GetChildrenMutable()[1]; // Check if it's a constant 1 (from count_star rewrite) - if (val_child->type == ExpressionType::VALUE_CONSTANT) { + if (val_child->GetExpressionType() == ExpressionType::VALUE_CONSTANT) { auto &const_expr = val_child->Cast(); - if (const_expr.value.IsNull() || const_expr.value == Value::BIGINT(1)) { + if (const_expr.GetValue().IsNull() || const_expr.GetValue() == Value::BIGINT(1)) { // count_star — no children } else { plain_children.push_back(val_child->Copy()); @@ -1305,14 +1330,14 @@ void RewriteClipAggregates(OptimizerExtensionInput &input, unique_ptr 1) { + } else if (orig_name == "count" && aggr.GetChildrenMutable().size() > 1) { // count with column reference - plain_children.push_back(aggr.children[1]->Copy()); + plain_children.push_back(aggr.GetChildrenMutable()[1]->Copy()); lower_expressions.push_back(BindPlainAggregate(input, "count", std::move(plain_children))); } else { // sum, min, max — extract the value child (children[1]) - if (aggr.children.size() >= 2) { - plain_children.push_back(aggr.children[1]->Copy()); + if (aggr.GetChildrenMutable().size() >= 2) { + plain_children.push_back(aggr.GetChildrenMutable()[1]->Copy()); } lower_expressions.push_back(BindPlainAggregate(input, orig_name, std::move(plain_children))); } @@ -1323,14 +1348,14 @@ void RewriteClipAggregates(OptimizerExtensionInput &input, unique_ptrexpressions.size(); i++) { auto &aggr = agg->expressions[i]->Cast(); - string pac_name = aggr.function.name; + string pac_name = aggr.FunctionMutable().GetName().GetIdentifierName(); bool noised = IsNoisedVariant(pac_name); string orig = GetOriginalAggregate(pac_name); // Reference to lower aggregate's result auto lower_type = pre_agg.lower_agg->types[pre_agg.num_original_groups + 1 + i]; - unique_ptr lower_ref = - make_uniq(lower_type, ColumnBinding(pre_agg.lower_agg_index, i)); + unique_ptr lower_ref = make_uniq( + lower_type, ColumnBinding(TableIndex(pre_agg.lower_agg_index), ProjectionIndex(i))); // as_clip_sum has integer + DECIMAL overloads but no FLOAT/DOUBLE. // Cast FLOAT/DOUBLE to BIGINT so binding succeeds. diff --git a/src/query_processing/privacy_fk_resolver.cpp b/src/query_processing/privacy_fk_resolver.cpp index 80e54cc5..6c23d0ab 100644 --- a/src/query_processing/privacy_fk_resolver.cpp +++ b/src/query_processing/privacy_fk_resolver.cpp @@ -45,7 +45,7 @@ AccessibleFKTableResult FindAccessibleFKTable(unique_ptr &plan, for (auto *node : table_nodes) { auto &node_get = node->get()->Cast(); - idx_t node_table_idx = node_get.table_index; + idx_t node_table_idx = node_get.table_index.index; // Check if this table is in the search root's subtree if (search_root && !HasTableIndexInSubtree(search_root, node_table_idx)) { @@ -55,7 +55,7 @@ AccessibleFKTableResult FindAccessibleFKTable(unique_ptr &plan, // Check if this table's columns are accessible if (AreTableColumnsAccessible(plan.get(), node_table_idx)) { result.get_node = &node_get; - result.table_index = node_table_idx; + result.table_index.index = node_table_idx; result.table_name = table_name; result.fk_columns = fk_result.fk_columns; result.found = true; diff --git a/src/query_processing/privacy_plan_traversal.cpp b/src/query_processing/privacy_plan_traversal.cpp index a41950ff..2c5bf218 100644 --- a/src/query_processing/privacy_plan_traversal.cpp +++ b/src/query_processing/privacy_plan_traversal.cpp @@ -41,7 +41,7 @@ unique_ptr *FindPrivacyUnitGetNode(unique_ptr if (!pu_table_name.empty()) { auto &get = cur->Cast(); auto tblptr = get.GetTable(); - if (tblptr && tblptr->name == pu_table_name) { + if (tblptr && tblptr->name.GetIdentifierName() == pu_table_name) { found_ptr = cur_ptr; break; } @@ -77,7 +77,7 @@ LogicalGet *FindTableScanInSubtree(LogicalOperator *subtree, const string &table if (subtree->type == LogicalOperatorType::LOGICAL_GET) { auto &get = subtree->Cast(); auto tblptr = get.GetTable(); - if (tblptr && StringUtil::CIEquals(tblptr->name, table_name)) { + if (tblptr && StringUtil::CIEquals(tblptr->name.GetIdentifierName(), table_name)) { return &get; } } @@ -161,7 +161,7 @@ unique_ptr *FindNodeRefByTable(unique_ptr *roo if (cur->type == LogicalOperatorType::LOGICAL_GET) { auto &get = cur->Cast(); auto tblptr = get.GetTable(); - if (tblptr && tblptr->name == table_name) { + if (tblptr && tblptr->name.GetIdentifierName() == table_name) { if (parent_out) { *parent_out = entry.parent; } @@ -240,7 +240,7 @@ bool HasTableInSubtree(LogicalOperator *op, const string &table_name) { if (op->type == LogicalOperatorType::LOGICAL_GET) { auto &get = op->Cast(); auto tblptr = get.GetTable(); - if (tblptr && tblptr->name == table_name) { + if (tblptr && tblptr->name.GetIdentifierName() == table_name) { return true; } } @@ -266,7 +266,7 @@ static void CollectTableNamesInSubtree(LogicalOperator *op, std::unordered_setCast(); auto tblptr = get.GetTable(); if (tblptr) { - table_names.insert(tblptr->name); + table_names.insert(tblptr->name.GetIdentifierName()); } } for (auto &child : op->children) { @@ -283,7 +283,7 @@ void BuildCTETableMap(LogicalOperator *op, CTETableMap &cte_map) { if (op->type == LogicalOperatorType::LOGICAL_MATERIALIZED_CTE) { auto &cte = op->Cast(); - idx_t cte_idx = cte.table_index; + idx_t cte_idx = cte.table_index.index; // children[0] = CTE definition, children[1] = consumer if (!cte.children.empty()) { @@ -306,7 +306,7 @@ static void MergeCTERefsIntoSet(LogicalOperator *op, idx_t target_cte_idx, CTETa } if (op->type == LogicalOperatorType::LOGICAL_CTE_REF) { auto &ref = op->Cast(); - auto it = cte_map.find(ref.cte_index); + auto it = cte_map.find(ref.cte_index.index); if (it != cte_map.end()) { cte_map[target_cte_idx].insert(it->second.begin(), it->second.end()); } @@ -325,7 +325,7 @@ static void ResolveCTERefsInDefinitions(LogicalOperator *op, CTETableMap &cte_ma if (op->type == LogicalOperatorType::LOGICAL_MATERIALIZED_CTE) { auto &cte = op->Cast(); if (!cte.children.empty()) { - MergeCTERefsIntoSet(cte.children[0].get(), cte.table_index, cte_map); + MergeCTERefsIntoSet(cte.children[0].get(), cte.table_index.index, cte_map); } } for (auto &child : op->children) { @@ -371,7 +371,7 @@ bool HasTableInSubtreeCTE(LogicalOperator *op, const string &table_name, const C if (op->type == LogicalOperatorType::LOGICAL_GET) { auto &get = op->Cast(); auto tblptr = get.GetTable(); - if (tblptr && tblptr->name == table_name) { + if (tblptr && tblptr->name.GetIdentifierName() == table_name) { return true; } } @@ -379,7 +379,7 @@ bool HasTableInSubtreeCTE(LogicalOperator *op, const string &table_name, const C // Follow CTE references through the map if (op->type == LogicalOperatorType::LOGICAL_CTE_REF) { auto &ref = op->Cast(); - auto it = cte_map.find(ref.cte_index); + auto it = cte_map.find(ref.cte_index.index); if (it != cte_map.end() && it->second.count(table_name) > 0) { return true; } @@ -407,7 +407,7 @@ void FindAllNodesByTable(unique_ptr *root, const string &table_ if (cur->type == LogicalOperatorType::LOGICAL_GET) { auto &get = cur->Cast(); auto tblptr = get.GetTable(); - if (tblptr && tblptr->name == table_name) { + if (tblptr && tblptr->name.GetIdentifierName() == table_name) { results.push_back(root); } } @@ -439,7 +439,7 @@ bool HasTableIndexInSubtree(LogicalOperator *op, idx_t table_index) { // Check if this is a LogicalGet with the target table index if (op->type == LogicalOperatorType::LOGICAL_GET) { auto &get = op->Cast(); - if (get.table_index == table_index) { + if (get.table_index.index == table_index) { return true; } } @@ -462,26 +462,26 @@ LogicalOperator *FindOperatorByTableIndex(LogicalOperator *op, idx_t table_index } if (op->type == LogicalOperatorType::LOGICAL_GET) { auto &get = op->Cast(); - if (get.table_index == table_index) { + if (get.table_index.index == table_index) { return op; } } else if (op->type == LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY) { auto &aggr = op->Cast(); - if (aggr.group_index == table_index || aggr.aggregate_index == table_index) { + if (aggr.group_index.index == table_index || aggr.aggregate_index.index == table_index) { return op; } } else if (op->type == LogicalOperatorType::LOGICAL_PROJECTION) { auto &proj = op->Cast(); - if (proj.table_index == table_index) { + if (proj.table_index.index == table_index) { return op; } } else if (op->type == LogicalOperatorType::LOGICAL_CTE_REF) { auto &cte_ref = op->Cast(); - if (cte_ref.table_index == table_index) { + if (cte_ref.table_index.index == table_index) { return op; } } else if (op->type == LogicalOperatorType::LOGICAL_UNION) { - if (op->Cast().table_index == table_index) { + if (op->Cast().table_index.index == table_index) { return op; } } @@ -501,7 +501,7 @@ LogicalMaterializedCTE *FindMaterializedCTE(LogicalOperator *op, idx_t cte_table } if (op->type == LogicalOperatorType::LOGICAL_MATERIALIZED_CTE) { auto &cte = op->Cast(); - if (cte.table_index == cte_table_index) { + if (cte.table_index.index == cte_table_index) { return &cte; } } @@ -526,7 +526,7 @@ void FindAllNodesByTableIndex(unique_ptr *root, idx_t table_ind // Check if this is a LogicalGet with the target table index if (cur->type == LogicalOperatorType::LOGICAL_GET) { auto &get = cur->Cast(); - if (get.table_index == table_index) { + if (get.table_index.index == table_index) { results.push_back(root); } } @@ -636,7 +636,7 @@ static bool CheckColumnsAccessible(LogicalOperator *op, idx_t table_index) { // If this is the target table, it's accessible from here if (op->type == LogicalOperatorType::LOGICAL_GET) { auto &get = op->Cast(); - if (get.table_index == table_index) { + if (get.table_index.index == table_index) { return true; } } @@ -737,7 +737,7 @@ static LogicalGet *FindLogicalGetByTableIndex(LogicalOperator *op, idx_t table_i } if (op->type == LogicalOperatorType::LOGICAL_GET) { auto &get = op->Cast(); - if (get.table_index == table_index) { + if (get.table_index.index == table_index) { return &get; } } @@ -756,7 +756,7 @@ static std::pair GetColumnInfoFromBinding(LogicalOperator *subtr return {"", ""}; } - auto *get = FindLogicalGetByTableIndex(subtree, binding.table_index); + auto *get = FindLogicalGetByTableIndex(subtree, binding.table_index.index); if (!get) { return {"", ""}; } @@ -766,7 +766,7 @@ static std::pair GetColumnInfoFromBinding(LogicalOperator *subtr return {"", ""}; } - string table_name = table_entry->name; + string table_name = table_entry->name.GetIdentifierName(); // Get column name from the binding const auto &column_ids = get->GetColumnIds(); @@ -774,7 +774,7 @@ static std::pair GetColumnInfoFromBinding(LogicalOperator *subtr return {table_name, ""}; } - string col_name = get->GetColumnName(column_ids[binding.column_index]); + string col_name = get->GetColumnName(column_ids[binding.column_index]).GetIdentifierName(); return {table_name, col_name}; } @@ -793,9 +793,9 @@ bool AggregateGroupsByPUKey(LogicalAggregate *agg, const PrivacyCompatibilityRes continue; } ExpressionIterator::EnumerateExpression(group_expr, [&](Expression &expr) { - if (expr.type == ExpressionType::BOUND_COLUMN_REF) { + if (expr.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = expr.Cast(); - group_bindings.push_back(col_ref.binding); + group_bindings.push_back(col_ref.BindingMutable()); } }); } @@ -810,7 +810,7 @@ bool AggregateGroupsByPUKey(LogicalAggregate *agg, const PrivacyCompatibilityRes auto &table_name = col_info.first; auto &col_name = col_info.second; #if PRIVACY_DEBUG - PRIVACY_DEBUG_PRINT("AggregateGroupsByPUKey: binding [" + std::to_string(binding.table_index) + "." + + PRIVACY_DEBUG_PRINT("AggregateGroupsByPUKey: binding [" + std::to_string(binding.table_index.index) + "." + std::to_string(binding.column_index) + "] -> table='" + table_name + "' col='" + col_name + "'"); #endif @@ -895,6 +895,11 @@ vector FilterTargetAggregatesWithPUKeyCheck(const vector include_outer_aggregates; for (auto *agg : all_aggregates) { + // Correlation-domain deduplication is a GROUP BY without aggregate + // expressions. It must not move privacy responsibility to another node. + if (agg->expressions.empty()) { + continue; + } // Check if this aggregate has a child aggregate (nested aggregate — direct subtree) for (auto &child : agg->children) { auto *inner_agg = FindFirstChildAggregate(child.get()); @@ -924,7 +929,7 @@ vector FilterTargetAggregatesWithPUKeyCheck(const vectorexpressions.empty()) { continue; } bool other_has_direct = false; @@ -953,6 +958,11 @@ vector FilterTargetAggregatesWithPUKeyCheck(const vectorexpressions.empty()) { + continue; + } // Skip aggregates that are marked to be skipped (inner aggregates that group by PU key) if (skip_aggregates.count(agg) > 0) { #if PRIVACY_DEBUG @@ -1057,12 +1067,12 @@ LogicalAggregate *FindInnerAggregateWithPUKeyGroup(LogicalAggregate *target_agg, // Now find which group column is the PU key and get its output binding for (idx_t group_idx = 0; group_idx < inner_agg->groups.size(); group_idx++) { auto &group_expr = inner_agg->groups[group_idx]; - if (!group_expr || group_expr->type != ExpressionType::BOUND_COLUMN_REF) { + if (!group_expr || group_expr->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { continue; } auto &col_ref = group_expr->Cast(); - auto col_info = GetColumnInfoFromBinding(inner_agg, col_ref.binding); + auto col_info = GetColumnInfoFromBinding(inner_agg, col_ref.BindingMutable()); auto &table_name = col_info.first; auto &col_name = col_info.second; if (table_name.empty() || col_name.empty()) { @@ -1081,7 +1091,7 @@ LogicalAggregate *FindInnerAggregateWithPUKeyGroup(LogicalAggregate *target_agg, // Found the PU key group column // The group columns are output as the first N columns of the aggregate // Their bindings use the aggregate's group_index - ColumnBinding current_binding(inner_agg->group_index, group_idx); + ColumnBinding current_binding {TableIndex(inner_agg->group_index.index), ProjectionIndex(group_idx)}; // Now propagate this binding through any projections between inner_agg and target_agg // The path_to_inner contains operators from outer's child down to (but not including) inner_agg @@ -1095,12 +1105,13 @@ LogicalAggregate *FindInnerAggregateWithPUKeyGroup(LogicalAggregate *target_agg, bool found = false; for (idx_t expr_idx = 0; expr_idx < proj.expressions.size(); expr_idx++) { auto &expr = proj.expressions[expr_idx]; - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { + if (expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &proj_col_ref = expr->Cast(); - if (proj_col_ref.binding == current_binding) { + if (proj_col_ref.BindingMutable() == current_binding) { // This projection expression references our binding // The new binding is the projection's output - current_binding = ColumnBinding(proj.table_index, expr_idx); + current_binding = + ColumnBinding(TableIndex(proj.table_index.index), ProjectionIndex(expr_idx)); found = true; break; } diff --git a/src/table_functions/dp_sass_stability_query.cpp b/src/table_functions/dp_sass_stability_query.cpp index 80a77210..a36911ce 100644 --- a/src/table_functions/dp_sass_stability_query.cpp +++ b/src/table_functions/dp_sass_stability_query.cpp @@ -32,13 +32,13 @@ struct DpSassStabilityQueryState : public GlobalTableFunctionState { idx_t offset = 0; }; -static void AddOutputColumn(vector &return_types, vector &names, const string &name, +static void AddOutputColumn(vector &return_types, vector &names, const string &name, const LogicalType &type) { - names.push_back(name); + names.emplace_back(name); return_types.push_back(type); } -static void AddStabilityOutputColumns(vector &return_types, vector &names) { +static void AddStabilityOutputColumns(vector &return_types, vector &names) { AddOutputColumn(return_types, names, "group_row", LogicalType::BIGINT); AddOutputColumn(return_types, names, "aggregate_index", LogicalType::INTEGER); AddOutputColumn(return_types, names, "aggregate_name", LogicalType::VARCHAR); @@ -59,7 +59,7 @@ static void AddStabilityOutputColumns(vector &return_types, vector< AddOutputColumn(return_types, names, "stable_mad_1", LogicalType::BOOLEAN); } -static void AddNoiseScaleOutputColumns(vector &return_types, vector &names) { +static void AddNoiseScaleOutputColumns(vector &return_types, vector &names) { AddOutputColumn(return_types, names, "group_row", LogicalType::BIGINT); AddOutputColumn(return_types, names, "aggregate_index", LogicalType::INTEGER); AddOutputColumn(return_types, names, "aggregate_name", LogicalType::VARCHAR); @@ -113,7 +113,7 @@ static void RunSetting(Connection &conn, const string &sql) { static void CopySetting(ClientContext &context, Connection &conn, const string &setting_name) { Value value; - if (!context.TryGetCurrentSetting(setting_name, value) || value.IsNull()) { + if (!context.TryGetCurrentSetting(Identifier(setting_name), value) || value.IsNull()) { return; } RunSetting(conn, "SET " + setting_name + " = " + value.ToSQLString()); @@ -179,7 +179,7 @@ static string BuildGroupKeyJson(QueryResult &result, const vector &row, i json += ","; } json += "\""; - json += EscapeJsonString(result.names[col]); + json += EscapeJsonString(result.GetNames()[col].GetIdentifierName()); json += "\":"; auto &value = row[col]; if (value.IsNull()) { @@ -266,8 +266,8 @@ ExecuteDiagnosticQuery(ClientContext &context, const string &sql, const string & values.push_back(Value::BIGINT(static_cast(group_row))); values.push_back(Value::INTEGER(record.aggregate_index)); idx_t result_name_idx = group_cols + aggregate_index; - if (result_columns_align && result_name_idx < result->names.size()) { - values.push_back(Value(result->names[result_name_idx])); + if (result_columns_align && result_name_idx < result->GetNames().size()) { + values.push_back(Value(result->GetNames()[result_name_idx].GetIdentifierName())); } else { values.push_back(Value("agg_" + std::to_string(record.aggregate_index))); } @@ -302,7 +302,7 @@ static vector ExecuteNoiseScaleQuery(ClientContext &con } static unique_ptr DpSassStabilityQueryBind(ClientContext &context, TableFunctionBindInput &input, - vector &return_types, vector &names) { + vector &return_types, vector &names) { if (input.inputs.empty() || input.inputs[0].IsNull()) { throw BinderException("dp_sass_stability_query requires a non-NULL SQL string"); } @@ -312,7 +312,8 @@ static unique_ptr DpSassStabilityQueryBind(ClientContext &context, } static unique_ptr DpSassNoiseScaleQueryBind(ClientContext &context, TableFunctionBindInput &input, - vector &return_types, vector &names) { + vector &return_types, + vector &names) { if (input.inputs.empty() || input.inputs[0].IsNull()) { throw BinderException("dp_sass_noise_scale_query requires a non-NULL SQL string"); } @@ -329,15 +330,14 @@ static unique_ptr DpSassStabilityQueryInit(ClientConte static void DpSassStabilityQueryFunction(ClientContext &context, TableFunctionInput &data_p, DataChunk &output) { auto &bind_data = data_p.bind_data->Cast(); auto &state = data_p.global_state->Cast(); - idx_t count = 0; - while (state.offset < bind_data.rows.size() && count < STANDARD_VECTOR_SIZE) { + idx_t count = MinValue(bind_data.rows.size() - state.offset, STANDARD_VECTOR_SIZE); + output.SetChildCardinality(count); + for (idx_t row_idx = 0; row_idx < count; row_idx++) { auto &row = bind_data.rows[state.offset++]; for (idx_t col = 0; col < row.values.size(); col++) { - output.SetValue(col, count, row.values[col]); + output.data[col].SetValue(row_idx, row.values[col]); } - count++; } - output.SetCardinality(count); } } // namespace diff --git a/src/test/test_compiler_functions.cpp b/src/test/test_compiler_functions.cpp index 250c241c..1fbe653b 100644 --- a/src/test/test_compiler_functions.cpp +++ b/src/test/test_compiler_functions.cpp @@ -48,7 +48,7 @@ static void CollectColumnBindingsRecursiveForTest(LogicalOperator *node, std::se } auto binds = node->GetColumnBindings(); for (auto &b : binds) { - out.insert({b.table_index, b.column_index}); + out.insert({b.table_index.index, b.column_index}); } for (auto &c : node->children) { CollectColumnBindingsRecursiveForTest(c.get(), out); @@ -63,11 +63,11 @@ static bool VerifyAllExprBindingsExist(LogicalOperator *node, const std::setexpressions) { ExpressionIterator::VisitExpression(*expr, [&](const BoundColumnRefExpression &bcr) { - auto key = std::make_pair(bcr.binding.table_index, bcr.binding.column_index); + auto key = std::make_pair(bcr.Binding().table_index.index, bcr.Binding().column_index); if (bindings.find(key) == bindings.end()) { std::cerr << "Verification failure: expression references missing binding (table=" - << bcr.binding.table_index << ", col=" << bcr.binding.column_index << ") in operator " - << node->GetName() << "\n"; + << bcr.Binding().table_index.index << ", col=" << bcr.Binding().column_index + << ") in operator " << node->GetName() << "\n"; ok = false; } }); @@ -98,7 +98,7 @@ static void CollectBindingProducers(LogicalOperator *node, } auto binds = node->GetColumnBindings(); for (auto &b : binds) { - out[{b.table_index, b.column_index}].push_back(node); + out[{b.table_index.index, b.column_index}].push_back(node); } for (auto &c : node->children) { CollectBindingProducers(c.get(), out); @@ -200,7 +200,7 @@ static bool VerifyExprTypesMatchProducers(LogicalOperator *root) { for (auto &expr : node->expressions) { ExpressionIterator::VisitExpression( *expr, [&](const BoundColumnRefExpression &bcr) { - auto key = std::make_pair(bcr.binding.table_index, bcr.binding.column_index); + auto key = std::make_pair(bcr.Binding().table_index.index, bcr.Binding().column_index); auto it = producers.find(key); if (it == producers.end() || it->second.empty()) { // missing producer already reported elsewhere; skip @@ -208,17 +208,18 @@ static bool VerifyExprTypesMatchProducers(LogicalOperator *root) { } auto *producer = it->second[0]; auto types = producer->types; - if (bcr.binding.column_index >= (idx_t)types.size()) { + if (bcr.Binding().column_index >= (idx_t)types.size()) { std::cerr << "Verification failure: expression references column index " - << bcr.binding.column_index << " beyond producer types size " << types.size() << "\n"; + << bcr.Binding().column_index << " beyond producer types size " << types.size() + << "\n"; ok = false; return; } - const auto &prod_type = types[bcr.binding.column_index]; - if (bcr.return_type.id() != prod_type.id()) { + const auto &prod_type = types[bcr.Binding().column_index]; + if (bcr.GetReturnType().id() != prod_type.id()) { std::cerr << "Verification failure: expression type mismatch for binding (table=" - << bcr.binding.table_index << ", col=" << bcr.binding.column_index - << "): expr_type=" << bcr.return_type.ToString() + << bcr.Binding().table_index.index << ", col=" << bcr.Binding().column_index + << "): expr_type=" << bcr.GetReturnType().ToString() << " producer_type=" << prod_type.ToString() << "\n"; ok = false; } @@ -261,7 +262,7 @@ static bool VerifyOperatorBindingsConsistentWithTypes(LogicalOperator *root) { for (auto &expr : node->expressions) { // Enumerate all expressions under this unique_ptr and check their return types ExpressionIterator::EnumerateExpression(expr, [&](const Expression &e) { - if (e.return_type.id() == LogicalTypeId::INVALID) { + if (e.GetReturnType().id() == LogicalTypeId::INVALID) { std::cerr << "Verification failure: expression in operator " << node->GetName() << " has INVALID return type\n"; ok = false; @@ -342,18 +343,20 @@ static int RunRobustPlanVerification(const unique_ptr &root) { // Test 1: basic ReplaceNode insertion and binding remap static int Test_BasicReplaceNode() { // Build a simple plan: root projection (table_index = 0) with a child DummyScan (table_index = 0) - auto child = make_uniq(0); + auto child = make_uniq(TableIndex(0)); vector> proj_exprs; // Projection expression referencing child's first column (binding {0,0}) - proj_exprs.emplace_back(make_uniq(LogicalType::INTEGER, ColumnBinding(0, 0))); + proj_exprs.emplace_back( + make_uniq(LogicalType::INTEGER, ColumnBinding(TableIndex(0), ProjectionIndex(0)))); // store root as unique_ptr so it can be passed to ReplaceNode - unique_ptr root = make_uniq(0, std::move(proj_exprs)); + unique_ptr root = make_uniq(TableIndex(0), std::move(proj_exprs)); root->children.emplace_back(std::move(child)); // Prepare a new parent projection that contains its own child (strict replace semantics) - unique_ptr new_parent = make_uniq(1, vector>()); - new_parent->children.emplace_back(make_uniq(5)); + unique_ptr new_parent = + make_uniq(TableIndex(1), vector>()); + new_parent->children.emplace_back(make_uniq(TableIndex(5))); ReplaceNode(root, root->children[0], new_parent); // After insertion, validate @@ -383,16 +386,18 @@ static int Test_BasicReplaceNode() { // Test 2: multiple consecutive ReplaceNode calls produce consistent remapping static int Test_ConsecutiveReplaceNode() { // Start with root projection (table_index = 0) and a DummyScan child (table_index = 0) - auto child = make_uniq(0); + auto child = make_uniq(TableIndex(0)); vector> proj_exprs; - proj_exprs.emplace_back(make_uniq(LogicalType::INTEGER, ColumnBinding(0, 0))); - unique_ptr root = make_uniq(0, std::move(proj_exprs)); + proj_exprs.emplace_back( + make_uniq(LogicalType::INTEGER, ColumnBinding(TableIndex(0), ProjectionIndex(0)))); + unique_ptr root = make_uniq(TableIndex(0), std::move(proj_exprs)); root->children.emplace_back(std::move(child)); // First insertion - unique_ptr new_parent1 = make_uniq(1, vector>()); - new_parent1->children.emplace_back(make_uniq(7)); + unique_ptr new_parent1 = + make_uniq(TableIndex(1), vector>()); + new_parent1->children.emplace_back(make_uniq(TableIndex(7))); ReplaceNode(root, root->children[0], new_parent1); // Second insertion above the same original child @@ -400,8 +405,9 @@ static int Test_ConsecutiveReplaceNode() { if (!root->children[0]) { return 7; } - unique_ptr new_parent2 = make_uniq(2, vector>()); - new_parent2->children.emplace_back(make_uniq(8)); + unique_ptr new_parent2 = + make_uniq(TableIndex(2), vector>()); + new_parent2->children.emplace_back(make_uniq(TableIndex(8))); ReplaceNode(root->children[0], root->children[0]->children[0], new_parent2); // Validate structure: there should be two projections between root and the DummyScan @@ -435,17 +441,19 @@ static int Test_ConsecutiveReplaceNode() { // Test 3: inserting parent when child has multiple columns and projection references later column static int Test_MultiColumnChildBinding() { // DummyScan with two columns: logical behavior is simulated via projections using column bindings - auto child = make_uniq(0); + auto child = make_uniq(TableIndex(0)); vector> proj_exprs; // Reference first column from child binding - proj_exprs.emplace_back(make_uniq(LogicalType::INTEGER, ColumnBinding(0, 0))); - unique_ptr root = make_uniq(0, std::move(proj_exprs)); + proj_exprs.emplace_back( + make_uniq(LogicalType::INTEGER, ColumnBinding(TableIndex(0), ProjectionIndex(0)))); + unique_ptr root = make_uniq(TableIndex(0), std::move(proj_exprs)); root->children.emplace_back(std::move(child)); // Create new_parent with a child scan - this matches the pattern used in other tests - unique_ptr new_parent = make_uniq(1, vector>()); - new_parent->children.emplace_back(make_uniq(17)); + unique_ptr new_parent = + make_uniq(TableIndex(1), vector>()); + new_parent->children.emplace_back(make_uniq(TableIndex(17))); ReplaceNode(root, root->children[0], new_parent); @@ -471,15 +479,17 @@ static int Test_MultiColumnChildBinding() { // Test 4: insert a projection between an aggregate and root static int Test_InsertProjectionBetweenAggregateAndRoot() { // Setup: root projection referencing an aggregate child - auto agg_child = make_uniq(0, 0, vector>()); - unique_ptr root = make_uniq(0, vector>()); + auto agg_child = make_uniq(TableIndex(0), TableIndex(0), vector>()); + unique_ptr root = make_uniq(TableIndex(0), vector>()); // projection referencing the aggregate's (would-be) column 0 - root->expressions.emplace_back(make_uniq(LogicalType::INTEGER, ColumnBinding(0, 0))); + root->expressions.emplace_back( + make_uniq(LogicalType::INTEGER, ColumnBinding(TableIndex(0), ProjectionIndex(0)))); root->children.emplace_back(std::move(agg_child)); // Insert a LogicalProjection as new parent above the aggregate - unique_ptr new_proj = make_uniq(12, vector>()); - new_proj->children.emplace_back(make_uniq(13)); + unique_ptr new_proj = + make_uniq(TableIndex(12), vector>()); + new_proj->children.emplace_back(make_uniq(TableIndex(13))); ReplaceNode(root, root->children[0], new_proj); @@ -505,14 +515,15 @@ static int Test_InsertProjectionBetweenAggregateAndRoot() { // Test 6: insert a filter parent above a scan and ensure child moved static int Test_InsertFilterParent() { - auto child = make_uniq(0); + auto child = make_uniq(TableIndex(0)); vector> proj_exprs; - proj_exprs.emplace_back(make_uniq(LogicalType::INTEGER, ColumnBinding(0, 0))); - unique_ptr root = make_uniq(0, std::move(proj_exprs)); + proj_exprs.emplace_back( + make_uniq(LogicalType::INTEGER, ColumnBinding(TableIndex(0), ProjectionIndex(0)))); + unique_ptr root = make_uniq(TableIndex(0), std::move(proj_exprs)); root->children.emplace_back(std::move(child)); unique_ptr new_filter = make_uniq(); - new_filter->children.emplace_back(make_uniq(16)); + new_filter->children.emplace_back(make_uniq(TableIndex(16))); ReplaceNode(root, root->children[0], new_filter); // Verify plan @@ -567,15 +578,18 @@ static int Test_InsertFilterParent() { // Test 7: ensure ReplaceNode preserves column positions when no remapping is required static int Test_PreserveColumnsNoRemap() { - auto child = make_uniq(0); + auto child = make_uniq(TableIndex(0)); vector> proj_exprs; - proj_exprs.emplace_back(make_uniq(LogicalType::INTEGER, ColumnBinding(0, 0))); - proj_exprs.emplace_back(make_uniq(LogicalType::INTEGER, ColumnBinding(0, 1))); - unique_ptr root = make_uniq(0, std::move(proj_exprs)); + proj_exprs.emplace_back( + make_uniq(LogicalType::INTEGER, ColumnBinding(TableIndex(0), ProjectionIndex(0)))); + proj_exprs.emplace_back( + make_uniq(LogicalType::INTEGER, ColumnBinding(TableIndex(0), ProjectionIndex(1)))); + unique_ptr root = make_uniq(TableIndex(0), std::move(proj_exprs)); root->children.emplace_back(std::move(child)); - unique_ptr new_parent = make_uniq(1, vector>()); - new_parent->children.emplace_back(make_uniq(17)); + unique_ptr new_parent = + make_uniq(TableIndex(1), vector>()); + new_parent->children.emplace_back(make_uniq(TableIndex(17))); ReplaceNode(root, root->children[0], new_parent); // Validate structure @@ -597,10 +611,10 @@ static int Test_PreserveColumnsNoRemap() { // Verify column indices are preserved (0 and 1) auto &bcr0 = root->expressions[0]->Cast(); auto &bcr1 = root->expressions[1]->Cast(); - if (bcr0.binding.column_index != 0) { + if (bcr0.Binding().column_index != 0) { return 31; } - if (bcr1.binding.column_index != 1) { + if (bcr1.Binding().column_index != 1) { return 32; } @@ -610,15 +624,16 @@ static int Test_PreserveColumnsNoRemap() { // New test: verifier should detect duplicate producers for the same ColumnBinding static int Test_VerifierDetectsDuplicateProducers() { // Create a join whose left and right are both DummyScans with the same table_index (0) - auto left = make_uniq(0); - auto right = make_uniq(0); + auto left = make_uniq(TableIndex(0)); + auto right = make_uniq(TableIndex(0)); auto base_join = make_uniq(JoinType::INNER); base_join->children.emplace_back(std::move(left)); base_join->children.emplace_back(std::move(right)); - unique_ptr root = make_uniq(0, vector>()); + unique_ptr root = make_uniq(TableIndex(0), vector>()); // Add a projection expression that references the duplicated binding (table=0,col=0) - root->expressions.emplace_back(make_uniq(LogicalType::INTEGER, ColumnBinding(0, 0))); + root->expressions.emplace_back( + make_uniq(LogicalType::INTEGER, ColumnBinding(TableIndex(0), ProjectionIndex(0)))); root->children.emplace_back(std::move(base_join)); int res = RunRobustPlanVerification(root); @@ -678,7 +693,7 @@ int RunCompilerFunctionTests() { // Test 1: Simple scan - check that columns from t1 belong to t1 { - Parser parser; + Parser parser(*con.context); parser.ParseQuery("SELECT a, b FROM t1;"); Planner planner(*con.context); @@ -744,7 +759,7 @@ int RunCompilerFunctionTests() { // Join key columns (t1.a, t2.x) belong to BOTH tables via equivalence. // Non-join-key columns (t1.b, t2.y) belong ONLY to their own table. { - Parser parser; + Parser parser(*con.context); parser.ParseQuery("SELECT t1.a, t1.b, t2.x, t2.y FROM t1 INNER JOIN t2 ON t1.a = t2.x;"); Planner planner(*con.context); planner.CreatePlan(std::move(parser.statements[0])); @@ -841,7 +856,7 @@ int RunCompilerFunctionTests() { // Join conditions: t1.a = t2.x AND t2.x = t3.m → transitive: t1.a = t2.x = t3.m // Join key columns belong to all three tables; non-join-key columns stay isolated. { - Parser parser; + Parser parser(*con.context); parser.ParseQuery("SELECT * FROM t1 INNER JOIN t2 ON t1.a = t2.x INNER JOIN t3 ON t2.x = t3.m;"); Planner planner(*con.context); planner.CreatePlan(std::move(parser.statements[0])); @@ -943,7 +958,7 @@ int RunCompilerFunctionTests() { // Test 4: Non-existent table { - Parser parser; + Parser parser(*con.context); parser.ParseQuery("SELECT a FROM t1;"); Planner planner(*con.context); planner.CreatePlan(std::move(parser.statements[0])); @@ -991,7 +1006,7 @@ int RunCompilerFunctionTests() { // Test 5: Subquery with multiple tables { - Parser parser; + Parser parser(*con.context); parser.ParseQuery("SELECT * FROM (SELECT t1.a, t2.x FROM t1, t2) AS sub;"); Planner planner(*con.context); planner.CreatePlan(std::move(parser.statements[0])); @@ -1070,7 +1085,7 @@ int RunCompilerFunctionTests() { "SELECT t2.x, SUM(t3.n) FROM t1 INNER JOIN t2 ON t1.a = t2.x INNER JOIN t3 ON t2.x = t3.m GROUP BY t2.x;"; // Create plan using parser, planner and optimizer - Parser parser; + Parser parser(*con.context); parser.ParseQuery(query); Planner planner(*con.context); planner.CreatePlan(std::move(parser.statements[0])); @@ -1110,7 +1125,7 @@ int RunCompilerFunctionTests() { } auto &group_expr = agg_node->groups[0]; - if (group_expr->type != ExpressionType::BOUND_COLUMN_REF) { + if (group_expr->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { std::cerr << "FAIL: ColumnBelongsToTable test 6: group expression is not a column ref" << "\n"; return 1; @@ -1148,8 +1163,8 @@ int RunCompilerFunctionTests() { // The GROUP BY column could belong to either t1 or t2 since we're joining on the key // (t1.a = t2.x), making them equivalent - bool belongs_to_t1 = ColumnBelongsToTable(*plan, "t1", col_ref.binding); - bool belongs_to_t2 = ColumnBelongsToTable(*plan, "t2", col_ref.binding); + bool belongs_to_t1 = ColumnBelongsToTable(*plan, "t1", col_ref.Binding()); + bool belongs_to_t2 = ColumnBelongsToTable(*plan, "t2", col_ref.Binding()); // The column should belong to at least one of the tables (t1 or t2) if (!belongs_to_t1 && !belongs_to_t2) { @@ -1176,7 +1191,7 @@ int RunCompilerFunctionTests() { const string query = "SELECT a, SUM(b) FROM t1 GROUP BY a ORDER BY a;"; // Create plan using parser, planner and optimizer - Parser parser; + Parser parser(*con.context); parser.ParseQuery(query); Planner planner(*con.context); planner.CreatePlan(std::move(parser.statements[0])); @@ -1239,7 +1254,7 @@ int RunCompilerFunctionTests() { } auto &group_expr = agg_node->groups[0]; - if (group_expr->type != ExpressionType::BOUND_COLUMN_REF) { + if (group_expr->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { std::cerr << "FAIL: ColumnBelongsToTable test 7: group expression is not a column ref" << "\n"; return 1; @@ -1248,7 +1263,7 @@ int RunCompilerFunctionTests() { auto &group_col_ref = group_expr->Cast(); // Verify that the GROUP BY column 'a' belongs to t1 - if (!ColumnBelongsToTable(*plan, "t1", group_col_ref.binding)) { + if (!ColumnBelongsToTable(*plan, "t1", group_col_ref.Binding())) { std::cerr << "FAIL: ColumnBelongsToTable test 7: GROUP BY column 'a' should belong to t1" << "\n"; return 1; @@ -1267,7 +1282,7 @@ int RunCompilerFunctionTests() { ExpressionIterator::VisitExpression( *agg_expr, [&](const BoundColumnRefExpression &col_ref) { // Verify that column 'b' belongs to t1 - if (!ColumnBelongsToTable(*plan, "t1", col_ref.binding)) { + if (!ColumnBelongsToTable(*plan, "t1", col_ref.Binding())) { std::cerr << "FAIL: ColumnBelongsToTable test 7: SUM(b) column 'b' should belong to t1" << "\n"; found_b_column = false; @@ -1321,7 +1336,7 @@ int RunCompilerFunctionTests() { "pu_table.a = non_pu_table.x GROUP BY non_pu_table.y;"; // Create plan using parser, planner and optimizer - Parser parser; + Parser parser(*con.context); parser.ParseQuery(query); Planner planner(*con.context); planner.CreatePlan(std::move(parser.statements[0])); @@ -1361,7 +1376,7 @@ int RunCompilerFunctionTests() { } auto &group_expr = agg_node->groups[0]; - if (group_expr->type != ExpressionType::BOUND_COLUMN_REF) { + if (group_expr->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { std::cerr << "FAIL: ColumnBelongsToTable test 8: group expression is not a column ref" << "\n"; return 1; @@ -1370,14 +1385,14 @@ int RunCompilerFunctionTests() { auto &col_ref = group_expr->Cast(); // The GROUP BY column should belong to non_pu_table - if (!ColumnBelongsToTable(*plan, "non_pu_table", col_ref.binding)) { + if (!ColumnBelongsToTable(*plan, "non_pu_table", col_ref.Binding())) { std::cerr << "FAIL: ColumnBelongsToTable test 8: GROUP BY column should belong to non_pu_table" << "\n"; return 1; } // Should NOT belong to pu_table - if (ColumnBelongsToTable(*plan, "pu_table", col_ref.binding)) { + if (ColumnBelongsToTable(*plan, "pu_table", col_ref.Binding())) { std::cerr << "FAIL: ColumnBelongsToTable test 8: GROUP BY column should not belong to pu_table" << "\n"; return 1; @@ -1407,7 +1422,7 @@ int RunCompilerFunctionTests() { "pu_table.a = non_pu_table.x GROUP BY non_pu_table.x;"; // Create plan using parser, planner and optimizer - Parser parser; + Parser parser(*con.context); parser.ParseQuery(query); Planner planner(*con.context); planner.CreatePlan(std::move(parser.statements[0])); @@ -1447,7 +1462,7 @@ int RunCompilerFunctionTests() { } auto &group_expr = agg_node->groups[0]; - if (group_expr->type != ExpressionType::BOUND_COLUMN_REF) { + if (group_expr->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { std::cerr << "FAIL: ColumnBelongsToTable test 9: group expression is not a column ref" << "\n"; return 1; @@ -1456,7 +1471,7 @@ int RunCompilerFunctionTests() { auto &col_ref = group_expr->Cast(); // The GROUP BY column physically comes from non_pu_table - if (!ColumnBelongsToTable(*plan, "non_pu_table", col_ref.binding)) { + if (!ColumnBelongsToTable(*plan, "non_pu_table", col_ref.Binding())) { std::cerr << "FAIL: ColumnBelongsToTable test 9: GROUP BY column should belong to non_pu_table" << "\n"; return 1; @@ -1464,7 +1479,7 @@ int RunCompilerFunctionTests() { // BUT because it's a join key (pu_table.a = non_pu_table.x), it should ALSO be detected // as belonging to pu_table due to join key equivalence - if (!ColumnBelongsToTable(*plan, "pu_table", col_ref.binding)) { + if (!ColumnBelongsToTable(*plan, "pu_table", col_ref.Binding())) { std::cerr << "FAIL: ColumnBelongsToTable test 9: JOIN KEY column should ALSO belong to pu_table " "(due to join key equivalence)" << "\n"; diff --git a/src/test/test_plan_traversal.cpp b/src/test/test_plan_traversal.cpp index 345f4cb3..e19246e8 100644 --- a/src/test/test_plan_traversal.cpp +++ b/src/test/test_plan_traversal.cpp @@ -28,7 +28,7 @@ int RunPlanTraversalTests() { // Test 1: FindPrivacyUnitGetNode - find specific table std::cerr << "\n--- Test FindPrivacyUnitGetNode ---\n"; - Parser parser; + Parser parser(*con.context); parser.ParseQuery("SELECT SUM(val) FROM t_test;"); Planner planner(*con.context); planner.CreatePlan(std::move(parser.statements[0])); @@ -75,7 +75,7 @@ int RunPlanTraversalTests() { // Test 4: FindTopAggregate - no aggregate should throw std::cerr << "\n--- Test FindTopAggregate (no aggregate) ---\n"; - Parser parser2; + Parser parser2(*con.context); parser2.ParseQuery("SELECT val FROM t_test;"); Planner planner2(*con.context); planner2.CreatePlan(std::move(parser2.statements[0])); diff --git a/src/utils/privacy_helpers.cpp b/src/utils/privacy_helpers.cpp index 38310eb0..803fbbfe 100644 --- a/src/utils/privacy_helpers.cpp +++ b/src/utils/privacy_helpers.cpp @@ -1,3 +1,7 @@ +#include "duckdb/main/database_manager.hpp" +#include "duckdb/catalog/catalog_entry/window_function_catalog_entry.hpp" +#include "duckdb/function/function_binder.hpp" +#include "duckdb/planner/expression/bound_window_expression.hpp" #include "utils/privacy_helpers.hpp" #include @@ -23,6 +27,7 @@ #include "duckdb/planner/operator/logical_distinct.hpp" #include "duckdb/planner/operator/logical_order.hpp" #include "duckdb/planner/expression/bound_columnref_expression.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" #include "duckdb/planner/operator/logical_comparison_join.hpp" #include "duckdb/planner/expression_iterator.hpp" #include @@ -39,6 +44,26 @@ using idx_set = std::unordered_set; namespace duckdb { +unique_ptr BindRankingWindow(ClientContext &context, ExpressionType type) { + Identifier name; + switch (type) { + case ExpressionType::WINDOW_ROW_NUMBER: + name = "row_number"; + break; + case ExpressionType::WINDOW_RANK_DENSE: + name = "dense_rank"; + break; + default: + throw InternalException("Unsupported privacy ranking window"); + } + auto &entry = Catalog::GetSystemCatalog(context).GetEntry( + context, QualifiedName(Catalog::GetSystemCatalog(context).GetName(), Identifier::DefaultSchema(), name)); + FunctionBinder binder(context); + vector order_types; + vector arg_order_types; + return binder.BindWindowFunction(entry.functions.GetFunctionByOffset(0), {}, order_types, arg_order_types); +} + string Sanitize(const string &in) { string out; for (char c : in) { @@ -95,7 +120,8 @@ idx_t GetNextTableIndex(unique_ptr &plan) { continue; } auto tbls = cur->GetTableIndex(); - for (auto t : tbls) { + for (auto typed_index : tbls) { + auto t = typed_index.index; if (t != DConstants::INVALID_INDEX && (max_index == DConstants::INVALID_INDEX || t > max_index)) { max_index = t; } @@ -112,7 +138,8 @@ void CollectTableIndicesRecursive(LogicalOperator *node, idx_set &out) { return; } auto tbls = node->GetTableIndex(); - for (auto t : tbls) { + for (auto typed_index : tbls) { + auto t = typed_index.index; if (t != DConstants::INVALID_INDEX) { out.insert(t); } @@ -130,7 +157,8 @@ static void CollectTableIndicesExcluding(LogicalOperator *node, LogicalOperator return; // skip this subtree entirely } auto tbls = node->GetTableIndex(); - for (auto t : tbls) { + for (auto typed_index : tbls) { + auto t = typed_index.index; if (t != DConstants::INVALID_INDEX) { out.insert(t); } @@ -146,73 +174,74 @@ void ApplyIndexMapToSubtree(LogicalOperator *node, const std::unordered_map(node)) { - if (map.find(get_ptr->table_index) != map.end()) { - get_ptr->table_index = map.at(get_ptr->table_index); + if (map.find(get_ptr->table_index.index) != map.end()) { + get_ptr->table_index.index = map.at(get_ptr->table_index.index); } } else if (auto proj_ptr = dynamic_cast(node)) { - if (map.find(proj_ptr->table_index) != map.end()) { - proj_ptr->table_index = map.at(proj_ptr->table_index); + if (map.find(proj_ptr->table_index.index) != map.end()) { + proj_ptr->table_index.index = map.at(proj_ptr->table_index.index); } } else if (auto setop_ptr = dynamic_cast(node)) { - if (map.find(setop_ptr->table_index) != map.end()) { - setop_ptr->table_index = map.at(setop_ptr->table_index); + if (map.find(setop_ptr->table_index.index) != map.end()) { + setop_ptr->table_index.index = map.at(setop_ptr->table_index.index); } } else if (auto insert_ptr = dynamic_cast(node)) { - if (map.find(insert_ptr->table_index) != map.end()) { - insert_ptr->table_index = map.at(insert_ptr->table_index); + if (map.find(insert_ptr->table_index.index) != map.end()) { + insert_ptr->table_index.index = map.at(insert_ptr->table_index.index); } } else if (auto dummy_ptr = dynamic_cast(node)) { - if (map.find(dummy_ptr->table_index) != map.end()) { - dummy_ptr->table_index = map.at(dummy_ptr->table_index); + if (map.find(dummy_ptr->table_index.index) != map.end()) { + dummy_ptr->table_index.index = map.at(dummy_ptr->table_index.index); } } else if (auto coldata_ptr = dynamic_cast(node)) { - if (map.find(coldata_ptr->table_index) != map.end()) { - coldata_ptr->table_index = map.at(coldata_ptr->table_index); + if (map.find(coldata_ptr->table_index.index) != map.end()) { + coldata_ptr->table_index.index = map.at(coldata_ptr->table_index.index); } } else if (auto upd_ptr = dynamic_cast(node)) { - if (map.find(upd_ptr->table_index) != map.end()) { - upd_ptr->table_index = map.at(upd_ptr->table_index); + if (map.find(upd_ptr->table_index.index) != map.end()) { + upd_ptr->table_index.index = map.at(upd_ptr->table_index.index); } } else if (auto del_ptr = dynamic_cast(node)) { - if (map.find(del_ptr->table_index) != map.end()) { - del_ptr->table_index = map.at(del_ptr->table_index); + if (map.find(del_ptr->table_index.index) != map.end()) { + del_ptr->table_index.index = map.at(del_ptr->table_index.index); } } else if (auto cte_ptr = dynamic_cast(node)) { - if (map.find(cte_ptr->table_index) != map.end()) { - cte_ptr->table_index = map.at(cte_ptr->table_index); + if (map.find(cte_ptr->table_index.index) != map.end()) { + cte_ptr->table_index.index = map.at(cte_ptr->table_index.index); } } else if (auto rcte_ptr = dynamic_cast(node)) { - if (map.find(rcte_ptr->table_index) != map.end()) { - rcte_ptr->table_index = map.at(rcte_ptr->table_index); + if (map.find(rcte_ptr->table_index.index) != map.end()) { + rcte_ptr->table_index.index = map.at(rcte_ptr->table_index.index); } } else if (auto eg_ptr = dynamic_cast(node)) { - if (map.find(eg_ptr->table_index) != map.end()) { - eg_ptr->table_index = map.at(eg_ptr->table_index); + if (map.find(eg_ptr->table_index.index) != map.end()) { + eg_ptr->table_index.index = map.at(eg_ptr->table_index.index); } } else if (auto cteref_ptr = dynamic_cast(node)) { - if (map.find(cteref_ptr->table_index) != map.end()) { - cteref_ptr->table_index = map.at(cteref_ptr->table_index); + if (map.find(cteref_ptr->table_index.index) != map.end()) { + cteref_ptr->table_index.index = map.at(cteref_ptr->table_index.index); } } else if (auto dg_ptr = dynamic_cast(node)) { - if (map.find(dg_ptr->table_index) != map.end()) { - dg_ptr->table_index = map.at(dg_ptr->table_index); + if (map.find(dg_ptr->table_index.index) != map.end()) { + dg_ptr->table_index.index = map.at(dg_ptr->table_index.index); } } else if (auto unnest_ptr = dynamic_cast(node)) { - if (map.find(unnest_ptr->unnest_index) != map.end()) { - unnest_ptr->unnest_index = map.at(unnest_ptr->unnest_index); + if (map.find(unnest_ptr->unnest_index.index) != map.end()) { + unnest_ptr->unnest_index.index = map.at(unnest_ptr->unnest_index.index); } } // Special handling for LogicalAggregate: multiple indices if (auto agg_ptr = dynamic_cast(node)) { - if (map.find(agg_ptr->aggregate_index) != map.end()) { - agg_ptr->aggregate_index = map.at(agg_ptr->aggregate_index); + if (map.find(agg_ptr->aggregate_index.index) != map.end()) { + agg_ptr->aggregate_index.index = map.at(agg_ptr->aggregate_index.index); } - if (map.find(agg_ptr->group_index) != map.end()) { - agg_ptr->group_index = map.at(agg_ptr->group_index); + if (map.find(agg_ptr->group_index.index) != map.end()) { + agg_ptr->group_index.index = map.at(agg_ptr->group_index.index); } - if (agg_ptr->groupings_index != DConstants::INVALID_INDEX && map.find(agg_ptr->groupings_index) != map.end()) { - agg_ptr->groupings_index = map.at(agg_ptr->groupings_index); + if (agg_ptr->groupings_index.index != DConstants::INVALID_INDEX && + map.find(agg_ptr->groupings_index.index) != map.end()) { + agg_ptr->groupings_index.index = map.at(agg_ptr->groupings_index.index); } } @@ -263,7 +292,7 @@ void ReplaceNode(unique_ptr &root, unique_ptr // Build index remapping for any new-subtree indices that collide with external indices std::unordered_map index_map; if (!new_subtree_indices.empty()) { - idx_t next_idx = binder ? binder->GenerateTableIndex() : GetNextTableIndex(root); + idx_t next_idx = binder ? binder->GenerateTableIndex().index : GetNextTableIndex(root); for (auto idx : new_subtree_indices) { if (idx == DConstants::INVALID_INDEX) { continue; @@ -272,11 +301,11 @@ void ReplaceNode(unique_ptr &root, unique_ptr // find a fresh index not in external_indices and not in new_subtree_indices while (external_indices.find(next_idx) != external_indices.end() || new_subtree_indices.find(next_idx) != new_subtree_indices.end()) { - next_idx = binder ? binder->GenerateTableIndex() : next_idx + 1; + next_idx = binder ? binder->GenerateTableIndex().index : next_idx + 1; } index_map[idx] = next_idx; external_indices.insert(next_idx); - next_idx = binder ? binder->GenerateTableIndex() : next_idx + 1; + next_idx = binder ? binder->GenerateTableIndex().index : next_idx + 1; } } } @@ -293,9 +322,9 @@ void ReplaceNode(unique_ptr &root, unique_ptr ColumnBindingReplacer replacer; if (!index_map.empty()) { for (auto &old_b : old_top_bindings) { - auto it = index_map.find(old_b.table_index); + auto it = index_map.find(old_b.table_index.index); if (it != index_map.end()) { - ColumnBinding new_b(it->second, old_b.column_index); + ColumnBinding new_b {TableIndex(it->second), ProjectionIndex(old_b.column_index)}; replacer.replacement_bindings.emplace_back(old_b, new_b); } } @@ -320,14 +349,14 @@ vector FindPacKey(ClientContext &context, const string &table_name) { vector out; out.reserve(cols.size()); for (auto &col_name : cols) { - string tmp = col_name; + Identifier tmp(col_name); try { auto col_idx = table_entry.GetColumnIndex(tmp); auto &col = table_entry.GetColumn(col_idx); if (!col.Type().IsNumeric()) { return {}; } - out.push_back(col.GetName()); + out.push_back(col.GetName().GetIdentifierName()); } catch (...) { return {}; } @@ -351,16 +380,16 @@ vector FindPacKey(ClientContext &context, const string &table_name) { // Validate columns exist and are numeric by looking up the table in the catalog if (dot_pos != string::npos) { string schema = table_name.substr(0, dot_pos); - auto entry = - catalog.GetEntry(context, CatalogType::TABLE_ENTRY, schema, unqualified, OnEntryNotFound::RETURN_NULL); + auto entry = catalog.GetEntry(context, CatalogType::TABLE_ENTRY, Identifier(schema), Identifier(unqualified), + OnEntryNotFound::RETURN_NULL); if (entry) { return check_and_return_numeric_columns(entry->Cast(), pac_meta->primary_key_columns); } } else { CatalogSearchPath path(context); for (auto &entry_path : path.Get()) { - auto entry = catalog.GetEntry(context, CatalogType::TABLE_ENTRY, entry_path.schema, table_name, - OnEntryNotFound::RETURN_NULL); + auto entry = catalog.GetEntry(context, CatalogType::TABLE_ENTRY, entry_path.GetSchema(), + Identifier(table_name), OnEntryNotFound::RETURN_NULL); if (entry) { return check_and_return_numeric_columns(entry->Cast(), pac_meta->primary_key_columns); @@ -437,7 +466,8 @@ std::unordered_map> FindPacLinkPath(ClientContext &contex if (dot_pos != string::npos) { string schema = tbl_name.substr(0, dot_pos); string tbl = tbl_name.substr(dot_pos + 1); - auto entry = catalog.GetEntry(context, CatalogType::TABLE_ENTRY, schema, tbl, OnEntryNotFound::RETURN_NULL); + auto entry = catalog.GetEntry(context, CatalogType::TABLE_ENTRY, Identifier(schema), Identifier(tbl), + OnEntryNotFound::RETURN_NULL); if (entry) { return StringUtil::Lower(tbl); // return unqualified table name, normalized to lowercase } @@ -446,8 +476,8 @@ std::unordered_map> FindPacLinkPath(ClientContext &contex // Non-qualified: walk the search path; if found return unqualified table name CatalogSearchPath path(context); for (auto &entry_path : path.Get()) { - auto entry = catalog.GetEntry(context, CatalogType::TABLE_ENTRY, entry_path.schema, tbl_name, - OnEntryNotFound::RETURN_NULL); + auto entry = catalog.GetEntry(context, CatalogType::TABLE_ENTRY, entry_path.GetSchema(), + Identifier(tbl_name), OnEntryNotFound::RETURN_NULL); if (entry) { return StringUtil::Lower(tbl_name); // already unqualified, normalized to lowercase } @@ -538,7 +568,7 @@ ReplanGuard::~ReplanGuard() { // Configuration helpers string GetPacCompiledPath(ClientContext &context, const string &default_path) { Value v; - context.TryGetCurrentSetting("priv_compiled_path", v); + context.TryGetCurrentSetting(Identifier("priv_compiled_path"), v); string path = v.IsNull() ? default_path : v.ToString(); if (!path.empty() && path.back() != '/') { path.push_back('/'); @@ -548,7 +578,7 @@ string GetPacCompiledPath(ClientContext &context, const string &default_path) { int64_t GetPacM(ClientContext &context, int64_t default_m) { Value v; - context.TryGetCurrentSetting("pac_m", v); + context.TryGetCurrentSetting(Identifier("pac_m"), v); if (!v.IsNull()) { try { int64_t m = v.GetValue(); @@ -563,7 +593,7 @@ int64_t GetPacM(ClientContext &context, int64_t default_m) { bool IsPacNoiseEnabled(ClientContext &context, bool default_value) { Value v; - context.TryGetCurrentSetting("privacy_noise", v); + context.TryGetCurrentSetting(Identifier("privacy_noise"), v); if (v.IsNull()) { return default_value; } @@ -577,7 +607,7 @@ bool IsPacNoiseEnabled(ClientContext &context, bool default_value) { static string GetNormalizedStringSetting(ClientContext &context, const string &setting_name, const string &default_value) { Value v; - if (!context.TryGetCurrentSetting(setting_name, v) || v.IsNull()) { + if (!context.TryGetCurrentSetting(Identifier(setting_name), v) || v.IsNull()) { return default_value; } string s = v.ToString(); @@ -596,7 +626,7 @@ string GetPrivacyMode(ClientContext &context) { static bool TryGetDoubleSetting(ClientContext &context, const char *setting_name, double &out) { Value v; - if (!context.TryGetCurrentSetting(setting_name, v) || v.IsNull()) { + if (!context.TryGetCurrentSetting(Identifier(setting_name), v) || v.IsNull()) { return false; } try { @@ -652,7 +682,7 @@ bool TryGetDpDelta(ClientContext &context, double &out) { // Add implementation for GetPacCompileMethod string GetPacCompileMethod(ClientContext &context, const string &default_method) { Value v; - context.TryGetCurrentSetting("pac_compile_method", v); + context.TryGetCurrentSetting(Identifier("pac_compile_method"), v); if (v.IsNull()) { return default_method; } @@ -672,7 +702,7 @@ string GetPacCompileMethod(ClientContext &context, const string &default_method) // Helper to safely retrieve boolean settings with defaults bool GetBooleanSetting(ClientContext &context, const string &setting_name, bool default_value) { Value val; - if (context.TryGetCurrentSetting(setting_name, val) && !val.IsNull()) { + if (context.TryGetCurrentSetting(Identifier(setting_name), val) && !val.IsNull()) { return val.GetValue(); } return default_value; @@ -683,7 +713,7 @@ static void CollectOperatorsByTableIndex(LogicalOperator &op, std::unordered_map &table_index_to_op) { auto bindings = op.GetColumnBindings(); if (!bindings.empty()) { - idx_t op_table_index = bindings[0].table_index; + idx_t op_table_index = bindings[0].table_index.index; table_index_to_op[op_table_index] = &op; } for (auto &child : op.children) { @@ -693,6 +723,26 @@ static void CollectOperatorsByTableIndex(LogicalOperator &op, } } +// Casts and the optimizer's reversible storage encodings retain column lineage. +static optional_ptr FindLineageColumn(const Expression &expression) { + auto *expr = &expression; + while (expr->GetExpressionClass() == ExpressionClass::BOUND_FUNCTION) { + auto &function = expr->Cast(); + auto name = function.Function().GetName().GetIdentifierName(); + if (expr->GetExpressionType() != ExpressionType::OPERATOR_CAST && + !StringUtil::StartsWith(name, "__internal_compress_") && + !StringUtil::StartsWith(name, "__internal_decompress_")) { + break; + } + if (function.GetChildren().empty()) { + return nullptr; + } + expr = function.GetChildren()[0].get(); + } + return expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF ? &expr->Cast() + : nullptr; +} + // Helper to trace a binding back through the plan to find which LogicalGet it originates from LogicalGet *TraceBindingToSource(LogicalOperator &plan, const ColumnBinding &binding, idx_t *out_column_index) { // Build a map from table_index to the operator that produces it @@ -700,7 +750,7 @@ LogicalGet *TraceBindingToSource(LogicalOperator &plan, const ColumnBinding &bin CollectOperatorsByTableIndex(plan, table_index_to_op); // Start from the binding's table_index and trace backwards - idx_t current_table_index = binding.table_index; + idx_t current_table_index = binding.table_index.index; idx_t current_column_index = binding.column_index; // Limit iterations to prevent infinite loops @@ -737,10 +787,10 @@ LogicalGet *TraceBindingToSource(LogicalOperator &plan, const ColumnBinding &bin auto &expr = proj.expressions[current_column_index]; // If it's a BoundColumnRefExpression, follow it - if (expr->type == ExpressionType::BOUND_COLUMN_REF) { - auto &col_ref = expr->Cast(); - current_table_index = col_ref.binding.table_index; - current_column_index = col_ref.binding.column_index; + auto col_ref = FindLineageColumn(*expr); + if (col_ref) { + current_table_index = col_ref->Binding().table_index.index; + current_column_index = col_ref->Binding().column_index; continue; } @@ -757,10 +807,10 @@ LogicalGet *TraceBindingToSource(LogicalOperator &plan, const ColumnBinding &bin // This is a group column, trace through the group expression auto &group_expr = agg.groups[current_column_index]; - if (group_expr->type == ExpressionType::BOUND_COLUMN_REF) { - auto &col_ref = group_expr->Cast(); - current_table_index = col_ref.binding.table_index; - current_column_index = col_ref.binding.column_index; + auto col_ref = FindLineageColumn(*group_expr); + if (col_ref) { + current_table_index = col_ref->Binding().table_index.index; + current_column_index = col_ref->Binding().column_index; continue; } return nullptr; @@ -785,7 +835,7 @@ LogicalGet *TraceBindingToSource(LogicalOperator &plan, const ColumnBinding &bin if (current_column_index < left_count && op->children[0]) { auto left_bindings = op->children[0]->GetColumnBindings(); if (current_column_index < left_bindings.size()) { - current_table_index = left_bindings[current_column_index].table_index; + current_table_index = left_bindings[current_column_index].table_index.index; current_column_index = left_bindings[current_column_index].column_index; continue; } @@ -796,7 +846,7 @@ LogicalGet *TraceBindingToSource(LogicalOperator &plan, const ColumnBinding &bin idx_t right_col_idx = current_column_index - left_count; auto right_bindings = op->children[1]->GetColumnBindings(); if (right_col_idx < right_bindings.size()) { - current_table_index = right_bindings[right_col_idx].table_index; + current_table_index = right_bindings[right_col_idx].table_index.index; current_column_index = right_bindings[right_col_idx].column_index; continue; } @@ -809,7 +859,7 @@ LogicalGet *TraceBindingToSource(LogicalOperator &plan, const ColumnBinding &bin if (!op->children.empty() && op->children[0]) { auto child_bindings = op->children[0]->GetColumnBindings(); if (current_column_index < child_bindings.size()) { - current_table_index = child_bindings[current_column_index].table_index; + current_table_index = child_bindings[current_column_index].table_index.index; current_column_index = child_bindings[current_column_index].column_index; } else { return nullptr; @@ -832,16 +882,29 @@ static void CollectJoinKeyEquivalences(LogicalOperator &op, for (auto &cond : join.conditions) { // Only consider equality conditions - if (cond.comparison != ExpressionType::COMPARE_EQUAL) { + if (cond.GetComparisonType() != ExpressionType::COMPARE_EQUAL) { continue; } // Check if both sides are column references - if (cond.left && cond.right && cond.left->type == ExpressionType::BOUND_COLUMN_REF && - cond.right->type == ExpressionType::BOUND_COLUMN_REF) { - auto &left_ref = cond.left->Cast(); - auto &right_ref = cond.right->Cast(); - equivalences.emplace_back(left_ref.binding, right_ref.binding); + if (cond.LeftReference() && cond.RightReference() && + cond.LeftReference()->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF && + cond.RightReference()->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { + auto &left_ref = cond.LeftReference()->Cast(); + auto &right_ref = cond.RightReference()->Cast(); + equivalences.emplace_back(left_ref.BindingMutable(), right_ref.BindingMutable()); + } + } + } + + // Optimizer-inserted projections can rename a join key or encode it for + // grouping. Include these aliases in the same equivalence graph. + if (op.type == LogicalOperatorType::LOGICAL_PROJECTION) { + auto &projection = op.Cast(); + for (idx_t i = 0; i < projection.expressions.size(); i++) { + auto ref = FindLineageColumn(*projection.expressions[i]); + if (ref) { + equivalences.emplace_back(ColumnBinding(projection.table_index, ProjectionIndex(i)), ref->Binding()); } } } @@ -859,7 +922,7 @@ static void FindEquivalentBindings(const ColumnBinding &binding, const vector> &equivalences, std::unordered_set &visited, vector &result) { // Create a unique key for the binding - uint64_t key = (static_cast(binding.table_index) << 32) | binding.column_index; + uint64_t key = (static_cast(binding.table_index.index) << 32) | binding.column_index; if (visited.find(key) != visited.end()) { return; } @@ -868,10 +931,11 @@ static void FindEquivalentBindings(const ColumnBinding &binding, // Find all equivalences involving this binding for (auto &eq : equivalences) { - if (eq.first.table_index == binding.table_index && eq.first.column_index == binding.column_index) { + if (eq.first.table_index.index == binding.table_index.index && eq.first.column_index == binding.column_index) { FindEquivalentBindings(eq.second, equivalences, visited, result); } - if (eq.second.table_index == binding.table_index && eq.second.column_index == binding.column_index) { + if (eq.second.table_index.index == binding.table_index.index && + eq.second.column_index == binding.column_index) { FindEquivalentBindings(eq.first, equivalences, visited, result); } } @@ -899,7 +963,7 @@ bool ColumnBelongsToTable(LogicalOperator &plan, const string &table_name, const continue; } - if (table_entry->name == table_name) { + if (table_entry->name.GetIdentifierName() == table_name) { return true; } } @@ -909,11 +973,11 @@ bool ColumnBelongsToTable(LogicalOperator &plan, const string &table_name, const // Recursively remap BoundColumnRefExpression bindings in an expression tree. static void RemapBindingsInExpr(Expression &e, const std::unordered_map &map) { - if (e.type == ExpressionType::BOUND_COLUMN_REF) { + if (e.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) { auto &col_ref = e.Cast(); - auto it = map.find(col_ref.binding.table_index); + auto it = map.find(col_ref.BindingMutable().table_index.index); if (it != map.end()) { - col_ref.binding.table_index = it->second; + col_ref.BindingMutable().table_index.index = it->second; } } ExpressionIterator::EnumerateChildren(e, [&map](Expression &child) { RemapBindingsInExpr(child, map); }); @@ -947,8 +1011,8 @@ void RemapBindingsInSubtree(LogicalOperator &op, const std::unordered_map(&op)) { for (auto &cond : cj->conditions) { - remap_expr(cond.left); - remap_expr(cond.right); + remap_expr(cond.LeftReference()); + remap_expr(cond.RightReference()); } for (auto &expr : cj->duplicate_eliminated_columns) { remap_expr(expr); @@ -996,9 +1060,9 @@ std::unordered_map RemapSubtreeIndices(LogicalOperator *subtree, B if (old_idx == DConstants::INVALID_INDEX) { continue; } - idx_t fresh = binder.GenerateTableIndex(); + idx_t fresh = binder.GenerateTableIndex().index; while (used.find(fresh) != used.end()) { - fresh = binder.GenerateTableIndex(); + fresh = binder.GenerateTableIndex().index; } map[old_idx] = fresh; used.insert(fresh); diff --git a/test/sql/dp_filterless.test b/test/sql/dp_filterless.test index a8e9ce3c..25d53b7b 100644 --- a/test/sql/dp_filterless.test +++ b/test/sql/dp_filterless.test @@ -367,7 +367,7 @@ EXPLAIN SELECT grp, SUM(amount) FROM filterless_group_cap_three GROUP BY grp; ---- -physical_plan :[\s\S]*UNNEST[\s\S]*arg_min_nulls_last[\s\S]* +physical_plan :[\s\S]*Unnest[\s\S]*arg_min_nulls_last[\s\S]* statement ok SET dp_max_groups_contributed = 1; @@ -541,7 +541,7 @@ EXPLAIN SELECT AVG(amount) FROM filterless_sampled_events WHERE qualifies; ---- -physical_plan :[\s\S]*UNGROUPED_AGGREGATE[\s\S]*priv_filterless_avg[\s\S]*priv_filterless_sum_pair[\s\S]*priv_filterless_count_pair[\s\S]* +physical_plan :[\s\S]*Ungrouped Aggregate[\s\S]*priv_filterless_avg[\s\S]*priv_filterless_sum_pair[\s\S]*priv_filterless_count_pair[\s\S]* # COUNT(expr) is zero when every qualifying value is NULL, so fused AVG keeps # SQL's NULL result instead of dividing by a synthetic denominator. diff --git a/test/sql/dp_sass_m.test b/test/sql/dp_sass_m.test index 78760a86..32e6cad1 100644 --- a/test/sql/dp_sass_m.test +++ b/test/sql/dp_sass_m.test @@ -874,7 +874,7 @@ EXPLAIN SELECT category AS g1, length(category) AS g2, COUNT(DISTINCT v) FROM sass_m_events GROUP BY g1, g2; ---- -logical_opt :[\s\S]*RANK_DENSE[\s\S]*ORDER BY[\s\S]*#0 ASC[\s\S]*#1 ASC[\s\S]* +logical_opt :[\s\S]*dense_rank[\s\S]*ORDER BY[\s\S]*#0 ASC[\s\S]*#1 ASC[\s\S]* # Noisy aggregate cells must use independent random streams even when their lane # lists and public domains are identical. Cover both the optimized m=64 path and diff --git a/test/sql/dp_standard.test b/test/sql/dp_standard.test index 22f9c780..16a1eaeb 100644 --- a/test/sql/dp_standard.test +++ b/test/sql/dp_standard.test @@ -502,7 +502,7 @@ WHERE EXISTS ( AND o2.store <> o1.store ); ---- -dp_standard: lateral or correlated subqueries are not supported by DP mechanisms +dp_standard: scalar or mark subqueries are not supported by DP mechanisms # --- MIN/MAX (dp_standard) --- # Values are clipped to the public domain [dp_minmax_lower_bound, dp_minmax_upper_bound], so global diff --git a/test/sql/pac_bitslice_compiler_uncorrelated_subqueries.test b/test/sql/pac_bitslice_compiler_uncorrelated_subqueries.test index edfb34d7..b61f4b6f 100644 --- a/test/sql/pac_bitslice_compiler_uncorrelated_subqueries.test +++ b/test/sql/pac_bitslice_compiler_uncorrelated_subqueries.test @@ -358,7 +358,9 @@ ORDER BY # Test 4: TPC-H Q9-style query with complex arithmetic in subquery # Tests: Subquery with multiple joins, complex arithmetic expression (subtraction of products), GROUP BY multiple columns -# This query computes profit by nation and year +# This query computes profit by nation and year. DuckDB 2.0 widens the arithmetic +# to DECIMAL(21,4), selecting the INT128/double SUM state rather than the +# approximate INT64 state used for DuckDB 1.5's DECIMAL(18,4). query III SELECT nation, @@ -389,10 +391,10 @@ ORDER BY nation, o_year DESC; ---- -FRANCE 1996 -619.3152 +FRANCE 1996 -620.0000 FRANCE 1995 0.0000 GERMANY 1996 0.0000 -GERMANY 1995 NULL +GERMANY 1995 -420.0000 # Clean up statement ok diff --git a/test/sql/pac_categorical_subqueries.test b/test/sql/pac_categorical_subqueries.test index 45e15655..5e2e7e82 100644 --- a/test/sql/pac_categorical_subqueries.test +++ b/test/sql/pac_categorical_subqueries.test @@ -524,12 +524,49 @@ SELECT count(*) FROM ( ) ); ---- -2 +1 + +# Independently evaluate all 64 sampled worlds with ordinary SQL. Only supplier +# 1 passes in at least half of the worlds. DuckDB 2.0's nullable-COUNT wrapper +# contains two COUNT states; both must remain world-specific through its CASE. +statement ok +SET priv_check = false; + +statement ok +SET priv_rewrite = false; + +query I +WITH sampled AS ( + SELECT l_suppkey, j AS world + FROM lineitem JOIN orders ON l_orderkey = o_orderkey + CROSS JOIN range(64) t(j) + WHERE ((priv_hash(hash(o_custkey)) >> j) & 1::UBIGINT) <> 0 +), counts AS ( + SELECT world, l_suppkey, COUNT(*) AS n FROM sampled GROUP BY ALL +), means AS ( + SELECT world, AVG(n) AS mean_count FROM counts GROUP BY world +), decisions AS ( + SELECT s_suppkey, SUM((COALESCE(n, 0) > mean_count)::INTEGER) AS true_worlds + FROM supplier CROSS JOIN range(64) t(j) + LEFT JOIN counts ON j = counts.world AND s_suppkey = l_suppkey + LEFT JOIN means ON j = means.world + GROUP BY s_suppkey +) +SELECT COUNT(*) FROM decisions WHERE true_worlds >= 32; +---- +1 + +statement ok +SET priv_rewrite = true; + +statement ok +SET priv_check = true; # ============================================================================ # Test 11: Categorical with AVG aggregate # Find parts where avg quantity exceeds threshold -# NOTE: Result is 2 with proper avg decomposition (as_noised_avg → sum/count) +# All four parts pass in a majority of worlds; both correlated and global +# averages must remain world-specific through the comparison. # ============================================================================ query I @@ -551,7 +588,42 @@ SELECT count(*) FROM ( ) ); ---- -2 +4 + +# Ordinary SQL oracle: report the number of passing worlds for each part. +statement ok +SET priv_check = false; + +statement ok +SET priv_rewrite = false; + +query II +WITH sampled AS ( + SELECT l_partkey, l_quantity, j AS world + FROM lineitem JOIN orders ON l_orderkey = o_orderkey + CROSS JOIN range(64) t(j) + WHERE ((priv_hash(hash(o_custkey)) >> j) & 1::UBIGINT) <> 0 +), avgs AS ( + SELECT world, l_partkey, AVG(l_quantity) AS a FROM sampled GROUP BY ALL +), means AS ( + SELECT world, AVG(l_quantity) AS m FROM sampled GROUP BY world +) +SELECT p_partkey, SUM((a > 0.8 * m)::INTEGER) AS true_worlds +FROM part CROSS JOIN range(64) t(j) +LEFT JOIN avgs ON j = avgs.world AND p_partkey = l_partkey +LEFT JOIN means ON j = means.world +GROUP BY p_partkey ORDER BY p_partkey; +---- +1 41 +2 37 +3 50 +4 50 + +statement ok +SET priv_rewrite = true; + +statement ok +SET priv_check = true; # ============================================================================ # Test 12: Correlated categorical subquery with multiple tables @@ -798,3 +870,79 @@ DROP TABLE scalar_case_reason; statement ok DROP TABLE scalar_case_pu; + + +# Right-anti membership filters must preserve ownership joins, duplicates, +# NULL-safe matches, and the output layout when the build/probe sides change. +# as_count scales each world by two for 50% sampling; divide out that scale. +statement ok +PRAGMA clear_privacy_metadata; + +statement ok +SET priv_rewrite=false; + +statement ok +SET pac_sample_diversity_check=false; + +statement ok +CREATE TABLE anti_owners AS SELECT i::BIGINT k FROM range(1000) r(i); + +statement ok +CREATE TABLE anti_facts AS +SELECT (i % 1000)::BIGINT k, (i % 7)::INTEGER v FROM range(10000) r(i); + +statement ok +CREATE TABLE anti_outer(k BIGINT, v INTEGER); + +statement ok +INSERT INTO anti_outer VALUES (10, 3), (10, 3), (11, 4), (2000, 1), (NULL, 1), (12, 0); + +query I +SELECT list_extract(as_count(3::UBIGINT), 1) / 2 +FROM anti_outer a +WHERE a.v > 0 AND NOT EXISTS ( + SELECT 1 FROM anti_owners o JOIN anti_facts f ON o.k = f.k + WHERE f.k = a.k AND f.v <> a.v AND f.v > 1 +); +---- +2 + +# NULL=NULL must remain a match when neither join key is proven non-null. +statement ok +CREATE TABLE anti_nullable(k BIGINT, v INTEGER); + +statement ok +INSERT INTO anti_nullable VALUES (NULL, 2), (10, 3), (11, 4); + +query I +SELECT list_extract(as_count(3::UBIGINT), 1) / 2 +FROM anti_outer a +WHERE a.v > 0 AND NOT EXISTS ( + SELECT 1 FROM anti_nullable n + WHERE n.k IS NOT DISTINCT FROM a.k AND n.v <> a.v +); +---- +4 + +# A LIMIT is a barrier: filtering below it can change which evidence row exists. +query I +SELECT list_extract(as_count(3::UBIGINT), 1) / 2 +FROM anti_outer a +WHERE a.v > 0 AND NOT EXISTS ( + SELECT 1 FROM (SELECT * FROM anti_nullable ORDER BY k NULLS FIRST LIMIT 1) n + WHERE n.k IS NOT DISTINCT FROM a.k AND n.v <> a.v +); +---- +4 + +statement ok +DROP TABLE anti_nullable; + +statement ok +DROP TABLE anti_outer; + +statement ok +DROP TABLE anti_facts; + +statement ok +DROP TABLE anti_owners; diff --git a/test/sql/pac_noise_independence.test b/test/sql/pac_noise_independence.test new file mode 100644 index 00000000..eb939366 --- /dev/null +++ b/test/sql/pac_noise_independence.test @@ -0,0 +1,310 @@ +# name: test/sql/pac_noise_independence.test +# description: Independent noise draws across groups and finalization batches +# group: [sql] + +require privacy + +statement ok +PRAGMA clear_privacy_metadata; + +statement ok +SET privacy_seed = 42 + +statement ok +SET threads = 1 + +statement ok +SET pac_mi = 0.0078125 + +statement ok +SET pac_ptracking = false + +statement ok +SET pac_sample_diversity_check = false + +# 1000 groups of one entity: each is NULL with probability popcount(~key_hash)/64 = 1/2 +statement ok +CREATE TABLE solo AS SELECT i AS g, i AS e, 1 + i % 5 AS v FROM range(1000) t(i) + +# 1000 groups of 20 entities (no world empty: never NULL) +statement ok +CREATE TABLE cells AS SELECT i // 20 AS g, i AS e, 1 + i % 7 AS v FROM range(20000) t(i) + +# ============================================================================ +# NULL decisions are independent across groups (not all-or-nothing) +# ============================================================================ + +query I +SELECT count(*) FILTER (WHERE s IS NULL) BETWEEN 400 AND 600 +FROM (SELECT g, as_noised_sum(priv_hash(hash(e)), v) AS s FROM solo GROUP BY g) +---- +true + +query I +SELECT count(*) FILTER (WHERE c IS NULL) BETWEEN 400 AND 600 +FROM (SELECT g, as_noised_count(priv_hash(hash(e))) AS c FROM solo GROUP BY g) +---- +true + +query II +SELECT count(*) FILTER (WHERE mx IS NULL) BETWEEN 400 AND 600, count(*) FILTER (WHERE mn IS NULL) BETWEEN 400 AND 600 +FROM (SELECT g, as_noised_max(priv_hash(hash(e)), v) AS mx, as_noised_min(priv_hash(hash(e)), v) AS mn + FROM solo GROUP BY g) +---- +true true + +query II +SELECT count(*) FILTER (WHERE s IS NULL) BETWEEN 400 AND 600, count(*) FILTER (WHERE mx IS NULL) BETWEEN 400 AND 600 +FROM (SELECT g, as_noised_clip_sum(priv_hash(hash(e)), v) AS s, as_noised_clip_max(priv_hash(hash(e)), v) AS mx + FROM solo GROUP BY g) +---- +true true + +# ============================================================================ +# Noise is independent across groups: released values fall on either side of the exact ones about equally +# often (500 +- 16 of 1000; with one draw shared by all groups, nearly all fall on the same side) +# ============================================================================ + +query II +SELECT count(*) FILTER (WHERE s > x) BETWEEN 400 AND 600, count(*) FILTER (WHERE s < x) BETWEEN 400 AND 600 +FROM (SELECT g, as_noised_sum(priv_hash(hash(e)), v) AS s, sum(v) AS x FROM cells GROUP BY g) +---- +true true + +query II +SELECT count(*) FILTER (WHERE c > x) BETWEEN 400 AND 600, count(*) FILTER (WHERE c < x) BETWEEN 400 AND 600 +FROM (SELECT g, as_noised_count(priv_hash(hash(e))) AS c, count(*) AS x FROM cells GROUP BY g) +---- +true true + +query II +SELECT count(*) FILTER (WHERE s > x) BETWEEN 400 AND 600, count(*) FILTER (WHERE s < x) BETWEEN 400 AND 600 +FROM (SELECT g, as_noised_clip_sum(priv_hash(hash(e)), v) AS s, sum(v) AS x FROM cells GROUP BY g) +---- +true true + +query II +SELECT count(*) FILTER (WHERE s > x) BETWEEN 400 AND 600, count(*) FILTER (WHERE s < x) BETWEEN 400 AND 600 +FROM (SELECT g, as_noised(as_sum(priv_hash(hash(e)), v)) AS s, sum(v) AS x FROM cells GROUP BY g) +---- +true true + +# ============================================================================ +# Equal counter values in different groups still receive independent draws. +# Reinitializing the generator from counters would make all 1000 NULL decisions identical. +# ============================================================================ + +query I +SELECT count(*) FILTER (WHERE c IS NULL) BETWEEN 400 AND 600 +FROM (SELECT i, as_noised_count(4294967295::UBIGINT) AS c FROM range(1000) t(i) GROUP BY i) +---- +true + +# Constant counter arrays across several vectors must not be constant-folded, +# or restart the same RNG stream in each vector. Check all scalar release forms. +query III +SELECT count(DISTINCT l) > 3000, count(DISTINCT f) > 3000, count(DISTINCT d) > 3000 +FROM (SELECT as_noised(range(64)::FLOAT[]) AS l, + priv_finalize(range(64)::FLOAT[]) AS f, + as_noised_div(range(64)::FLOAT[], list_transform(range(64), lambda x: 1::FLOAT)) AS d + FROM range(4096)) t +---- +true true true + +# A fixed seed repeats for the same execution setup, even after unrelated draws. +statement ok +CREATE TABLE seeded_first AS +SELECT g, as_noised_sum(priv_hash(hash(e)), v) AS s, + as_noised_count(priv_hash(hash(e))) AS c, + as_noised_min(priv_hash(hash(e)), v) AS mn, + as_noised_max(priv_hash(hash(e)), v) AS mx, + as_noised_clip_sum(priv_hash(hash(e)), v) AS cs, + as_noised_clip_max(priv_hash(hash(e)), v) AS cm +FROM solo GROUP BY g + +statement ok +CREATE TABLE unrelated_draws AS SELECT as_noised(range(64)::FLOAT[]) AS v FROM range(8192) + +statement ok +CREATE TABLE seeded_second AS +SELECT g, as_noised_sum(priv_hash(hash(e)), v) AS s, + as_noised_count(priv_hash(hash(e))) AS c, + as_noised_min(priv_hash(hash(e)), v) AS mn, + as_noised_max(priv_hash(hash(e)), v) AS mx, + as_noised_clip_sum(priv_hash(hash(e)), v) AS cs, + as_noised_clip_max(priv_hash(hash(e)), v) AS cm +FROM solo GROUP BY g + +query I +SELECT count(*) FROM seeded_first a JOIN seeded_second b USING (g) +WHERE a.s IS DISTINCT FROM b.s OR a.c IS DISTINCT FROM b.c + OR a.mn IS DISTINCT FROM b.mn OR a.mx IS DISTINCT FROM b.mx + OR a.cs IS DISTINCT FROM b.cs OR a.cm IS DISTINCT FROM b.cm +---- +0 + +# Prepared scalar releases must reset streams on every execution and advance +# between vectors; bind data can survive many executions of the statement. +statement ok +CREATE TABLE seeded_prepared(i BIGINT, l FLOAT, f FLOAT, d DOUBLE) + +statement ok +PREPARE seeded_release AS INSERT INTO seeded_prepared +SELECT i, as_noised(range(64)::FLOAT[]), priv_finalize(range(64)::FLOAT[]), + as_noised_div(range(64)::FLOAT[], list_transform(range(64), lambda x: 1::FLOAT)) +FROM range(4096) t(i) + +statement ok +EXECUTE seeded_release + +statement ok +EXECUTE seeded_release + +query I +SELECT count(*) FROM seeded_prepared a JOIN seeded_prepared b USING (i) +WHERE a.rowid < 4096 AND b.rowid >= 4096 +AND (a.l IS DISTINCT FROM b.l OR a.f IS DISTINCT FROM b.f OR a.d IS DISTINCT FROM b.d) +---- +0 + +query III +SELECT count(DISTINCT l) > 3000, count(DISTINCT f) > 3000, count(DISTINCT d) > 3000 +FROM seeded_prepared WHERE rowid < 4096 +---- +true true true + +statement ok +DEALLOCATE seeded_release + +statement ok +CREATE TABLE seeded_counts(g BIGINT, c BIGINT) + +statement ok +PREPARE seeded_count AS INSERT INTO seeded_counts +SELECT g, as_noised_count(priv_hash(hash(e))) FROM solo GROUP BY g + +statement ok +EXECUTE seeded_count + +statement ok +EXECUTE seeded_count + +query I +SELECT count(*) FROM seeded_counts a JOIN seeded_counts b USING (g) +WHERE a.rowid < 1000 AND b.rowid >= 1000 AND a.c IS DISTINCT FROM b.c +---- +0 + +statement ok +DEALLOCATE seeded_count + +# Correlated scalar releases must draw independently for repeated outer keys, +# even when their deterministic aggregate counters are shared. +statement ok +SET priv_rewrite=false + +statement ok +CREATE TABLE correlated_noise_input AS +SELECT 0::INTEGER AS k, i AS e, 1 + i % 7 AS v FROM range(128) t(i) + +statement ok +CREATE TABLE correlated_noise_outer AS +SELECT i AS id, 0::INTEGER AS k, CASE WHEN i % 2 = 0 THEN 'KEEP' ELSE 'DROP' END AS label +FROM range(4096) t(i) + +statement ok +CREATE TABLE correlated_noise_result AS +SELECT o.id, + (SELECT as_noised(as_sum(priv_hash(hash(i.e)), i.v)) + FROM correlated_noise_input i WHERE i.k = o.k) AS s, + (SELECT as_noised_div(as_sum(priv_hash(hash(i.e)), i.v), + as_count(priv_hash(hash(i.e)), i.v)) + FROM correlated_noise_input i WHERE i.k = o.k) AS d +FROM correlated_noise_outer o + +query II +SELECT count(DISTINCT s) > 3000, count(DISTINCT d) > 3000 +FROM correlated_noise_result +---- +true true + +# With noise suppressed, the rewrite must exactly preserve empty inputs, +# NULL correlation keys, duplicate keys, and multiple scalar subqueries. +statement ok +INSERT INTO correlated_noise_outer VALUES (4096, 1, 'KEEP'), (4097, NULL, 'DROP') + +statement ok +SET pac_mi=0 + +statement ok +SET disabled_optimizers='extension' + +statement ok +CREATE TABLE correlated_noise_reference AS +SELECT o.id, + (SELECT as_noised(as_sum(priv_hash(hash(i.e)), i.v)) + FROM correlated_noise_input i WHERE i.k = o.k) AS s, + (SELECT as_noised_div(as_sum(priv_hash(hash(i.e)), i.v), + as_count(priv_hash(hash(i.e)), i.v)) + FROM correlated_noise_input i WHERE i.k = o.k) AS d +FROM correlated_noise_outer o + +statement ok +SET disabled_optimizers='' + +statement ok +CREATE TABLE correlated_noise_actual AS +SELECT o.id, + (SELECT as_noised(as_sum(priv_hash(hash(i.e)), i.v)) + FROM correlated_noise_input i WHERE i.k = o.k) AS s, + (SELECT as_noised_div(as_sum(priv_hash(hash(i.e)), i.v), + as_count(priv_hash(hash(i.e)), i.v)) + FROM correlated_noise_input i WHERE i.k = o.k) AS d +FROM correlated_noise_outer o + +query I +SELECT count(*) FROM ( + (SELECT * FROM correlated_noise_actual EXCEPT ALL SELECT * FROM correlated_noise_reference) + UNION ALL + (SELECT * FROM correlated_noise_reference EXCEPT ALL SELECT * FROM correlated_noise_actual) +) differences +---- +0 + + +# Pure outer predicates can move before the shared counters, while the noisy +# comparison and its potentially throwing arithmetic retain their output layout. +statement ok +SET disabled_optimizers='extension' + +statement ok +CREATE TABLE correlated_filtered_reference AS +SELECT o.id FROM correlated_noise_outer o +WHERE o.label = 'KEEP' AND o.id * 5 < ( + SELECT as_noised_div(as_sum(priv_hash(hash(i.e)), i.v), + as_count(priv_hash(hash(i.e)), i.v)) + FROM correlated_noise_input i WHERE i.k = o.k +) + +statement ok +SET disabled_optimizers='' + +query I +SELECT count(*) FROM ( + (SELECT o.id FROM correlated_noise_outer o + WHERE o.label = 'KEEP' AND o.id * 5 < ( + SELECT as_noised_div(as_sum(priv_hash(hash(i.e)), i.v), + as_count(priv_hash(hash(i.e)), i.v)) + FROM correlated_noise_input i WHERE i.k = o.k + ) EXCEPT ALL SELECT * FROM correlated_filtered_reference) + UNION ALL + (SELECT * FROM correlated_filtered_reference EXCEPT ALL + SELECT o.id FROM correlated_noise_outer o + WHERE o.label = 'KEEP' AND o.id * 5 < ( + SELECT as_noised_div(as_sum(priv_hash(hash(i.e)), i.v), + as_count(priv_hash(hash(i.e)), i.v)) + FROM correlated_noise_input i WHERE i.k = o.k + )) +) differences +---- +0 diff --git a/test/sql/pac_plan.test b/test/sql/pac_plan.test index 3f5f4c41..ebbbfcc5 100644 --- a/test/sql/pac_plan.test +++ b/test/sql/pac_plan.test @@ -798,7 +798,7 @@ WHERE GROUP BY c_nationkey; ---- -logical_opt :.*COMPARISON_JOIN.*bit_or.*bit_or.* +logical_opt :.*Comparison Join.*bit_or.*bit_or.* # ============================================================ diff --git a/test/sql/pac_sum.test b/test/sql/pac_sum.test index 8b857ae0..fdcf0bb0 100644 --- a/test/sql/pac_sum.test +++ b/test/sql/pac_sum.test @@ -374,3 +374,322 @@ a AS (SELECT as_sum(mask, v) AS lanes FROM t) SELECT count(*) FILTER (WHERE lanes[j+1] / 2 = 196605) FROM a CROSS JOIN range(64) worlds(j) ---- 64 + + +# DuckDB 2.0 state export preserves buffers, positive/negative cascades, and +# pointer-backed counters. Values include a level above the inline-pointer area. +statement ok +SET priv_rewrite=false; + +statement ok +SET threads=1; + +statement ok +SET pac_mi=0; + +statement ok +SET pac_sample_diversity_check=false; + +statement ok +CREATE TABLE spill_sum_input AS +SELECT i%5 g,priv_hash(hash(i)) h, + CASE WHEN i%3=0 THEN -(1::BIGINT<<48) ELSE 1024::BIGINT END v +FROM range(10000)t(i); + +statement ok +CREATE TABLE spill_sum_states AS +SELECT g,as_sum(h,v) EXPORT_STATE s, + as_sum(h,v::DOUBLE) EXPORT_STATE d, + as_noised_sum(h,v::DECIMAL(18,2)) EXPORT_STATE dec, + as_sum(h,abs(v)::UBIGINT,0.75) EXPORT_STATE u +FROM spill_sum_input GROUP BY g; + +query I +SELECT count(*) FROM spill_sum_states a JOIN ( + SELECT g,as_sum(h,v) s,as_sum(h,v::DOUBLE) d, + as_noised_sum(h,v::DECIMAL(18,2)) dec, + as_sum(h,abs(v)::UBIGINT,0.75) u + FROM spill_sum_input GROUP BY g +)b USING(g) +WHERE finalize(a.s) IS DISTINCT FROM b.s OR finalize(a.d) IS DISTINCT FROM b.d + OR finalize(a.dec) IS DISTINCT FROM b.dec OR finalize(a.u) IS DISTINCT FROM b.u; +---- +0 + +# One or two values stay buffered; NULL-only/empty groups have no inner state. +query I +SELECT finalize(s) IS NOT DISTINCT FROM expected FROM ( + SELECT as_sum(h,NULL::BIGINT) EXPORT_STATE s,as_sum(h,NULL::BIGINT) expected + FROM spill_sum_input +); +---- +true + +query I +SELECT finalize(s) IS NOT DISTINCT FROM expected FROM ( + SELECT as_sum(h,v) EXPORT_STATE s,as_sum(h,v) expected + FROM spill_sum_input WHERE false +); +---- +true + +query I +SELECT finalize(s) IS NOT DISTINCT FROM expected FROM ( + SELECT as_sum(h,v) EXPORT_STATE s,as_sum(h,v) expected + FROM (SELECT h,v FROM spill_sum_input LIMIT 1) +); +---- +true + +query I +SELECT finalize(s) IS NOT DISTINCT FROM expected FROM ( + SELECT as_sum(h,v) EXPORT_STATE s,as_sum(h,v) expected + FROM (SELECT h,v FROM spill_sum_input LIMIT 2) +); +---- +true + +# Compare the complete result under arena pressure with an in-memory execution. +# Mixed integer, floating-point, and DECIMAL states share the same hash table. +statement ok +SET memory_limit='256MB'; + +statement ok +CREATE TABLE spill_sum_expected AS +SELECT count(*) n,sum(hash(a,b,c)) checksum FROM ( + SELECT i//8 g, + as_sum(18446744073709551615::UBIGINT,1::BIGINT) a, + as_sum(18446744073709551615::UBIGINT,2::DOUBLE) b, + as_noised_sum(18446744073709551615::UBIGINT,3::DECIMAL(18,2)) c + FROM range(800000)t(i) GROUP BY g +); + +statement ok +SET debug_force_external=true; + +statement ok +SET memory_limit='32MB'; + +query II +SELECT actual.n,actual.checksum=expected.checksum FROM ( + SELECT count(*) n,sum(hash(a,b,c)) checksum FROM ( + SELECT i//8 g, + as_sum(18446744073709551615::UBIGINT,1::BIGINT) a, + as_sum(18446744073709551615::UBIGINT,2::DOUBLE) b, + as_noised_sum(18446744073709551615::UBIGINT,3::DECIMAL(18,2)) c + FROM range(800000)t(i) GROUP BY g + ) +)actual,spill_sum_expected expected; +---- +100000 true + +statement ok +SET debug_force_external=false; + +statement ok +RESET memory_limit; + + +# Exported states remain combinable, including the high and negative levels. +query I +SELECT finalize(combine(a.s,b.s)) IS NOT DISTINCT FROM c.expected FROM +(SELECT as_sum(h,v) EXPORT_STATE s FROM spill_sum_input WHERE g<3)a, +(SELECT as_sum(h,v) EXPORT_STATE s FROM spill_sum_input WHERE g>=3)b, +(SELECT as_sum(h,v) expected FROM spill_sum_input)c; +---- +true + +statement error +SELECT finalize(to_aggregate_state(''::BLOB,'as_sum',['UBIGINT','BIGINT'])); +---- +Truncated stochastic aggregate state + + +# DECIMAL binding replaces its implementation; retain the logical third +# argument (sampling correction) when rebinding exported states. +query I +SELECT finalize(s) IS NOT DISTINCT FROM expected FROM ( + SELECT as_noised_sum(h,v::DECIMAL(18,2),0.75) EXPORT_STATE s, + as_noised_sum(h,v::DECIMAL(18,2),0.75) expected + FROM spill_sum_input +); +---- +true + + +# Q78-shaped failure: sparse local groups stay buffered in the sink. Combining +# them must not manufacture a permanent full state for each final group. +statement ok +SET debug_force_external=true; + +statement ok +SET memory_limit='32MB'; + +query II +SELECT count(*),count(*) FILTER (WHERE a=4 AND b=8 AND c=12) FROM ( + SELECT i%100000 g, + as_noised_sum(18446744073709551615::UBIGINT,1::BIGINT) a, + as_noised_sum(18446744073709551615::UBIGINT,2::DOUBLE) b, + as_noised_sum(18446744073709551615::UBIGINT,3::DECIMAL(18,2)) c + FROM range(200000)t(i) GROUP BY g +); +---- +100000 100000 + +statement ok +SET debug_force_external=false; + +statement ok +RESET memory_limit; + +# The two-value buffer boundary: combine two buffered states, or a buffered +# destination and a materialized source. Every world must retain every value. +query I +SELECT count(*) FROM ( + SELECT finalize(combine(a.s,b.s)) lanes FROM + (SELECT as_sum(18446744073709551615::UBIGINT,v) EXPORT_STATE s FROM (VALUES (13::BIGINT),(17))t(v))a, + (SELECT as_sum(18446744073709551615::UBIGINT,29::BIGINT) EXPORT_STATE s)b +),unnest(lanes)t(v) WHERE v=118; +---- +64 + +query I +SELECT count(*) FROM ( + SELECT finalize(combine(a.s,b.s)) lanes FROM + (SELECT as_sum(18446744073709551615::UBIGINT,v) EXPORT_STATE s FROM (VALUES (1::BIGINT),(2))t(v))a, + (SELECT as_sum(18446744073709551615::UBIGINT,v) EXPORT_STATE s FROM (VALUES (3::BIGINT),(4),(5))t(v))b +),unnest(lanes)t(v) WHERE v=30; +---- +64 + +# All stochastic aggregates in a hash table must provide a state layout. +# Exercise COUNT's SWAR overflow, MIN/MAX, AVG decomposition and list-input states. +statement ok +SET priv_rewrite=false + +statement ok +SET pac_mi=0 + +statement ok +SET pac_sample_diversity_check=false + +statement ok +SET threads=1 + +statement ok +SET memory_limit='512MB' + +statement ok +CREATE TABLE mixed_spill_expected AS +SELECT count(*) n,sum(hash(s,c,a,mn,mx)) checksum FROM ( + SELECT i//8 g,as_sum(18446744073709551615::UBIGINT,1::BIGINT) s, + as_count(18446744073709551615::UBIGINT) c, + as_noised_avg(18446744073709551615::UBIGINT,2::DOUBLE) a, + as_min(18446744073709551615::UBIGINT,i%8) mn, + as_max(18446744073709551615::UBIGINT,i%8) mx + FROM range(800000)t(i) GROUP BY g +) + +statement ok +SET memory_limit='32MB' + +query I +SELECT r.n=e.n AND r.checksum=e.checksum FROM mixed_spill_expected e,( + SELECT count(*) n,sum(hash(s,c,a,mn,mx)) checksum FROM ( + SELECT i//8 g,as_sum(18446744073709551615::UBIGINT,1::BIGINT) s, + as_count(18446744073709551615::UBIGINT) c, + as_noised_avg(18446744073709551615::UBIGINT,2::DOUBLE) a, + as_min(18446744073709551615::UBIGINT,i%8) mn, + as_max(18446744073709551615::UBIGINT,i%8) mx + FROM range(800000)t(i) GROUP BY g + ))r +---- +true + +statement ok +SET memory_limit='512MB' + +statement ok +DROP TABLE mixed_spill_expected + +# String states allocate variable byte counts in the shared import arena. +# Rebuilding the SUM state must realign its tagged inner-state pointer. +statement ok +SET memory_limit='32MB' + +statement ok +SET debug_force_external=true + +query II +SELECT count(*),count(*) FILTER (WHERE list_min(s)=6 AND list_max(s)=6 AND length(t)=39) FROM ( + SELECT i//3 g,string_agg(repeat('x',13),'') t, + as_sum(18446744073709551615::UBIGINT,1::BIGINT) s + FROM range(300000)t(i) GROUP BY g +) +---- +100000 100000 + +statement ok +SET debug_force_external=false + +statement ok +SET memory_limit='512MB' + +query II +SELECT finalize(s)=expected,finalize(c)=expected FROM ( + SELECT as_noised_count(18446744073709551615::UBIGINT,0.75) EXPORT_STATE s, + as_noised_count(18446744073709551615::UBIGINT,1::BIGINT,0.75) EXPORT_STATE c, + as_noised_count(18446744073709551615::UBIGINT,0.75) expected FROM range(1000) +) +---- +true true + +# NULL exported states occur naturally after an outer join. Native finalize +# propagates NULL, and combine keeps the non-NULL input unchanged. +query I +WITH s AS MATERIALIZED ( + SELECT as_sum(18446744073709551615::UBIGINT,1::BIGINT) EXPORT_STATE x, + as_sum(18446744073709551615::UBIGINT,1::BIGINT) expected +) +SELECT finalize(CASE WHEN i=0 THEN x END) IS NOT DISTINCT FROM CASE WHEN i=0 THEN expected END +FROM s,range(2)t(i) ORDER BY i +---- +true +true + +query I +WITH s AS MATERIALIZED ( + SELECT as_sum(18446744073709551615::UBIGINT,1::BIGINT) EXPORT_STATE x, + as_sum(18446744073709551615::UBIGINT,1::BIGINT) expected +) +SELECT finalize(combine(CASE WHEN i=0 THEN x END,CASE WHEN i=1 THEN x END)) IS NOT DISTINCT FROM expected +FROM s,range(2)t(i) ORDER BY i +---- +true +true + +# Clustered updates reuse their selected views across groups. Exercise a partial +# chunk, nullable values, dictionary selections, and constant hash/value inputs. +query II +SELECT count(*),count(*) FILTER (WHERE list_extract(s,1)=2*expected + AND list_min(k)=2*3*n AND list_max(k)=2*3*n AND list_min(z)=0 AND list_max(z)=0) +FROM ( + SELECT i%13 g, + as_sum(CASE WHEN i%2=0 THEN 18446744073709551615::UBIGINT ELSE 1::UBIGINT END, + CASE WHEN i%5=0 THEN NULL ELSE (i%101)::INTEGER END) s, + as_sum(18446744073709551615::UBIGINT,3::INTEGER) k, + as_sum(18446744073709551615::UBIGINT,NULL::INTEGER) z, + sum(CASE WHEN i%5=0 THEN NULL ELSE (i%101)::INTEGER END) expected,count(*) n + FROM range(10000)t(i) WHERE i%7<>0 GROUP BY g +) +---- +13 13 + +statement ok +SET pac_sample_diversity_check=true + +statement ok +SET memory_limit='80%' + +statement ok +SET priv_rewrite=true