diff --git a/.gitignore b/.gitignore index b9f264b..9d67701 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,15 @@ duckdb_unittest_tempdir/ testext test/python/__pycache__/ .Rhistory +*.csv +vcpkg/ +resources/* +duckdb/ +temp_dir/ +mydb.duckdb +plans/ +*.json +archived_tests/* +queries/synthetic_query/*.csv +queries/synthetic_query/*.db +*.zip diff --git a/CMakeLists.txt b/CMakeLists.txt index dbc450d..b00bbb8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,4 +1,6 @@ cmake_minimum_required(VERSION 3.5) +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) # Set extension name here set(TARGET_NAME cactusdb) @@ -7,8 +9,16 @@ set(TARGET_NAME cactusdb) # 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) find_package(Eigen3 REQUIRED) +if(DEFINED ENV{LIBTORCH_DIR}) + set(Torch_DIR "$ENV{LIBTORCH_DIR}/share/cmake/Torch") +elseif(DEFINED LIBTORCH_DIR) + set(Torch_DIR "${LIBTORCH_DIR}/share/cmake/Torch") +endif() +find_package(Torch REQUIRED) set(EXTENSION_NAME ${TARGET_NAME}_extension) @@ -19,18 +29,21 @@ set(LOADABLE_EXTENSION_NAME ${TARGET_NAME}_loadable_extension) project(${TARGET_NAME}) include_directories(src/include) include_directories(${EIGEN3_INCLUDE_DIR}) +include_directories(${TORCH_INCLUDE_DIRS}) set(EXTENSION_SOURCES src/cactusdb_extension.cpp) add_subdirectory(src/ml_functions) add_subdirectory(src/util) +add_subdirectory(src/optimization) +add_subdirectory(src/cost_model) message(STATUS "EXTENSION_SOURCES include ${EXTENSION_SOURCES}") build_static_extension(${TARGET_NAME} ${EXTENSION_SOURCES}) build_loadable_extension(${TARGET_NAME} " " ${EXTENSION_SOURCES}) # Link OpenSSL in both the static library as the loadable extension -target_link_libraries(${EXTENSION_NAME} OpenSSL::SSL OpenSSL::Crypto Eigen3::Eigen) -target_link_libraries(${LOADABLE_EXTENSION_NAME} OpenSSL::SSL OpenSSL::Crypto Eigen3::Eigen) +target_link_libraries(${EXTENSION_NAME} OpenSSL::SSL OpenSSL::Crypto Eigen3::Eigen ${TORCH_LIBRARIES}) +target_link_libraries(${LOADABLE_EXTENSION_NAME} OpenSSL::SSL OpenSSL::Crypto Eigen3::Eigen ${TORCH_LIBRARIES}) install( TARGETS ${EXTENSION_NAME} diff --git a/README.md b/README.md index 21029be..ac8265e 100644 --- a/README.md +++ b/README.md @@ -1,21 +1,60 @@ # OPTBench +# OPTBench This repository is based on https://github.com/duckdb/extension-template, check it out if you want to build and ship your own DuckDB extension. --- -This extension, OPTBench, allows you to ... . +This extension, OPTBench, is a DuckDB extension for benchmarking query optimizers. ## Building -### Managing dependencies -DuckDB extensions use VCPKG for dependency management. Enabling VCPKG is very simple: follow the [installation instructions](https://vcpkg.io/en/getting-started) or just run the following: -```shell -git clone https://github.com/Microsoft/vcpkg.git -./vcpkg/bootstrap-vcpkg.sh -export VCPKG_TOOLCHAIN_PATH=`pwd`/vcpkg/scripts/buildsystems/vcpkg.cmake +### Prerequisites +Install the following system dependencies before building: + +```sh +# Ubuntu / Debian +sudo apt install libssl-dev libeigen3-dev +``` + +### Clone with submodules +This repo uses git submodules (`duckdb` and `extension-ci-tools`). Always clone with: + +```sh +git clone --recurse-submodules +``` + +If you already cloned without submodules, run: + +```sh +git submodule update --init --recursive +``` + +### LibTorch dependency +This project requires LibTorch (the C++ PyTorch distribution). If you do not have CUDA available, download the CPU-only build: + +```sh +wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.1.2%2Bcpu.zip +unzip libtorch-cxx11-abi-shared-with-deps-2.1.2+cpu.zip -d ~/ ``` -Note: VCPKG is only required for extensions that want to rely on it for dependency management. If you want to develop an extension without dependencies or want to do your own dependency management, just skip this step. Note that the example extension uses VCPKG to build with a dependency for instructive purposes, so when skipping this step, the build may not work without removing the dependency. + +> **Note:** LibTorch 2.7.x prebuilt CPU binaries have internal header inconsistencies and will fail to compile. Use 2.1.2. + +Then set `LIBTORCH_DIR` before building: + +```sh +export LIBTORCH_DIR=~/libtorch +``` + +Alternatively, pass it directly to `make`: +```sh +LIBTORCH_DIR=~/libtorch make +``` + +> **Note:** If you switch `LIBTORCH_DIR` after a previous build, clear the CMake cache first to avoid picking up stale paths: +> ```sh +> rm -rf build/release/CMakeCache.txt build/release/CMakeFiles +> ``` ### Build steps Now to build the extension, run: @@ -35,6 +74,7 @@ The main binaries that will be built are: ## Running the extension To run the extension code, simply start the shell with `./build/release/duckdb`. +Now we can use the features from the extension directly in DuckDB. The template contains a single scalar function `cactusdb()` that takes a string argument and returns a string: Now we can use the features from the extension directly in DuckDB. The template contains a single scalar function `cactusdb()` that takes a string argument and returns a string: ``` D select cactusdb('Jane') as result; diff --git a/archived_tests/one_hot_encoder_rewrite_rule.test b/archived_tests/one_hot_encoder_rewrite_rule.test new file mode 100644 index 0000000..632467d --- /dev/null +++ b/archived_tests/one_hot_encoder_rewrite_rule.test @@ -0,0 +1,26 @@ +# name: test/sql/replace_binding_rule.test +# group: [sql] +require cactusdb + +statement ok +LOAD cactusdb; + +statement ok +PRAGMA disable_optimizer; + +# If you want to assert a value, use a query block with expected output: +# (adjust the expected list to match your CSV) +query I +SELECT one_hot_encode('red', '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/test/resources/month.csv', TRUE) AS y; +---- +[1.0, 0.0, 0.0, 0.0] + +statement ok +PRAGMA enable_optimizer; + +query I +SELECT one_hot_encode('red', '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/test/resources/month.csv', TRUE) AS y; +---- +[0.0, 0.0, 0.0, 1.0] + + diff --git a/expedia.json b/expedia.json new file mode 100644 index 0000000..cce837e --- /dev/null +++ b/expedia.json @@ -0,0 +1,185 @@ +{ + "total_bytes_written": 0, + "total_bytes_read": 0, + "rows_returned": 79756, + "latency": 0.619100106, + "result_set_size": 2871216, + "query_name": "SELECT\n s.prop_id,\n s.srch_id,\n tree_predict(\n list_transform(\n list_concat(\n /* 1) NUMERICALS — per-column scaling (exact order) */\n min_max_scaler(list_value(CAST(s.prop_location_score1 AS DOUBLE)),\n '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_location_score1.txt'),\n min_max_scaler(list_value(CAST(s.prop_location_score2 AS DOUBLE)),\n '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_location_score2.txt'),\n min_max_scaler(list_value(CAST(s.prop_log_historical_price AS DOUBLE)),\n '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_log_historical_price.txt'),\n min_max_scaler(list_value(CAST(s.price_usd AS DOUBLE)),\n '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/price_usd.txt'),\n min_max_scaler(list_value(CAST(s.orig_destination_distance AS DOUBLE)),\n '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/orig_destination_distance.txt'),\n min_max_scaler(list_value(CAST(h.prop_review_score AS DOUBLE)),\n '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_review_score.txt'),\n min_max_scaler(list_value(CAST(h.avg_bookings_usd AS DOUBLE)),\n '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/avg_bookings_usd.txt'),\n min_max_scaler(list_value(CAST(h.stdev_bookings_usd AS DOUBLE)),\n '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/stdev_bookings_usd.txt'),\n\n /* 2) CATEGORICALS — one-hot, exact order of your Python */\n -- stringEncoderSpecs (string→int): TRUE\n one_hot_encode(CAST(s.position AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/position.csv', TRUE),\n one_hot_encode(CAST(h.prop_country_id AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/prop_country_id.csv', TRUE),\n\n -- encoderSpecs (int→int): FALSE\n one_hot_encode(CAST(h.prop_starrating AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/prop_starrating.csv', FALSE),\n one_hot_encode(CAST(h.prop_brand_bool AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/prop_brand_bool.csv', FALSE),\n one_hot_encode(CAST(h.count_clicks AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/count_clicks.csv', FALSE),\n one_hot_encode(CAST(h.count_bookings AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/count_bookings.csv', FALSE),\n\n -- stringEncoderSpecs (string→int): TRUE\n one_hot_encode(CAST(r.year AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/year.csv', TRUE),\n one_hot_encode(CAST(r.month AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/month.csv', TRUE),\n one_hot_encode(CAST(r.weekofyear AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/weekofyear.csv', TRUE),\n one_hot_encode(CAST(r.time AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/time.csv', TRUE),\n one_hot_encode(CAST(r.site_id AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/site_id.csv', TRUE),\n one_hot_encode(CAST(r.visitor_location_country_id AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/visitor_location_country_id.csv', TRUE),\n one_hot_encode(CAST(r.srch_destination_id AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_destination_id.csv', TRUE),\n\n -- encoderSpecs (int→int): FALSE\n one_hot_encode(CAST(r.srch_length_of_stay AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_length_of_stay.csv', FALSE),\n one_hot_encode(CAST(r.srch_booking_window AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_booking_window.csv', FALSE),\n one_hot_encode(CAST(r.srch_adults_count AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_adults_count.csv', FALSE),\n one_hot_encode(CAST(r.srch_children_count AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_children_count.csv', FALSE),\n one_hot_encode(CAST(r.srch_room_count AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_room_count.csv', FALSE),\n one_hot_encode(CAST(r.srch_saturday_night_bool AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_saturday_night_bool.csv', FALSE),\n one_hot_encode(CAST(r.random_bool AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/random_bool.csv', FALSE)\n ),\n x -> CAST(x AS FLOAT) -- tree_predict expects LIST\n ),\n '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/tree/0.txt',\n FALSE\n )\nFROM Expedia_S_listings_extension2 AS s\nJOIN Expedia_R1_hotels2 AS h ON s.prop_id = h.prop_id\nJOIN Expedia_R2_searches AS r ON s.srch_id = r.srch_id\nWHERE s.prop_location_score1 > 1\n AND s.prop_location_score2 > 0.1\n AND s.prop_log_historical_price > 4\n AND h.count_bookings > 5\n AND r.srch_booking_window > 10\n AND r.srch_length_of_stay > 1;", + "blocked_thread_time": 0.0, + "system_peak_buffer_memory": 41269248, + "system_peak_temp_dir_size": 0, + "cpu_time": 4.15678951, + "extra_info": {}, + "cumulative_cardinality": 637172, + "cumulative_rows_scanned": 7586096, + "children": [ + { + "total_bytes_written": 0, + "total_bytes_read": 0, + "result_set_size": 2871216, + "operator_name": "PROJECTION", + "cpu_time": 4.15678951, + "extra_info": { + "Projections": [ + "prop_id", + "srch_id", + "tree_predict(list_transform(list_concat(min_max_scaler(list_value(prop_location_score1), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_location_score1.txt'), min_max_scaler(list_value(prop_location_score2), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_location_score2.txt'), min_max_scaler(list_value(prop_log_historical_price), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_log_historical_price.txt'), min_max_scaler(list_value(price_usd), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/price_usd.txt'), min_max_scaler(list_value(orig_destination_distance), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/orig_destination_distance.txt'), min_max_scaler(list_value(CAST(prop_review_score AS DOUBLE)), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_review_score.txt'), min_max_scaler(list_value(avg_bookings_usd), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/avg_bookings_usd.txt'), min_max_scaler(list_value(stdev_bookings_usd), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/stdev_bookings_usd.txt'), one_hot_encode(position, '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/position.csv', true), one_hot_encode(prop_country_id, '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/prop_country_id.csv', true), one_hot_encode(CAST(prop_starrating AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/prop_starrating.csv', false), one_hot_encode(CAST(prop_brand_bool AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/prop_brand_bool.csv', false), one_hot_encode(CAST(count_clicks AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/count_clicks.csv', false), one_hot_encode(CAST(count_bookings AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/count_bookings.csv', false), one_hot_encode(year, '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/year.csv', true), one_hot_encode(month, '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/month.csv', true), one_hot_encode(weekofyear, '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/weekofyear.csv', true), one_hot_encode(time, '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/time.csv', true), one_hot_encode(site_id, '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/site_id.csv', true), one_hot_encode(visitor_location_country_id, '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/visitor_location_country_id.csv', true), one_hot_encode(srch_destination_id, '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_destination_id.csv', true), one_hot_encode(CAST(srch_length_of_stay AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_length_of_stay.csv', false), one_hot_encode(CAST(srch_booking_window AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_booking_window.csv', false), one_hot_encode(CAST(srch_adults_count AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_adults_count.csv', false), one_hot_encode(CAST(srch_children_count AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_children_count.csv', false), one_hot_encode(CAST(srch_room_count AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_room_count.csv', false), one_hot_encode(CAST(srch_saturday_night_bool AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_saturday_night_bool.csv', false), one_hot_encode(CAST(random_bool AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/random_bool.csv', false))), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/tree/0.txt', false)" + ], + "Estimated Cardinality": "7220" + }, + "cumulative_cardinality": 637172, + "operator_type": "PROJECTION", + "operator_cardinality": 79756, + "cumulative_rows_scanned": 7586096, + "operator_rows_scanned": 0, + "operator_timing": 3.9706238020000004, + "children": [ + { + "total_bytes_written": 0, + "total_bytes_read": 0, + "result_set_size": 26798016, + "operator_name": "HASH_JOIN", + "cpu_time": 0.18616570799999999, + "extra_info": { + "Join Type": "INNER", + "Conditions": "srch_id = srch_id", + "Estimated Cardinality": "7220" + }, + "cumulative_cardinality": 557416, + "operator_type": "HASH_JOIN", + "operator_cardinality": 79756, + "cumulative_rows_scanned": 7586096, + "operator_rows_scanned": 0, + "operator_timing": 0.027968030999999997, + "children": [ + { + "total_bytes_written": 0, + "total_bytes_read": 0, + "result_set_size": 35619696, + "operator_name": "HASH_JOIN", + "cpu_time": 0.153371983, + "extra_info": { + "Join Type": "INNER", + "Conditions": "prop_id = prop_id", + "Estimated Cardinality": "36103" + }, + "cumulative_cardinality": 466608, + "operator_type": "HASH_JOIN", + "operator_cardinality": 212022, + "cumulative_rows_scanned": 7549075, + "operator_rows_scanned": 0, + "operator_timing": 0.040581439999999996, + "children": [ + { + "total_bytes_written": 0, + "total_bytes_read": 0, + "result_set_size": 21772608, + "operator_name": "SEQ_SCAN ", + "cpu_time": 0.11156367, + "extra_info": { + "Table": "Expedia_S_listings_extension2", + "Type": "Sequential Scan", + "Projections": [ + "prop_id", + "srch_id", + "prop_location_score1", + "prop_location_score2", + "prop_log_historical_price", + "price_usd", + "orig_destination_distance", + "position" + ], + "Filters": [ + "prop_location_score1>1.0", + "prop_location_score2>0.1", + "prop_log_historical_price>4.0" + ], + "Estimated Cardinality": "188428" + }, + "cumulative_cardinality": 247416, + "operator_type": "TABLE_SCAN", + "operator_cardinality": 247416, + "cumulative_rows_scanned": 7537136, + "operator_rows_scanned": 7537136, + "operator_timing": 0.11156367, + "children": [] + }, + { + "total_bytes_written": 0, + "total_bytes_read": 0, + "result_set_size": 688320, + "operator_name": "SEQ_SCAN ", + "cpu_time": 0.0012268730000000003, + "extra_info": { + "Table": "Expedia_R1_hotels2", + "Type": "Sequential Scan", + "Projections": [ + "prop_id", + "count_bookings", + "prop_review_score", + "avg_bookings_usd", + "stdev_bookings_usd", + "prop_country_id", + "prop_starrating", + "prop_brand_bool", + "count_clicks" + ], + "Filters": "count_bookings>5", + "Estimated Cardinality": "2387" + }, + "cumulative_cardinality": 7170, + "operator_type": "TABLE_SCAN", + "operator_cardinality": 7170, + "cumulative_rows_scanned": 11939, + "operator_rows_scanned": 11939, + "operator_timing": 0.0012268730000000003, + "children": [] + } + ] + }, + { + "total_bytes_written": 0, + "total_bytes_read": 0, + "result_set_size": 2033568, + "operator_name": "SEQ_SCAN ", + "cpu_time": 0.004825693999999999, + "extra_info": { + "Table": "Expedia_R2_searches", + "Type": "Sequential Scan", + "Projections": [ + "srch_id", + "srch_booking_window", + "srch_length_of_stay", + "year", + "month", + "weekofyear", + "time", + "site_id", + "visitor_location_country_id", + "srch_destination_id", + "srch_adults_count", + "srch_children_count", + "srch_room_count", + "srch_saturday_night_bool", + "random_bool" + ], + "Filters": [ + "srch_booking_window>10", + "srch_length_of_stay>1" + ], + "Estimated Cardinality": "7404" + }, + "cumulative_cardinality": 11052, + "operator_type": "TABLE_SCAN", + "operator_cardinality": 11052, + "cumulative_rows_scanned": 37021, + "operator_rows_scanned": 37021, + "operator_timing": 0.004825693999999999, + "children": [] + } + ] + } + ] + } + ] +} \ No newline at end of file diff --git a/log.txt b/log.txt new file mode 100644 index 0000000..baa85ee --- /dev/null +++ b/log.txt @@ -0,0 +1,233 @@ + +[CactusDynamicOpt] Initial Plan: +┌───────────────────────────┐ +│ ATTACH │ +│ ──────────────────── │ +└───────────────────────────┘ + + +[CactusDynamicOpt] Final Plan: +┌───────────────────────────┐ +│ ATTACH │ +│ ──────────────────── │ +└───────────────────────────┘ + + +[CactusDynamicOpt] Initial Plan: +┌───────────────────────────┐ +│ SET │ +│ ──────────────────── │ +└───────────────────────────┘ + + +[CactusDynamicOpt] Final Plan: +┌───────────────────────────┐ +│ SET │ +│ ──────────────────── │ +└───────────────────────────┘ + + +[CactusDynamicOpt] Initial Plan: +┌───────────────────────────┐ +│ SET │ +│ ──────────────────── │ +└───────────────────────────┘ + + +[CactusDynamicOpt] Final Plan: +┌───────────────────────────┐ +│ SET │ +│ ──────────────────── │ +└───────────────────────────┘ + + +[CactusDynamicOpt] Initial Plan: +┌───────────────────────────┐ +│ SET │ +│ ──────────────────── │ +└───────────────────────────┘ + + +[CactusDynamicOpt] Final Plan: +┌───────────────────────────┐ +│ SET │ +│ ──────────────────── │ +└───────────────────────────┘ + + +[CactusDynamicOpt] Initial Plan: +┌───────────────────────────┐ +│ ORDER_BY │ +│ ──────────────────── │ +│ s.impression_id │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ +│ PROJECTION │ +│ ──────────────────── │ +│ Expressions: │ +│ impression_id │ +│ listing_row_id │ +│ property_type │ +│ ml_argmax(cdb_softmax │ +│ (mat_add(mat_mul(mat_add │ +│ (mat_mul(list_concat │ +│ (user_profile_vec, │ +│ search_context_vec, │ +│ listing_feature_vec), ' │ +│ /home/local/ASUAD/jrtandel│ +│ /cactusdb-duckdb-extension│ +│ /queries/synthetic_query │ +│ /nn_params_150_512_2/nn_W1│ +│ .csv', 'dense'), '/home │ +│ /local/ASUAD/jrtandel │ +│ /cactusdb-duckdb-extension│ +│ /queries/synthetic_query │ +│ /nn_params_150_512_2/nn_b1│ +│ .csv'), '/home/local/ASUAD│ +│ /jrtandel/cactusdb-duckdb │ +│ -extension/queries │ +│ /synthetic_query │ +│ /nn_params_150_512_2/nn_W2│ +│ .csv', 'dense'), '/home │ +│ /local/ASUAD/jrtandel │ +│ /cactusdb-duckdb-extension│ +│ /queries/synthetic_query │ +│ /nn_params_150_512_2/nn_b2│ +│ .csv'))) │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ +│ COMPARISON_JOIN │ +│ ──────────────────── │ +│ Join Type: INNER │ +│ ├──────────────┐ +│ Conditions: │ │ +│ (listing_id = listing_id) │ │ +└─────────────┬─────────────┘ │ +┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +│ SEQ_SCAN ││ SEQ_SCAN │ +│ ──────────────────── ││ ──────────────────── │ +│ Table: ││ Table: │ +│ search_impressions ││ listing_metadata │ +│ ││ │ +│ Type: Sequential Scan ││ Type: Sequential Scan │ +└───────────────────────────┘└───────────────────────────┘ + + +[CactusDynamicOpt] Final Plan: +┌───────────────────────────┐ +│ ORDER_BY │ +│ ──────────────────── │ +│ s.impression_id │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ +│ PROJECTION │ +│ ──────────────────── │ +│ Expressions: │ +│ impression_id │ +│ listing_row_id │ +│ property_type │ +│ #5 │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ +│ COMPARISON_JOIN │ +│ ──────────────────── │ +│ Join Type: INNER │ +│ ├──────────────┐ +│ Conditions: │ │ +│ (listing_id = listing_id) │ │ +└─────────────┬─────────────┘ │ +┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +│ PROJECTION ││ SEQ_SCAN │ +│ ──────────────────── ││ ──────────────────── │ +│ Expressions: ││ Table: │ +│ #0 ││ listing_metadata │ +│ #1 ││ │ +│ #2 ││ Type: Sequential Scan │ +│ #3 ││ │ +│ #4 ││ │ +│ ml_argmax(cdb_softmax ││ │ +│ (mat_add(mat_mul(mat_add ││ │ +│ (mat_mul(list_concat ││ │ +│ (user_profile_vec, ││ │ +│ search_context_vec, ││ │ +│ listing_feature_vec), ' ││ │ +│ /home/local/ASUAD/jrtandel││ │ +│ /cactusdb-duckdb-extension││ │ +│ /queries/synthetic_query ││ │ +│ /nn_params_150_512_2/nn_W1││ │ +│ .csv', 'sparse'), '/home ││ │ +│ /local/ASUAD/jrtandel ││ │ +│ /cactusdb-duckdb-extension││ │ +│ /queries/synthetic_query ││ │ +│ /nn_params_150_512_2/nn_b1││ │ +│ .csv'), '/home/local/ASUAD││ │ +│ /jrtandel/cactusdb-duckdb ││ │ +│ -extension/queries ││ │ +│ /synthetic_query ││ │ +│ /nn_params_150_512_2/nn_W2││ │ +│ .csv', 'sparse'), '/home ││ │ +│ /local/ASUAD/jrtandel ││ │ +│ /cactusdb-duckdb-extension││ │ +│ /queries/synthetic_query ││ │ +│ /nn_params_150_512_2/nn_b2││ │ +│ .csv'))) ││ │ +└─────────────┬─────────────┘└───────────────────────────┘ +┌─────────────┴─────────────┐ +│ SEQ_SCAN │ +│ ──────────────────── │ +│ Table: │ +│ search_impressions │ +│ │ +│ Type: Sequential Scan │ +└───────────────────────────┘ + +┌───────────────┬────────────────┬───────────────┬────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐ +│ impression_id │ listing_row_id │ property_type │ ml_argmax(cdb_softmax(mat_add(mat_mul(mat_add(mat_mul(list_concat(s.user_profile_vec, s.search_context_vec, s.listing_feature_vec), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/queries/synthetic_query/nn_params_150_512_2/nn_W1.csv', 'dense'), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/queries/synthetic_query/nn_params_150_512_2/nn_b1.csv'), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/queries/synthetic_query/nn_params_150_512_2/nn_W2.csv', 'dense'), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/queries/synthetic_query/nn_params_150_512_2/nn_b2.csv'))) │ +│ int64 │ int64 │ int64 │ int32 │ +├───────────────┼────────────────┼───────────────┼────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤ +│ 0 │ 74026 │ 34 │ 1 │ +│ 0 │ 146251 │ 9 │ 1 │ +│ 0 │ 120994 │ 41 │ 1 │ +│ 0 │ 138687 │ 31 │ 1 │ +│ 0 │ 124436 │ 45 │ 1 │ +│ 0 │ 148069 │ 19 │ 1 │ +│ 0 │ 122560 │ 36 │ 1 │ +│ 0 │ 163553 │ 22 │ 1 │ +│ 0 │ 191611 │ 6 │ 1 │ +│ 0 │ 140026 │ 8 │ 1 │ +│ 0 │ 112396 │ 32 │ 1 │ +│ 0 │ 91592 │ 6 │ 1 │ +│ 0 │ 182195 │ 41 │ 1 │ +│ 0 │ 141028 │ 7 │ 1 │ +│ 0 │ 105511 │ 40 │ 1 │ +│ 0 │ 63623 │ 0 │ 1 │ +│ 0 │ 192415 │ 37 │ 1 │ +│ 0 │ 97952 │ 8 │ 1 │ +│ 0 │ 181131 │ 31 │ 1 │ +│ 0 │ 172611 │ 39 │ 1 │ +│ · │ · │ · │ · │ +│ · │ · │ · │ · │ +│ · │ · │ · │ · │ +│ 499999 │ 144233 │ 0 │ 1 │ +│ 499999 │ 121415 │ 5 │ 1 │ +│ 499999 │ 69362 │ 9 │ 1 │ +│ 499999 │ 157211 │ 8 │ 1 │ +│ 499999 │ 153869 │ 29 │ 1 │ +│ 499999 │ 128390 │ 17 │ 1 │ +│ 499999 │ 7499 │ 6 │ 1 │ +│ 499999 │ 35630 │ 35 │ 1 │ +│ 499999 │ 68698 │ 1 │ 1 │ +│ 499999 │ 179187 │ 4 │ 1 │ +│ 499999 │ 12742 │ 45 │ 1 │ +│ 499999 │ 136334 │ 19 │ 1 │ +│ 499999 │ 66967 │ 42 │ 1 │ +│ 499999 │ 132220 │ 26 │ 1 │ +│ 499999 │ 63892 │ 19 │ 1 │ +│ 499999 │ 12500 │ 14 │ 1 │ +│ 499999 │ 134088 │ 42 │ 1 │ +│ 499999 │ 181000 │ 36 │ 1 │ +│ 499999 │ 63209 │ 2 │ 1 │ +│ 499999 │ 172703 │ 31 │ 1 │ +├───────────────┴────────────────┴───────────────┴────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤ +│ 20006670 rows (40 shown) 4 columns │ +└─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘ diff --git a/queries/expedia_query.sql b/queries/expedia_query.sql new file mode 100644 index 0000000..0918b3c --- /dev/null +++ b/queries/expedia_query.sql @@ -0,0 +1,90 @@ +-- Attach the database file and use its schema +ATTACH '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/data/expedia_imbridge2.db' + AS expedia_db (READ_ONLY); +SET schema 'expedia_db.main'; + +-- Optimizer Config +PRAGMA enable_profiling; +SET enable_profiling = 'json'; +SET profiling_output = '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/expedia.json'; +-- PRAGMA profiling_output='/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/plans/profile.json'; +-- PRAGMA disable_optimizer; +-- PRAGMA explain_output='json'; +-- SET enable_profiling = 'query_tree_optimizer'; +-- SET profiling_mode = 'detailed'; + +LOAD 'cactusdb'; + +set threads = 8; + +-- Now query the tables +-- SELECT 'Expedia_S_listings_extension2' AS which, * FROM Expedia_S_listings_extension2 LIMIT 10; +-- SELECT 'Expedia_R1_hotels2' AS which, * FROM Expedia_R1_hotels2 LIMIT 10; +-- SELECT 'Expedia_R2_searches' AS which, * FROM Expedia_R2_searches LIMIT 10; +SELECT + s.prop_id, + s.srch_id, + tree_predict( + list_transform( + list_concat( + /* 1) NUMERICALS — per-column scaling (exact order) */ + min_max_scaler(list_value(CAST(s.prop_location_score1 AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_location_score1.txt'), + min_max_scaler(list_value(CAST(s.prop_location_score2 AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_location_score2.txt'), + min_max_scaler(list_value(CAST(s.prop_log_historical_price AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_log_historical_price.txt'), + min_max_scaler(list_value(CAST(s.price_usd AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/price_usd.txt'), + min_max_scaler(list_value(CAST(s.orig_destination_distance AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/orig_destination_distance.txt'), + min_max_scaler(list_value(CAST(h.prop_review_score AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_review_score.txt'), + min_max_scaler(list_value(CAST(h.avg_bookings_usd AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/avg_bookings_usd.txt'), + min_max_scaler(list_value(CAST(h.stdev_bookings_usd AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/stdev_bookings_usd.txt'), + + /* 2) CATEGORICALS — one-hot, exact order of your Python */ + -- stringEncoderSpecs (string→int): TRUE + one_hot_encode(CAST(s.position AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/position.csv', TRUE), + one_hot_encode(CAST(h.prop_country_id AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/prop_country_id.csv', TRUE), + + -- encoderSpecs (int→int): FALSE + one_hot_encode(CAST(h.prop_starrating AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/prop_starrating.csv', FALSE), + one_hot_encode(CAST(h.prop_brand_bool AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/prop_brand_bool.csv', FALSE), + one_hot_encode(CAST(h.count_clicks AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/count_clicks.csv', FALSE), + one_hot_encode(CAST(h.count_bookings AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/count_bookings.csv', FALSE), + + -- stringEncoderSpecs (string→int): TRUE + one_hot_encode(CAST(r.year AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/year.csv', TRUE), + one_hot_encode(CAST(r.month AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/month.csv', TRUE), + one_hot_encode(CAST(r.weekofyear AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/weekofyear.csv', TRUE), + one_hot_encode(CAST(r.time AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/time.csv', TRUE), + one_hot_encode(CAST(r.site_id AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/site_id.csv', TRUE), + one_hot_encode(CAST(r.visitor_location_country_id AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/visitor_location_country_id.csv', TRUE), + one_hot_encode(CAST(r.srch_destination_id AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_destination_id.csv', TRUE), + + -- encoderSpecs (int→int): FALSE + one_hot_encode(CAST(r.srch_length_of_stay AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_length_of_stay.csv', FALSE), + one_hot_encode(CAST(r.srch_booking_window AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_booking_window.csv', FALSE), + one_hot_encode(CAST(r.srch_adults_count AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_adults_count.csv', FALSE), + one_hot_encode(CAST(r.srch_children_count AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_children_count.csv', FALSE), + one_hot_encode(CAST(r.srch_room_count AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_room_count.csv', FALSE), + one_hot_encode(CAST(r.srch_saturday_night_bool AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_saturday_night_bool.csv', FALSE), + one_hot_encode(CAST(r.random_bool AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/random_bool.csv', FALSE) + ), + x -> CAST(x AS FLOAT) -- tree_predict expects LIST + ), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/tree/0.txt', + FALSE + ) +FROM Expedia_S_listings_extension2 AS s +JOIN Expedia_R1_hotels2 AS h ON s.prop_id = h.prop_id +JOIN Expedia_R2_searches AS r ON s.srch_id = r.srch_id +WHERE s.prop_location_score1 > 1 + AND s.prop_location_score2 > 0.1 + AND s.prop_log_historical_price > 4 + AND h.count_bookings > 5 + AND r.srch_booking_window > 10 + AND r.srch_length_of_stay > 1; diff --git a/queries/expedia_udf_filter_query.sql b/queries/expedia_udf_filter_query.sql new file mode 100644 index 0000000..b960ad4 --- /dev/null +++ b/queries/expedia_udf_filter_query.sql @@ -0,0 +1,99 @@ +-- Attach the database file and use its schema +ATTACH '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/data/expedia_imbridge2.db' + AS expedia_db (READ_ONLY); +SET schema 'expedia_db.main'; + +-- Optimizer Config +PRAGMA disable_optimizer; +-- PRAGMA explain_output='optimized_only'; +-- SET enable_profiling = 'query_tree_optimizer'; +-- SET profiling_mode = 'detailed'; + +LOAD 'cactusdb'; + +set threads = 8; + +EXPLAIN ANALYZE +SELECT + b.prop_id, + b.srch_id, + b.prediction, + r2.weekofyear, + r2.srch_booking_window +FROM ( + -- Inner: compute prediction after joining S+H+R and doing basic selective filters + SELECT + s.prop_id, + s.srch_id, + tree_predict( + list_transform( + list_concat( + /* 1) NUMERICALS — per-column scaling (exact order) */ + min_max_scaler(list_value(CAST(s.prop_location_score1 AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_location_score1.txt'), + min_max_scaler(list_value(CAST(s.prop_location_score2 AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_location_score2.txt'), + min_max_scaler(list_value(CAST(s.prop_log_historical_price AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_log_historical_price.txt'), + min_max_scaler(list_value(CAST(s.price_usd AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/price_usd.txt'), + min_max_scaler(list_value(CAST(s.orig_destination_distance AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/orig_destination_distance.txt'), + min_max_scaler(list_value(CAST(h.prop_review_score AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/prop_review_score.txt'), + min_max_scaler(list_value(CAST(h.avg_bookings_usd AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/avg_bookings_usd.txt'), + min_max_scaler(list_value(CAST(h.stdev_bookings_usd AS DOUBLE)), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/scaler_files/stdev_bookings_usd.txt'), + + /* 2) CATEGORICALS — one-hot, exact order of your Python */ + -- stringEncoderSpecs (string→int): TRUE + one_hot_encode(CAST(s.position AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/position.csv', TRUE), + one_hot_encode(CAST(h.prop_country_id AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/prop_country_id.csv', TRUE), + + -- encoderSpecs (int→int): FALSE + one_hot_encode(CAST(h.prop_starrating AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/prop_starrating.csv', FALSE), + one_hot_encode(CAST(h.prop_brand_bool AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/prop_brand_bool.csv', FALSE), + one_hot_encode(CAST(h.count_clicks AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/count_clicks.csv', FALSE), + one_hot_encode(CAST(h.count_bookings AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/count_bookings.csv', FALSE), + + -- stringEncoderSpecs (string→int): TRUE + one_hot_encode(CAST(r.year AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/year.csv', TRUE), + one_hot_encode(CAST(r.month AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/month.csv', TRUE), + one_hot_encode(CAST(r.weekofyear AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/weekofyear.csv', TRUE), + one_hot_encode(CAST(r.time AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/time.csv', TRUE), + one_hot_encode(CAST(r.site_id AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/site_id.csv', TRUE), + one_hot_encode(CAST(r.visitor_location_country_id AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/visitor_location_country_id.csv', TRUE), + one_hot_encode(CAST(r.srch_destination_id AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_destination_id.csv', TRUE), + + -- encoderSpecs (int→int): FALSE + one_hot_encode(CAST(r.srch_length_of_stay AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_length_of_stay.csv', FALSE), + one_hot_encode(CAST(r.srch_booking_window AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_booking_window.csv', FALSE), + one_hot_encode(CAST(r.srch_adults_count AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_adults_count.csv', FALSE), + one_hot_encode(CAST(r.srch_children_count AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_children_count.csv', FALSE), + one_hot_encode(CAST(r.srch_room_count AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_room_count.csv', FALSE), + one_hot_encode(CAST(r.srch_saturday_night_bool AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/srch_saturday_night_bool.csv', FALSE), + one_hot_encode(CAST(r.random_bool AS VARCHAR), '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/onehotencoders/random_bool.csv', FALSE) + ), + x -> CAST(x AS FLOAT) + ), + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/expedia/tree/0.txt', + FALSE + ) AS prediction + FROM Expedia_S_listings_extension2 AS s + JOIN Expedia_R1_hotels2 AS h ON s.prop_id = h.prop_id + JOIN Expedia_R2_searches AS r ON s.srch_id = r.srch_id + WHERE s.prop_location_score1 > 1 + AND s.prop_location_score2 > 0.1 + AND s.prop_log_historical_price > 4 + AND h.count_bookings > 5 + AND r.srch_booking_window > 10 + AND r.srch_length_of_stay > 1 +) AS b +-- Outer: extra join layer (gives optimizer a place to push the prediction filter below) +JOIN Expedia_R2_searches AS r2 + ON b.srch_id = r2.srch_id +WHERE + -- This is the UDF in the WHERE filter (not just as a projected column) + b.prediction > 0.01 + AND r2.weekofyear IN ('10','11'); -- an extra selective filter on the outer side diff --git a/queries/forest_query.sql b/queries/forest_query.sql new file mode 100644 index 0000000..a54da06 --- /dev/null +++ b/queries/forest_query.sql @@ -0,0 +1,30 @@ +-- queries/test_decision_forest.sql + + + +Load cactusdb; + +-- Build 3 rows of 6756-dimensional feature vectors +WITH feature_rows AS ( + SELECT + id, + list(CAST(i AS FLOAT)) AS features + FROM ( + -- three ids + SELECT * FROM (VALUES (1), (2), (3)) AS ids(id), + -- 0 .. 6755 + range(0, 6756) AS r(i) + ) + GROUP BY id +) +SELECT + id, + features, + decision_forest_predict( + features, + '/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/imbridgextenx/flights/rf_dot_trees_custom', + 6756, + TRUE + ) +FROM feature_rows +ORDER BY id; diff --git a/queries/q1.sql b/queries/q1.sql new file mode 100644 index 0000000..217b128 --- /dev/null +++ b/queries/q1.sql @@ -0,0 +1,24 @@ +-- q1.sql + +LOAD cactusdb; + +-- create the inputs table +CREATE TABLE inputs( + id INTEGER, + a DOUBLE[] +); + +-- data +INSERT INTO inputs VALUES + (1, [2.0, 1.0]), + (2, [1.0, 3.0]), + (3, [5.0, 5.0]); + +-- run the NN pipeline +SELECT + mat_add( + mat_mul(a, 'w_nn.csv', 'dense'), + 'b_zero.csv' + ) , id +FROM inputs +ORDER BY id; diff --git a/queries/results.txt b/queries/results.txt new file mode 100644 index 0000000..4a5c3ed --- /dev/null +++ b/queries/results.txt @@ -0,0 +1,183 @@ +Expedia query + +Settings : +thread count = 1 +with default optimizer | no default optimizer +3.66s | 3.70s + +thread count = 8 +with default optimizer | no default optimizer +0.622s | 0.637s + + +Query plan without optimizer : +┌────────────────────────────────────────────────┐ +│┌──────────────────────────────────────────────┐│ +││ Total Time: 0.637s ││ +│└──────────────────────────────────────────────┘│ +└────────────────────────────────────────────────┘ +┌───────────────────────────┐ +│ QUERY │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ +│ EXPLAIN_ANALYZE │ +│ ──────────────────── │ +│ 0 rows │ +│ (0.00s) │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ +│ PROJECTION │ +│ ──────────────────── │ +│ prop_id │ +│ srch_id │ +│ prediction │ +│ │ +│ 79,756 rows │ +│ (4.22s) │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ +│ FILTER │ +│ ──────────────────── │ +│ ((prop_location_score1 > │ +│ CAST(1 AS DOUBLE)) AND │ +│ (prop_location_score2 > │ +│ CAST(0.1 AS DOUBLE)) AND │ +│ (prop_log_historical_price│ +│ > CAST(4 AS DOUBLE)) AND │ +│ (count_bookings > CAST(5 │ +│ AS BIGINT)) AND │ +│ (srch_booking_window > │ +│ CAST(10 AS BIGINT)) AND │ +│ (srch_length_of_stay > │ +│ CAST(1 AS BIGINT))) │ +│ │ +│ 79,756 rows │ +│ (0.04s) │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ +│ HASH_JOIN │ +│ ──────────────────── │ +│ Join Type: INNER │ +│ │ +│ Conditions: ├───────────────────────────────────────────┐ +│ srch_id = srch_id │ │ +│ │ │ +│ 942,142 rows │ │ +│ (0.10s) │ │ +└─────────────┬─────────────┘ │ +┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ +│ HASH_JOIN │ │ TABLE_SCAN │ +│ ──────────────────── │ │ ──────────────────── │ +│ Join Type: INNER │ │ Table: │ +│ │ │ Expedia_R2_searches │ +│ Conditions: ├──────────────┐ │ │ +│ prop_id = prop_id │ │ │ Type: Sequential Scan │ +│ │ │ │ │ +│ 942,142 rows │ │ │ 37,021 rows │ +│ (0.12s) │ │ │ (0.00s) │ +└─────────────┬─────────────┘ │ └───────────────────────────┘ +┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +│ TABLE_SCAN ││ TABLE_SCAN │ +│ ──────────────────── ││ ──────────────────── │ +│ Table: ││ Table: │ +│Expedia_S_listings_extensio││ Expedia_R1_hotels2 │ +│ n2 ││ │ +│ ││ Type: Sequential Scan │ +│ Type: Sequential Scan ││ │ +│ ││ │ +│ 942,142 rows ││ 11,939 rows │ +│ (0.08s) ││ (0.00s) │ +└───────────────────────────┘└───────────────────────────┘ + +Query Plan with optimizer : + +┌───────────────────────────┐ +│ QUERY │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ +│ EXPLAIN_ANALYZE │ +│ ──────────────────── │ +│ 0 rows │ +│ (0.00s) │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ +│ PROJECTION │ +│ ──────────────────── │ +│ prop_id │ +│ srch_id │ +│ prediction │ +│ │ +│ 79,756 rows │ +│ (4.26s) │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ +│ HASH_JOIN │ +│ ──────────────────── │ +│ Join Type: INNER │ +│ │ +│ Conditions: ├───────────────────────────────────────────┐ +│ srch_id = srch_id │ │ +│ │ │ +│ 79,756 rows │ │ +│ (0.03s) │ │ +└─────────────┬─────────────┘ │ +┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ +│ HASH_JOIN │ │ TABLE_SCAN │ +│ ──────────────────── │ │ ──────────────────── │ +│ Join Type: INNER │ │ Table: │ +│ │ │ Expedia_R2_searches │ +│ Conditions: │ │ │ +│ prop_id = prop_id │ │ Type: Sequential Scan │ +│ │ │ │ +│ │ │ Projections: │ +│ │ │ srch_id │ +│ │ │ srch_booking_window │ +│ │ │ srch_length_of_stay │ +│ │ │ year │ +│ │ │ month │ +│ │ │ weekofyear │ +│ │ │ time │ +│ ├──────────────┐ │ site_id │ +│ │ │ │visitor_location_country_id│ +│ │ │ │ srch_destination_id │ +│ │ │ │ srch_adults_count │ +│ │ │ │ srch_children_count │ +│ │ │ │ srch_room_count │ +│ │ │ │ srch_saturday_night_bool │ +│ │ │ │ random_bool │ +│ │ │ │ │ +│ │ │ │ Filters: │ +│ │ │ │ srch_booking_window>10 │ +│ │ │ │ srch_length_of_stay>1 │ +│ │ │ │ │ +│ 212,022 rows │ │ │ 11,052 rows │ +│ (0.04s) │ │ │ (0.00s) │ +└─────────────┬─────────────┘ │ └───────────────────────────┘ +┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +│ TABLE_SCAN ││ TABLE_SCAN │ +│ ──────────────────── ││ ──────────────────── │ +│ Table: ││ Table: │ +│Expedia_S_listings_extensio││ Expedia_R1_hotels2 │ +│ n2 ││ │ +│ ││ Type: Sequential Scan │ +│ Type: Sequential Scan ││ │ +│ ││ Projections: │ +│ Projections: ││ prop_id │ +│ prop_id ││ count_bookings │ +│ srch_id ││ prop_review_score │ +│ prop_location_score1 ││ avg_bookings_usd │ +│ prop_location_score2 ││ stdev_bookings_usd │ +│ prop_log_historical_price ││ prop_country_id │ +│ price_usd ││ prop_starrating │ +│ orig_destination_distance ││ prop_brand_bool │ +│ position ││ count_clicks │ +│ ││ │ +│ Filters: ││ Filters: │ +│ prop_location_score1>1.0 ││ count_bookings>5 │ +│ prop_location_score2>0.1 ││ │ +│ prop_log_historical_price ││ │ +│ >4.0 ││ │ +│ ││ │ +│ 247,416 rows ││ 7,170 rows │ +│ (0.11s) ││ (0.00s) │ +└───────────────────────────┘└───────────────────────────┘ \ No newline at end of file diff --git a/queries/synthetic_query/model_generation.py b/queries/synthetic_query/model_generation.py new file mode 100644 index 0000000..3273480 --- /dev/null +++ b/queries/synthetic_query/model_generation.py @@ -0,0 +1,192 @@ +#!/usr/bin/env python3 +import argparse +import os +from pathlib import Path + +import numpy as np + + +def parse_layers(s: str): + """ + Parse a comma-separated list of integers into layer sizes. + Example: "150,16,2" -> [150, 16, 2] + """ + parts = [p.strip() for p in s.split(",") if p.strip()] + if len(parts) < 2: + raise ValueError( + "You must specify at least input and output layer, e.g. '150,2' or '150,16,2'" + ) + try: + layers = [int(p) for p in parts] + except ValueError: + raise ValueError(f"Invalid layer sizes in --layers: {s}") + if any(d <= 0 for d in layers): + raise ValueError("All layer sizes must be positive integers") + return layers + + +def init_weights(in_dim, out_dim, init: str, scale: float): + """ + Initialize weight matrix of shape (in_dim, out_dim). + """ + if init == "normal": + # N(0, scale^2) + return np.random.normal(loc=0.0, scale=scale, size=(in_dim, out_dim)) + elif init == "uniform": + # U(-scale, scale) + return np.random.uniform(low=-scale, high=scale, size=(in_dim, out_dim)) + else: + raise ValueError(f"Unknown weight init '{init}'") + + +def init_bias(out_dim, bias_mode: str, scale: float): + """ + Initialize bias vector of shape (1, out_dim). + """ + if bias_mode == "zeros": + return np.zeros((1, out_dim), dtype=float) + elif bias_mode == "normal": + return np.random.normal(loc=0.0, scale=scale, size=(1, out_dim)) + elif bias_mode == "uniform": + return np.random.uniform(low=-scale, high=scale, size=(1, out_dim)) + else: + raise ValueError(f"Unknown bias mode '{bias_mode}'") + + +def save_matrix_csv(matrix: np.ndarray, path: Path, precision: int = 10): + """ + Save matrix to CSV with given precision, no header. + Each row = one line; values comma-separated. + """ + path.parent.mkdir(parents=True, exist_ok=True) + fmt = f"%.{precision}g" # good balance for DuckDB + not too huge + np.savetxt(path, matrix, delimiter=",", fmt=fmt) + + +def generate_params( + layers, + out_dir: Path, + prefix: str, + weight_init: str, + weight_scale: float, + bias_mode: str, + bias_scale: float, + precision: int, +): + """ + layers = [input_dim, hidden1, ..., output_dim] + For each consecutive pair (d_in, d_out) generate: + - weights: (d_in x d_out) -> {prefix}_W{idx}.csv + - bias: (1 x d_out) -> {prefix}_b{idx}.csv + where idx starts at 1. + """ + out_dir.mkdir(parents=True, exist_ok=True) + + print(f"Generating params in '{out_dir}' with prefix '{prefix}'") + print(f"Layers: {layers}") + for i in range(len(layers) - 1): + d_in = layers[i] + d_out = layers[i + 1] + layer_idx = i + 1 + + print(f" Layer {layer_idx}: {d_in} -> {d_out}") + + W = init_weights(d_in, d_out, init=weight_init, scale=weight_scale) + b = init_bias(d_out, bias_mode=bias_mode, scale=bias_scale) + + w_path = out_dir / f"{prefix}_W{layer_idx}.csv" + b_path = out_dir / f"{prefix}_b{layer_idx}.csv" + + save_matrix_csv(W, w_path, precision=precision) + save_matrix_csv(b, b_path, precision=precision) + + print(f" Weights: {w_path} (shape {W.shape})") + print(f" Biases : {b_path} (shape {b.shape})") + + print("Done.") + + +def main(): + parser = argparse.ArgumentParser( + description="Generate random NN weights/biases as CSVs for DuckDB mat_mul/mat_add." + ) + parser.add_argument( + "--layers", + type=str, + required=True, + help="Comma-separated layer sizes including input and output. " + "Example: '150,16,2' for 150→16→2.", + ) + parser.add_argument( + "--out-dir", + type=str, + default="nn_params", + help="Output directory for CSV files (default: nn_params).", + ) + parser.add_argument( + "--prefix", + type=str, + default="nn", + help="Prefix for filenames (default: 'nn').", + ) + parser.add_argument( + "--weight-init", + type=str, + choices=["normal", "uniform"], + default="normal", + help="Weight initialization distribution (default: normal).", + ) + parser.add_argument( + "--weight-scale", + type=float, + default=0.1, + help="Scale (stddev or range) for weight init (default: 0.1).", + ) + parser.add_argument( + "--bias-mode", + type=str, + choices=["zeros", "normal", "uniform"], + default="normal", + help="Bias initialization mode (default: zeros).", + ) + parser.add_argument( + "--bias-scale", + type=float, + default=0.1, + help="Scale for bias init when bias-mode is normal/uniform (default: 0.1).", + ) + parser.add_argument( + "--precision", + type=int, + default=10, + help="Number of significant digits in CSV output (default: 10).", + ) + parser.add_argument( + "--seed", + type=int, + default=None, + help="Random seed for reproducibility (default: None).", + ) + + args = parser.parse_args() + + if args.seed is not None: + np.random.seed(args.seed) + + layers = parse_layers(args.layers) + out_dir = Path(args.out_dir) + + generate_params( + layers=layers, + out_dir=out_dir, + prefix=args.prefix, + weight_init=args.weight_init, + weight_scale=args.weight_scale, + bias_mode=args.bias_mode, + bias_scale=args.bias_scale, + precision=args.precision, + ) + + +if __name__ == "__main__": + main() diff --git a/queries/synthetic_query/synthetic_query.sql b/queries/synthetic_query/synthetic_query.sql new file mode 100644 index 0000000..b9205d5 --- /dev/null +++ b/queries/synthetic_query/synthetic_query.sql @@ -0,0 +1,42 @@ +ATTACH 'queries/synthetic_query/synthetic_database.db' + AS synthetic_db (READ_ONLY); + +SET schema 'synthetic_db.synthetic'; + +SET disabled_optimizers='expression_rewriter, empty_result_pullup, cte_filter_pusher, regex_range, in_clause, join_order, deliminator, unnest_rewriter, unused_columns, statistics_propagation, common_subexpressions, common_aggregate, column_lifetime, limit_pushdown, top_n, build_side_probe_side, compressed_materialization, duplicate_groups, reorder_filter, sampling_pushdown, join_filter_pushdown, materialized_cte, sum_rewriter, late_materialization, cte_inlining'; + +SET enable_extended_optimizer = '2'; + +SELECT + s.impression_id AS impression_id, + m.listing_row_id AS listing_row_id, + m.property_type AS property_type, + ml_argmax( + cdb_softmax( + -- layer 2: (h * W2) + b2 -> 1 x 2 logits + mat_add( + mat_mul( + -- layer 1: (x * W1) + b1 -> 1 x 16 hidden + mat_add( + mat_mul( + list_concat( + s.user_profile_vec, + s.search_context_vec, + s.listing_feature_vec + ), -- 150-dim feature vector + 'queries/synthetic_query/nn_params_150_512_2/nn_W1.csv', + 'dense' + ), + 'queries/synthetic_query/nn_params_150_512_2/nn_b1.csv' + ), + 'queries/synthetic_query/nn_params_150_512_2/nn_W2.csv', + 'dense' + ), + 'queries/synthetic_query/nn_params_150_512_2/nn_b2.csv' + ) + ) + ) +FROM search_impressions AS s +JOIN listing_metadata AS m + USING (listing_id) +ORDER BY s.impression_id; diff --git a/queries/synthetic_query/table_generation.py b/queries/synthetic_query/table_generation.py new file mode 100644 index 0000000..db5acdc --- /dev/null +++ b/queries/synthetic_query/table_generation.py @@ -0,0 +1,254 @@ +import numpy as np +import pandas as pd +import duckdb + + +def generate_synthetic_tables_vector( + n_left=10_000, + n_right=5_000, + join_key_distinct=1_000, + feature_cols=("user_profile_vec", "search_context_vec", "listing_feature_vec"), + feature_dim=50, + feature_nnz_ratio=0.3, + zero_row_ratio=0.9, + left_other_distinct=None, + right_other_distinct=None, + overlap_fraction=1.0, # controls join cardinality + random_state=None, +): + """ + Generate two tables for join + sparse feature-vector experiments. + + Table search_impressions: + - impression_id: unique row id + - listing_id: join column + - feature_cols: each is a vector of length `feature_dim` (Python list) + so the full feature vector has len = len(feature_cols) * feature_dim + + Table listing_metadata: + - listing_row_id: unique row id + - listing_id: join column (same domain size as left, but with + controllable overlap) + - optional extra columns with controlled distinct counts. + + Parameters + ---------- + n_left, n_right : int + Number of rows in left / right tables. + join_key_distinct : int + Number of distinct values per table for the join_key domain (K). + overlap_fraction : float in [0,1] + Fraction of those K keys that are shared between left and right. + - 1.0 => full overlap (max join size, original behavior) + - 0.0 => disjoint domains (join cardinality ~ 0) + - in between => approximately linear scaling of join cardinality: + E[|L ⋈ R|] ≈ overlap_fraction * (n_left * n_right / K) + feature_cols : tuple[str] + Names of the feature columns in search_impressions. + feature_dim : int + Length of each feature column vector (e.g., 50). + feature_nnz_ratio : float in [0,1] + Target fraction of non-zero entries across ALL feature positions, + i.e. across shape (n_left, len(feature_cols), feature_dim), + including zero rows. + zero_row_ratio : float in [0,1] + Fraction of rows whose entire feature vector (all dims of all + feature_cols) is zero. + left_other_distinct, right_other_distinct : dict | None + Optional dict: column_name -> #distinct values for extra columns. + random_state : int | None + RNG seed. + + Returns + ------- + search_impressions : pd.DataFrame + listing_metadata : pd.DataFrame + summary : dict + Contains realized stats: + - left_nnz_ratio + - zero_rows_left + - join_key_distinct_left/right + - approx_inner_join_cardinality + - etc. + """ + rng = np.random.default_rng(random_state) + + # --- Left/right ids --- + left_ids = np.arange(n_left) + right_ids = np.arange(n_right) + + # --- Construct join domains with controllable overlap --- + K = int(join_key_distinct) + overlap_fraction = float(np.clip(overlap_fraction, 0.0, 1.0)) + + if overlap_fraction == 1.0: + # Full overlap: original behavior + left_domain = np.arange(K) + right_domain = np.arange(K) + elif overlap_fraction == 0.0: + # Disjoint domains: no key intersection => join cardinality ~ 0 + left_domain = np.arange(K) + right_domain = np.arange(K, 2 * K) + else: + K_overlap = max(1, int(round(overlap_fraction * K))) + # Shared keys: 0 .. K_overlap-1 + overlap_keys = np.arange(K_overlap) + + # Left-only keys: K_overlap .. K-1 + left_exclusive = np.arange(K_overlap, K) + + # Right-only keys: K .. K + (K - K_overlap) - 1 + right_exclusive = np.arange(K, K + (K - K_overlap)) + + left_domain = np.concatenate([overlap_keys, left_exclusive]) + right_domain = np.concatenate([overlap_keys, right_exclusive]) + + left_join = rng.choice(left_domain, size=n_left, replace=True) + right_join = rng.choice(right_domain, size=n_right, replace=True) + + # Real-world column names + left_data = {"impression_id": left_ids, "listing_id": left_join} + right_data = {"listing_row_id": right_ids, "listing_id": right_join} + + # --- Extra non-feature columns with controlled distinct counts --- + def add_extras(data_dict, n_rows, distinct_cfg): + if not distinct_cfg: + return + for col, n_dist in distinct_cfg.items(): + domain = np.arange(n_dist) + data_dict[col] = rng.choice(domain, size=n_rows, replace=True) + + add_extras(left_data, n_left, left_other_distinct) + add_extras(right_data, n_right, right_other_distinct) + + # --- Feature tensor: shape (n_left, num_feature_cols, feature_dim) --- + m = len(feature_cols) + d = feature_dim + total_positions = n_left * m * d + + feature_nnz_ratio = float(np.clip(feature_nnz_ratio, 0.0, 1.0)) + zero_row_ratio = float(np.clip(zero_row_ratio, 0.0, 1.0)) + + # Decide which rows are fully zero + zero_row_mask = rng.random(n_left) < zero_row_ratio + n_zero_rows = int(zero_row_mask.sum()) + n_active_rows = n_left - n_zero_rows + + active_positions = n_active_rows * m * d + if active_positions == 0: + per_pos_prob = 0.0 + else: + target_nnz = feature_nnz_ratio * total_positions + per_pos_prob = float(np.clip(target_nnz / active_positions, 0.0, 1.0)) + + # Initialize feature tensor with zeros + feat_tensor = np.zeros((n_left, m, d), dtype=float) + + # Fill active rows with Bernoulli mask * N(0,1) values + if n_active_rows > 0 and per_pos_prob > 0.0: + active_idx = np.where(~zero_row_mask)[0] + mask = rng.random((n_active_rows, m, d)) < per_pos_prob + values = rng.normal(loc=0.0, scale=1.0, size=(n_active_rows, m, d)) + feat_tensor[active_idx] = mask * values + + # Attach feature columns to left_data as Python lists of length `feature_dim` + for j, col in enumerate(feature_cols): + left_data[col] = [feat_tensor[i, j, :].tolist() for i in range(n_left)] + + search_impressions = pd.DataFrame(left_data) + listing_metadata = pd.DataFrame(right_data) + + # --- Realized stats --- + nnz = int(np.count_nonzero(feat_tensor)) + left_nnz_ratio = nnz / total_positions if total_positions > 0 else 0.0 + + row_sums = np.abs(feat_tensor).sum(axis=(1, 2)) + zero_rows_left = float((row_sums == 0).mean()) + + # still using old summary keys for backwards-compat; now refer to listing_id + join_key_distinct_left = int(search_impressions["listing_id"].nunique()) + join_key_distinct_right = int(listing_metadata["listing_id"].nunique()) + + join_counts_left = search_impressions["listing_id"].value_counts() + join_counts_right = listing_metadata["listing_id"].value_counts() + approx_inner_join_cardinality = int( + ( + join_counts_left.to_frame("l") + .join(join_counts_right.to_frame("r"), how="inner") + .fillna(0) + .eval("l * r") + .sum() + ) + ) + + # theoretical max join cardinality with full overlap, for reference + max_expected_join = (n_left * n_right) / max(1, K) + expected_join_given_overlap = overlap_fraction * max_expected_join + + summary = { + "left_nnz_ratio": left_nnz_ratio, # global nnz over 3×50 vector + "zero_rows_left": zero_rows_left, # fraction of fully-zero rows + "per_pos_prob_active_rows": per_pos_prob, + "n_zero_rows_left": n_zero_rows, + "feature_dim": d, + "join_key_distinct_left": join_key_distinct_left, + "join_key_distinct_right": join_key_distinct_right, + "approx_inner_join_cardinality": approx_inner_join_cardinality, + "max_expected_join_full_overlap": max_expected_join, + "expected_join_given_overlap": expected_join_given_overlap, + "overlap_fraction": overlap_fraction, + } + + return search_impressions, listing_metadata, summary + + +if __name__ == "__main__": + # 1) Generate synthetic tables + search_impressions, listing_metadata, stats = generate_synthetic_tables_vector( + n_left=50_0000, # 500k + n_right=20_0000, # 200k + join_key_distinct=5000, + feature_cols=("user_profile_vec", "search_context_vec", "listing_feature_vec"), + feature_dim=50, + feature_nnz_ratio=0.0002, + zero_row_ratio=0.98, + left_other_distinct={"search_segment": 10}, + right_other_distinct={"property_type": 50}, + overlap_fraction=1000, # will be clipped to 1.0 by np.clip + random_state=42, + ) + + print(stats) + + # Optionally keep CSV names aligned with table names + search_impressions.to_csv("search_impressions.csv", index=False) + listing_metadata.to_csv("listing_metadata.csv", index=False) + + # 2) Path where you want your .db file + db_path = "/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/queries/synthetic_query/synthetic_database.db" + + # 3) Create DuckDB file and store tables + con = duckdb.connect(db_path) # this creates the file if it doesn't exist + + # optional: create a separate schema for synthetic data + con.execute("CREATE SCHEMA IF NOT EXISTS synthetic;") + con.execute("SET schema 'synthetic';") + + # register DataFrames as temporary views + con.register("search_impressions_df", search_impressions) + con.register("listing_metadata_df", listing_metadata) + + # create persistent tables inside the .db file + con.execute("DROP TABLE IF EXISTS search_impressions;") + con.execute("DROP TABLE IF EXISTS listing_metadata;") + + con.execute( + "CREATE TABLE search_impressions AS SELECT * FROM search_impressions_df;" + ) + con.execute("CREATE TABLE listing_metadata AS SELECT * FROM listing_metadata_df;") + + # collect stats for better cardinality estimation + con.execute("ANALYZE;") + + con.close() + print("Created DuckDB file at:", db_path) diff --git a/src/cactusdb_extension.cpp b/src/cactusdb_extension.cpp index f872cba..186be78 100644 --- a/src/cactusdb_extension.cpp +++ b/src/cactusdb_extension.cpp @@ -1,6 +1,8 @@ #include "cactusdb_extension.hpp" #include "duckdb.hpp" #include "functions.hpp" +#include "optimization/register_optimizers.hpp" +#include "duckdb/main/database.hpp" // OpenSSL linked through vcpkg #include @@ -9,7 +11,18 @@ namespace duckdb { static void LoadInternal(ExtensionLoader &loader) { + auto &db = loader.GetDatabaseInstance(); + auto &config = DBConfig::GetConfig(db); + + // Add New Configs + // + config.AddExtensionOption("cactusdb_enable_optimizers", "Enable All Optimizers", LogicalType::BOOLEAN, + Value::BOOLEAN(true)); + config.AddExtensionOption("enable_extended_optimizer", "Enable Specific Optimizer", LogicalType::VARCHAR, ""); + + RegisterMLScalarFunctions(loader); + RegisterCactusOptimizers(db); } diff --git a/src/cost_model/CMakeLists.txt b/src/cost_model/CMakeLists.txt new file mode 100644 index 0000000..029bffd --- /dev/null +++ b/src/cost_model/CMakeLists.txt @@ -0,0 +1,9 @@ +# Append optimizer sources to the extension's global list +list(APPEND EXTENSION_SOURCES + ${CMAKE_CURRENT_SOURCE_DIR}/cost_model.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/estimator.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/udf_cost_coeff.cpp +) + +# Expose back to parent +set(EXTENSION_SOURCES "${EXTENSION_SOURCES}" PARENT_SCOPE) diff --git a/src/cost_model/cost_model.cpp b/src/cost_model/cost_model.cpp new file mode 100644 index 0000000..053cd2d --- /dev/null +++ b/src/cost_model/cost_model.cpp @@ -0,0 +1,130 @@ +#include "duckdb.hpp" +#include "cost_model/cost_model.hpp" +#include "duckdb/planner/operator/logical_projection.hpp" +#include "duckdb/planner/operator/logical_get.hpp" +#include "duckdb/planner/operator/logical_filter.hpp" +#include "duckdb/planner/operator/logical_join.hpp" +#include "duckdb/planner/operator/logical_aggregate.hpp" +#include "duckdb/planner/operator/logical_order.hpp" +#include "duckdb/common/types/vector.hpp" +#include + +// using namespace duckdb; +namespace duckdb { + +CostModel::Node CostModel::MapNode(const LogicalOperator &op) { + switch (op.type) { + case LogicalOperatorType::LOGICAL_FILTER: return Node::Filter; + case LogicalOperatorType::LOGICAL_PROJECTION: return Node::Projection; + case LogicalOperatorType::LOGICAL_GET: return Node::Get; + case LogicalOperatorType::LOGICAL_COMPARISON_JOIN: + case LogicalOperatorType::LOGICAL_ANY_JOIN: return Node::HashJoin; + case LogicalOperatorType::LOGICAL_ORDER_BY: return Node::Order; + case LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY: return Node::Aggregate; + case LogicalOperatorType::LOGICAL_DELIM_JOIN: + case LogicalOperatorType::LOGICAL_CROSS_PRODUCT: return Node::NestedLoopJoin; + default: return Node::Unknown; + } +} + +CostEstimate SimpleCostModel::GetCost(LogicalOperator &op, const std::vector &inputs) { + if (op.children.empty()) return HandleLeaf(op); + return HandleInternal(op, inputs); +} + +CostEstimate SimpleCostModel::HandleLeaf(LogicalOperator &op) { + // LOGICAL_GET: use table stats if available (row_count); cols from types.size() + if (op.type == LogicalOperatorType::LOGICAL_GET) { + auto &get = op.Cast(); + idx_t rows = get.estimated_cardinality; // DuckDB keeps this updated + if (rows == DConstants::INVALID_INDEX) rows = 0; + idx_t cols = get.types.size(); + return CostEstimate(0.0, rows, cols); + } + // Default leaf + return CostEstimate(0.0, 1, op.types.size()); +} + +CostEstimate SimpleCostModel::HandleInternal(LogicalOperator &op, const std::vector &inputs) { + auto node = MapNode(op); + switch (node) { + case Node::Aggregate: + case Node::Order: { + // cost = c0 * rows_in ; rows_out ≈ λ * rows_in + auto stat = std::static_pointer_cast(inputs[0].stats); + auto coeff = UdfCostCoeff::Get("Aggregation"); // or "Order" + double c0 = coeff.empty() ? 1.0 : coeff[0]; + double lambda = 0.2; + return CostEstimate(c0 * stat->rows, (idx_t)(lambda * stat->rows), stat->cols); + } + case Node::Filter: { + auto stat = std::static_pointer_cast(inputs[0].stats); + double lambda = 0.5; // configurable selectivity + return CostEstimate(lambda * stat->rows, (idx_t)(lambda * stat->rows), stat->cols); + } + case Node::HashJoin: { + // Left (M×N), Right (O×P) + auto l = std::static_pointer_cast(inputs[0].stats); + auto r = std::static_pointer_cast(inputs[1].stats); + auto coeff = UdfCostCoeff::Get("HashJoin"); + double a = coeff.size() > 0 ? coeff[0] : 1.0; + double b = coeff.size() > 1 ? coeff[1] : 0.0; + double g = coeff.size() > 2 ? coeff[2] : 0.0; + double cost = a * l->rows + b * l->rows * l->cols + g * r->rows * r->cols; + return CostEstimate(cost, std::max(l->rows, r->rows), l->cols + r->cols); + } + case Node::NestedLoopJoin: { + auto l = std::static_pointer_cast(inputs[0].stats); + auto r = std::static_pointer_cast(inputs[1].stats); + auto coeff = UdfCostCoeff::Get("NestedLoopJoin"); + double a = coeff.size() > 0 ? coeff[0] : 1.0; + double b = coeff.size() > 1 ? coeff[1] : 0.0; + double g = coeff.size() > 2 ? coeff[2] : 0.0; + double cost = a * l->rows + b * l->rows * l->cols + g * r->rows * r->cols; + return CostEstimate(cost, std::max(l->rows, r->rows), r->cols); + } + case Node::Projection: + return HandleProjection(op, inputs); + default: + // passthrough + auto stat = std::static_pointer_cast(inputs[0].stats); + return CostEstimate(0.0, stat->rows, stat->cols); + } +} + +CostEstimate SimpleCostModel::HandleProjection(LogicalOperator &op, const std::vector &inputs) { + auto &proj = op.Cast(); + // Very similar to your Velox logic: parse UDFs from expressions and sum costs + std::vector udfs; udfs.reserve(proj.expressions.size()); + for (auto &expr : proj.expressions) { + std::cout << "\nProjection Expr: " << expr->ToString() << "\n"; + udfs.push_back(expr->ToString()); // parse later inside UDFChainCost + } + return UDFChainCost(udfs, inputs); +} + +CostEstimate SimpleCostModel::UDFChainCost(const std::vector &udfs_raw, + const std::vector &inputs) { + auto stat = std::static_pointer_cast(inputs[0].stats); + CostEstimate out(0.0, stat->rows, stat->cols); + + std::vector udf_ids; + udf_ids.reserve(udfs_raw.size()); + for (auto &s : udfs_raw) { + // very simple tokenization: split by '(' and drop last token + auto pos = s.find('('); + if (pos != std::string::npos) udf_ids.push_back(s.substr(0, pos)); + } + std::reverse(udf_ids.begin(), udf_ids.end()); + + for (auto &u : udf_ids) { + std::cout << "\nUDF detected: " << u << "\n"; + auto coeff = UdfCostCoeff::Get(u); + // Plug your learned model here; as placeholder, linear in rows + double c = coeff.empty() ? 0.0 : coeff[0] * out.output_rows; + out.cost += c; + } + return out; +} + +} // namespace duckdb diff --git a/src/cost_model/estimator.cpp b/src/cost_model/estimator.cpp new file mode 100644 index 0000000..07bfe9d --- /dev/null +++ b/src/cost_model/estimator.cpp @@ -0,0 +1,32 @@ +#include "duckdb.hpp" +#include "cost_model/estimator.hpp" +#include "cost_model/stat_source.hpp" +#include "duckdb/planner/logical_operator.hpp" + +namespace duckdb { + +CostEstimate Estimator::Estimate(LogicalOperator &root) { + return SimpleEstimator(std::move(model)).Estimate(root); +} + +CostEstimate SimpleEstimator::Estimate(LogicalOperator &node) { + // 1) visit children first, accumulate their costs + double child_cost_sum = 0.0; + std::vector inputs; inputs.reserve(node.children.size()); + for (auto &child_ptr : node.children) { + auto &child = *child_ptr; + auto est = Estimate(child); + child_cost_sum += est.cost; + inputs.emplace_back(node.GetName(), Source::Type::NODE, + std::make_shared(est.output_rows, est.output_cols)); + } + + // 2) ask model for current op's cost given inputs + auto cur = model->GetCost(node, inputs); + + // 3) add child costs + cur.cost += child_cost_sum; + return cur; +} + +} // namespace duckdb diff --git a/src/cost_model/udf_cost_coeff.cpp b/src/cost_model/udf_cost_coeff.cpp new file mode 100644 index 0000000..d6b9f09 --- /dev/null +++ b/src/cost_model/udf_cost_coeff.cpp @@ -0,0 +1,47 @@ +// .cpp +#include "cost_model/udf_cost_coeff.hpp" +#include +#include +#include "duckdb.hpp" + +namespace duckdb { +static std::unordered_map> *g_map = nullptr; +static bool g_init = false; + +std::unordered_map> &UdfCostCoeff::Map() { + if (!g_map) g_map = new std::unordered_map>(); + return *g_map; +} +bool &UdfCostCoeff::Initialized() { return g_init; } + +void UdfCostCoeff::InitFromFile(const std::string &path) { + if (Initialized()) return; + auto p = path; + if (p.empty()) { + + p = "/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/resources/cost_model/udf_coefficient_aws.txt"; + } + std::ifstream in(p); + if (!in.is_open()) { Initialized() = true; return; } // empty map ok + std::string line; + while (std::getline(in, line)) { + auto pos = line.find(':'); + if (pos == std::string::npos) continue; + auto k = line.substr(0, pos); + auto v = std::stod(line.substr(pos + 1)); + Map()[k].push_back(v); + } + Initialized() = true; +} + +const std::vector& UdfCostCoeff::Get(const std::string &key) { + if (!Initialized()) InitFromFile(); + auto &m = Map(); + auto it = m.find(key); + static std::vector empty; + if (it == m.end()){ + throw std::runtime_error("UDF Cost Coefficient not found for key: " + key); + } + return it == m.end() ? empty : it->second; +} +} // namespace duckdb diff --git a/src/include/cost_model/cost_estimate.hpp b/src/include/cost_model/cost_estimate.hpp new file mode 100644 index 0000000..395ce7f --- /dev/null +++ b/src/include/cost_model/cost_estimate.hpp @@ -0,0 +1,11 @@ +#include "duckdb.hpp" + +namespace duckdb { +struct CostEstimate { + double cost; + idx_t output_rows; + idx_t output_cols; + CostEstimate(double c=0, idx_t r=0, idx_t c2=0) + : cost(c), output_rows(r), output_cols(c2) {} +}; +} // namespace duckdb diff --git a/src/include/cost_model/cost_model.hpp b/src/include/cost_model/cost_model.hpp new file mode 100644 index 0000000..0a5e910 --- /dev/null +++ b/src/include/cost_model/cost_model.hpp @@ -0,0 +1,38 @@ +#pragma once +#include "cost_model/stat_source.hpp" +#include "cost_model/cost_estimate.hpp" +#include "cost_model/udf_cost_coeff.hpp" +#include "duckdb/planner/logical_operator.hpp" +#include "duckdb.hpp" + +namespace duckdb { + +class CostModel { +public: + enum class Node { + Filter, Projection, Aggregate, HashJoin, Order, Get, LocalPartition, + Unnest, RowNumber, NestedLoopJoin, Unknown + }; + virtual ~CostModel() = default; + virtual CostEstimate GetCost(duckdb::LogicalOperator &op, + const std::vector &inputs) = 0; + + static Node MapNode(const duckdb::LogicalOperator &op); +}; + +class SimpleCostModel : public CostModel { +public: + CostEstimate GetCost(duckdb::LogicalOperator &op, + const std::vector &inputs) override; + +private: + CostEstimate HandleLeaf(duckdb::LogicalOperator &op); + CostEstimate HandleInternal(duckdb::LogicalOperator &op, + const std::vector &inputs); + CostEstimate HandleProjection(duckdb::LogicalOperator &op, + const std::vector &inputs); + CostEstimate UDFChainCost(const std::vector &udfs, + const std::vector &inputs); +}; + +} // namespace cactusdb diff --git a/src/include/cost_model/estimator.hpp b/src/include/cost_model/estimator.hpp new file mode 100644 index 0000000..0a74d31 --- /dev/null +++ b/src/include/cost_model/estimator.hpp @@ -0,0 +1,26 @@ +#pragma once +#include "cost_model/cost_model.hpp" +#include "duckdb/planner/logical_operator.hpp" +#include "duckdb.hpp" + +namespace duckdb { + +class Estimator { +protected: + std::unique_ptr model; + +public: + explicit Estimator(std::unique_ptr m) : model(std::move(m)) {} + virtual ~Estimator() = default; + + // Recursively compute cost from a logical root + virtual CostEstimate Estimate(duckdb::LogicalOperator &root); +}; + +class SimpleEstimator : public Estimator { +public: + using Estimator::Estimator; + CostEstimate Estimate(duckdb::LogicalOperator &root) override; +}; + +} // namespace duckdb diff --git a/src/include/cost_model/stat_source.hpp b/src/include/cost_model/stat_source.hpp new file mode 100644 index 0000000..f13d698 --- /dev/null +++ b/src/include/cost_model/stat_source.hpp @@ -0,0 +1,29 @@ +#pragma once +#include "duckdb.hpp" +#include "duckdb/common/types.hpp" +#include +#include + +namespace duckdb { + +struct Stat { + explicit Stat(std::string n) : name(std::move(n)) {} + std::string name; + virtual ~Stat() = default; +}; + +struct OutputStat : public Stat { + duckdb::idx_t rows, cols; + OutputStat(duckdb::idx_t r, duckdb::idx_t c) : Stat(""), rows(r), cols(c) {} +}; + +struct Source { + enum class Type { NODE, FILE, VECTOR, DATABASE }; + std::string name; + Type type; + std::shared_ptr stats; + Source(std::string n, Type t, std::shared_ptr s) + : name(std::move(n)), type(t), stats(std::move(s)) {} +}; + +} // namespace duckdb diff --git a/src/include/cost_model/udf_cost_coeff.hpp b/src/include/cost_model/udf_cost_coeff.hpp new file mode 100644 index 0000000..6377d09 --- /dev/null +++ b/src/include/cost_model/udf_cost_coeff.hpp @@ -0,0 +1,17 @@ +// .hpp +#pragma once +#include +#include +#include +#include "duckdb.hpp" + +namespace duckdb { +class UdfCostCoeff { +public: + static const std::vector& Get(const std::string &key); + static void InitFromFile(const std::string &path = ""); +private: + static std::unordered_map> &Map(); + static bool &Initialized(); +}; +} // namespace cactusdb diff --git a/src/include/ml_functions/argmax.hpp b/src/include/ml_functions/argmax.hpp new file mode 100644 index 0000000..498e75b --- /dev/null +++ b/src/include/ml_functions/argmax.hpp @@ -0,0 +1,25 @@ +#pragma once + +#include +#include +#include + +#include "duckdb.hpp" +#include "duckdb/function/function.hpp" + +namespace duckdb { + +struct ML_Argmax { + // We don't need bind state; this just enforces arity (optionally types) + static unique_ptr Bind(ClientContext &context, + ScalarFunction &bound_function, + vector> &args); + + // Executor: argmax over each LIST row + static void Execute(DataChunk &args, ExpressionState &state, Vector &result); + + // Registration + static void Register(ExtensionLoader &loader); +}; + +} // namespace duckdb diff --git a/src/include/ml_functions/decision_forest.hpp b/src/include/ml_functions/decision_forest.hpp new file mode 100644 index 0000000..361d7c6 --- /dev/null +++ b/src/include/ml_functions/decision_forest.hpp @@ -0,0 +1,78 @@ +#pragma once + +#include "duckdb.hpp" +#include "duckdb/function/function.hpp" +#include "ml_functions/decision_tree.hpp" + +#include +#include +#include + +namespace duckdb { + +struct ML_DecisionForest { + // Max number of trees (same spirit as Velox code) + static constexpr idx_t MAX_NUM_TREES = 1600; + + // ------------------------- + // Forest implementation + // ------------------------- + class Forest { + public: + // We reuse your Tree implementation from ML_DecisionTree + std::vector trees; + bool is_classification; + + Forest() : is_classification(false) { + } + + Forest(const std::string &folder_path, bool is_cls); + + // Build forest from a folder of .dot tree files + void ConstructFromFolder(const std::string &folder_path); + + // For explicit list of tree paths (if ever needed) + void ConstructFromPaths(const std::vector &tree_paths); + + // Predict for a single row (features is a pointer to length num_features) + float PredictRow(const float *row_base, idx_t num_features) const; + }; + + using ForestPtr = std::shared_ptr; + + // ------------------------- + // Bind data + // ------------------------- + struct BindData final : public FunctionData { + ForestPtr forest; + idx_t num_features; + std::string forest_path; + bool is_classification; + + BindData(); + BindData(const ForestPtr &f, + idx_t num_features_p, + const std::string &forest_path_p, + bool is_classification_p); + + unique_ptr Copy() const override; + bool Equals(const FunctionData &other_p) const override; + }; + + // ------------------------- + // Scalar function hooks + // ------------------------- + static unique_ptr Bind( + ClientContext &context, + ScalarFunction &bound_function, + vector> &args); + + static void Execute( + DataChunk &args, + ExpressionState &state, + Vector &result); + + static void Register(ExtensionLoader &loader); +}; + +} // namespace duckdb diff --git a/src/include/ml_functions/decision_tree.hpp b/src/include/ml_functions/decision_tree.hpp new file mode 100644 index 0000000..d3623da --- /dev/null +++ b/src/include/ml_functions/decision_tree.hpp @@ -0,0 +1,91 @@ +/* + * Decision tree inference UDF for DuckDB (tree_predict) + * + * Layout: + * - ML_DecisionTree::Node / Tree: in-memory decision tree + * - ML_DecisionTree::BindData: per-binding state + tree cache + * - ML_DecisionTree::Bind / Execute / Register: UDF wiring + */ + +#pragma once + +#include "duckdb.hpp" +#include +#include +#include +#include + +namespace duckdb { + +class ML_DecisionTree { +public: + // Match Velox side: enough for typical XGBoost dumps + static constexpr idx_t MAX_NUM_NODES_PER_TREE = 512; + + struct Node { + union { + float threshold; + float leaf_value; + }; + int32_t index_id; + int32_t left_child; + int32_t right_child; + bool is_leaf; + bool is_miss_track_left; + + Node() + : threshold(0.0f), index_id(-1), left_child(-1), right_child(-1), + is_leaf(true), is_miss_track_left(true) { + } + }; + + class Tree { + public: + Node nodes[MAX_NUM_NODES_PER_TREE]; + int32_t tree_id; + + Tree(); + Tree(int32_t id, const std::string &tree_path); + + static void ConstructTreeFromPath(const std::string &tree_path, Node *nodes); + + static void ConstructTreeFromPathHelper(const std::string &tree_path, + std::vector &relationships, + std::vector &inner_nodes, + std::vector &leaf_nodes); + + static void ProcessInnerNodes(std::vector &inner_nodes, Node *nodes); + static void ProcessLeafNodes(std::vector &leaf_nodes, Node *nodes); + static void ProcessRelationships(std::vector &relationships, Node *nodes); + + float PredictSingle(const float *row_base) const; + float PredictSingleMissing(const float *row_base) const; + }; + + using TreePtr = std::shared_ptr; + + struct BindData : public FunctionData { + bool has_missing_default; + std::unordered_map tree_cache; + + BindData(); + explicit BindData(bool hm); + + unique_ptr Copy() const override; + bool Equals(const FunctionData &other) const override; + }; + +public: + // Bind: validates argument types + captures constant default for has_missing if present. + static unique_ptr Bind(ClientContext &context, + ScalarFunction &bound_function, + vector> &args); + + // Execute: evaluates tree_predict over a LIST(FLOAT) + path + has_missing. + static void Execute(DataChunk &args, ExpressionState &state, Vector &result); + + // Register: registers "tree_predict" into the extension. + static void Register(ExtensionLoader &loader); +}; + +} // namespace duckdb diff --git a/src/include/ml_functions/matrix_addition.hpp b/src/include/ml_functions/matrix_addition.hpp new file mode 100644 index 0000000..5fc25de --- /dev/null +++ b/src/include/ml_functions/matrix_addition.hpp @@ -0,0 +1,107 @@ +#pragma once + +#include +#include +#include + +#include "duckdb.hpp" // umbrella header (brings in FunctionData, ScalarFunction, DataChunk, etc.) +#include "duckdb/function/function.hpp" + +#include +#include +#include + +namespace duckdb { + struct ML_MatrixAddition { + struct BindData final : public FunctionData { + vector biases; // length n + idx_t n = 0; // length of biases + std::string biases_file; + bool has_bound_biases = false; + BindData() = default; + BindData(const vector &b, idx_t n_, + const std::string &file, bool bound) + : biases(b), n(n_), biases_file(file), has_bound_biases(bound) {} + unique_ptr Copy() const override { + // deep copy (vector/string copy is deep) + return make_uniq(biases, n, biases_file, has_bound_biases); + } + bool Equals(const FunctionData &other_p) const override { + auto &o = other_p.Cast(); + return has_bound_biases == o.has_bound_biases && + n == o.n && + biases_file == o.biases_file && + biases == o.biases; + } + }; + + static unique_ptr Bind(ClientContext &context, + ScalarFunction &bound_function, + vector> &args); + static void Execute(DataChunk &args, ExpressionState &state, Vector &result); + static void Register(ExtensionLoader &loader); + + static void ParseCSVBiases(const std::string &path, + std::vector &b, + idx_t &n) { + std::ifstream f(path); + if (!f.is_open()) { + throw InvalidInputException("mat_add: cannot open biases file: %s", path.c_str()); + } + + b.clear(); + n = 0; + + std::string line; + // read first non-empty line + while (std::getline(f, line)) { + // trim whitespace + auto l = line.find_first_not_of(" \t\r\n"); + if (l == std::string::npos) continue; // skip blank line + auto r = line.find_last_not_of(" \t\r\n"); + line = line.substr(l, r - l + 1); + break; + } + + if (line.empty()) { + throw InvalidInputException("mat_add: empty biases file: %s", path.c_str()); + } + + std::stringstream ss(line); + std::string cell; + while (std::getline(ss, cell, ',')) { + // trim token + auto l = cell.find_first_not_of(" \t\r\n"); + auto r = cell.find_last_not_of(" \t\r\n"); + std::string tok = (l == std::string::npos) ? "" : cell.substr(l, r - l + 1); + if (tok.empty()) { + throw InvalidInputException("mat_add: empty field in biases file: %s", path.c_str()); + } + try { + b.emplace_back(std::stod(tok)); + } catch (...) { + throw InvalidInputException("mat_add: non-numeric bias '%s' in %s", + tok.c_str(), path.c_str()); + } + } + + if (b.empty()) { + throw InvalidInputException("mat_add: no values found in %s", path.c_str()); + } + + n = static_cast(b.size()); + + // if there are more non-empty lines, treat as error (we expect exactly one row) + std::string extra; + while (std::getline(f, extra)) { + auto l2 = extra.find_first_not_of(" \t\r\n"); + if (l2 != std::string::npos) { + throw InvalidInputException("mat_add: multiple rows found in %s; expected a single 1xN row", + path.c_str()); + } + } + } + + + }; +} // namespace duckdb end \ No newline at end of file diff --git a/src/include/ml_functions/matrix_multiplication.hpp b/src/include/ml_functions/matrix_multiplication.hpp new file mode 100644 index 0000000..7e7e1cc --- /dev/null +++ b/src/include/ml_functions/matrix_multiplication.hpp @@ -0,0 +1,131 @@ +#pragma once + +#include +#include +#include + +#include "duckdb.hpp" // umbrella header (brings in FunctionData, ScalarFunction, DataChunk, etc.) +#include "duckdb/function/function.hpp" + +#include +#include +#include + +namespace duckdb { + +struct ML_MatrixMultiply { + // TODO + // 2 and 3, only Dense and InputSparse supported + enum class MulMode : uint8_t { Dense = 0, InputSparse = 1, WeightSparse = 2, BothSparse = 3 }; + + // ------------------------- + // Bind Data + // ------------------------- + struct BindData final : public FunctionData { + std::vector weights; // row-major [k * n] + idx_t k = 0; // rows in weights (input length) + idx_t n = 0; // cols in weights (output length) + std::string weights_file; + bool has_bound_weights = false; + + // MatMul Mode + MulMode mode = MulMode::Dense; + + BindData() = default; + BindData(const std::vector &w, idx_t k_, idx_t n_, + const std::string &file, bool bound, MulMode mode_); + + unique_ptr Copy() const override; + bool Equals(const FunctionData &other) const override; + }; + + // ------------------------- + // Registration + // ------------------------- + static void Register(ExtensionLoader &loader); + + // ------------------------- + // DuckDB hooks + // ------------------------- + static unique_ptr Bind(ClientContext &context, + ScalarFunction &bound_function, + vector> &args); + + static void Execute(DataChunk &args, ExpressionState &state, Vector &result); + + // ------------------------- + // Helpers + // ------------------------- + static void ParseCSVWeights(const std::string &path, + std::vector &W, + idx_t &k, + idx_t &n){ + std::ifstream f(path); + if (!f.is_open()) { + throw InvalidInputException("matmul: cannot open weights file: %s", path.c_str()); + } + W.clear(); + k = 0; + + bool have_n = false; // <- replace the sentinel + n = 0; + + std::string line; + while (std::getline(f, line)) { + if (line.empty()) continue; + + std::stringstream ss(line); + std::string cell; + idx_t col_count = 0; + + while (std::getline(ss, cell, ',')) { + auto l = cell.find_first_not_of(" \t\r\n"); + auto r = cell.find_last_not_of(" \t\r\n"); + std::string trimmed = (l == std::string::npos) ? "" : cell.substr(l, r - l + 1); + + if (trimmed.empty()) { + throw InvalidInputException("matmul: empty field in weights file: %s", path.c_str()); + } + try { + W.emplace_back(std::stod(trimmed)); + } catch (...) { + throw InvalidInputException("matmul: non-numeric weight '%s' in %s", + trimmed.c_str(), path.c_str()); + } + col_count++; + } + + if (col_count == 0) continue; + + if (!have_n) { + n = col_count; + have_n = true; + } else if (n != col_count) { + throw InvalidInputException( + "matmul: ragged row in weights file %s (got %llu cols, expected %llu)", + path.c_str(), + (unsigned long long)col_count, + (unsigned long long)n + ); + } + k++; + } + + if (k == 0 || !have_n) { + throw InvalidInputException("matmul: empty weights file: %s", path.c_str()); + } + if (W.size() != (size_t)(k * n)) { + throw InternalException("matmul: internal size mismatch after parsing %s", path.c_str()); + } + } // end of ParseCSVWeights + + static ML_MatrixMultiply::MulMode ParseModeValue(const Value &v) { + auto s = StringUtil::Lower(v.ToString()); + if (s == "dense") return ML_MatrixMultiply::MulMode::Dense; + if (s == "sparse") return ML_MatrixMultiply::MulMode::InputSparse; + throw BinderException("matmul: unknown mode '%s' (expected 'dense' or 'input_sparse')", s.c_str()); + } + +}; + +} // namespace duckdb diff --git a/src/include/ml_functions/ml_decision_forest_table.hpp b/src/include/ml_functions/ml_decision_forest_table.hpp new file mode 100644 index 0000000..a24f0c0 --- /dev/null +++ b/src/include/ml_functions/ml_decision_forest_table.hpp @@ -0,0 +1,73 @@ +#pragma once + +#include "duckdb.hpp" +#include "duckdb/function/table_function.hpp" +#include "duckdb/common/vector.hpp" + +#include +#include + +namespace duckdb { + +struct ML_ForestTable { + // ------------------------- + // Global scan state + // ------------------------- + struct GlobalState : public GlobalTableFunctionState { + vector tree_paths; + idx_t current_row; + idx_t total_rows; + + GlobalState() : current_row(0), total_rows(0) { + } + + idx_t MaxThreads() const override { + // Single threaded scan is fine for this helper + return 1; + } + }; + + // ------------------------- + // Bind data + // ------------------------- + struct BindData final : public FunctionData { + std::string forest_path; + vector tree_paths; + + BindData() = default; + BindData(std::string forest_path_p, vector tree_paths_p) + : forest_path(std::move(forest_path_p)), tree_paths(std::move(tree_paths_p)) { + } + + unique_ptr Copy() const override { + return make_uniq(forest_path, tree_paths); + } + + bool Equals(const FunctionData &other_p) const override { + auto &o = other_p.Cast(); + return forest_path == o.forest_path && tree_paths == o.tree_paths; + } + }; + + // ------------------------- + // Hook signatures + // ------------------------- + static unique_ptr Bind( + ClientContext &context, + TableFunctionBindInput &input, + duckdb::vector &return_types, + duckdb::vector &names); + + static unique_ptr InitGlobal( + ClientContext &context, + TableFunctionInitInput &input); + + static void Function( + ClientContext &context, + TableFunctionInput &data, + DataChunk &output); + + static void Register(ExtensionLoader &loader); +}; + +} // namespace duckdb diff --git a/src/include/ml_functions/one_hot_encoder.hpp b/src/include/ml_functions/one_hot_encoder.hpp new file mode 100644 index 0000000..17b8c8e --- /dev/null +++ b/src/include/ml_functions/one_hot_encoder.hpp @@ -0,0 +1,82 @@ +#pragma once + +#include "duckdb.hpp" +#include "duckdb/function/function.hpp" +#include +#include +#include +#include + +namespace duckdb { + +// --------------------------- +// Helper Definition +// --------------------------- +static void one_hot_encoder_bind_constructor( + const std::string path, + bool is_string, + std::map &str_map, + std::map &int_map +) { + std::ifstream in(path); + if (!in.is_open()) { + throw std::runtime_error("OneHot mapping file not found: " + path); + } + + std::string line; + std::getline(in, line); // optional header skip + + if (is_string) { + while (std::getline(in, line)) { + std::istringstream ss(line); + std::string valStr, posStr; + if (!std::getline(ss, valStr, ',')) continue; + if (!std::getline(ss, posStr, ',')) continue; + int position = std::stoi(posStr); + str_map[valStr] = position; + } + } else { + while (std::getline(in, line)) { + std::istringstream ss(line); + std::string valStr, posStr; + if (!std::getline(ss, valStr, ',')) continue; + if (!std::getline(ss, posStr, ',')) continue; + long long value = std::stoll(valStr); + int position = std::stoi(posStr); + int_map[(duckdb::idx_t)value] = position; + } + } +} + +struct ML_OneHotEncoder { + // ------------------------- + // BindData + // ------------------------- + struct BindData final : public FunctionData { + std::map str_to_index; + std::map int_to_index; + bool is_string; + + BindData() = default; + BindData(const std::map &s2i, + const std::map &i2i, + bool str) + : str_to_index(s2i), int_to_index(i2i), is_string(str) {} + + unique_ptr Copy() const override; + bool Equals(const FunctionData &other_p) const override; + }; + + // ------------------------- + // Lifecycle + // ------------------------- + static unique_ptr Bind(ClientContext &context, + ScalarFunction &bound_function, + vector> &args); + + static void Execute(DataChunk &args, ExpressionState &state, Vector &result); + + static void Register(ExtensionLoader &loader); +}; + +} // namespace duckdb diff --git a/src/include/ml_functions/sigmoid.hpp b/src/include/ml_functions/sigmoid.hpp new file mode 100644 index 0000000..46c4a92 --- /dev/null +++ b/src/include/ml_functions/sigmoid.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include +#include +#include + +#include "duckdb.hpp" +#include "duckdb/function/function.hpp" + +namespace duckdb { + +struct ML_Sigmoid { + static unique_ptr Bind(ClientContext &context, + ScalarFunction &bound_function, + vector> &args); + + static void Execute(DataChunk &args, ExpressionState &state, Vector &result); + + static void Register(ExtensionLoader &loader); +}; + +} // namespace duckdb diff --git a/src/include/ml_functions/softmax.hpp b/src/include/ml_functions/softmax.hpp new file mode 100644 index 0000000..9f66e02 --- /dev/null +++ b/src/include/ml_functions/softmax.hpp @@ -0,0 +1,25 @@ +#pragma once + +#include +#include +#include + +#include "duckdb.hpp" +#include "duckdb/function/function.hpp" + +namespace duckdb { + +struct ML_Softmax { + // No bind state required + static unique_ptr Bind(ClientContext &context, + ScalarFunction &bound_function, + vector> &args); + + // Execute softmax over each LIST row + static void Execute(DataChunk &args, ExpressionState &state, Vector &result); + + // Register the function with a unique name to avoid collisions + static void Register(ExtensionLoader &loader); +}; + +} // namespace duckdb diff --git a/src/include/ml_functions/torch_dnn.hpp b/src/include/ml_functions/torch_dnn.hpp new file mode 100644 index 0000000..dc7af54 --- /dev/null +++ b/src/include/ml_functions/torch_dnn.hpp @@ -0,0 +1,91 @@ +#pragma once + +#include "duckdb.hpp" +#include "duckdb/function/function.hpp" + +// libtorch +#include + +#include +#include +#include + +namespace duckdb { + +struct ML_TorchDNN { + // Must match what the .cpp expects + enum class KernelType { + MatMul, + MatAdd, + BatchNorm, + ReLU, + Sigmoid, + Softmax, + Argmax + }; + + struct BindData final : public FunctionData { + // Fully materialized model (ModuleHolder) + torch::nn::Sequential model; + // dims[0] = input_dim, dims.back() = output_dim + std::vector dims; + bool has_argmax; + // Optional metadata-only field; may be empty for inline/fused cases + std::string model_file; + + BindData(torch::nn::Sequential model_p, + std::vector dims_p, + bool has_argmax_p, + std::string file = std::string()) + : model(std::move(model_p)), + dims(std::move(dims_p)), + has_argmax(has_argmax_p), + model_file(std::move(file)) { + } + + unique_ptr Copy() const override { + // torch::nn::Sequential (ModuleHolder) is copyable + return make_uniq(model, dims, has_argmax, model_file); + } + + bool Equals(const FunctionData &other_p) const override { + auto &o = other_p.Cast(); + // We do not attempt deep model-graph equality; config is enough. + return has_argmax == o.has_argmax && + dims == o.dims && + model_file == o.model_file; + } + }; + + // NOTE: For the "file-based" mode, LoadModelFromFile maps a model spec + // string to kernel_types + dims + weight_tensors. + // For fused/inline mode, a rewrite rule can skip this entirely and build + // an ML_TorchDNN::BindData directly from UDF bind_data. + static void LoadModelFromFile( + const std::string &path, + std::vector &kernel_types, + std::vector &dims, + std::vector &weight_tensors, + bool &has_argmax); + + // Default binder: expects last argument to be a constant model_file. + static unique_ptr Bind( + ClientContext &context, + ScalarFunction &bound_function, + vector> &args); + + // Execute on LIST input, returning LIST or INTEGER + // depending on has_argmax + return type. + static void Execute( + DataChunk &args, + ExpressionState &state, + Vector &result); + + // Register three overloads: + // torchdnn(LIST, model_file VARCHAR) -> LIST + // torchdnn(LIST, k INTEGER, model_file VARCHAR) -> INTEGER + // torchdnn(LIST, k BIGINT, model_file VARCHAR) -> INTEGER + static void Register(ExtensionLoader &loader); +}; + +} // namespace duckdb diff --git a/src/include/optimization/actions/DecisionForestUDF2RelationRewriteAction.hpp b/src/include/optimization/actions/DecisionForestUDF2RelationRewriteAction.hpp new file mode 100644 index 0000000..2e506f7 --- /dev/null +++ b/src/include/optimization/actions/DecisionForestUDF2RelationRewriteAction.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include "duckdb.hpp" + +namespace duckdb { + +struct DecisionForestUDF2RelationRewriteAction { + // 1) Check if this rule *could* apply on this plan (pattern exists) + static bool check(OptimizerExtensionInput &input, + unique_ptr &plan); + + // 2) Actually apply the rewrite + static bool apply(OptimizerExtensionInput &input, + unique_ptr &plan); +}; + +} // namespace duckdb diff --git a/src/include/optimization/actions/MLDecompositionPushdownRewriteAction.hpp b/src/include/optimization/actions/MLDecompositionPushdownRewriteAction.hpp new file mode 100644 index 0000000..956ab63 --- /dev/null +++ b/src/include/optimization/actions/MLDecompositionPushdownRewriteAction.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include "duckdb.hpp" + +namespace duckdb { + +// Push mat_mul(...) closer to the base table: +// - Simple case: PROJECTION(mat_mul(...)) over a single child -> insert a +// projection below that projection that computes mat_mul, and replace the +// expression with a BoundReferenceExpression. +// - Join case: PROJECTION(mat_mul(...), other_cols) over a COMPARISON_JOIN: +// if mat_mul depends on exactly one table (say the left input), insert a +// projection on that join child that computes mat_mul, and replace the +// expression in the top projection with a reference to the new join output +// column. +class MLDecompositionPushdownRewriteAction { +public: + static bool check(OptimizerExtensionInput &input, unique_ptr &plan); + static bool apply(OptimizerExtensionInput &input, unique_ptr &plan); +}; + +} // namespace duckdb diff --git a/src/include/optimization/actions/MatMulDense2SparseRewriteAction.hpp b/src/include/optimization/actions/MatMulDense2SparseRewriteAction.hpp new file mode 100644 index 0000000..a08a6cf --- /dev/null +++ b/src/include/optimization/actions/MatMulDense2SparseRewriteAction.hpp @@ -0,0 +1,10 @@ +#include "duckdb.hpp" + +namespace duckdb { + struct MatMulDense2SparseRewriteAction { + // It needs to have a check method and a rewrite method + static bool check(OptimizerExtensionInput &input, unique_ptr &plan); + + static bool apply(OptimizerExtensionInput &input, unique_ptr &plan); + }; +} \ No newline at end of file diff --git a/src/include/optimization/actions/MultiLayerUDF2TorchNNRewriteAction.hpp b/src/include/optimization/actions/MultiLayerUDF2TorchNNRewriteAction.hpp new file mode 100644 index 0000000..0e81809 --- /dev/null +++ b/src/include/optimization/actions/MultiLayerUDF2TorchNNRewriteAction.hpp @@ -0,0 +1,18 @@ +// optimization/actions/MultiLayerUDF2TorchNNRewriteAction.hpp +#pragma once + +#include "duckdb.hpp" + +namespace duckdb { + +struct MultiLayerUDF2TorchNNRewriteAction { + // 1) Check if this rule *could* apply on this plan (pattern exists) + static bool check(OptimizerExtensionInput &input, + unique_ptr &plan); + + // 2) Actually apply the rewrite (replace mat_add(mat_mul(...)) with torchdnn) + static bool apply(OptimizerExtensionInput &input, + unique_ptr &plan); +}; + +} // namespace duckdb diff --git a/src/include/optimization/add5_rule.hpp b/src/include/optimization/add5_rule.hpp new file mode 100644 index 0000000..9dc6e14 --- /dev/null +++ b/src/include/optimization/add5_rule.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "duckdb.hpp" + +struct Add5Rule { + static void Apply(duckdb::ClientContext &, + duckdb::unique_ptr &plan); +}; \ No newline at end of file diff --git a/src/include/optimization/catalog.hpp b/src/include/optimization/catalog.hpp new file mode 100644 index 0000000..0ef21a3 --- /dev/null +++ b/src/include/optimization/catalog.hpp @@ -0,0 +1,19 @@ +#include "duckdb.hpp" +#include + +namespace duckdb { + + class Catalog_optimization { + public: + Catalog_optimization() = default; + idx_t cacheQueryPlan(ClientContext &ctx,unique_ptr &plan); + unique_ptr resetQueryPlanFromCache(ClientContext &ctx,idx_t cacheId); + unique_ptr resetQueryPlan(ClientContext &ctx); + + + private: + std::map > queryPlanCaches_; + idx_t initQueryPlanCacheId = 0; + idx_t currentCacheId = 0; + }; //end of class Catalog +} //end of duckdb namespace \ No newline at end of file diff --git a/src/include/optimization/dynamic_optimizer.hpp b/src/include/optimization/dynamic_optimizer.hpp new file mode 100644 index 0000000..60beb25 --- /dev/null +++ b/src/include/optimization/dynamic_optimizer.hpp @@ -0,0 +1,7 @@ +#include "duckdb.hpp" + +namespace duckdb{ + struct DynamicOpt { + static void Optimize(OptimizerExtensionInput &input, duckdb::unique_ptr &plan); + }; +} \ No newline at end of file diff --git a/src/include/optimization/feature_extraction/query2vec/query_annotation.hpp b/src/include/optimization/feature_extraction/query2vec/query_annotation.hpp new file mode 100644 index 0000000..9a6c4d1 --- /dev/null +++ b/src/include/optimization/feature_extraction/query2vec/query_annotation.hpp @@ -0,0 +1,24 @@ +#pragma once +#include +#include +#include +#include "duckdb.hpp" +#include "duckdb/planner/logical_operator.hpp" + +namespace duckdb { + +struct AnnotatedPlanNode { + idx_t node_id{0}; + const LogicalOperator *plan_node{nullptr}; // non-owning pointer + std::vector> children; + vector features; // placeholder for extracted features +}; + +// Call with: AnnotatePlanTree(*logical_plan_unique_ptr); +std::unique_ptr AnnotatePlanTree(const LogicalOperator &root); + +void DisplayAnnotatedPlan(const AnnotatedPlanNode &node, idx_t depth = 0); + +AnnotatedPlanNode* FindNodeById(AnnotatedPlanNode &node, idx_t target_id); + +} // namespace duckdb diff --git a/src/include/optimization/feature_extraction/query2vec/query_feature_extraction.hpp b/src/include/optimization/feature_extraction/query2vec/query_feature_extraction.hpp new file mode 100644 index 0000000..737f0c6 --- /dev/null +++ b/src/include/optimization/feature_extraction/query2vec/query_feature_extraction.hpp @@ -0,0 +1,16 @@ +#include "duckdb.hpp" +#include "optimization/feature_extraction/query2vec/query_annotation.hpp" + +namespace duckdb{ + + std::string filter_feature_extract(const AnnotatedPlanNode &node); + std::string get_filters_feature_extract(const AnnotatedPlanNode &node); + std::string get_table_feature_extract(const AnnotatedPlanNode &node); + std::string aggregate_feature_extract(const AnnotatedPlanNode &node); + std::string join_features_extract(const AnnotatedPlanNode &node); + + void traverse_and_extract(const AnnotatedPlanNode &node); + + void extract_features(const AnnotatedPlanNode &node); + +} \ No newline at end of file diff --git a/src/include/optimization/metric_catalog.hpp b/src/include/optimization/metric_catalog.hpp new file mode 100644 index 0000000..6f90436 --- /dev/null +++ b/src/include/optimization/metric_catalog.hpp @@ -0,0 +1,103 @@ +#pragma once + +#include "duckdb/common/types.hpp" +#include "duckdb/common/unordered_map.hpp" +#include "duckdb/common/string.hpp" + +namespace duckdb { + +// ================================ +// Schema-level metric container +// ================================ + +struct SchemaMetricSet { + // metric_name -> value + unordered_map values; + + bool Has(const string &metric_name) const { + return values.find(metric_name) != values.end(); + } + + // Get metric value or default if not present + double Get(const string &metric_name, double default_value = 0.0) const { + auto it = values.find(metric_name); + if (it == values.end()) { + return default_value; + } + return it->second; + } +}; + +// ================================ +// MetricCatalog: schema -> metrics +// ================================ + +// Simple singleton catalog hard-coded in the extension. +// Granularity: schema. Each schema has a map. +class MetricCatalog { +public: + // Global singleton instance + static const MetricCatalog &Get() { + static MetricCatalog instance; + return instance; + } + + // Get metrics for a specific schema (nullptr if not found) + const SchemaMetricSet *GetSchemaMetrics(const string &schema_name) const { + auto it = schema_metrics.find(schema_name); + if (it == schema_metrics.end()) { + return nullptr; + } + return &it->second; + } + + // Convenience: get single metric for a schema, with default + double GetMetric(const string &schema_name, + const string &metric_name, + double default_value = 0.0) const { + auto schema = GetSchemaMetrics(schema_name); + if (!schema) { + return default_value; + } + return schema->Get(metric_name, default_value); + } + +private: + MetricCatalog() { + Initialize(); + } + + void Initialize() { + // ---- Define metrics for each synthetic schema here ---- + // Example schema 1: synthetic dense features + { + SchemaMetricSet s; + s.values["join_ratio"] = 5e-5; // big joins + s.values["sparsity"] = 0.20; // 95% non-zero -> dense + schema_metrics["synthetic"] = std::move(s); + } + + // Example schema 2: synthetic sparse features + { + SchemaMetricSet s; + s.values["join_ratio"] = 5e-5; + s.values["sparsity"] = 0.10; // 10% non-zero -> sparse + schema_metrics["syn_sparse"] = std::move(s); + } + + // Add more schemas as needed: + // { + // SchemaMetricSet s; + // s.values["join_ratio"] = ...; + // s.values["sparsity"] = ...; + // s.values["some_other_metric"] = ...; + // schema_metrics["my_schema_name"] = std::move(s); + // } + } + +private: + // schema_name -> SchemaMetricSet + unordered_map schema_metrics; +}; + +} // namespace duckdb diff --git a/src/include/optimization/plan_dump.hpp b/src/include/optimization/plan_dump.hpp new file mode 100644 index 0000000..4df9071 --- /dev/null +++ b/src/include/optimization/plan_dump.hpp @@ -0,0 +1,67 @@ +#pragma once +#include "duckdb/common/serializer/binary_serializer.hpp" +#include "duckdb/common/serializer/buffered_file_writer.hpp" +#include "duckdb/common/file_system.hpp" +#include "duckdb/planner/logical_operator.hpp" +#include "duckdb/main/client_context.hpp" +// #include "duckdb/common/base64.hpp" +#include "duckdb/common/serializer/memory_stream.hpp" +#include "duckdb/common/types/blob.hpp" +#include "duckdb/common/types/string_type.hpp" + + +#include +#include + +namespace duckdb { + +inline void DumpPlanToFile(duckdb::ClientContext &ctx, + LogicalOperator &op, + const std::string &path) { + + // open a buffered writer to the target file + auto &fs = FileSystem::GetFileSystem(ctx); + BufferedFileWriter writer(fs, path, FileFlags::FILE_FLAGS_WRITE | + FileFlags::FILE_FLAGS_FILE_CREATE | + FileFlags::FILE_FLAGS_APPEND); + + SerializationOptions options; + options.serialization_compatibility = SerializationCompatibility::Latest(); + + BinarySerializer serializer(writer, options); + serializer.Begin(); + op.Serialize(serializer); + serializer.End(); + + writer.Sync(); + writer.Close(); +} + +inline void DumpPlanBase64(duckdb::ClientContext &ctx, + duckdb::LogicalOperator &plan, + const std::string &txt_path) { + using namespace duckdb; + + // 1) serialize plan into memory + MemoryStream ms(Allocator::Get(ctx)); + SerializationOptions opts; + opts.serialization_compatibility = SerializationCompatibility::Latest(); + BinarySerializer ser(ms, opts); + ser.Begin(); + plan.Serialize(ser); + ser.End(); + + // 2) wrap the raw bytes in a string_t (no copy) + const_data_ptr_t bytes = ms.GetData(); + idx_t nbytes = ms.GetPosition(); // current size written + string_t blob_view(const_char_ptr_cast(bytes), nbytes); + + // 3) Base64 encode using Blob helpers + std::string b64 = Blob::ToBase64(blob_view); + + // 4) write text-safe output + std::ofstream out(txt_path, std::ios::out | std::ios::trunc); + out << b64; +} + +} // namespace duckdb diff --git a/src/include/optimization/register_optimizers.hpp b/src/include/optimization/register_optimizers.hpp new file mode 100644 index 0000000..c71f4f0 --- /dev/null +++ b/src/include/optimization/register_optimizers.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "duckdb.hpp" +#include "duckdb/main/database.hpp" + +namespace duckdb { +void RegisterCactusOptimizers(DatabaseInstance &db); +} // namespace duckdb diff --git a/src/include/optimization/replace_binding_rule.hpp b/src/include/optimization/replace_binding_rule.hpp new file mode 100644 index 0000000..d7308de --- /dev/null +++ b/src/include/optimization/replace_binding_rule.hpp @@ -0,0 +1,26 @@ +#pragma once +#include "duckdb.hpp" +#include +#include + +using namespace duckdb; + +struct ReplaceBindingRule { + static constexpr const char *kFuncName = "one_hot_encode"; + + std::string from_path; + std::string to_path; + std::optional require_third_bool; // std::nullopt = don't care + + ReplaceBindingRule(std::string fromp, std::string top, std::optional third = std::nullopt) + : from_path(std::move(fromp)), to_path(std::move(top)), require_third_bool(third) {} + + static void Optimize(OptimizerExtensionInput &input, + duckdb::unique_ptr &plan, + const ReplaceBindingRule &cfg); + +private: + static duckdb::unique_ptr RewriteExpr(duckdb::unique_ptr expr, + const ReplaceBindingRule &cfg); + static void RewriteOp(LogicalOperator &op, const ReplaceBindingRule &cfg); +}; diff --git a/src/include/optimization/rule_manager.hpp b/src/include/optimization/rule_manager.hpp new file mode 100644 index 0000000..c8da6cc --- /dev/null +++ b/src/include/optimization/rule_manager.hpp @@ -0,0 +1,16 @@ +#include "duckdb.hpp" + +namespace duckdb{ + enum class RewriteRule : uint8_t { + MatMulDense2Sparse = 0, + // NextRuleGoesHere, + }; + + // Human-readable name + inline const char *ToString(RewriteRule r) { + switch (r) { + case RewriteRule::MatMulDense2Sparse: return "MatMulDense2Sparse"; + default: return "Unknown"; + } + } +} //end of duckdb namespace \ No newline at end of file diff --git a/src/include/optimization/translator.hpp b/src/include/optimization/translator.hpp new file mode 100644 index 0000000..f4b714c --- /dev/null +++ b/src/include/optimization/translator.hpp @@ -0,0 +1,9 @@ +#include "duckdb.hpp" +#include "rule_manager.hpp" + +namespace duckdb{ + class translator { + public : + private : + }; +} \ No newline at end of file diff --git a/src/ml_functions/CMakeLists.txt b/src/ml_functions/CMakeLists.txt index 5d5c838..3d9987f 100644 --- a/src/ml_functions/CMakeLists.txt +++ b/src/ml_functions/CMakeLists.txt @@ -1,4 +1,19 @@ +set(ML_FUNCTIONS_SRCS + ${CMAKE_CURRENT_SOURCE_DIR}/ml_functions_scalar.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/matrix_multiplication.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/min_max_scaler.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/one_hot_encoder.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/decision_tree.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/matrix_addition.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/argmax.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/sigmoid.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/torch_dnn.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/decision_forest.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/ml_decision_forest_table.cpp +) + set(EXTENSION_SOURCES ${EXTENSION_SOURCES} - ${CMAKE_CURRENT_SOURCE_DIR}/ml_functions_scalar.cpp + ${ML_FUNCTIONS_SRCS} PARENT_SCOPE) diff --git a/src/ml_functions/argmax.cpp b/src/ml_functions/argmax.cpp new file mode 100644 index 0000000..a58751a --- /dev/null +++ b/src/ml_functions/argmax.cpp @@ -0,0 +1,103 @@ +#include "ml_functions/argmax.hpp" + +#include + +#include "duckdb/common/exception.hpp" +#include "duckdb/common/types.hpp" +#include "duckdb/common/vector.hpp" +#include "duckdb/execution/expression_executor.hpp" +#include "duckdb/function/scalar_function.hpp" +#include "duckdb/parser/parsed_data/create_scalar_function_info.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" + +namespace duckdb { + +// ---------------- Bind ---------------- +unique_ptr ML_Argmax::Bind(ClientContext &context, + ScalarFunction &bound_function, + vector> &args) { + if (args.size() != 1) { + throw BinderException("argmax: expects exactly 1 argument: LIST"); + } + // If you want to enforce the element type strictly at bind time: + // if (args[0]->return_type != LogicalType::LIST(LogicalType::DOUBLE)) { ... } + return nullptr; // no bind info needed +} + +// ---------------- Execute ---------------- +void ML_Argmax::Execute(DataChunk &args, ExpressionState &state, Vector &result) { + const idx_t count = args.size(); + + // Input: LIST + Vector &in = args.data[0]; + UnifiedVectorFormat in_uvf; + in.ToUnifiedFormat(count, in_uvf); + const auto *entries = UnifiedVectorFormat::GetData(in_uvf); + + // Child payload for the list (flat array of doubles) + auto &child = ListVector::GetEntry(in); + const idx_t total_elems = ListVector::GetListSize(in); + UnifiedVectorFormat child_uvf; + child.ToUnifiedFormat(total_elems, child_uvf); + const double *child_data = UnifiedVectorFormat::GetData(child_uvf); + const auto &child_valid = child_uvf.validity; + + // Output: INTEGER + result.SetVectorType(VectorType::FLAT_VECTOR); + auto *out = FlatVector::GetData(result); + auto &out_valid = FlatVector::Validity(result); + + for (idx_t i = 0; i < count; ++i) { + const idx_t sel = in_uvf.sel->get_index(i); + + // NULL row → NULL result + if (!in_uvf.validity.RowIsValid(sel)) { + out_valid.Set(i, false); + continue; + } + + const auto le = entries[sel]; + const idx_t len = le.length; + const idx_t off = le.offset; + + if (len == 0) { + throw InvalidInputException("argmax: empty list at row %llu", + (unsigned long long)i); + } + + // Disallow NULL elements inside the list (strict mode) + if (!child_valid.AllValid()) { + for (idx_t j = 0; j < len; ++j) { + if (!child_valid.RowIsValid(off + j)) { + throw InvalidInputException( + "argmax: NULL element in list at row %llu (position %llu)", + (unsigned long long)i, (unsigned long long)j + ); + } + } + } + + // Map this row as an Eigen row-vector and get argmax + using RowVec = Eigen::Matrix; + Eigen::Map row(child_data + off, (Eigen::Index)len); + Eigen::Index max_row, max_col; + row.maxCoeff(&max_row, &max_col); // max_row will be 0 for a row-vector + + out[i] = static_cast(max_col); // 0-based index + } +} + +// ---------------- Register ---------------- +void ML_Argmax::Register(ExtensionLoader &loader) { + ScalarFunction fn( + "ml_argmax", + { LogicalType::LIST(LogicalType::DOUBLE) }, // x + LogicalType::INTEGER, // returns index + Execute + ); + fn.bind = Bind; + fn.stability = FunctionStability::VOLATILE; + loader.RegisterFunction(fn); +} + +} // namespace duckdb diff --git a/src/ml_functions/decision_forest.cpp b/src/ml_functions/decision_forest.cpp new file mode 100644 index 0000000..fbc8cde --- /dev/null +++ b/src/ml_functions/decision_forest.cpp @@ -0,0 +1,269 @@ +#include "ml_functions/decision_forest.hpp" + +#include "duckdb/common/exception.hpp" +#include "duckdb/common/list.hpp" +#include "duckdb/execution/expression_executor.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" + +#include +#include +#include +#include + +namespace duckdb { + +// ========================= +// Forest methods +// ========================= + +ML_DecisionForest::Forest::Forest(const std::string &folder_path, bool is_cls) + : is_classification(is_cls) { + ConstructFromFolder(folder_path); +} + +void ML_DecisionForest::Forest::ConstructFromFolder(const std::string &folder_path_in) { + std::filesystem::path folder(folder_path_in); + if (!std::filesystem::exists(folder) || !std::filesystem::is_directory(folder)) { + throw InvalidInputException( + "decision_forest_predict: forest folder '%s' does not exist or is not a directory", + folder_path_in.c_str()); + } + + std::vector tree_paths; + tree_paths.reserve(256); + + for (auto const &entry : std::filesystem::directory_iterator(folder)) { + if (entry.is_regular_file()) { + tree_paths.push_back(entry.path().string()); + } + } + + if (tree_paths.empty()) { + throw InvalidInputException( + "decision_forest_predict: forest folder '%s' contains no tree files", + folder_path_in.c_str()); + } + + // Sort for deterministic loading order + std::sort(tree_paths.begin(), tree_paths.end()); + + ConstructFromPaths(tree_paths); +} + +void ML_DecisionForest::Forest::ConstructFromPaths(const std::vector &tree_paths) { + if (tree_paths.size() > MAX_NUM_TREES) { + throw InvalidInputException( + "decision_forest_predict: forest has %llu trees, exceeds MAX_NUM_TREES=%llu", + (unsigned long long)tree_paths.size(), + (unsigned long long)MAX_NUM_TREES); + } + + trees.clear(); + trees.reserve(tree_paths.size()); + + // Reuse ML_DecisionTree::Tree for each file + for (idx_t i = 0; i < tree_paths.size(); ++i) { + const auto &path = tree_paths[i]; + trees.emplace_back((int)i, path); + } + + // Optional logging hook if you want + // std::cout << "[DecisionForest] Loaded " << trees.size() << " trees\n"; +} + +float ML_DecisionForest::Forest::PredictRow(const float *row_base, idx_t num_features) const { + if (trees.empty()) { + return 0.0f; + } + + double accumulated = 0.0; + for (auto const &tree : trees) { + // We use the non-missing variant (Velox forest didn't handle missing explicitly) + accumulated += (double)tree.PredictSingle(row_base); + } + + double avg = accumulated / (double)trees.size(); + + if (is_classification) { + return (avg > 0.0) ? 1.0f : 0.0f; + } + return static_cast(avg); +} + +// ========================= +// BindData methods +// ========================= + +ML_DecisionForest::BindData::BindData() + : num_features(0), forest_path(), is_classification(false) { +} + +ML_DecisionForest::BindData::BindData(const ForestPtr &f, + idx_t num_features_p, + const std::string &forest_path_p, + bool is_classification_p) + : forest(f), + num_features(num_features_p), + forest_path(forest_path_p), + is_classification(is_classification_p) { +} + +unique_ptr ML_DecisionForest::BindData::Copy() const { + // ForestPtr is shared_ptr; copying is cheap and safe (forest is immutable after load) + return make_uniq(forest, num_features, forest_path, is_classification); +} + +bool ML_DecisionForest::BindData::Equals(const FunctionData &other_p) const { + auto &o = other_p.Cast(); + // pointer equality is fine for forest; plus config fields + return forest.get() == o.forest.get() && + num_features == o.num_features && + forest_path == o.forest_path && + is_classification == o.is_classification; +} + +// ========================= +// Bind +// ========================= +// +// decision_forest_predict( +// features LIST(FLOAT), +// forest_folder VARCHAR, +// num_features INTEGER, +// is_classification BOOLEAN +// ) -> FLOAT +// +unique_ptr ML_DecisionForest::Bind( + ClientContext &context, + ScalarFunction &bound_function, + vector> &args) { + + if (args.size() != 4) { + throw InvalidInputException( + "decision_forest_predict expects 4 arguments: " + "features LIST(FLOAT), forest_folder VARCHAR, num_features INTEGER, is_classification BOOLEAN"); + } + + // 1) Check features type + const auto &feat_type = args[0]->return_type; + if (feat_type.id() != LogicalTypeId::LIST || + ListType::GetChildType(feat_type).id() != LogicalTypeId::FLOAT) { + throw InvalidInputException( + "decision_forest_predict: first argument must be LIST(FLOAT)"); + } + + // 2) Evaluate forest_folder, num_features, is_classification as constants + auto vForestFolder = ExpressionExecutor::EvaluateScalar(context, *args[1]); + auto vNumFeatures = ExpressionExecutor::EvaluateScalar(context, *args[2]); + auto vIsClass = ExpressionExecutor::EvaluateScalar(context, *args[3]); + + if (vForestFolder.IsNull()) { + throw InvalidInputException("decision_forest_predict: forest_folder cannot be NULL"); + } + if (vNumFeatures.IsNull()) { + throw InvalidInputException("decision_forest_predict: num_features cannot be NULL"); + } + if (vIsClass.IsNull()) { + throw InvalidInputException("decision_forest_predict: is_classification cannot be NULL"); + } + + const std::string forest_folder = vForestFolder.ToString(); + const idx_t num_features = (idx_t)vNumFeatures.GetValue(); + const bool is_classification = vIsClass.GetValue(); + + if (num_features == 0) { + throw InvalidInputException("decision_forest_predict: num_features must be > 0"); + } + + // 3) Load forest once at bind time + ForestPtr forest_ptr; + try { + forest_ptr = std::make_shared(forest_folder, is_classification); + } catch (const std::exception &e) { + throw InvalidInputException( + "decision_forest_predict: failed to load forest from '%s': %s", + forest_folder.c_str(), e.what()); + } + + return make_uniq(forest_ptr, num_features, forest_folder, is_classification); +} + +// ========================= +// Execute +// ========================= + +void ML_DecisionForest::Execute( + DataChunk &args, + ExpressionState &state, + Vector &result) { + + const idx_t count = args.size(); + if (count == 0) { + return; + } + + // Access bind data + const auto &func_expr = state.expr.Cast(); + auto &bd = func_expr.bind_info->Cast(); + const auto &forest = *bd.forest; + const idx_t num_features = bd.num_features; + + // features: LIST(FLOAT) + Vector &features_vec = args.data[0]; + auto &features_valid = FlatVector::Validity(features_vec); + auto *entries = FlatVector::GetData(features_vec); + + Vector &child = ListVector::GetEntry(features_vec); + auto &child_valid = FlatVector::Validity(child); + const float *vals = FlatVector::GetData(child); + + // Prepare result vector + result.SetVectorType(VectorType::FLAT_VECTOR); + auto out = FlatVector::GetData(result); + auto &out_valid = FlatVector::Validity(result); + + for (idx_t i = 0; i < count; ++i) { + if (!features_valid.RowIsValid(i)) { + out_valid.SetInvalid(i); + continue; + } + + const auto ent = entries[i]; + if (ent.length != num_features) { + throw InvalidInputException( + "decision_forest_predict: row %llu feature length %llu != num_features %llu", + (unsigned long long)i, + (unsigned long long)ent.length, + (unsigned long long)num_features); + } + + const float *row_base = vals + ent.offset; + + // (Velox forest didn’t implement explicit missing handling; we do the same here) + float pred = forest.PredictRow(row_base, num_features); + out[i] = pred; + } +} + +// ========================= +// Register +// ========================= + +void ML_DecisionForest::Register(ExtensionLoader &loader) { + ScalarFunction fn( + "decision_forest_predict", + { + LogicalType::LIST(LogicalType::FLOAT), // features + LogicalType::VARCHAR, // forest folder path + LogicalType::INTEGER, // num_features + LogicalType::BOOLEAN // is_classification + }, + LogicalType::FLOAT, // prediction + Execute); + + fn.bind = Bind; + fn.stability = FunctionStability::VOLATILE; + loader.RegisterFunction(fn); +} + +} // namespace duckdb diff --git a/src/ml_functions/decision_tree.cpp b/src/ml_functions/decision_tree.cpp new file mode 100644 index 0000000..80bddf5 --- /dev/null +++ b/src/ml_functions/decision_tree.cpp @@ -0,0 +1,420 @@ +//===----------------------------------------------------------------------===// +// decision_tree.cpp +// Implementation of tree_predict(features, tree_path, has_missing) +//===----------------------------------------------------------------------===// + +#include "ml_functions/decision_tree.hpp" + +#include "duckdb/common/exception.hpp" +#include "duckdb/common/list.hpp" +// #include "duckdb/common/types/unified_vector_format.hpp" +#include "duckdb/execution/expression_executor.hpp" +#include "duckdb/planner/expression/bound_constant_expression.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" + +#include +#include +#include +#include +#include + +namespace duckdb { + +// ========================= +// Tree implementation +// ========================= + +ML_DecisionTree::Tree::Tree() : tree_id(0) { + for (idx_t i = 0; i < MAX_NUM_NODES_PER_TREE; ++i) { + nodes[i] = Node(); + } +} + +ML_DecisionTree::Tree::Tree(int32_t id, const std::string &tree_path) : tree_id(id) { + for (idx_t i = 0; i < MAX_NUM_NODES_PER_TREE; ++i) { + nodes[i] = Node(); + } + ConstructTreeFromPath(tree_path, nodes); +} + +void ML_DecisionTree::Tree::ConstructTreeFromPath(const std::string &tree_path, Node *nodes) { + std::vector relationships; + std::vector inner_nodes; + std::vector leaf_nodes; + ConstructTreeFromPathHelper(tree_path, relationships, inner_nodes, leaf_nodes); + ProcessInnerNodes(inner_nodes, nodes); + ProcessLeafNodes(leaf_nodes, nodes); + ProcessRelationships(relationships, nodes); +} + +void ML_DecisionTree::Tree::ConstructTreeFromPathHelper( + const std::string &tree_path, + std::vector &relationships, + std::vector &inner_nodes, + std::vector &leaf_nodes) { + + std::ifstream input_file(tree_path); + if (!input_file.is_open()) { + throw InvalidInputException("DecisionTree: failed to open tree file '%s'", tree_path.c_str()); + } + + std::string line; + while (std::getline(input_file, line)) { + if (line.empty() || line.find("graph") != std::string::npos || line.find('}') != std::string::npos) { + continue; // skip header/footer/empty + } + if (line.find("->") != std::string::npos) { + relationships.push_back(line); + } else if (line.find("leaf") != std::string::npos) { + leaf_nodes.push_back(line); + } else if (line.find("label") != std::string::npos) { + inner_nodes.push_back(line); + } else { + // ignore + } + } + input_file.close(); +} + +void ML_DecisionTree::Tree::ProcessInnerNodes(std::vector &inner_nodes, Node *nodes) { + for (idx_t i = 0; i < inner_nodes.size(); ++i) { + const std::string ¤t = inner_nodes[i]; + int node_id; + int index_id; + float threshold; + + // node_id + auto pos_label = current.find("[ label"); + if (pos_label == std::string::npos) { + throw std::runtime_error("[DecisionTree] Error extracting inner node id"); + } + node_id = std::stoi(current.substr(4, pos_label - 1 - 4)); + if (node_id < 0 || (idx_t)node_id >= MAX_NUM_NODES_PER_TREE) { + throw std::runtime_error("[DecisionTree] Node id out of range in inner node"); + } + + // index_id + auto pos_f = current.find('f'); + auto pos_lt = current.find('<'); + if (pos_f == std::string::npos || pos_lt == std::string::npos || pos_f >= pos_lt) { + throw std::runtime_error("[DecisionTree] Error extracting inner node indexID"); + } + index_id = std::stoi(current.substr(pos_f + 1, pos_lt - pos_f - 1)); + + // threshold + auto pos_lt2 = current.find('<'); + auto pos_end = current.find("\" ]"); + if (pos_lt2 == std::string::npos || pos_end == std::string::npos || pos_lt2 >= pos_end) { + throw std::runtime_error("[DecisionTree] Error extracting inner node threshold"); + } + threshold = std::stof(current.substr(pos_lt2 + 1, pos_end - pos_lt2 - 1)); + + Node &n = nodes[node_id]; + n.index_id = index_id; + n.is_leaf = false; + n.left_child = -1; + n.right_child = -1; + n.threshold = threshold; + n.is_miss_track_left = false; // default: noMissing/right + } +} + +void ML_DecisionTree::Tree::ProcessLeafNodes(std::vector &leaf_nodes, Node *nodes) { + for (idx_t i = 0; i < leaf_nodes.size(); ++i) { + const std::string ¤t = leaf_nodes[i]; + int node_id; + float leaf_value = -1.0f; + + auto pos_bracket = current.find('['); + if (pos_bracket == std::string::npos) { + throw std::runtime_error("[DecisionTree] Error extracting leaf node id"); + } + node_id = std::stoi(current.substr(4, pos_bracket - 1 - 4)); + if (node_id < 0 || (idx_t)node_id >= MAX_NUM_NODES_PER_TREE) { + throw std::runtime_error("[DecisionTree] Node id out of range in leaf node"); + } + + auto pos_leaf = current.find("leaf="); + auto pos_end = current.find("\" ]"); + if (pos_leaf == std::string::npos || pos_end == std::string::npos || pos_leaf >= pos_end) { + throw std::runtime_error("[DecisionTree] Error extracting leaf node value"); + } + leaf_value = std::stof(current.substr(pos_leaf + 5, pos_end - 3 - (pos_leaf + 5))); + + Node &n = nodes[node_id]; + n.index_id = -1; + n.is_leaf = true; + n.left_child = -1; + n.right_child = -1; + n.leaf_value = leaf_value; + n.is_miss_track_left = true; // doesn't matter for leaf nodes + } +} + +void ML_DecisionTree::Tree::ProcessRelationships(std::vector &relationships, Node *nodes) { + for (idx_t i = 0; i < relationships.size(); ++i) { + const std::string ¤t = relationships[i]; + int parent_id; + int child_id; + + auto pos_arrow = current.find("->"); + if (pos_arrow == std::string::npos) { + throw std::runtime_error("[DecisionTree] Error extracting parentNodeID"); + } + parent_id = std::stoi(current.substr(4, pos_arrow - 1 - 4)); + if (parent_id < 0 || (idx_t)parent_id >= MAX_NUM_NODES_PER_TREE) { + throw std::runtime_error("[DecisionTree] Parent id out of range"); + } + + auto pos_bracket = current.find('['); + if (pos_bracket == std::string::npos || pos_bracket <= pos_arrow + 2) { + throw std::runtime_error("[DecisionTree] Error extracting childNodeID"); + } + child_id = std::stoi(current.substr(pos_arrow + 3, pos_bracket - 1 - (pos_arrow + 3))); + if (child_id < 0 || (idx_t)child_id >= MAX_NUM_NODES_PER_TREE) { + throw std::runtime_error("[DecisionTree] Child id out of range"); + } + + if (current.find("yes, missing") != std::string::npos) { + nodes[parent_id].is_miss_track_left = true; + } + + Node &p = nodes[parent_id]; + if (p.left_child == -1) { + p.left_child = child_id; + } else if (p.right_child == -1) { + p.right_child = child_id; + } else { + throw std::runtime_error("[DecisionTree] children updated more than twice"); + } + } +} + +float ML_DecisionTree::Tree::PredictSingle(const float *row_base) const { + int cur = 0; + while (!nodes[cur].is_leaf) { + const Node &n = nodes[cur]; + const float fv = row_base[n.index_id]; + cur = (fv < n.threshold) ? n.left_child : n.right_child; + } + return nodes[cur].leaf_value; +} + +float ML_DecisionTree::Tree::PredictSingleMissing(const float *row_base) const { + int cur = 0; + while (!nodes[cur].is_leaf) { + const Node &n = nodes[cur]; + const float fv = row_base[n.index_id]; + if (std::isnan(fv)) { + cur = n.is_miss_track_left ? n.left_child : n.right_child; + } else { + cur = (fv < n.threshold) ? n.left_child : n.right_child; + } + } + return nodes[cur].leaf_value; +} + +// ========================= +// BindData implementation +// ========================= + +ML_DecisionTree::BindData::BindData() : has_missing_default(false) { +} + +ML_DecisionTree::BindData::BindData(bool hm) : has_missing_default(hm) { +} + +unique_ptr ML_DecisionTree::BindData::Copy() const { + auto copy = make_uniq(has_missing_default); + // share the cache; trees are immutable & read-only + copy->tree_cache = tree_cache; + return std::move(copy); +} + +bool ML_DecisionTree::BindData::Equals(const FunctionData &other_p) const { + auto &other = other_p.Cast(); + // Only semantics-relevant field is the default flag; cache does not matter. + return has_missing_default == other.has_missing_default; +} + +// ========================= +// Bind +// ========================= + +unique_ptr ML_DecisionTree::Bind( + ClientContext & /*context*/, + ScalarFunction & /*bound_function*/, + vector> &args) { + + if (args.size() != 3) { + throw InvalidInputException( + "tree_predict expects 3 arguments: " + "features LIST(FLOAT), tree_file VARCHAR, has_missing BOOLEAN"); + } + + const auto &feat_type = args[0]->return_type; + if (feat_type.id() != LogicalTypeId::LIST || + ListType::GetChildType(feat_type).id() != LogicalTypeId::FLOAT) { + throw InvalidInputException("tree_predict: first argument must be LIST(FLOAT)"); + } + if (args[1]->return_type.id() != LogicalTypeId::VARCHAR) { + throw InvalidInputException("tree_predict: second argument must be VARCHAR (tree file path)"); + } + if (args[2]->return_type.id() != LogicalTypeId::BOOLEAN) { + throw InvalidInputException("tree_predict: third argument must be BOOLEAN (has_missing)"); + } + + // Default has_missing value: only when the 3rd arg is a constant. + bool has_missing_default = false; + if (args[2]->expression_class == ExpressionClass::BOUND_CONSTANT) { + auto &const_expr = args[2]->Cast(); + if (!const_expr.value.IsNull()) { + has_missing_default = const_expr.value.GetValue(); + } + } + + return make_uniq(has_missing_default); +} + +// ========================= +// Execute +// ========================= + +void ML_DecisionTree::Execute( + DataChunk &args, + ExpressionState &state, + Vector &result) { + + const idx_t count = args.size(); + if (count == 0) { + return; + } + + const auto &func_expr = state.expr.Cast(); + auto &bd = func_expr.bind_info->Cast(); + + // ----------------------------------------------------------------- + // 1) features: LIST(FLOAT) in unified format + // ----------------------------------------------------------------- + Vector &features_vec = args.data[0]; + UnifiedVectorFormat features_uf; + features_vec.ToUnifiedFormat(count, features_uf); + auto list_entries = UnifiedVectorFormat::GetData(features_uf); + + // Child vector contains the actual FLOAT elements of all lists + Vector &child_vec = ListVector::GetEntry(features_vec); + const idx_t child_count = ListVector::GetListSize(features_vec); + UnifiedVectorFormat child_uf; + child_vec.ToUnifiedFormat(child_count, child_uf); + auto child_vals = UnifiedVectorFormat::GetData(child_uf); + + // ----------------------------------------------------------------- + // 2) tree_path: VARCHAR in unified format + // ----------------------------------------------------------------- + Vector &path_vec = args.data[1]; + UnifiedVectorFormat path_uf; + path_vec.ToUnifiedFormat(count, path_uf); + auto path_data = UnifiedVectorFormat::GetData(path_uf); + + // ----------------------------------------------------------------- + // 3) has_missing: BOOLEAN in unified format + // ----------------------------------------------------------------- + Vector &miss_vec = args.data[2]; + UnifiedVectorFormat miss_uf; + miss_vec.ToUnifiedFormat(count, miss_uf); + auto miss_data = UnifiedVectorFormat::GetData(miss_uf); + + // ----------------------------------------------------------------- + // 4) Output setup + // ----------------------------------------------------------------- + result.SetVectorType(VectorType::FLAT_VECTOR); + auto out = FlatVector::GetData(result); + auto &out_valid = FlatVector::Validity(result); + + for (idx_t i = 0; i < count; ++i) { + const idx_t feat_idx = features_uf.sel->get_index(i); + const idx_t path_idx = path_uf.sel->get_index(i); + const idx_t miss_idx = miss_uf.sel->get_index(i); + + // NULL list or NULL path -> NULL result + if (!features_uf.validity.RowIsValid(feat_idx) || + !path_uf.validity.RowIsValid(path_idx)) { + out_valid.SetInvalid(i); + continue; + } + + // list entry for this row + const auto ent = list_entries[feat_idx]; + + // Build a contiguous row buffer of features using child_uf + std::vector row_buf(ent.length); + for (idx_t k = 0; k < ent.length; ++k) { + const idx_t child_logical = ent.offset + k; + const idx_t child_phys = child_uf.sel->get_index(child_logical); + if (!child_uf.validity.RowIsValid(child_phys)) { + row_buf[k] = std::numeric_limits::quiet_NaN(); + } else { + row_buf[k] = child_vals[child_phys]; + } + } + + // Resolve tree_path + const std::string tree_path = path_data[path_idx].GetString(); + + // has_missing flag for this row (default + per-row override) + bool row_has_missing = bd.has_missing_default; + if (miss_uf.validity.RowIsValid(miss_idx)) { + row_has_missing = miss_data[miss_idx]; + } + + // Per-path tree cache lookup + TreePtr tree_ptr; + auto it = bd.tree_cache.find(tree_path); + if (it == bd.tree_cache.end()) { + TreePtr new_tree; + try { + new_tree = std::make_shared(0, tree_path); + } catch (const std::exception &e) { + throw InvalidInputException( + "tree_predict: failed to load/parse tree file '%s': %s", + tree_path.c_str(), e.what()); + } + it = bd.tree_cache.emplace(tree_path, new_tree).first; + } + tree_ptr = it->second; + const Tree &tree = *tree_ptr; + + // Prediction + float pred; + if (row_has_missing) { + pred = tree.PredictSingleMissing(row_buf.data()); + } else { + pred = tree.PredictSingle(row_buf.data()); + } + out[i] = pred; + } +} + +// ========================= +// Register +// ========================= + +void ML_DecisionTree::Register(ExtensionLoader &loader) { + ScalarFunction fn( + "tree_predict", + { + LogicalType::LIST(LogicalType::FLOAT), // features + LogicalType::VARCHAR, // tree file path + LogicalType::BOOLEAN // has_missing + }, + LogicalType::FLOAT, + Execute); + + fn.bind = Bind; + // Reads external files and depends on path contents; treat as volatile. + fn.stability = FunctionStability::VOLATILE; + + loader.RegisterFunction(fn); +} + +} // namespace duckdb diff --git a/src/ml_functions/matrix_addition.cpp b/src/ml_functions/matrix_addition.cpp new file mode 100644 index 0000000..8851324 --- /dev/null +++ b/src/ml_functions/matrix_addition.cpp @@ -0,0 +1,127 @@ +#include "ml_functions/matrix_addition.hpp" + +#include +#include +#include + +#include +#include + +#include "duckdb/common/constants.hpp" +#include "duckdb/common/enum_util.hpp" +#include "duckdb/common/exception.hpp" +#include "duckdb/common/string_util.hpp" +#include "duckdb/common/types.hpp" +#include "duckdb/common/types/blob.hpp" +#include "duckdb/common/types/value.hpp" +#include "duckdb/execution/expression_executor.hpp" +#include "duckdb/function/scalar_function.hpp" +#include "duckdb/parser/parsed_data/create_scalar_function_info.hpp" +#include "duckdb/planner/expression/bound_constant_expression.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/common/vector_operations/generic_executor.hpp" +#include "duckdb/common/vector_operations/septenary_executor.hpp" + +namespace duckdb { +// ========================= BindData impl ========================= +unique_ptr ML_MatrixAddition::Bind(ClientContext &context, + ScalarFunction &bound_function, + vector> &args) { + if (args.size() != 2) { + throw InvalidInputException("mat_add: expects 2 arguments: x, biases_file"); + } + + // (Optional) enforce constant filename at bind time + if (!args[1]->IsFoldable()) { + throw BinderException("mat_add: biases_file must be a constant string"); + } + + auto vFile = ExpressionExecutor::EvaluateScalar(context, *args[1]); + if (vFile.IsNull()) { + throw BinderException("mat_add: biases_file is NULL"); + } + const std::string path = vFile.ToString(); + + // Parse into std::vector because ParseCSVBiases takes std::vector& + std::vector biases_std; + idx_t n = 0; + ParseCSVBiases(path, biases_std, n); + + // Convert to duckdb::vector expected by BindData + vector biases_duck(biases_std.begin(), biases_std.end()); + + return make_uniq(biases_duck, n, path, /*bound*/ true); +} + +// ========================= Execute ========================= +void ML_MatrixAddition::Execute(DataChunk &args, ExpressionState &state, Vector &result) { + const idx_t count = args.size(); // tuples in this chunk + + // Input LIST + Vector &lhs = args.data[0]; + UnifiedVectorFormat lhs_uvf; + lhs.ToUnifiedFormat(count, lhs_uvf); + const auto *entries = UnifiedVectorFormat::GetData(lhs_uvf); + + auto &child = ListVector::GetEntry(lhs); + const idx_t total_elems = ListVector::GetListSize(lhs); + UnifiedVectorFormat child_uvf; + child.ToUnifiedFormat(total_elems, child_uvf); + const double *child_data = UnifiedVectorFormat::GetData(child_uvf); + + //read the bound data + const auto &func_expr = state.expr.Cast(); + const auto &info = func_expr.bind_info->Cast(); + if (!info.has_bound_biases) { + throw InvalidInputException("mat_add: biases are not bound"); + } + const idx_t n = info.n; + const std::vector &biases = info.biases; + + // Prepare result in a FlatVector + result.SetVectorType(VectorType::FLAT_VECTOR); + auto *res_entries = FlatVector::GetData(result); + auto &res_child = ListVector::GetEntry(result); + const idx_t total_out = count * n; + ListVector::Reserve(result, total_out); + ListVector::SetListSize(result, total_out); + double *out = FlatVector::GetData(res_child); + auto &res_valid = FlatVector::Validity(result); + + //main loop + idx_t cursor = 0; + for(idx_t i = 0;iget_index(i); + if(lhs_uvf.validity.RowIsValid(idx)){ + const auto ent = entries[idx]; + if(ent.length != n){ + throw InvalidInputException("mat_add: Input vector length %llu does not match the bound biases length %llu", (unsigned long long)ent.length, (unsigned long long)n); + } + res_entries[i] = {cursor, n}; + const double *row = child_data + ent.offset; + for (idx_t j = 0;j + LogicalType::VARCHAR // biases_file VARCHAR + }, + LogicalType::LIST(LogicalType::DOUBLE), // RETURNS LIST + Execute + ); + fn.bind = Bind; + loader.RegisterFunction(fn); + fn.stability = FunctionStability::VOLATILE; // reads a file; avoid constant folding if desired +} // register end +}//duckdb namespace end \ No newline at end of file diff --git a/src/ml_functions/matrix_multiplication.cpp b/src/ml_functions/matrix_multiplication.cpp new file mode 100644 index 0000000..d72bde9 --- /dev/null +++ b/src/ml_functions/matrix_multiplication.cpp @@ -0,0 +1,389 @@ +#include "ml_functions/matrix_multiplication.hpp" + +#include +#include +#include + +#include +#include + +#include "duckdb/common/constants.hpp" +#include "duckdb/common/enum_util.hpp" +#include "duckdb/common/exception.hpp" +#include "duckdb/common/string_util.hpp" +#include "duckdb/common/types.hpp" +#include "duckdb/common/types/blob.hpp" +#include "duckdb/common/types/value.hpp" +#include "duckdb/execution/expression_executor.hpp" +#include "duckdb/function/scalar_function.hpp" +#include "duckdb/parser/parsed_data/create_scalar_function_info.hpp" +#include "duckdb/planner/expression/bound_constant_expression.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/common/vector_operations/generic_executor.hpp" +#include "duckdb/common/vector_operations/septenary_executor.hpp" + +namespace duckdb { + +// ========================= +// BindData impl +// ========================= + +ML_MatrixMultiply::BindData::BindData(const std::vector &w, + idx_t k_, + idx_t n_, + const std::string &file, + bool bound, + MulMode mode_) + : weights(w), k(k_), n(n_), weights_file(file), has_bound_weights(bound), mode(mode_) {} + +unique_ptr ML_MatrixMultiply::BindData::Copy() const { + return std::make_unique(weights, k, n, weights_file, has_bound_weights,mode); +} + +bool ML_MatrixMultiply::BindData::Equals(const FunctionData &other_p) const { + auto &o = other_p.Cast(); + return has_bound_weights == o.has_bound_weights && + k == o.k && n == o.n && + weights_file == o.weights_file && + weights == o.weights && + mode == o.mode; +} + +// ========================= +// Bind +// ========================= + +unique_ptr ML_MatrixMultiply::Bind(ClientContext &context, + ScalarFunction &bound_function, + vector> &args) { + // Expect: x LIST, filename VARCHAR + if (args.size() != 3) { + throw BinderException("matmul expects 3 arguments: (x LIST, weights_file VARCHAR, mode VARCHAR)"); + } + + // Force filename to be constant; fold it at bind time + if (!args[1]->IsFoldable()) { + throw BinderException("matmul: weights_file must be a constant string"); + } + + auto vFile = ExpressionExecutor::EvaluateScalar(context, *args[1]); + if (vFile.IsNull()) { + throw BinderException("matmul: weights_file is NULL"); + } + const std::string path = vFile.ToString(); + + std::vector W; + idx_t K = 0, N = 0; + ParseCSVWeights(path, W, K, N); + + auto vMode = ExpressionExecutor::EvaluateScalar(context, *args[2]); + MulMode mode = ParseModeValue(vMode); + + // If you want execution to only receive the input vector, uncomment + // the following line to erase the filename argument at bind-time. + // Function::EraseArgument(bound_function, args, 1); + + return std::make_unique(W, K, N, path, /*bound*/ true, mode); +} + +// ========================= +// Execute +// ========================= + +// void ML_MatrixMultiply::Execute(DataChunk &args, ExpressionState &state, Vector &result) { +// idx_t count = args.size(); // number of tuples in a DataChunk + +// // Read left-hand side matrix from the argument +// Vector &lhs = args.data[0]; +// UnifiedVectorFormat lhs_uvf; +// lhs.ToUnifiedFormat(count, lhs_uvf); +// const auto *entries = UnifiedVectorFormat::GetData(lhs_uvf); // {offset, length} per row + +// auto &child = ListVector::GetEntry(lhs); // the child that holds the list payload +// const idx_t total_elems = ListVector::GetListSize(lhs); +// UnifiedVectorFormat child_uvf; +// child.ToUnifiedFormat(total_elems, child_uvf); +// const double *child_data = UnifiedVectorFormat::GetData(child_uvf); + +// // Build a dense row-major Eigen matrix X from the selected rows +// Eigen::Matrix X; +// idx_t columns = (idx_t)-1; + +// for (idx_t i = 0; i < count; i++) { +// const idx_t idx = lhs_uvf.sel->get_index(i); +// if (lhs_uvf.validity.RowIsValid(idx)) { +// const idx_t row_len = entries[idx].length; +// if (i == 0) { +// columns = row_len; +// X.resize((Eigen::Index)count, (Eigen::Index)columns); +// } else { +// if (columns != row_len) { +// throw InvalidInputException("Input Matrix has uneven rows. Error at row %llu", +// (unsigned long long)i); +// } +// } +// const double *row_ptr = child_data + entries[idx].offset; +// std::memcpy(X.data() + i * columns, row_ptr, sizeof(double) * (size_t)columns); +// } +// } + +// // Prepare weight matrix using the bound data +// const auto &func_expr = state.expr.Cast(); +// const auto &bd = func_expr.bind_info->Cast(); +// using RowMat = Eigen::Matrix; +// Eigen::Map W(bd.weights.data(), (Eigen::Index)bd.k, (Eigen::Index)bd.n); + +// // X (count x k) * W (k x n) = Y (count x n) +// Eigen::Matrix Y = X * W; + +// // Materialize result as LIST with one list per input row (length = n) +// result.SetVectorType(VectorType::FLAT_VECTOR); +// auto *res_entries = FlatVector::GetData(result); +// auto &res_child = ListVector::GetEntry(result); + +// const idx_t n = bd.n; +// const idx_t total = count * n; +// ListVector::Reserve(result, total); +// ListVector::SetListSize(result, total); +// double *out = FlatVector::GetData(res_child); +// auto &res_valid = FlatVector::Validity(result); + +// idx_t cursor = 0; +// for (idx_t i = 0; i < count; ++i) { +// const idx_t idx = lhs_uvf.sel->get_index(i); +// if (!lhs_uvf.validity.RowIsValid(idx)) { +// res_entries[i] = {0, 0}; +// res_valid.Set(i, false); +// continue; +// } +// res_entries[i].offset = cursor; +// res_entries[i].length = n; +// std::memcpy(out + cursor, Y.data() + i * n, sizeof(double) * (size_t)n); +// cursor += n; +// } +// } + +void ML_MatrixMultiply::Execute(DataChunk &args, ExpressionState &state, Vector &result) { + const idx_t count = args.size(); // tuples in this chunk + result.SetVectorType(VectorType::FLAT_VECTOR); + + // ---------------------------- + // 1) Read input LIST + // ---------------------------- + Vector &lhs = args.data[0]; + UnifiedVectorFormat lhs_uvf; + lhs.ToUnifiedFormat(count, lhs_uvf); + const auto *entries = UnifiedVectorFormat::GetData(lhs_uvf); + + auto &child = ListVector::GetEntry(lhs); + const idx_t total_elems = ListVector::GetListSize(lhs); + UnifiedVectorFormat child_uvf; + child.ToUnifiedFormat(total_elems, child_uvf); + const double *child_data = UnifiedVectorFormat::GetData(child_uvf); + + // ---------------------------- + // 2) Bind data + weight matrix + // ---------------------------- + const auto &func_expr = state.expr.Cast(); + const auto &bd = func_expr.bind_info->Cast(); + using RowMat = Eigen::Matrix; + Eigen::Map W(bd.weights.data(), (Eigen::Index)bd.k, (Eigen::Index)bd.n); + + // ---------------------------- + // 3) Common pre-pass: + // - determine columns (k) + // - detect all-zero rows + // ---------------------------- + struct RowInfo { + idx_t sel = DConstants::INVALID_INDEX; // index into base vector + idx_t offset = 0; + idx_t len = 0; + bool valid = false; + bool all_zero = true; // assume zero until we see a non-zero + }; + + std::vector rows(count); + idx_t columns = (idx_t)-1; + + for (idx_t i = 0; i < count; ++i) { + const idx_t sel = lhs_uvf.sel->get_index(i); + rows[i].sel = sel; + + if (!lhs_uvf.validity.RowIsValid(sel)) { + rows[i].valid = false; + rows[i].all_zero = true; + continue; + } + + rows[i].valid = true; + + const auto &e = entries[sel]; + rows[i].offset = e.offset; + rows[i].len = e.length; + + if (columns == (idx_t)-1) { + columns = e.length; + } else if (columns != e.length) { + throw InvalidInputException("matmul: input matrix has uneven rows. Error at row %llu", + (unsigned long long)i); + } + + // Scan for first non-zero to detect all-zero rows + const double *row_ptr = child_data + e.offset; + bool any_nz = false; + for (idx_t c = 0; c < e.length; ++c) { + if (row_ptr[c] != 0.0) { + any_nz = true; + break; + } + } + rows[i].all_zero = !any_nz; + } + + if (columns == (idx_t)-1) { + // all rows NULL → return NULL lists + auto *res_entries = FlatVector::GetData(result); + auto &res_child = ListVector::GetEntry(result); + ListVector::Reserve(result, 0); + ListVector::SetListSize(result, 0); + auto &res_valid = FlatVector::Validity(result); + for (idx_t i = 0; i < count; ++i) { + res_entries[i] = {0, 0}; + res_valid.Set(i, false); + } + return; + } + + if (columns != bd.k) { + throw InvalidInputException("matmul: input width (%llu) != weights rows k (%llu)", + (unsigned long long)columns, + (unsigned long long)bd.k); + } + + const idx_t n = bd.n; + const idx_t total_vals = count * n; + + auto *res_entries = FlatVector::GetData(result); + auto &res_child = ListVector::GetEntry(result); + auto &res_valid = FlatVector::Validity(result); + + ListVector::Reserve(result, total_vals); + ListVector::SetListSize(result, total_vals); + double *out = FlatVector::GetData(res_child); + + // We will give every non-NULL row a list of length n at offset i*n. + // Initialize child buffer to zeros so zero rows are basically free. + std::memset(out, 0, sizeof(double) * (size_t)total_vals); + + for (idx_t i = 0; i < count; ++i) { + if (!rows[i].valid) { + res_entries[i] = {0, 0}; + res_valid.Set(i, false); + } else { + res_entries[i].offset = i * n; + res_entries[i].length = n; + res_valid.Set(i, true); + } + } + + // -------------------------------------- + // 4) Dense mode: use all rows (like before) + // -------------------------------------- + if (bd.mode == MulMode::Dense) { + RowMat X_dense((Eigen::Index)count, (Eigen::Index)columns); + + for (idx_t i = 0; i < count; ++i) { + if (!rows[i].valid) { + // fill row with zeros (already zero-initialized by Eigen) + X_dense.row((Eigen::Index)i).setZero(); + continue; + } + const double *row_ptr = child_data + rows[i].offset; + std::memcpy(X_dense.data() + (ptrdiff_t)i * columns, + row_ptr, + sizeof(double) * (size_t)columns); + } + + RowMat Y = X_dense * W; // (count x k) * (k x n) + + // write results; we've already set offsets to i*n + for (idx_t i = 0; i < count; ++i) { + if (!rows[i].valid) { + continue; // NULL row already handled + } + std::memcpy(out + (ptrdiff_t)i * n, + Y.data() + (ptrdiff_t)i * n, + sizeof(double) * (size_t)n); + } + return; + } + + // -------------------------------------- + // 5) Sparse mode: row-sparse dense matmul + // - only build X for non-zero rows + // -------------------------------------- + idx_t nz_rows = 0; + for (idx_t i = 0; i < count; ++i) { + if (rows[i].valid && !rows[i].all_zero) { + ++nz_rows; + } + } + + if (nz_rows == 0) { + // all rows are zero (or NULL) → result is already all zeros + // (offsets/validity already set above) + return; + } + + RowMat X_compact((Eigen::Index)nz_rows, (Eigen::Index)columns); + + std::vector logical_to_compact(count, (idx_t)-1); + idx_t compact_idx = 0; + for (idx_t i = 0; i < count; ++i) { + if (!rows[i].valid || rows[i].all_zero) { + continue; + } + const double *row_ptr = child_data + rows[i].offset; + std::memcpy(X_compact.data() + (ptrdiff_t)compact_idx * columns, + row_ptr, + sizeof(double) * (size_t)columns); + logical_to_compact[i] = compact_idx; + ++compact_idx; + } + + RowMat Y_compact = X_compact * W; // (nz_rows x n) + + // scatter back into full result (zeros already there for zero rows) + for (idx_t i = 0; i < count; ++i) { + if (!rows[i].valid || rows[i].all_zero) { + // NULL rows: handled above; zero rows: already zeroed in 'out' + continue; + } + const idx_t cidx = logical_to_compact[i]; + std::memcpy(out + (ptrdiff_t)i * n, + Y_compact.data() + (ptrdiff_t)cidx * n, + sizeof(double) * (size_t)n); + } +} + + +// ========================= +// Register +// ========================= + +void ML_MatrixMultiply::Register(ExtensionLoader &loader) { + // Logical signature with filename; Bind may erase it so Exec sees only x. + ScalarFunction fn( + "mat_mul", + { LogicalType::LIST(LogicalType::DOUBLE), // x + LogicalType::VARCHAR, // weights_file + LogicalType::VARCHAR }, // MODE + LogicalType::LIST(LogicalType::DOUBLE), + Execute + ); + fn.bind = Bind; + loader.RegisterFunction(fn); + fn.stability = FunctionStability::VOLATILE; // reads a file; avoid constant folding if desired +} + +} // namespace duckdb diff --git a/src/ml_functions/min_max_scaler.cpp b/src/ml_functions/min_max_scaler.cpp new file mode 100644 index 0000000..0020947 --- /dev/null +++ b/src/ml_functions/min_max_scaler.cpp @@ -0,0 +1,230 @@ +#include "functions.hpp" +#include "duckdb.hpp" +#include "duckdb/function/function.hpp" +#include "duckdb/execution/expression_executor.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include +#include +#include +#include +#include +#include + +namespace duckdb { + struct ML_MinMaxScaler { + public: + // ------------------------- + // Helpers + // ------------------------- + + // ------------------------- + // Bind Data + // ------------------------- + struct BindData final : public FunctionData { + std::vector min; // min values for each column + std::vector max; // max values for each column + idx_t cols = 0; // number of columns + std::string min_max_file; + bool has_bound_min_max = false; + + BindData() = default; + + BindData(const std::vector &min_, const std::vector &max_, idx_t cols_, + const std::string &file, bool bound) + : min(min_), max(max_), cols(cols_), min_max_file(file), has_bound_min_max(bound) {} + + unique_ptr Copy() const override { + // deep copy (vector/string copy is deep) + return make_uniq(min, max, cols, min_max_file, has_bound_min_max); + } + + bool Equals(const FunctionData &other_p) const override { + auto &o = other_p.Cast(); + return has_bound_min_max == o.has_bound_min_max && + cols == o.cols && + min_max_file == o.min_max_file && + min == o.min && + max == o.max; + } + }; + + // ------------------------- + // Bind Functions + // ------------------------- + static unique_ptr Bind_List(ClientContext &context, + ScalarFunction &bound_function, + vector> &args) { + // Expecting 4 arguments: x, min_list, max_list, no_of_columns + if (args.size() != 4) { + throw InvalidInputException("min_max_scaler with list arguments expects 4 arguments: x, min_list, max_list, no_of_columns"); + } + //TODO + auto vmin_list = ExpressionExecutor::EvaluateScalar(context, *args[1]); + auto vmax_list = ExpressionExecutor::EvaluateScalar(context, *args[2]); + auto vcols = ExpressionExecutor::EvaluateScalar(context, *args[3]); + + std::vector min; + std::vector max; + + const auto &min_children = ListValue::GetChildren(vmin_list); + const auto &max_children = ListValue::GetChildren(vmax_list); + min.reserve(min_children.size()); + max.reserve(max_children.size()); + + for (const auto &child : min_children) min.push_back(child.GetValue()); + for (const auto &child : max_children) max.push_back(child.GetValue()); + + std::string path = "from_arguments"; + return make_uniq(min, max, min.size(), path, /*bound*/ true); + } + + static unique_ptr Bind_File(ClientContext &context, + ScalarFunction &bound_function, + vector> &args) { + // Expecting 2 arguments: x, min_max_file + if (args.size() != 2) { + throw InvalidInputException("min_max_scaler with file argument expects 2 arguments: x, min_max_file"); + } + auto vFile = ExpressionExecutor::EvaluateScalar(context, *args[1]); + if (vFile.IsNull()) { + throw BinderException("matmul: weights_file is NULL"); + } + const std::string path = vFile.ToString(); + // const std::string path = args.data[1]->ToString(); + // // std::cout << "Min-Max file: " << path << std::endl; + + if(!std::filesystem::exists(path)){ + throw InvalidInputException("min_max_scaler: File does not exist - " + path); + } + + std::vector min; + std::vector max; + std::ifstream file(path); + + std::string line; + int line_count = 0; + + while(std::getline(file,line)){ + std::istringstream iss(line); + double value; + while (iss>>value){ + // std::cout << "Read value: " << value << std::endl; + if(line_count == 0){ + min.push_back(value); + } + else if(line_count == 1){ + max.push_back(value); + } + else { + throw InvalidInputException("min_max_scaler: File should contain only two lines - first line for min values and second line for max values"); + } + } + line_count++; + // std::cout << "Read line " << line_count << ": " << line << std::endl; + } + file.close(); + + // std::cout << "Min values: "; + // for(auto v : min) // std::cout << v << " "; + // std::cout << std::endl; + // std::cout << "Max values: "; + // for(auto v : max) // std::cout << v << " "; + // std::cout << std::endl; + + return make_uniq(min, max, min.size(), path, /*bound*/ true); + } + + // ------------------------- + // Execute + // ------------------------- + static void Execute(DataChunk &args, ExpressionState &state, Vector &result) { + idx_t count = args.size(); //number of tuples in a DataChunk + + // Read the feature vector from the argument + Vector &f = args.data[0]; + UnifiedVectorFormat f_uvf; + f.ToUnifiedFormat(count, f_uvf); + const auto *entries = UnifiedVectorFormat::GetData(f_uvf); + auto &child = ListVector::GetEntry(f); + UnifiedVectorFormat child_uvf; + child.ToUnifiedFormat(count, child_uvf); + const double *x_data = UnifiedVectorFormat::GetData(child_uvf); + + //read the bound data + const auto &func_expr = state.expr.Cast(); + const auto &info = func_expr.bind_info->Cast(); + if (!info.has_bound_min_max) { + throw InvalidInputException("min_max_scaler: min and max values are not bound"); + } + const idx_t cols = info.cols; + const std::vector &min = info.min; + const std::vector &max = info.max; + + // Prepare result in a FlatVector + result.SetVectorType(VectorType::FLAT_VECTOR); + auto *res_entries = FlatVector::GetData(result); + auto &res_child = ListVector::GetEntry(result); + const idx_t total_out = count * cols; + ListVector::Reserve(result, total_out); + ListVector::SetListSize(result, total_out); + double *out = FlatVector::GetData(res_child); + auto &res_valid = FlatVector::Validity(result); + + //main loop + idx_t cursor = 0; + for(idx_t i = 0;iget_index(i); + if(f_uvf.validity.RowIsValid(idx)){ + const auto ent = entries[idx]; + if(ent.length != cols){ + throw InvalidInputException("min_max_scaler: Input vector length %llu does not match the bound min/max length %llu, %llu", (unsigned long long)ent.length, (unsigned long long)cols, (unsigned long long)min.size()); + } + // std::cout << " Input length: " << ent.length << ", cols: " << cols << std::endl; + res_entries[i] = {cursor, cols}; + const double *row = x_data + ent.offset; + for (idx_t j = 0;j + +namespace duckdb { + +using std::string; + +// ------------------------- +// Bind +// ------------------------- +unique_ptr ML_ForestTable::Bind( + ClientContext &context, + TableFunctionBindInput &input, + duckdb::vector &return_types, + duckdb::vector &names) { + + if (input.inputs.size() != 1) { + throw InvalidInputException("forest_trees(path) expects exactly 1 VARCHAR argument"); + } + + auto path_val = input.inputs[0]; + if (path_val.IsNull()) { + throw InvalidInputException("forest_trees: path cannot be NULL"); + } + + string path = path_val.ToString(); + if (!std::filesystem::exists(path)) { + throw InvalidInputException("forest_trees: path '%s' does not exist", path.c_str()); + } + + // Normalize: ensure it's a directory + if (!std::filesystem::is_directory(path)) { + throw InvalidInputException("forest_trees: '%s' is not a directory", path.c_str()); + } + + // Collect tree file paths at bind time + vector tree_paths; + for (auto &entry : std::filesystem::directory_iterator(path)) { + if (!entry.is_regular_file()) { + continue; + } + tree_paths.push_back(entry.path().string()); + } + + // Sort for deterministic ordering (optional but nice) + std::sort(tree_paths.begin(), tree_paths.end()); + + // Define output schema: (tree_id INT, tree_path VARCHAR) + return_types.clear(); + names.clear(); + + return_types.push_back(LogicalType::INTEGER); + names.push_back("tree_id"); + + return_types.push_back(LogicalType::VARCHAR); + names.push_back("tree_path"); + + return make_uniq(std::move(path), std::move(tree_paths)); +} + +// ------------------------- +// InitGlobal +// ------------------------- +unique_ptr ML_ForestTable::InitGlobal( + ClientContext &context, + TableFunctionInitInput &input) { + + auto &bind = input.bind_data->Cast(); + + auto state = make_uniq(); + state->tree_paths = bind.tree_paths; + state->total_rows = state->tree_paths.size(); + state->current_row = 0; + + return std::move(state); +} + +// ------------------------- +// Function +// ------------------------- +void ML_ForestTable::Function( + ClientContext &context, + TableFunctionInput &data, + DataChunk &output) { + + auto &gst = data.global_state->Cast(); + + if (gst.current_row >= gst.total_rows) { + output.SetCardinality(0); + return; + } + + idx_t remaining = gst.total_rows - gst.current_row; + idx_t max_rows = std::min(STANDARD_VECTOR_SIZE, remaining); + output.SetCardinality(max_rows); + + auto tree_id_vec = FlatVector::GetData(output.data[0]); + auto tree_path_vec = FlatVector::GetData(output.data[1]); + + for (idx_t i = 0; i < max_rows; i++) { + idx_t row_idx = gst.current_row + i; + + tree_id_vec[i] = static_cast(row_idx); + + const string &path = gst.tree_paths[row_idx]; + tree_path_vec[i] = StringVector::AddString(output.data[1], path); + } + + gst.current_row += max_rows; +} + +// ------------------------- +// Register +// ------------------------- +void ML_ForestTable::Register(ExtensionLoader &loader) { + TableFunction tf( + "forest_trees", + {LogicalType::VARCHAR}, // forest folder path + Function // table_function_t + ); + + tf.bind = Bind; + tf.init_global = InitGlobal; + + loader.RegisterFunction(tf); +} + +} // namespace duckdb diff --git a/src/ml_functions/ml_functions_scalar.cpp b/src/ml_functions/ml_functions_scalar.cpp index 7a00bcd..7c9d539 100644 --- a/src/ml_functions/ml_functions_scalar.cpp +++ b/src/ml_functions/ml_functions_scalar.cpp @@ -1,224 +1,33 @@ -#include -#include -#include "functions.hpp" -#include "function_builder.hpp" #include "duckdb.hpp" -#include "duckdb/common/types/value.hpp" -#include "duckdb/common/types.hpp" -#include "duckdb/common/enum_util.hpp" -#include "duckdb/common/exception.hpp" -#include "duckdb/common/string_util.hpp" -#include "duckdb/function/scalar_function.hpp" -#include -#include "duckdb/common/constants.hpp" -#include "duckdb/common/types/blob.hpp" -#include "duckdb/common/vector_operations/generic_executor.hpp" -#include "duckdb/execution/expression_executor.hpp" -#include "duckdb/planner/expression/bound_function_expression.hpp" -#include "duckdb/common/vector_operations/septenary_executor.hpp" -#include "duckdb/planner/expression/bound_constant_expression.hpp" +#include "ml_functions/matrix_multiplication.hpp" +#include "min_max_scaler.cpp" +#include "ml_functions/one_hot_encoder.hpp" +#include "ml_functions/decision_tree.hpp" +#include "ml_functions/matrix_addition.hpp" +#include "ml_functions/argmax.hpp" +#include "ml_functions/softmax.hpp" +#include "ml_functions/sigmoid.hpp" +#include "ml_functions/torch_dnn.hpp" +#include "ml_functions/decision_forest.hpp" +#include "ml_functions/ml_decision_forest_table.hpp" namespace duckdb { -class MatrixData final : public FunctionData { - -public: - unsigned int num_rows_; - unsigned int num_columns_; - - bool use_gpu_; - - explicit MatrixData(int num_rows, int num_columns, bool use_gpu) { - num_rows_ = num_rows; - num_columns_ = num_columns; - use_gpu_ = use_gpu; - } - - unique_ptr Copy() const override { - return make_uniq(num_rows_, num_columns_, use_gpu_); - } - - bool Equals(const FunctionData &other) const override { - auto &other_data = other.Cast(); - return num_rows_ == other_data.num_rows_ && num_columns_ == other_data.num_columns_ && - use_gpu_ == other_data.use_gpu_; - } -}; - -class MatrixDataByPointer final : public FunctionData { - -public: - float *weights_; - unsigned int num_rows_; - unsigned int num_columns_; - bool use_gpu_; - - explicit MatrixDataByPointer(float *weights, int num_rows, int num_columns, bool use_gpu) { - weights_ = new float[num_rows * num_columns]; - std::memcpy(weights_, weights, num_rows * num_columns * sizeof(float)); - num_rows_ = num_rows; - num_columns_ = num_columns; - use_gpu_ = use_gpu; - } - - unique_ptr Copy() const override { - return make_uniq(weights_, num_rows_, num_columns_, use_gpu_); - } - - bool Equals(const FunctionData &other) const override { - auto &other_data = other.Cast(); - return num_rows_ == other_data.num_rows_ && num_columns_ == other_data.num_columns_ && - weights_ == other_data.weights_ && use_gpu_ == other_data.use_gpu_; - } -}; - -/* -static unique_ptr MatrixBindByPointer(ClientContext &context, ScalarFunction &bound_function, - vector> &arguments) { - - const auto expr0 = ExpressionExecutor::EvaluateScalar(context, *arguments[0]); - const auto weights_value = expr0.GetPointer(); - const auto expr1 = ExpressionExecutor::EvaluateScalar(context, *arguments[1]); - const auto num_rows_value = expr1.GetValue(); - const auto expr2 = ExpressionExecutor::EvaluateScalar(context, *arguments[2]); - const auto num_columns_value = expr2.GetValue(); - const auto expr3 = ExpressionExecutor::EvaluateScalar(context, *arguments[3]); - const auto use_gpu = expr3.GetValue(); - return make_uniq((float *)weights_value, num_rows_value, num_columns_value, use_gpu); - -} - - -static unique_ptr MatrixBind(ClientContext &context, ScalarFunction &bound_function, - vector> &arguments) { - return make_uniq(8, 2, false); -} - -*/ - -struct ML_MatrixMultiply { - -public: - static void Execute(DataChunk &args, ExpressionState &state, Vector &result) { - - const auto count = args.size(); - - auto &lhs = ArrayVector::GetEntry(args.data[0]); - auto &rhs = ArrayVector::GetEntry(args.data[1]); - auto &res = ArrayVector::GetEntry(result); - - const auto &lhs_validity = FlatVector::Validity(lhs); - const auto &rhs_validity = FlatVector::Validity(rhs); - - UnifiedVectorFormat lhs_format; - UnifiedVectorFormat rhs_format; - - args.data[0].ToUnifiedFormat(count, lhs_format); - args.data[1].ToUnifiedFormat(count, rhs_format); - - auto lhs_data = FlatVector::GetData(lhs); - auto rhs_data = FlatVector::GetData(rhs); - auto res_data = FlatVector::GetData(res); - - /* - for (int i = 0; i < count; i++) { - std::cout << "lhs-"<< i << ":" << lhs.GetValue(i) << std::endl; - std::cout << "rhs-"<< i << ":"<< rhs.GetValue(i) << std::endl; - } - std::cout << "lhs type:" << int(lhs.GetVectorType()) << std::endl; - std::cout << "lhs type:" << lhs.GetType().ToString() << std::endl; - std::cout << "lhs type:" << EnumUtil::ToChars(lhs.GetType().InternalType()) << std::endl; - std::cout << "rhs type:" << int(GetTypeIdSize(lhs.GetType().InternalType())) << std::endl; - std::cout << "rhs type:" << int(rhs.GetVectorType()) << std::endl; - std::cout << "rhs type:" << EnumUtil::ToChars(rhs.GetType().InternalType()) << std::endl; - std::cout << "rhs type:" << int(GetTypeIdSize(rhs.GetType().InternalType())) << std::endl; - - auto lhs_buffer = lhs.GetBuffer(); - auto rhs_buffer = rhs.GetBuffer(); - - std::cout << "lhs buffer type:" << EnumUtil::ToChars(lhs_buffer->GetBufferType()) << - std::endl; std::cout << "rhs buffer type:" << EnumUtil::ToChars(rhs_buffer->GetBufferType()) - << std::endl; - */ - - Value v2 = args.data[2].GetValue(0); - Value v3 = args.data[3].GetValue(0); - int num_rows = IntegerValue::Get(v2); - int num_columns = IntegerValue::Get(v3); - std::cout << "num rows: " << num_rows << "; num columns: " << num_columns << std::endl; - - const auto rhs_idx = rhs_format.sel->get_index(0); - const auto right_offset = rhs_idx * num_rows; - if (!rhs_validity.CheckAllValid(right_offset + num_rows, right_offset)) { - throw InvalidInputException(StringUtil::Format("right argument can not contain NULL values")); - } - const auto rhs_data_ptr = rhs_data + right_offset; - Eigen::Map> weight_matrix( - (float *)rhs_data_ptr, num_rows, num_columns); - - for (idx_t i = 0; i < count; i++) { - const auto lhs_idx = lhs_format.sel->get_index(i); - - if (!lhs_format.validity.RowIsValid(lhs_idx)) { - FlatVector::SetNull(result, i, true); - continue; - } - - const auto left_offset = lhs_idx * num_rows; - if (!lhs_validity.CheckAllValid(left_offset + num_rows, left_offset)) { - throw InvalidInputException(StringUtil::Format("left argument can not contain NULL values")); - } - - const auto result_offset = i * num_rows; - - const auto lhs_data_ptr = lhs_data + left_offset; - const auto res_data_ptr = res_data + result_offset; - - bool use_gpu = false; - - if (use_gpu) { - // TODO: implementation of matrix multiplication in GPU - throw std::runtime_error("GPU implementation of Matrix Multiple is not implemented."); - } else { - - Eigen::Map> input_matrix( - (float *)lhs_data_ptr, 1, num_rows); - - Eigen::Matrix result_matrix = - input_matrix * weight_matrix; - - // dispatch results to multiple rows - - std::memcpy((float *)res_data_ptr, result_matrix.row(0).data(), num_columns); - } - } - } - - static void RegisterByArray(ExtensionLoader &loader, int num_features, int num_neurons) { - - FunctionBuilder::RegisterScalar(loader, "matmul", [&](ScalarFunctionBuilder &func) { - func.AddVariant([&](ScalarFunctionVariantBuilder &variant) { - variant.AddParameter("input", LogicalType::ARRAY(LogicalType::FLOAT, num_features)); - variant.AddParameter("weight", LogicalType::ARRAY(LogicalType::FLOAT, num_features * num_neurons)); - variant.AddParameter("num_rows", LogicalType::INTEGER); - variant.AddParameter("num_columns", LogicalType::INTEGER); - variant.SetReturnType(LogicalType::ARRAY(LogicalType::FLOAT, num_neurons)); - variant.SetFunction(Execute); - }); - - func.SetDescription(R"( - Returns matrix results multiplying each input with a pre-defined weight matrix. - )"); - - func.SetTag("ext", "cactusdb"); - func.SetTag("category", "relation"); - }); - } -}; - +// This is the symbol your extension’s init code should call. void RegisterMLScalarFunctions(ExtensionLoader &loader) { - - ML_MatrixMultiply::RegisterByArray(loader, 8, 2); + ML_MatrixMultiply::Register(loader); + ML_MinMaxScaler::Register(loader); + ML_OneHotEncoder::Register(loader); + ML_DecisionTree::Register(loader); + ML_MatrixAddition::Register(loader); + ML_Argmax::Register(loader); + ML_Softmax::Register(loader); + ML_Sigmoid::Register(loader); + ML_TorchDNN::Register(loader); + ML_DecisionForest::Register(loader); + ML_ForestTable::Register(loader); + // RegisterSoftmax(loader); + // RegisterSigmoid(loader); } } // namespace duckdb diff --git a/src/ml_functions/one_hot_encoder.cpp b/src/ml_functions/one_hot_encoder.cpp new file mode 100644 index 0000000..760d8a1 --- /dev/null +++ b/src/ml_functions/one_hot_encoder.cpp @@ -0,0 +1,122 @@ +#include "ml_functions/one_hot_encoder.hpp" +#include "duckdb/execution/expression_executor.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/common/types/value.hpp" +#include "duckdb/common/types/validity_mask.hpp" +#include +#include +#include +#include +#include + +namespace duckdb { + + +// ---------- BindData impl ---------- +unique_ptr ML_OneHotEncoder::BindData::Copy() const { + return make_uniq(str_to_index, int_to_index, is_string); +} +bool ML_OneHotEncoder::BindData::Equals(const FunctionData &other_p) const { + auto &o = other_p.Cast(); + return str_to_index == o.str_to_index && + int_to_index == o.int_to_index && + is_string == o.is_string; +} + +// ---------- Bind ---------- +unique_ptr ML_OneHotEncoder::Bind(ClientContext &context, + ScalarFunction & /*bound_function*/, + vector> &args) { + // Expecting 3 arguments: x, categories_file, is_string + if (args.size() != 3) { + throw InvalidInputException("one_hot_encode expects 3 arguments: x, categories_file, is_string"); + } + + auto vFile = ExpressionExecutor::EvaluateScalar(context, *args[1]); + auto vIsString = ExpressionExecutor::EvaluateScalar(context, *args[2]); + + const std::string path = vFile.ToString(); + const bool is_string = vIsString.GetValue(); + + std::map str_map; + std::map int_map; + + // Load mapping (single helper function defined elsewhere) + one_hot_encoder_bind_constructor(path, is_string, str_map, int_map); + + return make_uniq(str_map, int_map, is_string); +} + +// ---------- Execute ---------- +void ML_OneHotEncoder::Execute(DataChunk &args, ExpressionState &state, Vector &result) { + const idx_t count = args.size(); + + const auto &func_expr = state.expr.Cast(); + const auto &bd = func_expr.bind_info->Cast(); + const idx_t k = bd.is_string ? (idx_t)bd.str_to_index.size() + : (idx_t)bd.int_to_index.size(); + + Vector &f = args.data[0]; + UnifiedVectorFormat f_uvf; + f.ToUnifiedFormat(count, f_uvf); + + result.SetVectorType(VectorType::FLAT_VECTOR); + auto *res_entries = FlatVector::GetData(result); + auto &res_child = ListVector::GetEntry(result); + const idx_t total = count * k; + ListVector::Reserve(result, total); + ListVector::SetListSize(result, total); + + auto *out = FlatVector::GetData(res_child); + auto &res_valid = FlatVector::Validity(result); + + idx_t cursor = 0; + for (idx_t i = 0; i < count; i++) { + const idx_t sel_idx = f_uvf.sel->get_index(i); + if (f_uvf.validity.RowIsValid(sel_idx)) { + res_entries[i] = {cursor, k}; + const auto sref = UnifiedVectorFormat::GetData(f_uvf)[sel_idx]; + + idx_t pos = DConstants::INVALID_INDEX; + if (bd.is_string) { + std::string key = sref.GetString(); + auto it = bd.str_to_index.find(key); + if (it != bd.str_to_index.end()) pos = it->second; + } else { + idx_t key = (idx_t)std::stoll(sref.GetString()); + auto it = bd.int_to_index.find(key); + if (it != bd.int_to_index.end()) pos = it->second; + } + + if (pos != DConstants::INVALID_INDEX) { + for (idx_t j = 0; j < k; j++) { + out[cursor++] = (j == pos) ? 1.0 : 0.0; + } + } else { + for (idx_t j = 0; j < k; j++) { + out[cursor++] = 0.0; + } + } + } else { + // null input → null list + res_entries[i] = {0, 0}; + res_valid.SetInvalid(i); + } + } +} + +// ---------- Register ---------- +void ML_OneHotEncoder::Register(ExtensionLoader &loader) { + // (x VARCHAR, categories_file VARCHAR, is_string BOOLEAN) -> DOUBLE[] + ScalarFunction fn( + "one_hot_encode", + { LogicalType::VARCHAR, LogicalType::VARCHAR, LogicalType::BOOLEAN }, + LogicalType::LIST(LogicalType::DOUBLE), + Execute + ); + fn.bind = Bind; + fn.stability = FunctionStability::VOLATILE; // reads a file; avoid constant folding if desired + loader.RegisterFunction(fn); +} + +} // namespace duckdb \ No newline at end of file diff --git a/src/ml_functions/sigmoid.cpp b/src/ml_functions/sigmoid.cpp new file mode 100644 index 0000000..22adb82 --- /dev/null +++ b/src/ml_functions/sigmoid.cpp @@ -0,0 +1,113 @@ +#include "ml_functions/sigmoid.hpp" + +#include +#include + +#include "duckdb/common/exception.hpp" +#include "duckdb/common/types.hpp" +#include "duckdb/common/vector.hpp" +#include "duckdb/function/scalar_function.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" + +namespace duckdb { + +// -------- Bind -------- +unique_ptr ML_Sigmoid::Bind(ClientContext &context, + ScalarFunction &bound_function, + vector> &args) { + if (args.size() != 1) { + throw BinderException("cdb_sigmoid: expects exactly 1 argument: LIST"); + } + return nullptr; +} + +// -------- Execute -------- +void ML_Sigmoid::Execute(DataChunk &args, ExpressionState &state, Vector &result) { + const idx_t count = args.size(); + + // Input: LIST + Vector &in = args.data[0]; + UnifiedVectorFormat in_uvf; + in.ToUnifiedFormat(count, in_uvf); + const auto *entries = UnifiedVectorFormat::GetData(in_uvf); + + // Child payload + auto &child = ListVector::GetEntry(in); + const idx_t total_elems = ListVector::GetListSize(in); + UnifiedVectorFormat child_uvf; + child.ToUnifiedFormat(total_elems, child_uvf); + const double *child_data = UnifiedVectorFormat::GetData(child_uvf); + const auto &child_valid = child_uvf.validity; + + // Output LIST + result.SetVectorType(VectorType::FLAT_VECTOR); + auto *res_entries = FlatVector::GetData(result); + auto &res_child = ListVector::GetEntry(result); + auto &res_valid = FlatVector::Validity(result); + + idx_t out_cursor = 0; + ListVector::Reserve(result, total_elems); + + for (idx_t i = 0; i < count; ++i) { + const idx_t sel = in_uvf.sel->get_index(i); + + if (!in_uvf.validity.RowIsValid(sel)) { + res_entries[i] = {0, 0}; + res_valid.Set(i, false); + continue; + } + + const auto le = entries[sel]; + const idx_t len = le.length; + const idx_t off = le.offset; + + if (len == 0) { + throw InvalidInputException("cdb_sigmoid: empty list at row %llu", + (unsigned long long)i); + } + if (!child_valid.AllValid()) { + for (idx_t j = 0; j < len; ++j) { + if (!child_valid.RowIsValid(off + j)) { + throw InvalidInputException( + "cdb_sigmoid: NULL element in list at row %llu (position %llu)", + (unsigned long long)i, (unsigned long long)j + ); + } + } + } + + using RowArr = Eigen::Array; + using RowVec = Eigen::Matrix; + + Eigen::Map row(child_data + off, (Eigen::Index)len); + + // Sigmoid: 1 / (1 + exp(-x)), computed vectorized with Eigen + RowArr probs = 1.0 / (1.0 + (-row.array()).exp()); + + // Materialize + const idx_t prev_total = ListVector::GetListSize(result); + ListVector::SetListSize(result, prev_total + len); + double *out = FlatVector::GetData(res_child); + + res_entries[i].offset = out_cursor; + res_entries[i].length = len; + std::memcpy(out + out_cursor, probs.data(), sizeof(double) * (size_t)len); + out_cursor += len; + } +} + +// -------- Register -------- +void ML_Sigmoid::Register(ExtensionLoader &loader) { + // Unique name to avoid built-in collisions + ScalarFunction fn( + "sigmoid", + { LogicalType::LIST(LogicalType::DOUBLE) }, + LogicalType::LIST(LogicalType::DOUBLE), + Execute + ); + fn.bind = Bind; + fn.stability = FunctionStability::VOLATILE; + loader.RegisterFunction(fn); +} + +} // namespace duckdb diff --git a/src/ml_functions/softmax.cpp b/src/ml_functions/softmax.cpp new file mode 100644 index 0000000..182e9c3 --- /dev/null +++ b/src/ml_functions/softmax.cpp @@ -0,0 +1,132 @@ +#include "ml_functions/softmax.hpp" + +#include + +#include "duckdb/common/exception.hpp" +#include "duckdb/common/types.hpp" +#include "duckdb/common/vector.hpp" +#include "duckdb/function/scalar_function.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" + +namespace duckdb { + +// -------- Bind -------- +unique_ptr ML_Softmax::Bind(ClientContext &context, + ScalarFunction &bound_function, + vector> &args) { + if (args.size() != 1) { + throw BinderException("cdb_softmax: expects exactly 1 argument: LIST"); + } + // Type checking can be left to the binder; no bind info needed + return nullptr; +} + +// -------- Execute -------- +void ML_Softmax::Execute(DataChunk &args, ExpressionState &state, Vector &result) { + const idx_t count = args.size(); + + // Input: LIST + Vector &in = args.data[0]; + UnifiedVectorFormat in_uvf; + in.ToUnifiedFormat(count, in_uvf); + const auto *entries = UnifiedVectorFormat::GetData(in_uvf); + + // Child payload vector for the list + auto &child = ListVector::GetEntry(in); + const idx_t total_elems = ListVector::GetListSize(in); + UnifiedVectorFormat child_uvf; + child.ToUnifiedFormat(total_elems, child_uvf); + const double *child_data = UnifiedVectorFormat::GetData(child_uvf); + const auto &child_valid = child_uvf.validity; + + // Prepare result LIST + result.SetVectorType(VectorType::FLAT_VECTOR); + auto *res_entries = FlatVector::GetData(result); + auto &res_child = ListVector::GetEntry(result); + auto &res_valid = FlatVector::Validity(result); + + // We will append all outputs contiguously + idx_t out_cursor = 0; + + // We don’t yet know total size; reserve a heuristic to reduce reallocs + ListVector::Reserve(result, total_elems); + + for (idx_t i = 0; i < count; ++i) { + const idx_t sel = in_uvf.sel->get_index(i); + + if (!in_uvf.validity.RowIsValid(sel)) { + // NULL row -> NULL result + res_entries[i] = {0, 0}; + res_valid.Set(i, false); + continue; + } + + const auto le = entries[sel]; + const idx_t len = le.length; + const idx_t off = le.offset; + + if (len == 0) { + throw InvalidInputException("cdb_softmax: empty list at row %llu", + (unsigned long long)i); + } + + // Disallow NULL elements inside list (strict). If you prefer to treat + // NULLs as -inf or error differently, we can change this. + if (!child_valid.AllValid()) { + for (idx_t j = 0; j < len; ++j) { + if (!child_valid.RowIsValid(off + j)) { + throw InvalidInputException( + "cdb_softmax: NULL element in list at row %llu (position %llu)", + (unsigned long long)i, (unsigned long long)j + ); + } + } + } + + using RowVec = Eigen::Matrix; + using RowArr = Eigen::Array; + + // Map current row as Eigen vector + Eigen::Map row(child_data + off, (Eigen::Index)len); + + // Numerical stability: subtract row max before exp + const double row_max = row.maxCoeff(); + RowArr exps = (row.array() - row_max).exp(); + const double denom = exps.sum(); + if (denom == 0.0 || !std::isfinite(denom)) { + throw InvalidInputException("cdb_softmax: invalid denominator (sum(exp(.))=%lf) at row %llu", + denom, (unsigned long long)i); + } + RowArr probs = exps / denom; + + // Materialize into result child + const idx_t prev_total = ListVector::GetListSize(result); + ListVector::SetListSize(result, prev_total + len); + double *out = FlatVector::GetData(res_child); + + res_entries[i].offset = out_cursor; + res_entries[i].length = len; + + std::memcpy(out + out_cursor, + probs.data(), + sizeof(double) * (size_t)len); + out_cursor += len; + } +} + +// -------- Register -------- +void ML_Softmax::Register(ExtensionLoader &loader) { + // Unique name to avoid built-in name collisions + ScalarFunction fn( + "cdb_softmax", + { LogicalType::LIST(LogicalType::DOUBLE) }, + LogicalType::LIST(LogicalType::DOUBLE), + Execute + ); + fn.bind = Bind; + fn.stability = FunctionStability::VOLATILE; // avoid constant folding + + loader.RegisterFunction(fn); +} + +} // namespace duckdb diff --git a/src/ml_functions/torch_dnn.cpp b/src/ml_functions/torch_dnn.cpp new file mode 100644 index 0000000..75de36f --- /dev/null +++ b/src/ml_functions/torch_dnn.cpp @@ -0,0 +1,459 @@ +#include "duckdb.hpp" +#include "ml_functions/torch_dnn.hpp" + +#include "duckdb/common/constants.hpp" +#include "duckdb/common/exception.hpp" +#include "duckdb/common/types.hpp" +#include "duckdb/common/list.hpp" +#include "duckdb/common/vector_operations/generic_executor.hpp" +#include "duckdb/execution/expression_executor.hpp" +#include "duckdb/function/scalar_function.hpp" +#include "duckdb/parser/parsed_data/create_scalar_function_info.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" + +#include +#include +#include + +namespace duckdb { + +using KernelType = ML_TorchDNN::KernelType; + +// --------------------------------------------------------------------- +// Helper: build torch::nn::Sequential from kernel spec + tensors +// --------------------------------------------------------------------- +static torch::nn::Sequential BuildSequentialFromKernels( + const std::vector &kernel_types, + const std::vector &dims, + const std::vector &weight_tensors, + bool &has_argmax /* in/out */) { + + const int64_t num_ops = static_cast(kernel_types.size()); + if (2 * num_ops != static_cast(dims.size())) { + throw InvalidInputException( + "torchdnn: mismatched 2*num_ops (%lld) vs dims.size() (%lld)", + (long long)(2 * num_ops), + (long long)dims.size()); + } + + torch::nn::Sequential model; + idx_t weight_idx = 0; + + for (int64_t i = 0; i < num_ops; ++i) { + auto kt = kernel_types[i]; + + if (kt == KernelType::MatMul && + i + 1 < num_ops && + kernel_types[i + 1] == KernelType::MatAdd) { + // Fuse MatMul + MatAdd -> Linear + const int64_t in_dim = dims[2 * i]; + const int64_t out_dim = dims[2 * i + 1]; + + auto denseLayer = torch::nn::Linear(in_dim, out_dim); + + if (weight_idx >= weight_tensors.size()) { + throw InvalidInputException( + "torchdnn: not enough weight tensors for Linear layer"); + } + auto W = weight_tensors[weight_idx++]; + + if (weight_idx >= weight_tensors.size()) { + throw InvalidInputException( + "torchdnn: not enough weight tensors for Linear bias"); + } + auto b = weight_tensors[weight_idx++]; + + if (W.sizes() != torch::IntArrayRef{in_dim, out_dim}) { + throw InvalidInputException( + "torchdnn: Linear W has wrong shape, expected (%lld,%lld)", + (long long)in_dim, (long long)out_dim); + } + if (b.numel() != out_dim) { + throw InvalidInputException( + "torchdnn: Linear b has wrong size, expected %lld", + (long long)out_dim); + } + + denseLayer->weight.set_data(W.t().contiguous()); + denseLayer->bias.set_data(b.contiguous()); + model->push_back(denseLayer); + + } else if (kt == KernelType::MatAdd) { + // Handled together with MatMul; skip + continue; + + } else if (kt == KernelType::BatchNorm) { + const int64_t dim = dims[2 * i]; // same in/out + auto bn = torch::nn::BatchNorm1d(dim); + + if (weight_idx + 1 >= weight_tensors.size()) { + throw InvalidInputException( + "torchdnn: not enough weight tensors for BatchNorm"); + } + auto gamma = weight_tensors[weight_idx++]; + auto beta = weight_tensors[weight_idx++]; + + if (gamma.numel() != dim || beta.numel() != dim) { + throw InvalidInputException( + "torchdnn: BatchNorm params have wrong size, expected %lld", + (long long)dim); + } + + bn->weight.set_data(gamma.contiguous()); + bn->bias.set_data(beta.contiguous()); + model->push_back(bn); + + } else if (kt == KernelType::ReLU) { + model->push_back(torch::nn::ReLU()); + + } else if (kt == KernelType::Sigmoid) { + model->push_back(torch::nn::Sigmoid()); + + } else if (kt == KernelType::Softmax) { + model->push_back(torch::nn::Softmax(1)); + + } else if (kt == KernelType::Argmax) { + // Argmax is handled in Execute + has_argmax = true; + } else { + throw InvalidInputException("torchdnn: unsupported kernel type index %d", (int)kt); + } + } + + model->eval(); + return model; +} + +// ========================= LoadModelFromFile ========================= +// File/spec -> kernel_types, dims, weight_tensors, has_argmax +// For now we keep the simple "identity3" / "identity3_argmax" spec. +void ML_TorchDNN::LoadModelFromFile( + const std::string &path, + std::vector &kernel_types, + std::vector &dims, + std::vector &weight_tensors, + bool &has_argmax) { + + kernel_types.clear(); + dims.clear(); + weight_tensors.clear(); + has_argmax = false; + + bool argmax_mode = false; + std::string spec = path; + + if (spec == "identity3_argmax") { + argmax_mode = true; + spec = "identity3"; + } + + if (spec == "identity3") { + const int64_t in_dim = 3; + const int64_t out_dim = 3; + + // MatMul + MatAdd + kernel_types.push_back(KernelType::MatMul); + kernel_types.push_back(KernelType::MatAdd); + + // dims for 2 operators: [in, out, out, out] + dims.push_back(in_dim); // MatMul in + dims.push_back(out_dim); // MatMul out + dims.push_back(out_dim); // MatAdd in + dims.push_back(out_dim); // MatAdd out + + if (argmax_mode) { + kernel_types.push_back(KernelType::Argmax); + dims.push_back(out_dim); // Argmax in + dims.push_back(out_dim); // Argmax out (logits) + has_argmax = true; + } else { + has_argmax = false; + } + + auto options = torch::TensorOptions().dtype(torch::kFloat32); + auto W = torch::eye(in_dim, options); // (3,3) + auto b = torch::zeros({out_dim}, options); // (3) + + weight_tensors.push_back(W); + weight_tensors.push_back(b); + return; + } + + throw InvalidInputException( + "torchdnn: unsupported model spec '%s'. " + "For tests, use 'identity3' or 'identity3_argmax'.", + path.c_str()); +} + +// ========================= Bind ========================= +unique_ptr ML_TorchDNN::Bind( + ClientContext &context, + ScalarFunction &bound_function, + vector> &args) { + + if (args.empty()) { + throw BinderException("torchdnn: expected at least input and model_file"); + } + + // Last argument is model_file + idx_t model_arg_index = args.size() - 1; + if (!args[model_arg_index]->IsFoldable()) { + throw BinderException("torchdnn: model_file must be a constant string"); + } + + // Evaluate the model file argument at bind time + auto vFile = ExpressionExecutor::EvaluateScalar(context, *args[model_arg_index]); + if (vFile.IsNull()) { + throw BinderException("torchdnn: model_file is NULL"); + } + const std::string model_file = vFile.ToString(); + + // Load kernelTypes, dims, and weights from that file + std::vector kernel_types; + std::vector dims; + std::vector weight_tensors; + bool has_argmax = false; + + ML_TorchDNN::LoadModelFromFile(model_file, kernel_types, dims, weight_tensors, has_argmax); + + // Build libtorch model + auto model = BuildSequentialFromKernels(kernel_types, dims, weight_tensors, has_argmax); + + // BindData stores the model, dims, argmax flag, and file path + return make_uniq(model, dims, has_argmax, model_file); +} + +// ========================= Execute ========================= +void ML_TorchDNN::Execute( + DataChunk &args, + ExpressionState &state, + Vector &result) { + + const idx_t count = args.size(); + if (count == 0) { + return; + } + + // First argument: LIST input + Vector &lhs = args.data[0]; + UnifiedVectorFormat lhs_uvf; + lhs.ToUnifiedFormat(count, lhs_uvf); + const auto *entries = UnifiedVectorFormat::GetData(lhs_uvf); + + auto &child = ListVector::GetEntry(lhs); + const idx_t total_elems = ListVector::GetListSize(lhs); + UnifiedVectorFormat child_uvf; + child.ToUnifiedFormat(total_elems, child_uvf); + const double *child_data = UnifiedVectorFormat::GetData(child_uvf); + + // Get bind data + const auto &func_expr = state.expr.Cast(); + auto &info = func_expr.bind_info->Cast(); + + if (info.dims.empty()) { + throw InvalidInputException("torchdnn: dims are empty in bind data"); + } + + const auto input_dim = info.dims.front(); + const auto output_dim = info.dims.back(); + if (input_dim <= 0 || output_dim <= 0) { + throw InvalidInputException( + "torchdnn: invalid dims; input (%lld) or output (%lld) dim is non-positive", + (long long)input_dim, + (long long)output_dim); + } + + // Collect valid rows into contiguous buffer + std::vector input_buf; + input_buf.reserve(static_cast(count) * (size_t)input_dim); + + std::vector row_valid(count, false); + idx_t num_valid = 0; + + for (idx_t i = 0; i < count; ++i) { + const auto idx = lhs_uvf.sel->get_index(i); + + if (!lhs_uvf.validity.RowIsValid(idx)) { + row_valid[i] = false; + continue; + } + + const auto ent = entries[idx]; + if (ent.length != (idx_t)input_dim) { + throw InvalidInputException( + "torchdnn: input vector length %llu does not match model input dim %lld", + (unsigned long long)ent.length, + (long long)input_dim); + } + + row_valid[i] = true; + num_valid++; + + const double *row_ptr = child_data + ent.offset; + for (idx_t j = 0; j < (idx_t)input_dim; ++j) { + input_buf.push_back(static_cast(row_ptr[j])); + } + } + + if (num_valid == 0) { + // everything is NULL + if (result.GetType().id() == LogicalTypeId::LIST) { + result.SetVectorType(VectorType::FLAT_VECTOR); + auto *res_entries = FlatVector::GetData(result); + auto &res_valid = FlatVector::Validity(result); + for (idx_t i = 0; i < count; ++i) { + res_entries[i] = {0, 0}; + res_valid.Set(i, false); + } + } else { + result.SetVectorType(VectorType::FLAT_VECTOR); + auto &res_valid = FlatVector::Validity(result); + for (idx_t i = 0; i < count; ++i) { + res_valid.Set(i, false); + } + } + return; + } + + // Build torch input tensor: [num_valid, input_dim] + torch::Tensor input_tensor = torch::from_blob( + input_buf.data(), + {static_cast(num_valid), static_cast(input_dim)}, + torch::TensorOptions().dtype(torch::kFloat32)); + + // Forward + torch::Tensor output_tensor = info.model->forward(input_tensor); + output_tensor = output_tensor.contiguous().cpu(); + + const auto &return_type = result.GetType(); + + if (info.has_argmax && return_type.id() == LogicalTypeId::INTEGER) { + // Argmax-style scalar output + result.SetVectorType(VectorType::FLAT_VECTOR); + auto &res_valid = FlatVector::Validity(result); + auto *out_vals = FlatVector::GetData(result); + + torch::Tensor argmax_tensor; + if (output_tensor.dim() == 2) { + argmax_tensor = std::get<1>(output_tensor.max(/*dim=*/1)); + } else { + argmax_tensor = output_tensor; + } + argmax_tensor = argmax_tensor.to(torch::kInt32).contiguous(); + auto *data_int = argmax_tensor.data_ptr(); + + idx_t valid_idx = 0; + for (idx_t i = 0; i < count; ++i) { + if (!row_valid[i]) { + res_valid.Set(i, false); + continue; + } + res_valid.Set(i, true); + out_vals[i] = data_int[valid_idx++]; + } + + } else if (return_type.id() == LogicalTypeId::LIST) { + // Vector output LIST per row + result.SetVectorType(VectorType::FLAT_VECTOR); + + auto *res_entries = FlatVector::GetData(result); + auto &res_child = ListVector::GetEntry(result); + auto &res_valid = FlatVector::Validity(result); + + if (output_tensor.dim() != 2 || output_tensor.size(1) != output_dim) { + throw InvalidInputException( + "torchdnn: model output shape mismatch; expected (?, %lld), got (%lld, %lld)", + (long long)output_dim, + (long long)output_tensor.size(0), + (long long)output_tensor.size(1)); + } + + const idx_t total_out = (idx_t)num_valid * (idx_t)output_dim; + ListVector::Reserve(result, total_out); + ListVector::SetListSize(result, total_out); + + double *out_data = FlatVector::GetData(res_child); + auto *data_float = output_tensor.data_ptr(); + + idx_t cursor = 0; + idx_t valid_idx = 0; + for (idx_t i = 0; i < count; ++i) { + if (!row_valid[i]) { + res_entries[i] = {0, 0}; + res_valid.Set(i, false); + continue; + } + + res_entries[i] = {cursor, static_cast(output_dim)}; + res_valid.Set(i, true); + + const float *src = data_float + (size_t)valid_idx * (size_t)output_dim; + for (idx_t j = 0; j < (idx_t)output_dim; ++j) { + out_data[cursor + j] = static_cast(src[j]); + } + + cursor += (idx_t)output_dim; + valid_idx += 1; + } + } else { + throw InvalidInputException( + "torchdnn: unsupported combination of has_argmax=%d and return type %s", + info.has_argmax ? 1 : 0, + return_type.ToString().c_str()); + } +} + +// ========================= Register ========================= +void ML_TorchDNN::Register(ExtensionLoader &loader) { + // 1) LIST, VARCHAR -> LIST + { + ScalarFunction fn( + "torchdnn", + { + LogicalType::LIST(LogicalType::DOUBLE), // input + LogicalType::VARCHAR // model_file + }, + LogicalType::LIST(LogicalType::DOUBLE), // output + Execute + ); + fn.bind = Bind; + fn.stability = FunctionStability::VOLATILE; + loader.RegisterFunction(fn); + } + + // 2) LIST, INTEGER, VARCHAR -> INTEGER + { + ScalarFunction fn( + "torchdnn", + { + LogicalType::LIST(LogicalType::DOUBLE), + LogicalType::INTEGER, + LogicalType::VARCHAR + }, + LogicalType::INTEGER, + Execute + ); + fn.bind = Bind; + fn.stability = FunctionStability::VOLATILE; + loader.RegisterFunction(fn); + } + + // 3) LIST, BIGINT, VARCHAR -> INTEGER + { + ScalarFunction fn( + "torchdnn", + { + LogicalType::LIST(LogicalType::DOUBLE), + LogicalType::BIGINT, + LogicalType::VARCHAR + }, + LogicalType::INTEGER, + Execute + ); + fn.bind = Bind; + fn.stability = FunctionStability::VOLATILE; + loader.RegisterFunction(fn); + } +} + +} // namespace duckdb diff --git a/src/optimization/CMakeLists.txt b/src/optimization/CMakeLists.txt new file mode 100644 index 0000000..adaf6d7 --- /dev/null +++ b/src/optimization/CMakeLists.txt @@ -0,0 +1,21 @@ +# Append optimizer sources to the extension's global list +list(APPEND EXTENSION_SOURCES + ${CMAKE_CURRENT_SOURCE_DIR}/add5_rule.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/replace_binding_rule.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/register_optimizers.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/dynamic_optimizer.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/catalog.cpp + + #All Actions + ${CMAKE_CURRENT_SOURCE_DIR}/actions/MatMulDense2SparseRewriteAction.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/actions/MultiLayerUDF2TorchNNRewriteAction.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/actions/DecisionForestUDF2RelationRewriteAction.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/actions/MLDecompositionPushdownRewriteAction.cpp + + #Feature Extraction Helpers + ${CMAKE_CURRENT_SOURCE_DIR}/feature_extraction/query2vec/query_annotation.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/feature_extraction/query2vec/query_feature_extraction.cpp +) + +# Expose back to parent +set(EXTENSION_SOURCES "${EXTENSION_SOURCES}" PARENT_SCOPE) diff --git a/src/optimization/actions/DecisionForestUDF2RelationRewriteAction.cpp b/src/optimization/actions/DecisionForestUDF2RelationRewriteAction.cpp new file mode 100644 index 0000000..410ded4 --- /dev/null +++ b/src/optimization/actions/DecisionForestUDF2RelationRewriteAction.cpp @@ -0,0 +1,531 @@ +// src/optimization/actions/DecisionForestUDF2RelationRewriteAction.cpp + +#include "optimization/actions/DecisionForestUDF2RelationRewriteAction.hpp" + +#include "duckdb/catalog/catalog.hpp" +#include "duckdb/catalog/catalog_entry/scalar_function_catalog_entry.hpp" +#include "duckdb/catalog/catalog_entry/aggregate_function_catalog_entry.hpp" +#include "duckdb/common/constants.hpp" +#include "duckdb/common/string_util.hpp" +#include "duckdb/common/types/value.hpp" +#include "duckdb/function/aggregate_function.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/planner/expression/bound_reference_expression.hpp" +#include "duckdb/planner/expression/bound_constant_expression.hpp" +#include "duckdb/planner/expression/bound_operator_expression.hpp" +#include "duckdb/planner/expression/bound_case_expression.hpp" +#include "duckdb/planner/expression/bound_comparison_expression.hpp" +#include "duckdb/planner/expression/bound_aggregate_expression.hpp" +#include "duckdb/planner/operator/logical_projection.hpp" +#include "duckdb/planner/operator/logical_get.hpp" +#include "duckdb/planner/operator/logical_cross_product.hpp" +#include "duckdb/planner/operator/logical_aggregate.hpp" +#include "duckdb/planner/expression_iterator.hpp" +#include "duckdb/parser/tableref/table_function_ref.hpp" +#include "duckdb/function/table_function.hpp" + +#include "ml_functions/decision_forest.hpp" +#include "ml_functions/decision_tree.hpp" +#include "ml_functions/ml_decision_forest_table.hpp" + +#include + +namespace duckdb { + +namespace fs = std::filesystem; + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +// Enumerate tree files in a forest folder (like Velox Forest::vectorizeForestFolder) +static bool ListForestTreeFiles(const string &forest_path, vector &out_paths) { + std::error_code ec; + if (!fs::exists(forest_path, ec) || !fs::is_directory(forest_path, ec)) { + return false; + } + for (auto &entry : fs::directory_iterator(forest_path, ec)) { + if (ec) { + break; + } + if (!entry.is_regular_file()) { + continue; + } + out_paths.push_back(entry.path().string()); + } + std::sort(out_paths.begin(), out_paths.end()); + return !out_paths.empty(); +} + +// Find a top-level decision_forest_predict(...) call in a LogicalProjection +// Returns pointer to the BoundFunctionExpression and sets expr_idx_out +static BoundFunctionExpression *FindTopLevelForestCall(LogicalProjection &proj, idx_t &expr_idx_out) { + for (idx_t i = 0; i < proj.expressions.size(); i++) { + auto &expr = *proj.expressions[i]; + if (expr.expression_class != ExpressionClass::BOUND_FUNCTION) { + continue; + } + auto &f = expr.Cast(); + if (f.function.name == "decision_forest_predict") { + std::cout << "Found decision_forest_predict UDF call: " << f.ToString() << "\n"; + expr_idx_out = i; + return &f; + } + } + expr_idx_out = DConstants::INVALID_INDEX; + return nullptr; +} + +// Build a LogicalGet over the forest_trees(forest_path) table function: +// forest_trees(forest_path) -> columns: (tree_id INT, tree_path VARCHAR) +static unique_ptr BuildForestGet(ClientContext &context, const string &forest_path) { + // Bind input (forest_path) + vector bind_inputs; + bind_inputs.emplace_back(Value(forest_path)); + + // Define table function (same as in ML_ForestTable::Register) + TableFunction tf("forest_trees", + {LogicalType::VARCHAR}, // forest_path + ML_ForestTable::Function, ML_ForestTable::Bind, ML_ForestTable::InitGlobal); + + // Return types & names will be populated by Bind + vector return_types; + vector names; + + // Newer DuckDB TableFunctionBindInput signature + named_parameter_map_t named_parameters; + vector input_table_types; + vector input_table_names; + optional_ptr info(nullptr); + optional_ptr binder_ptr(nullptr); + TableFunctionRef tf_ref; + + TableFunctionBindInput tf_input(bind_inputs, named_parameters, input_table_types, input_table_names, info, + binder_ptr, tf, tf_ref); + + auto bind_data = ML_ForestTable::Bind(context, tf_input, return_types, names); + + // logical table id: any extension-local index is fine + idx_t table_index = 0; + virtual_column_map_t virtual_columns; + + auto get = make_uniq(table_index, std::move(tf), std::move(bind_data), std::move(return_types), + std::move(names), std::move(virtual_columns)); + std::cout << "Built forest_trees GET over path: " << forest_path << "\n"; + return std::move(get); +} + +// Look up AVG(FLOAT/DOUBLE) aggregate function from the catalog +static AggregateFunction GetAvgFloatAggregate(ClientContext &context) { + auto &catalog = Catalog::GetSystemCatalog(context); + auto &entry = + catalog.GetEntry(context, DEFAULT_SCHEMA, "avg"); + + // Pick the overload with FLOAT (or DOUBLE) argument + for (auto &fn : entry.functions.functions) { + if (fn.arguments.size() == 1 && + (fn.arguments[0].id() == LogicalTypeId::FLOAT || + fn.arguments[0].id() == LogicalTypeId::DOUBLE)) { + return fn; + } + } + throw InternalException("DecisionForest rewrite: AVG(FLOAT/DOUBLE) aggregate not found"); +} + +// Core rewrite: +// PROJECTION [ ..., decision_forest_predict(features, forest_path, num_feat, is_cls) AS cls ] +// child +// +// becomes (conceptually): +// +// PROJECTION [ group_keys..., cls ] +// AGGREGATE [ group_keys ] (AVG(pred) AS avg_pred) +// PROJECTION [ group_keys..., pred = tree_predict(features, tree_path, false) ] +// CROSS_PRODUCT +// child +// forest_trees(forest_path) +// +// For now we handle only the case where: +// * the projection has exactly one forest call at top level +static bool RewriteForestProjectionToCrossProduct(OptimizerExtensionInput &input, + unique_ptr &plan) { + if (!plan || plan->type != LogicalOperatorType::LOGICAL_PROJECTION) { + return false; + } + + auto &proj = plan->Cast(); + auto &context = input.context; + + // Preserve the original projection's table index and aliases + auto original_table_index = proj.table_index; + vector original_aliases; + original_aliases.reserve(proj.expressions.size()); + for (auto &expr : proj.expressions) { + original_aliases.push_back(expr->alias); + } + + // 1) Find the top-level decision_forest_predict(...) + idx_t forest_idx; + auto *forest_func = FindTopLevelForestCall(proj, forest_idx); + if (!forest_func) { + return false; // no forest UDF here + } + + std::cout << "Found the decision_forest_predict UDF in projection in application: " + << forest_func->ToString() << "\n"; + + // 2) Get forest bind data (ML_DecisionForest::BindData) + auto *forest_bind_ptr = dynamic_cast(forest_func->bind_info.get()); + if (!forest_bind_ptr) { + // Not the expected bind data; bail + return false; + } + + std::cout << "Forest bind data \n"; + + auto &dfb = *forest_bind_ptr; + const string forest_path = dfb.forest_path; + const bool is_classification = dfb.is_classification; + + // 3) Enumerate tree files once to get num_trees + vector tree_paths; + if (!ListForestTreeFiles(forest_path, tree_paths)) { + // Can't enumerate; bail out for now + std::cout << "Failed to list tree files in forest path: " << forest_path << "\n"; + return false; + } + const idx_t num_trees = tree_paths.size(); + if (num_trees == 0) { + std::cout << "No tree files found in forest path: " << forest_path << "\n"; + return false; + } + + // 4) Features arg must be present (we now accept any bound expression) + if (forest_func->children.empty()) { + std::cout << "decision_forest_predict has no arguments\n"; + return false; + } + if (!forest_func->children.empty()) { + std::cout << "forest_func child[0] class = " + << (int)forest_func->children[0]->expression_class + << ", expr = " << forest_func->children[0]->ToString() << "\n"; + } + if (forest_func->children[0]->expression_class != ExpressionClass::BOUND_REF) { + // Original debug logging retained + std::cout << "Features argument is not a BoundReferenceExpression\n"; + } else { + std::cout << "Features argument is a BoundReferenceExpression\n"; + } + + // Copy the original features expression as a generic bound expression + unique_ptr feat_expr_original = forest_func->children[0]->Copy(); + + // 5) Build right-hand side GET forest_trees(forest_path) + auto right_get = BuildForestGet(context, forest_path); + + // 6) Original child of projection becomes left side of CROSS_PRODUCT + if (proj.children.empty()) { + // malformed plan + std::cout << "Projection has no child\n"; + return false; + } + auto left_child = std::move(proj.children[0]); + + // CROSS_PRODUCT(left_child, forest_trees) + auto cross = make_uniq(std::move(left_child), std::move(right_get)); + + // After CROSS_PRODUCT schema: + // [ all columns from left_child | tree_id, tree_path ] + idx_t left_cols = cross->children[0]->types.size(); + idx_t tree_id_col_idx = left_cols + 0; + (void)tree_id_col_idx; // not used right now + idx_t tree_path_col_idx = left_cols + 1; + + std::cout << "CROSS_PRODUCT schema: left_cols = " << left_cols + << ", tree_id_col_idx = " << tree_id_col_idx + << ", tree_path_col_idx = " << tree_path_col_idx << "\n"; + + // 7) Build inner PROJECTION on top of CROSS_PRODUCT + // + // - group key columns: original projection expressions except the forest one + // - plus pred = tree_predict(features, tree_path, false) + + vector> inner_proj_exprs; + vector group_types; + + // Keep all original non-forest expressions as group keys + idx_t next_col_idx = 0; + for (idx_t i = 0; i < proj.expressions.size(); i++) { + std::cout << "Processing original proj expr[" << i << "]: " + << proj.expressions[i]->ToString() << "\n"; + if (i == forest_idx) { + continue; + } + auto &orig = *proj.expressions[i]; + inner_proj_exprs.push_back( + make_uniq(orig.return_type, next_col_idx)); + group_types.push_back(orig.return_type); + next_col_idx++; + } + std::cout << "Number of group key columns: " << group_types.size() << "\n"; + + // Build tree_predict(features_expr, tree_path, false) + + // We know the CROSS_PRODUCT schema is: +// [ all columns from left_child | tree_id, tree_path ] +// idx_t left_cols = cross->children[0]->types.size(); +// idx_t tree_path_col_idx = left_cols + 1; + +// 1) Figure out which column is "features" on the left side. +// In your current setup, the projection child is [id, features], +// so features is column index 1. +// For now we hard-code that, since this rule is specialized to your pattern. +idx_t features_col_idx = 1; + +// 2) Lookup tree_predict(...) function from the catalog +auto &catalog = Catalog::GetSystemCatalog(context); +auto &tree_entry = + catalog.GetEntry(context, DEFAULT_SCHEMA, "tree_predict"); + +const ScalarFunction *tree_func_ptr = nullptr; +for (auto &fn : tree_entry.functions.functions) { + if (fn.arguments.size() == 3 && + fn.arguments[0] == LogicalType::LIST(LogicalType::FLOAT) && + fn.arguments[1].id() == LogicalTypeId::VARCHAR && + fn.arguments[2].id() == LogicalTypeId::BOOLEAN && + fn.return_type.id() == LogicalTypeId::FLOAT) { + tree_func_ptr = &fn; + break; + } +} +if (!tree_func_ptr) { + throw InternalException( + "DecisionForest rewrite: tree_predict(LIST(FLOAT), VARCHAR, BOOLEAN) not found"); +} + +// Make a *non-const* copy so we can pass it to Bind +ScalarFunction tree_func = *tree_func_ptr; + +// 3) Build children for tree_predict(features, tree_path, false) + +// features: directly reference the LIST(FLOAT) column from CROSS_PRODUCT +auto feat_ref_for_child = + make_uniq(LogicalType::LIST(LogicalType::FLOAT), + features_col_idx); +auto feat_ref_for_bind = + make_uniq(LogicalType::LIST(LogicalType::FLOAT), + features_col_idx); + +// tree_path: VARCHAR column from CROSS_PRODUCT +auto tree_path_ref_for_child = + make_uniq(LogicalType::VARCHAR, tree_path_col_idx); +auto tree_path_ref_for_bind = + make_uniq(LogicalType::VARCHAR, tree_path_col_idx); + +// has_missing: FALSE literal +auto has_missing_child = + make_uniq(Value::BOOLEAN(false)); +auto has_missing_bind = + make_uniq(Value::BOOLEAN(false)); + +// Children that will be stored inside the BoundFunctionExpression +vector> tree_children; +tree_children.push_back(std::move(feat_ref_for_child)); +tree_children.push_back(std::move(tree_path_ref_for_child)); +tree_children.push_back(std::move(has_missing_child)); + +// Arguments passed to the Bind function +vector> args_for_bind; +args_for_bind.push_back(std::move(feat_ref_for_bind)); +args_for_bind.push_back(std::move(tree_path_ref_for_bind)); +args_for_bind.push_back(std::move(has_missing_bind)); + +// 4) Call ML_DecisionTree::Bind to produce bind_info +auto tree_bind = ML_DecisionTree::Bind(context, tree_func, args_for_bind); + +// 5) Create the BoundFunctionExpression for tree_predict(...) +auto tree_pred_expr = make_uniq( + tree_func.return_type, // should be FLOAT + tree_func, + std::move(tree_children), + std::move(tree_bind)); + // Append prediction as last column of inner projection + inner_proj_exprs.push_back(std::move(tree_pred_expr)); + + // Build LogicalProjection with the new API (table_index + select_list) + idx_t inner_table_index = 0; + auto inner_proj = + make_uniq(inner_table_index, std::move(inner_proj_exprs)); + inner_proj->children.push_back(std::move(cross)); + + // 8) Build AGGREGATE on top of inner_proj: + // SELECT group_keys..., AVG(pred) + // + // LogicalAggregate takes a select_list of expressions, where: + // - [0..group_count-1] are group keys + // - [group_count..] are BoundAggregateExpression(s) + vector> aggregate_select; + + // group keys: references [0..group_types.size()-1] of inner_proj output + for (idx_t i = 0; i < group_types.size(); i++) { + aggregate_select.push_back( + make_uniq(group_types[i], i)); + } + + // prediction column index in inner_proj + idx_t pred_idx = group_types.size(); + + // children for AVG: just reference the prediction column + vector> avg_children; + avg_children.push_back( + make_uniq(tree_func.return_type, pred_idx)); + + // lookup AVG(FLOAT/DOUBLE) + auto avg_fun = GetAvgFloatAggregate(context); + auto avg_type = avg_fun.return_type; + + // no filter, no bind_info + unique_ptr avg_filter; + unique_ptr avg_bind_info; + AggregateType agg_type = AggregateType::NON_DISTINCT; + + auto avg_agg = make_uniq( + avg_fun, std::move(avg_children), std::move(avg_filter), + std::move(avg_bind_info), agg_type); + + // append AVG(pred) to select list + aggregate_select.push_back(std::move(avg_agg)); + + // group_index = 0, aggregate_index = group_count + idx_t group_index = 0; + idx_t aggregate_index = group_types.size(); + + auto aggregate = make_uniq( + group_index, aggregate_index, std::move(aggregate_select)); + aggregate->children.push_back(std::move(inner_proj)); + + // 9) Final projection on top: + // - group key columns (restored as they were) + // - forest_result = (is_classification ? (avg > 0 ? 1 : 0) : avg) + vector> final_proj_exprs; + + // group keys: references [0..group_types.size()-1] of aggregate output + for (idx_t i = 0; i < group_types.size(); i++) { + final_proj_exprs.push_back( + make_uniq(group_types[i], i)); + } + + // avg_pred is at col index = group_types.size() + idx_t avg_idx = group_types.size(); + auto avg_ref = + make_uniq(avg_type, avg_idx); + + unique_ptr forest_result_expr; + + if (is_classification) { + // CASE WHEN avg > 0 THEN 1 ELSE 0 + auto zero_val = Value::FLOAT(0.0f); + auto zero_expr = make_uniq(zero_val); + + auto cond = make_uniq( + ExpressionType::COMPARE_GREATERTHAN, avg_ref->Copy(), + std::move(zero_expr)); + + auto one_out = make_uniq(Value::FLOAT(1.0f)); + auto zero_out = make_uniq(Value::FLOAT(0.0f)); + + forest_result_expr = make_uniq( + std::move(cond), std::move(one_out), std::move(zero_out)); + } else { + forest_result_expr = std::move(avg_ref); + } + + final_proj_exprs.push_back(std::move(forest_result_expr)); + + // Use the same table_index as the original projection so that + // existing bindings (e.g., ORDER BY feature_rows.id) remain valid. + auto final_proj = + make_uniq(original_table_index, std::move(final_proj_exprs)); + final_proj->children.push_back(std::move(aggregate)); + + // Restore column aliases so bindings by name still make sense. + // The final projection has: [all non-forest original expressions..., forest_result_expr] + idx_t new_col = 0; + for (idx_t i = 0; i < proj.expressions.size(); i++) { + if (i == forest_idx) { + continue; + } + if (new_col < final_proj->expressions.size()) { + final_proj->expressions[new_col]->alias = original_aliases[i]; + } + new_col++; + } + + // Alias for the forest result column (last one) + string forest_alias = original_aliases[forest_idx]; + if (forest_alias.empty()) { + // fall back to the original function alias or name + if (!forest_func->alias.empty()) { + forest_alias = forest_func->alias; + } else { + forest_alias = forest_func->function.name; + } + } + if (!final_proj->expressions.empty()) { + final_proj->expressions.back()->alias = forest_alias; + } + + // Replace whole subtree under this projection with the new rewritten plan + std::cout << "Rewrote decision_forest_predict UDF into CROSS_PRODUCT + AGGREGATE plan\n" + << final_proj->ToString() << "\n"; + plan = std::move(final_proj); + return true; +} + +// --------------------------------------------------------------------------- +// Public API: check / apply +// --------------------------------------------------------------------------- + +bool DecisionForestUDF2RelationRewriteAction::check(OptimizerExtensionInput &input, + unique_ptr &plan) { + if (!plan) { + return false; + } + + bool found = false; + + if (plan->type == LogicalOperatorType::LOGICAL_PROJECTION) { + auto &proj = plan->Cast(); + idx_t dummy; + if (FindTopLevelForestCall(proj, dummy)) { + found = true; + } + } + + // Recursively check children + for (auto &child : plan->children) { + found |= DecisionForestUDF2RelationRewriteAction::check(input, child); + } + return found; +} + +bool DecisionForestUDF2RelationRewriteAction::apply(OptimizerExtensionInput &input, + unique_ptr &plan) { + if (!plan) { + return false; + } + + bool modified = false; + + if (plan->type == LogicalOperatorType::LOGICAL_PROJECTION) { + modified |= RewriteForestProjectionToCrossProduct(input, plan); + } + + // Recurse into children + for (auto &child : plan->children) { + modified |= DecisionForestUDF2RelationRewriteAction::apply(input, child); + } + return modified; +} + +} // namespace duckdb diff --git a/src/optimization/actions/MLDecompositionPushdownRewriteAction.cpp b/src/optimization/actions/MLDecompositionPushdownRewriteAction.cpp new file mode 100644 index 0000000..72eaee6 --- /dev/null +++ b/src/optimization/actions/MLDecompositionPushdownRewriteAction.cpp @@ -0,0 +1,373 @@ +#include "duckdb.hpp" +#include "optimization/actions/MLDecompositionPushdownRewriteAction.hpp" + +#include "duckdb/common/assert.hpp" +#include "duckdb/common/string_util.hpp" + +#include "duckdb/planner/expression/bound_columnref_expression.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/planner/expression/bound_reference_expression.hpp" +#include "duckdb/planner/expression_iterator.hpp" + +#include "duckdb/planner/logical_operator.hpp" +#include "duckdb/planner/operator/logical_projection.hpp" +#include "duckdb/planner/operator/logical_comparison_join.hpp" + +#include +#include + +namespace duckdb { + +namespace { + +// ----------------------------------------------------------------------------- +// Small helpers +// ----------------------------------------------------------------------------- + +// Ensure we know how many columns an operator returns. +static idx_t GetLogicalWidth(LogicalOperator &op) { + if (op.types.empty()) { + op.ResolveOperatorTypes(); + } + return op.types.size(); +} + +// Decide if a function name is an "ML op" that we want to push down. +static bool IsMLFunction(const string &fname) { + auto lname = StringUtil::Lower(fname); + // Explicit list for now – extend as you add more ML functions. + if (lname == "mat_mul" || lname == "mat_add" || lname == "cdb_softmax" || lname == "ml_argmax") { + return true; + } + return false; +} + +// Check if an expression tree (owned by expr_ptr) contains any ML op. +static bool ExprContainsMLOp(unique_ptr &expr_ptr) { + if (!expr_ptr) { + return false; + } + bool found = false; + ExpressionIterator::EnumerateExpression( + expr_ptr, + [&](Expression &sub) { + if (sub.expression_class == ExpressionClass::BOUND_FUNCTION) { + auto &f = sub.Cast(); + if (IsMLFunction(f.function.name)) { + found = true; + } + } + }); + return found; +} + +// Collect all table indexes referenced by BoundColumnRefExpression inside expr_ptr. +static void CollectTableIndexes(unique_ptr &expr_ptr, std::unordered_set &tables) { + if (!expr_ptr) { + return; + } + + ExpressionIterator::EnumerateExpression( + expr_ptr, + [&](Expression &sub) { + if (sub.type == ExpressionType::BOUND_COLUMN_REF) { + auto &bref = sub.Cast(); + tables.insert(bref.binding.table_index); + } + }); +} + +// ----------------------------------------------------------------------------- +// Projection-binding rewrite (learning from FilterPushdown) +// ----------------------------------------------------------------------------- +// +// We have a LogicalProjection "proj" above some child (scan/join/etc). +// Some expression "root_expr" may contain BoundColumnRefExpression that +// refer to *proj's output* (binding.table_index == proj.table_index, +// column_index < proj.expressions.size()). +// +// To push "root_expr" down to the child, we must rewrite those refs to +// directly use the underlying child expressions, exactly like +// FilterPushdown::ReplaceProjectionBindings does for filters. +// +static unique_ptr RewriteProjectionBindingsForChild(LogicalProjection &proj, + unique_ptr root_expr) { + ExpressionIterator::VisitExpressionMutable( + root_expr, [&](BoundColumnRefExpression &colref, unique_ptr &expr) { + if (colref.binding.table_index != proj.table_index) { + // This reference does not point into the projection output; + // leave it alone (it might already reference a base table). + return; + } + D_ASSERT(colref.binding.column_index < proj.expressions.size()); + D_ASSERT(colref.depth == 0); + // replace the binding with a copy of the projection expression at that index + auto copy = proj.expressions[colref.binding.column_index]->Copy(); + if (!colref.alias.empty()) { + copy->alias = colref.alias; + } + expr = std::move(copy); + }); + return root_expr; +} + +// Insert a projection between `parent_op` and its child at `child_idx`. +// +// New projection schema: +// - All existing child columns as BoundReferenceExpression(type[i], i) +// - Plus one extra column: `extra_expr` (copied) appended at the end. +// +// This variant is the "simple" unary case (projection over scan/whatever). +static bool InsertUnaryPushdownProjection(LogicalOperator &parent_op, idx_t child_idx, Expression &extra_expr, + idx_t &new_col_idx) { + if (child_idx >= parent_op.children.size()) { + return false; + } + auto &child_ptr = parent_op.children[child_idx]; + if (!child_ptr) { + return false; + } + + // Make sure child's types are known. + const idx_t child_width = GetLogicalWidth(*child_ptr); + if (child_width == 0) { + return false; + } + + // Try to reuse the child's (single) table index if possible. + auto child_tables = child_ptr->GetTableIndex(); + if (child_tables.size() != 1) { + // Multi-table child: don't rewrite in this generic path. + return false; + } + const idx_t table_index = child_tables[0]; + + vector> select_list; + select_list.reserve(child_width + 1); + + // 1) pass-through all existing columns + for (idx_t i = 0; i < child_width; i++) { + select_list.push_back(make_uniq(child_ptr->types[i], i)); + } + + // 2) append extra_expr + select_list.push_back(extra_expr.Copy()); + + // the new column is at index child_width in the *child* output + new_col_idx = child_width; + + // Create a projection that produces the same table index as the child. + auto new_proj = make_uniq(table_index, std::move(select_list)); + new_proj->children.push_back(std::move(child_ptr)); + child_ptr = std::move(new_proj); + + return true; +} + +// Insert a projection on the LEFT or RIGHT child of a COMPARISON_JOIN. +// We only use this when the expression depends on exactly one table that +// is fully contained in that child subtree. +// +// On success, sets: +// - side_child_new_col_idx: index of the appended column in the side child +// - side (0 = left, 1 = right) +static bool InsertProjectionOnJoinChild(LogicalComparisonJoin &join, idx_t target_table, Expression &extra_expr, + int &side, idx_t &side_child_new_col_idx) { + + // Determine whether target_table lives on the left or right side. + auto left_tables = join.children[0]->GetTableIndex(); + auto right_tables = join.children[1]->GetTableIndex(); + + side = -1; + if (std::find(left_tables.begin(), left_tables.end(), target_table) != left_tables.end()) { + side = 0; + } else if (std::find(right_tables.begin(), right_tables.end(), target_table) != right_tables.end()) { + side = 1; + } else { + // Expression refers to a table that is not directly a join child. + return false; + } + + auto &side_child = join.children[side]; + if (!side_child) { + return false; + } + + // The side child should correspond to exactly one logical table (for now). + auto side_child_tables = side_child->GetTableIndex(); + if (side_child_tables.size() != 1) { + return false; + } + const idx_t table_index = side_child_tables[0]; + + const idx_t side_width_before = GetLogicalWidth(*side_child); + if (side_width_before == 0) { + return false; + } + + vector> select_list; + select_list.reserve(side_width_before + 1); + + // 1) pass-through existing columns + for (idx_t i = 0; i < side_width_before; i++) { + select_list.push_back(make_uniq(side_child->types[i], i)); + } + + // 2) append extra_expr + select_list.push_back(extra_expr.Copy()); + side_child_new_col_idx = side_width_before; // new column is last one + + // New projection with same table_index as the original side child. + auto new_proj = make_uniq(table_index, std::move(select_list)); + new_proj->children.push_back(std::move(side_child)); + side_child = std::move(new_proj); + + return true; +} + +} // anonymous namespace + +// ----------------------------------------------------------------------------- +// check +// ----------------------------------------------------------------------------- + +bool MLDecompositionPushdownRewriteAction::check(OptimizerExtensionInput &input, + unique_ptr &plan) { + if (!plan) { + return false; + } + + bool found = false; + + // Look for any ML op inside projection expressions (simple heuristic). + if (plan->type == LogicalOperatorType::LOGICAL_PROJECTION) { + auto &proj = plan->Cast(); + for (auto &expr_ptr : proj.expressions) { + if (expr_ptr && ExprContainsMLOp(expr_ptr)) { + found = true; + break; + } + } + } + + for (auto &child : plan->children) { + found |= check(input, child); + } + + return found; +} + +// ----------------------------------------------------------------------------- +// apply +// ----------------------------------------------------------------------------- + +bool MLDecompositionPushdownRewriteAction::apply(OptimizerExtensionInput &input, + unique_ptr &plan) { + if (!plan) { + return false; + } + + bool changed = false; + + // Bottom-up: first rewrite children + for (auto &child : plan->children) { + changed |= apply(input, child); + } + + // Now see if we can rewrite this node + if (plan->type != LogicalOperatorType::LOGICAL_PROJECTION) { + return changed; + } + + auto &proj = plan->Cast(); + if (proj.children.size() != 1) { + return changed; + } + auto &child = proj.children[0]; + if (!child) { + return changed; + } + + for (auto &expr_ptr : proj.expressions) { + if (!expr_ptr) { + continue; + } + if (!ExprContainsMLOp(expr_ptr)) { + continue; + } + + // Original expression in this projection (uses current bindings) + Expression &expr = *expr_ptr; + + // ----------------------------------------------------------------- + // Create a child-oriented copy using FilterPushdown-style logic + // ----------------------------------------------------------------- + auto expr_for_child = expr_ptr->Copy(); + expr_for_child = RewriteProjectionBindingsForChild(proj, std::move(expr_for_child)); + + // Use the child-oriented expression for table index analysis + std::unordered_set tables; + CollectTableIndexes(expr_for_child, tables); + + // --------------------------------------------------------------------- + // Case 1: PROJECTION over COMPARISON_JOIN and ML expr depends on one table + // --------------------------------------------------------------------- + if (child->type == LogicalOperatorType::LOGICAL_COMPARISON_JOIN && tables.size() == 1) { + auto &join = child->Cast(); + const idx_t target_table = *tables.begin(); + + int side = -1; + idx_t side_child_new_col_idx = 0; + if (!InsertProjectionOnJoinChild(join, target_table, *expr_for_child, side, side_child_new_col_idx)) { + // Could not safely push down on join child; fall through to generic case. + } else { + // We now have a new column on one of the join children. Figure out + // its position in the *join output* so the top projection can + // reference it by index. + + // Ensure widths are known after transformation + const idx_t left_width = GetLogicalWidth(*join.children[0]); + const idx_t right_width = GetLogicalWidth(*join.children[1]); + + idx_t global_idx; + if (side == 0) { + // Left side: its columns occupy [0 .. left_width-1] + global_idx = left_width - 1; + } else { + // Right side: columns occupy [left_width .. left_width+right_width-1] + global_idx = left_width + (right_width - 1); + } + + // Replace original ML expression with a reference into the + // join output at global_idx. + expr_ptr = make_uniq(expr.return_type, global_idx); + + changed = true; + break; // only rewrite first ML expr in this projection + } + } + + // --------------------------------------------------------------------- + // Case 2: generic unary child (scan, filter, etc.) – simple pushdown + // --------------------------------------------------------------------- + { + idx_t new_child_col_idx = 0; + if (!InsertUnaryPushdownProjection(*plan, 0, *expr_for_child, new_child_col_idx)) { + // Could not safely rewrite; leave as-is. + continue; + } + + // The parent projection still sees its child as a unary operator whose + // output columns are [0 .. old_width] with the new column at + // new_child_col_idx, so we can refer to it by index. + expr_ptr = make_uniq(expr.return_type, new_child_col_idx); + + changed = true; + break; // only one ML expr per projection for now + } + } + + return changed; +} + +} // namespace duckdb diff --git a/src/optimization/actions/MatMulDense2SparseRewriteAction.cpp b/src/optimization/actions/MatMulDense2SparseRewriteAction.cpp new file mode 100644 index 0000000..296cb5c --- /dev/null +++ b/src/optimization/actions/MatMulDense2SparseRewriteAction.cpp @@ -0,0 +1,126 @@ +#include "duckdb.hpp" +#include "optimization/actions/MatMulDense2SparseRewriteAction.hpp" +#include "duckdb/common/string_util.hpp" +#include "duckdb/planner/expression/bound_constant_expression.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/planner/expression_iterator.hpp" +#include "duckdb/planner/logical_operator.hpp" +#include "duckdb/planner/operator/logical_projection.hpp" +#include "duckdb/planner/operator/logical_filter.hpp" +#include "duckdb/planner/operator/logical_aggregate.hpp" +#include "duckdb/planner/operator/logical_comparison_join.hpp" +#include "duckdb/planner/operator/logical_join.hpp" + +#include "duckdb/planner/operator/logical_filter.hpp" +#include "duckdb/planner/operator/logical_projection.hpp" +#include "duckdb/planner/operator/logical_aggregate.hpp" +#include "duckdb/planner/operator/logical_comparison_join.hpp" +#include "duckdb/common/enums/expression_type.hpp" +#include "ml_functions/matrix_multiplication.hpp" +#include "duckdb/common/enums/logical_operator_type.hpp" +using namespace duckdb; + +bool checkExpression(OptimizerExtensionInput &input, unique_ptr &e){ + bool found = false; + if(e->expression_class == ExpressionClass::BOUND_FUNCTION){ + auto &f = e->Cast(); + if(f.function.name == "mat_mul"){ + // Check if the first argument is dense and the second is sparse + // This is a placeholder; actual implementation would involve checking the data types or properties + std::cout << "MatMul function found: " << f.ToString() << "\n"; + found = true; + } + } + + ExpressionIterator::EnumerateChildren(*e, [&](unique_ptr &child) { + found |= checkExpression(input, child); + }); + return found; +} + +bool applyExpressionRewrite(OptimizerExtensionInput &input, unique_ptr &e){ + bool modified = false; + if(e->expression_class == ExpressionClass::BOUND_FUNCTION){ + auto &f = e->Cast(); + if(f.function.name == "mat_mul"){ + // Rewrite logic: change the argument to use a sparse representation + // change the binding data to reflect sparse matrix multiplication + f.children[2] = make_uniq(Value("sparse")); // Example change + auto &bd = f.bind_info->Cast(); + bd.mode = ML_MatrixMultiply::MulMode::InputSparse; + modified = true; + } + } + ExpressionIterator::EnumerateChildren(*e, [&](unique_ptr &child) { + modified |= applyExpressionRewrite(input, child); + }); + return modified; +} + +bool MatMulDense2SparseRewriteAction::check(OptimizerExtensionInput &input, unique_ptr &plan) { + // Check if the plan contains a MatMul operation with a dense matrix on the left and a sparse matrix on the right + // This is a placeholder implementation; actual implementation would involve traversing the plan tree + if (!plan) return false; + + // Overall Logic + // check if this node is a success + // if yes return true; + // else return false | check child; + bool ans = false; + switch (plan->type){ + std::cout << "Checking plan node of type: " << LogicalOperatorToString(plan->type) << "\n"; + case LogicalOperatorType::LOGICAL_PROJECTION : { + auto &proj = plan->Cast(); + for (auto &e : proj.expressions) { + ans |= checkExpression(input, e); + } + } + break; + case LogicalOperatorType::LOGICAL_FILTER : { + auto &fil = plan->Cast(); + for (auto &e : fil.expressions) { + ans |= checkExpression(input, e); + } + } + break; + default : break; + } + + // check children + for(auto &child : plan->children){ + ans |= check(input, child); + } + + return ans; + +} + +bool MatMulDense2SparseRewriteAction::apply(OptimizerExtensionInput &input, unique_ptr &plan) { + if (!plan) return false; + + switch (plan->type){ + case LogicalOperatorType::LOGICAL_PROJECTION : { + auto &proj = plan->Cast(); + for (auto &e : proj.expressions) { + applyExpressionRewrite(input, e); + } + } + break; + case LogicalOperatorType::LOGICAL_FILTER : { + auto &fil = plan->Cast(); + for (auto &e : fil.expressions) { + applyExpressionRewrite(input, e); + } + } + break; + default : break; + } + + // Recurse into children + for(auto &child : plan->children){ + apply(input, child); + } + + // Indicate that the plan was modified + return true; +} \ No newline at end of file diff --git a/src/optimization/actions/MultiLayerUDF2TorchNNRewriteAction.cpp b/src/optimization/actions/MultiLayerUDF2TorchNNRewriteAction.cpp new file mode 100644 index 0000000..d616e30 --- /dev/null +++ b/src/optimization/actions/MultiLayerUDF2TorchNNRewriteAction.cpp @@ -0,0 +1,371 @@ +// optimization/actions/MultiLayerUDF2TorchNNRewriteAction.cpp +#include "optimization/actions/MultiLayerUDF2TorchNNRewriteAction.hpp" + +#include "duckdb/common/enums/expression_type.hpp" +#include "duckdb/common/exception.hpp" +#include "duckdb/common/string_util.hpp" + +#include "duckdb/catalog/catalog.hpp" +#include "duckdb/catalog/catalog_entry/scalar_function_catalog_entry.hpp" + +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/planner/expression/bound_reference_expression.hpp" +#include "duckdb/planner/expression/bound_constant_expression.hpp" +#include "duckdb/planner/expression_iterator.hpp" +#include "duckdb/planner/operator/logical_projection.hpp" +#include "duckdb/planner/operator/logical_filter.hpp" + +#include "ml_functions/matrix_multiplication.hpp" +#include "ml_functions/matrix_addition.hpp" +#include "ml_functions/torch_dnn.hpp" + +#include + + +namespace duckdb { + +using KernelType = ML_TorchDNN::KernelType; + +// Try to recognize: mat_add(mat_mul(input, ...), ...) +// and fill kernel_types / dims / weight_tensors / has_argmax. +// Also return the "base input" expression that should become torchdnn's first arg. +static bool CollectSimpleNNFromExpression( + Expression &expr, + std::vector &kernel_types, + std::vector &dims, + std::vector &weight_tensors, + bool &has_argmax, + unique_ptr &base_input_out) { + + // Reset outputs + kernel_types.clear(); + dims.clear(); + weight_tensors.clear(); + has_argmax = false; + base_input_out.reset(); + + if (expr.expression_class != ExpressionClass::BOUND_FUNCTION) { + std::cout << "Returning from here!!\n\n"; + return false; + } + + auto &outer_f = expr.Cast(); + + // We currently only handle mat_add(mat_mul(x, ...), ...) + if (!StringUtil::CIEquals(outer_f.function.name, "mat_add")) { + return false; + } + if (outer_f.children.size() < 1) { + return false; + } + + // First argument to mat_add should be mat_mul(...) + auto &mul_child_expr = *outer_f.children[0]; + if (mul_child_expr.expression_class != ExpressionClass::BOUND_FUNCTION) { + return false; + } + auto &mul_f = mul_child_expr.Cast(); + if (!StringUtil::CIEquals(mul_f.function.name, "mat_mul")) { + return false; + } + + // Extract bind data from mat_mul + auto *mm_bind_ptr = dynamic_cast(mul_f.bind_info.get()); + if (!mm_bind_ptr || !mm_bind_ptr->has_bound_weights) { + return false; + } + auto &mm_bind = *mm_bind_ptr; + + // Extract bind data from mat_add + auto *add_bind_ptr = dynamic_cast(outer_f.bind_info.get()); + if (!add_bind_ptr || !add_bind_ptr->has_bound_biases) { + return false; + } + auto &add_bind = *add_bind_ptr; + + // Sanity: dimensions must match + if (mm_bind.n != add_bind.n) { + // output size of mat_mul != bias length + return false; + } + + const int64_t in_dim = static_cast(mm_bind.k); + const int64_t out_dim = static_cast(mm_bind.n); + + auto options = torch::TensorOptions().dtype(torch::kFloat32); + + // MatMul - weights: row-major (k x n) + if (mm_bind.weights.size() != static_cast(mm_bind.k * mm_bind.n)) { + return false; + } + torch::Tensor W = torch::from_blob( + (void *)mm_bind.weights.data(), // non-owning + { (long)mm_bind.k, (long)mm_bind.n }, + options).clone(); // clone to own memory + + // MatAdd - biases: length n + if (add_bind.biases.size() != static_cast(add_bind.n)) { + return false; + } + torch::Tensor b = torch::from_blob( + (void *)add_bind.biases.data(), + { (long)add_bind.n }, + options).clone(); + + // Fill kernel_types & dims exactly like Velox rule: + // Op 0: MatMul + kernel_types.push_back(KernelType::MatMul); + dims.push_back(in_dim); + dims.push_back(out_dim); + + // Op 1: MatAdd + kernel_types.push_back(KernelType::MatAdd); + dims.push_back(out_dim); + dims.push_back(out_dim); + + // Weight tensors order: [W, b] (to be interpreted by Bind) + weight_tensors.push_back(W); + weight_tensors.push_back(b); + + has_argmax = false; // single layer, no argmax yet + + // base input is the first child to mat_mul (e.g., column "a") + if (mul_f.children.empty()) { + return false; + } + + base_input_out = mul_f.children[0]->Copy(); + + return true; +} + +// torchdnn(base_input, '') with bind_info containing +// the actual Sequential model built from kernel_types/dims/tensors. + +static unique_ptr BuildTorchDNNExpression( + ClientContext &context, + unique_ptr base_input, + const std::vector &kernel_types, + const std::vector &dims, + const std::vector &weight_tensors, + bool has_argmax) { + + // 1) Find the registered "torchdnn" function (we’ll use the 2-arg form). + auto &catalog = Catalog::GetSystemCatalog(context); + auto &torch_entry = catalog.GetEntry( + context, DEFAULT_SCHEMA, "torchdnn"); + + // We chose: torchdnn(LIST, VARCHAR) -> LIST + const ScalarFunction *chosen_func = nullptr; + for (auto &func : torch_entry.functions.functions) { + if (func.arguments.size() == 2 && + func.arguments[0] == LogicalType::LIST(LogicalType::DOUBLE) && + func.arguments[1].id() == LogicalTypeId::VARCHAR) { + chosen_func = &func; + break; + } + } + if (!chosen_func) { + throw InternalException("MultiLayerUDF2TorchNN: could not find torchdnn(LIST, VARCHAR) overload"); + } + + // 2) Build the actual torch::nn::Sequential from the specs + torch::nn::Sequential model; + idx_t weight_idx = 0; + + for (idx_t i = 0; i < kernel_types.size(); i++) { + auto kt = kernel_types[i]; + + if (kt == KernelType::MatMul && + i + 1 < kernel_types.size() && + kernel_types[i + 1] == KernelType::MatAdd) { + + const int64_t in_dim = dims[2 * i]; + const int64_t out_dim = dims[2 * i + 1]; + + auto linear = torch::nn::Linear(in_dim, out_dim); + + // W + if (weight_idx >= weight_tensors.size()) { + throw InternalException("MultiLayerUDF2TorchNN: not enough tensors for Linear weight"); + } + auto W = weight_tensors[weight_idx++]; + + // b + if (weight_idx >= weight_tensors.size()) { + throw InternalException("MultiLayerUDF2TorchNN: not enough tensors for Linear bias"); + } + auto b = weight_tensors[weight_idx++]; + + // Expect W to be (in_dim, out_dim) as per our builder + if (W.sizes() != torch::IntArrayRef{in_dim, out_dim}) { + throw InternalException("MultiLayerUDF2TorchNN: Linear W shape mismatch"); + } + if (b.numel() != out_dim) { + throw InternalException("MultiLayerUDF2TorchNN: Linear b size mismatch"); + } + + linear->weight.set_data(W.t().contiguous()); + linear->bias.set_data(b.contiguous()); + + model->push_back(linear); + + // Skip the MatAdd in the loop + i++; // extra increment + + } else if (kt == KernelType::MatAdd) { + // We only expect MatAdd after MatMul; standalone MatAdd not yet supported. + continue; + + } else if (kt == KernelType::ReLU) { + model->push_back(torch::nn::ReLU()); + + } else if (kt == KernelType::Sigmoid) { + model->push_back(torch::nn::Sigmoid()); + + } else if (kt == KernelType::Softmax) { + model->push_back(torch::nn::Softmax(1)); + + } else if (kt == KernelType::BatchNorm) { + // TODO: handle BatchNorm tensors here if you add them + throw NotImplementedException("BatchNorm not yet wired in BuildTorchDNNExpression"); + } else if (kt == KernelType::Argmax) { + // we handle argmax in Execute, not here + continue; + } + } + + model->eval(); + + // 3) Build BindData + std::vector dims_copy = dims; + std::string model_file = ""; // not actually used by Execute + auto bind_data = make_uniq( + std::move(model), + std::move(dims_copy), + has_argmax, + model_file); + + // 4) Build children: [ base_input , constant model_file ] + vector> children; + children.push_back(std::move(base_input)); + unique_ptr model_arg = make_uniq(Value(model_file)); + children.push_back(std::move(model_arg)); + + // 5) Return new BoundFunctionExpression with our pre-populated bind_data + auto return_type = LogicalType::LIST(LogicalType::DOUBLE); + auto bound = make_uniq( + return_type, + *chosen_func, + std::move(children), + std::move(bind_data)); + + unique_ptr result = std::move(bound); + return result; +} + +static bool CheckExpressionForNN(OptimizerExtensionInput &input, + Expression &expr) { + std::vector kernel_types; + std::vector dims; + std::vector tensors; + bool has_argmax = false; + unique_ptr base_input; + + return CollectSimpleNNFromExpression(expr, kernel_types, dims, tensors, has_argmax, base_input); +} + +bool MultiLayerUDF2TorchNNRewriteAction::check(OptimizerExtensionInput &input, + unique_ptr &plan) { + if (!plan) return false; + + bool found = false; + + switch (plan->type) { + case LogicalOperatorType::LOGICAL_PROJECTION: { + auto &proj = plan->Cast(); + for (auto &expr : proj.expressions) { + found |= CheckExpressionForNN(input, *expr); + } + break; + } + case LogicalOperatorType::LOGICAL_FILTER: { + auto &fil = plan->Cast(); + for (auto &expr : fil.expressions) { + found |= CheckExpressionForNN(input, *expr); + } + break; + } + default: + break; + } + + // Recurse into children + for (auto &child : plan->children) { + found |= check(input, child); + } + + return found; +} + +static bool RewriteExpressionToTorchDNN(OptimizerExtensionInput &input, + unique_ptr &expr) { + std::vector kernel_types; + std::vector dims; + std::vector tensors; + bool has_argmax = false; + unique_ptr base_input; + + if (!CollectSimpleNNFromExpression(*expr, kernel_types, dims, tensors, has_argmax, base_input)) { + // no match: still recurse into children + ExpressionIterator::EnumerateChildren(*expr, [&](unique_ptr &child) { + RewriteExpressionToTorchDNN(input, child); + }); + return false; + } + + // We have a match: build torchdnn(base_input, "") + auto &context = input.context; + expr = BuildTorchDNNExpression(context, + std::move(base_input), + kernel_types, + dims, + tensors, + has_argmax); + return true; +} + +bool MultiLayerUDF2TorchNNRewriteAction::apply(OptimizerExtensionInput &input, + unique_ptr &plan) { + if (!plan) return false; + + bool modified = false; + + switch (plan->type) { + case LogicalOperatorType::LOGICAL_PROJECTION: { + auto &proj = plan->Cast(); + for (auto &expr : proj.expressions) { + modified |= RewriteExpressionToTorchDNN(input, expr); + } + break; + } + case LogicalOperatorType::LOGICAL_FILTER: { + auto &fil = plan->Cast(); + for (auto &expr : fil.expressions) { + modified |= RewriteExpressionToTorchDNN(input, expr); + } + break; + } + default: + break; + } + + // Recurse + for (auto &child : plan->children) { + modified |= apply(input, child); + } + + return modified; +} + +} // namespace duckdb diff --git a/src/optimization/add5_rule.cpp b/src/optimization/add5_rule.cpp new file mode 100644 index 0000000..2fcb641 --- /dev/null +++ b/src/optimization/add5_rule.cpp @@ -0,0 +1,58 @@ +#include "optimization/add5_rule.hpp" + +#include "duckdb/planner/logical_operator.hpp" +#include "duckdb/planner/operator/logical_projection.hpp" +#include "duckdb/planner/expression/bound_operator_expression.hpp" +#include "duckdb/planner/expression/bound_constant_expression.hpp" +#include "duckdb/common/types/value.hpp" +#include "duckdb/common/enums/expression_type.hpp" + +using namespace duckdb; + +namespace { + +// Wrap FLOAT/DOUBLE expr as COALESCE(expr, default_value) +static unique_ptr WrapCoalesceDefault(unique_ptr expr) { + const auto &rt = expr->return_type; + + if (rt.id() == LogicalTypeId::FLOAT) { + auto defv = make_uniq(Value::FLOAT(0.0f)); // pick any default you want + auto op = make_uniq(ExpressionType::OPERATOR_COALESCE, LogicalType::FLOAT); + op->children.push_back(std::move(expr)); + op->children.push_back(std::move(defv)); + return std::move(op); + } + if (rt.id() == LogicalTypeId::DOUBLE) { + auto defv = make_uniq(Value::DOUBLE(0.0)); + auto op = make_uniq(ExpressionType::OPERATOR_COALESCE, LogicalType::DOUBLE); + op->children.push_back(std::move(expr)); + op->children.push_back(std::move(defv)); + return std::move(op); + } + + // leave other types untouched + return expr; +} + +static void RewriteOp(LogicalOperator &op) { + if (op.type == LogicalOperatorType::LOGICAL_PROJECTION) { + auto &proj = op.Cast(); + for (auto &e : proj.expressions) { + std::cout << "Before: " << e->ToString() << std::endl; + e = WrapCoalesceDefault(std::move(e)); + std::cout << " After: " << e->ToString() << std::endl; + } + } + for (auto &child : op.children) { + RewriteOp(*child); + } +} + +} // namespace + +void Add5Rule::Apply(ClientContext &, unique_ptr &plan) { + std::cout << plan->ToString() << "\n\n\n"; + if (!plan) return; + RewriteOp(*plan); + // optional: plan->ResolveOperatorTypes(); // helpful after edits +} diff --git a/src/optimization/catalog.cpp b/src/optimization/catalog.cpp new file mode 100644 index 0000000..fa7764f --- /dev/null +++ b/src/optimization/catalog.cpp @@ -0,0 +1,24 @@ +#include "duckdb.hpp" +#include "optimization/catalog.hpp" + +namespace duckdb { + idx_t Catalog_optimization::cacheQueryPlan(ClientContext &ctx,unique_ptr &plan) { + auto plan_to_cache = plan->Copy(ctx); + queryPlanCaches_[currentCacheId++] = std::move(plan_to_cache); + return currentCacheId-1; + } + + unique_ptr Catalog_optimization::resetQueryPlanFromCache(ClientContext &ctx,idx_t cacheId) { + auto it = queryPlanCaches_.find(cacheId); + if (it == queryPlanCaches_.end() || !it->second) { + throw std::runtime_error("Cache ID not found"); + } + // Return a fresh deep copy so caller owns an independent tree + return it->second->Copy(ctx); + } + + unique_ptr Catalog_optimization::resetQueryPlan(ClientContext &ctx) { + auto it = queryPlanCaches_.find(initQueryPlanCacheId); + return it->second->Copy(ctx); + } +} //end of duckdb namespace \ No newline at end of file diff --git a/src/optimization/dynamic_optimizer.cpp b/src/optimization/dynamic_optimizer.cpp new file mode 100644 index 0000000..08cd29f --- /dev/null +++ b/src/optimization/dynamic_optimizer.cpp @@ -0,0 +1,75 @@ +#include "duckdb.hpp" +#include "optimization/dynamic_optimizer.hpp" +#include "optimization/actions/MatMulDense2SparseRewriteAction.hpp" +#include "optimization/actions/MultiLayerUDF2TorchNNRewriteAction.hpp" +#include "optimization/actions/DecisionForestUDF2RelationRewriteAction.hpp" +#include "optimization/actions/MLDecompositionPushdownRewriteAction.hpp" +#include "optimization/catalog.hpp" +#include "cost_model/estimator.hpp" + +#include "optimization/feature_extraction/query2vec/query_annotation.hpp" +#include "optimization/feature_extraction/query2vec/query_feature_extraction.hpp" + +namespace duckdb{ + void DynamicOpt::Optimize(OptimizerExtensionInput &in, unique_ptr &plan) { + if (!plan) return; + // else return; + // Catalog_optimization cat; + // // // first step : check all rules and see which ones apply + // // // TODO Decompose this properly later with Catalog + // idx_t id = cat.cacheQueryPlan(in.context,plan); + + // // unique_ptr initial_plan = plan->Copy(in.context); + + // // std::cout << "\nCached the initial plan with cacheId: " << id << "\n"; + + // if(MatMulDense2SparseRewriteAction::check(in, plan)){ + // std::cout << "\nMatMulDense2SparseRewriteAction is applicable \n"; + // MatMulDense2SparseRewriteAction::apply(in, plan); + // } + + // if(MultiLayerUDF2TorchNNRewriteAction::check(in, plan)){ + // std::cout << "\nMultiLayerUDF2TorchNNRewriteAction is applicable \n"; + // MultiLayerUDF2TorchNNRewriteAction::apply(in, plan); + // } + + // if(DecisionForestUDF2RelationRewriteAction::check(in, plan)){ + // std::cout << "\nDecisionForestUDF2RelationRewriteAction is applicable \n"; + // DecisionForestUDF2RelationRewriteAction::apply(in, plan); + // } + + if(MLDecompositionPushdownRewriteAction::check(in, plan)){ + std::cout << "\nMLDecompositionPushdownRewriteAction is applicable \n"; + MLDecompositionPushdownRewriteAction::apply(in, plan); + } + + // SimpleEstimator est(std::make_unique()); + // auto total = est.Estimate(*plan); + // std::cout << "\n Final plan cost after Dynamic Optimization: " << total.cost; + + // auto annotated = AnnotatePlanTree(*plan); // pass by const-ref + // DisplayAnnotatedPlan(*annotated); + + // traverse_and_extract(*annotated); + + + + + + // plan = std::move(initial_plan); + + // plan = cat.resetQueryPlanFromCache(in.context,0); + + + // translate the rules in language that ReuableMCTS understands + // The language that we're using for the velox. + + // send all applicable rules to ReusableMCTS python runtime + + // recieve the rule that is applicable here + + // apply the rule on the plan + + // No type changes; ResolveOperatorTypes() not required. + }; +} \ No newline at end of file diff --git a/src/optimization/feature_extraction/query2vec/query_annotation.cpp b/src/optimization/feature_extraction/query2vec/query_annotation.cpp new file mode 100644 index 0000000..4e02ce5 --- /dev/null +++ b/src/optimization/feature_extraction/query2vec/query_annotation.cpp @@ -0,0 +1,47 @@ +#include "optimization/feature_extraction/query2vec/query_annotation.hpp" +#include "duckdb/planner/logical_operator.hpp" +#include "duckdb/common/to_string.hpp" + +namespace duckdb { + +static std::unique_ptr +AnnotateRec(const LogicalOperator &op, idx_t &counter) { + auto node = std::make_unique(); + node->node_id = ++counter; // pre-order id + node->plan_node = &op; // non-owning raw pointer + + node->children.reserve(op.children.size()); + for (idx_t i = 0; i < op.children.size(); ++i) { + const auto &child_op = *op.children[i]; // LogicalOperator& + auto child_node = AnnotateRec(child_op, counter); + node->children.push_back(std::move(child_node)); + } + return node; +} + +std::unique_ptr AnnotatePlanTree(const LogicalOperator &root) { + idx_t counter = 0; + return AnnotateRec(root, counter); +} + +void DisplayAnnotatedPlan(const AnnotatedPlanNode &node, idx_t depth) { + // guard + std::string name = LogicalOperatorToString(node.plan_node->type); + for (idx_t i = 0; i < depth; ++i) std::cout << " "; + std::cout << "Node ID: " << node.node_id << " Type: " << name << "\n"; + for (const auto &ch : node.children) { + DisplayAnnotatedPlan(*ch, depth + 1); + } +} + +AnnotatedPlanNode* FindNodeById(AnnotatedPlanNode &node, idx_t target_id) { + if (node.node_id == target_id) return &node; + for (auto &ch : node.children) { + if (auto *res = FindNodeById(*ch, target_id)) { + return res; + } + } + return nullptr; +} + +} // namespace duckdb diff --git a/src/optimization/feature_extraction/query2vec/query_feature_extraction.cpp b/src/optimization/feature_extraction/query2vec/query_feature_extraction.cpp new file mode 100644 index 0000000..453d12c --- /dev/null +++ b/src/optimization/feature_extraction/query2vec/query_feature_extraction.cpp @@ -0,0 +1,156 @@ +#include "duckdb.hpp" +#include "duckdb/common/string_util.hpp" +#include "duckdb/planner/operator/logical_filter.hpp" +#include "duckdb/planner/expression/bound_cast_expression.hpp" +#include "duckdb/planner/expression/bound_columnref_expression.hpp" +#include "duckdb/planner/expression/bound_constant_expression.hpp" +#include "duckdb/planner/expression/bound_comparison_expression.hpp" +#include "duckdb/planner/expression/bound_between_expression.hpp" +#include "duckdb/planner/operator/logical_get.hpp" +#include "duckdb/planner/operator/logical_join.hpp" +#include "duckdb/common/enums/expression_type.hpp" + +#include "optimization/feature_extraction/query2vec/query_annotation.hpp" +#include "optimization/feature_extraction/query2vec/query_feature_extraction.hpp" + +namespace duckdb{ + + void extract_features(const AnnotatedPlanNode &node){ + switch (node.plan_node->type){ + case LogicalOperatorType::LOGICAL_PROJECTION : { + //ml features will go here! + std::cout << "Extracting features from PROJECTION node ID: " << node.node_id << "\n"; + } + break; + + case LogicalOperatorType::LOGICAL_FILTER : { + std::string ffe = filter_feature_extract(node); + std::cout << "Extracting features from FILTER node ID: " << node.node_id << "\n"; + } + break; + + case LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY : { + std::string afe = aggregate_feature_extract(node); + std::cout << "Extracting features from AGGREGATE node ID: " << node.node_id << "\n"; + } + break; + + case LogicalOperatorType::LOGICAL_JOIN : + case LogicalOperatorType::LOGICAL_COMPARISON_JOIN : + case LogicalOperatorType::LOGICAL_ANY_JOIN : { + std::string jfe = join_features_extract(node); + std::cout << "Extracting features from JOIN node ID: " << node.node_id << "\n"; + std::cout << " Join Features: " << jfe << "\n"; + } + break; + + case LogicalOperatorType::LOGICAL_GET : { + std::string gffe = get_filters_feature_extract(node); + std::string gtfe = get_table_feature_extract(node); + std::cout << "Extracting features from GET node ID: " << node.node_id << "\n"; + std::cout << " Filters Features: " << gffe << "\n"; + std::cout << " Table Features: " << gtfe << "\n"; + } + break; + default : { + // extract generic features + std::cout << "Extracting generic features from node ID: " << node.node_id << "\n"; + } + break; + } + } + + std::string filter_feature_extract(const AnnotatedPlanNode &node){ + // it should have "input_column operator constant" structure + std::string filter_features = ""; + if(!node.plan_node->expressions.empty()){ + for(auto &expr : node.plan_node->expressions){ + if(expr->expression_class == ExpressionClass::BOUND_COMPARISON){ + auto &comp = expr->Cast(); + std::string left, right, op; + // left + if(comp.left->expression_class == ExpressionClass::BOUND_COLUMN_REF){ + auto &colref = comp.left->Cast(); + left = colref.GetName(); + } + // right + if(comp.right->expression_class == ExpressionClass::BOUND_CONSTANT){ + auto &const_expr = comp.right->Cast(); + right = const_expr.value.ToString(); + } + // operator + op = EnumUtil::ToString(comp.type); + // Combine features + filter_features += left + " " + op + " " + right + ","; + } + } + } + // #TODO Change later to proper edge case handling + return "empty empty empty,"; + } //end of filter_feature_extract + + std::string get_filters_feature_extract(const AnnotatedPlanNode &node) { + // Only for LOGICAL_GET + if (!node.plan_node || node.plan_node->type != LogicalOperatorType::LOGICAL_GET) { + return "No Node"; + } + + const auto &get = node.plan_node->Cast(); + if (get.table_filters.filters.empty()) { + return ""; + } + + std::vector atoms; + atoms.reserve(get.table_filters.filters.size()); + + // Each entry in the filter set is (column_idx -> root filter tree) + for (const auto &kv : get.table_filters.filters) { + const idx_t col_idx = kv.first; + const auto &filter = kv.second; // unique_ptr + if (!filter) continue; + + // Resolve the bound column name for this logical index + const std::string col_name = get.GetColumnName(ColumnIndex(col_idx)); + + // Let the filter render itself as "col OP value" (and/or nested trees) + atoms.push_back(filter->ToString(col_name)); + } + + // Join with ", " + std::string result; + for (idx_t i = 0; i < atoms.size(); ++i) { + if (i) result += ", "; + result += atoms[i]; + } + return result; + } //end of get_filters_feature_extract + + + std::string get_table_feature_extract(const AnnotatedPlanNode &node){ + auto &get = node.plan_node->Cast(); + idx_t rows = get.estimated_cardinality; // DuckDB keeps this updated + if (rows == DConstants::INVALID_INDEX) rows = 0; + idx_t cols = get.types.size(); + return "Rows:" + std::to_string(rows) + ",Cols:" + std::to_string(cols); + } + + std::string aggregate_feature_extract(const AnnotatedPlanNode &node){ + // Placeholder implementation + return "aggregate_features"; + } + + std::string join_features_extract(const AnnotatedPlanNode &node){ + // Placeholder implementation + const auto &lj = node.plan_node->Cast(); + JoinType jt = lj.join_type; + std::string jt_str = EnumUtil::ToString(jt); + return jt_str; + } + + void traverse_and_extract(const AnnotatedPlanNode &node){ + extract_features(node); + for (const auto &child : node.children) { + traverse_and_extract(*child); + } + } +}// end of namespace duckdb \ No newline at end of file diff --git a/src/optimization/register_optimizers.cpp b/src/optimization/register_optimizers.cpp new file mode 100644 index 0000000..594db1d --- /dev/null +++ b/src/optimization/register_optimizers.cpp @@ -0,0 +1,184 @@ +// register_optimizers.cpp +#include "optimization/register_optimizers.hpp" + +#include "duckdb/main/database.hpp" +#include "duckdb/optimizer/optimizer_extension.hpp" + +#include "optimization/add5_rule.hpp" +#include "optimization/replace_binding_rule.hpp" +#include "optimization/dynamic_optimizer.hpp" +#include "optimization/metric_catalog.hpp" + +#include "duckdb/planner/operator/logical_get.hpp" +#include "duckdb/catalog/catalog_entry/table_catalog_entry.hpp" +#include "duckdb/catalog/catalog_entry/schema_catalog_entry.hpp" + +// #include "optimization/plan_dump.hpp" + +#include +#include +#include + +//delete later +#include "optimization/actions/MatMulDense2SparseRewriteAction.hpp" +#include "optimization/actions/MultiLayerUDF2TorchNNRewriteAction.hpp" +#include "optimization/actions/MLDecompositionPushdownRewriteAction.hpp" + +namespace duckdb { + +//helper function to get schema name from LogicalOperator +static string GetQuerySchema(LogicalOperator &op) { + if (op.type == LogicalOperatorType::LOGICAL_GET) { + auto &get = op.Cast(); + auto table_entry = get.GetTable(); + if (table_entry) { + // This line may or may not compile depending on visibility: + return table_entry->schema.name; + + // If that complains about access, just fall back: + // return table_entry->GetName(); // as described above + } + return string(); + } + for (auto &child : op.children) { + auto s = GetQuerySchema(*child); + if (!s.empty()) return s; + } + return string(); +} + + +// ----------------------------- +// Pass 1: your existing Add5Rule +// ----------------------------- +struct CactusAdd5Opt : public OptimizerExtension { + CactusAdd5Opt() { optimize_function = Run; } + + static void Run(OptimizerExtensionInput &in, unique_ptr &plan) { + Add5Rule::Apply(in.context, plan); // always-on + } +}; + +// ----------------------------- +// Pass 2: binding-file replacement for one_hot_encode +// ----------------------------- +struct CactusReplaceBindingOpt : public OptimizerExtension { + CactusReplaceBindingOpt() { optimize_function = Run; } + + // Adjust these mappings however you like. You can add many. + // Each entry: { from_path, to_path, require_third_bool } + using MapEntry = std::tuple>; + + static void Run(OptimizerExtensionInput &in, unique_ptr &plan) { + if (!plan) return; + + // Example mappings (fill with your real paths): + // If you don't want to constrain the 3rd argument, use std::nullopt. + const std::vector mappings = { + { + "/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/test/resources/month.csv", + "/home/local/ASUAD/jrtandel/cactusdb-duckdb-extension/test/resources/month_alt.csv", + /*require_third_bool=*/std::optional{true} // only when the 3rd arg is TRUE + }, + // Add more file swaps here: + // { "", "", std::nullopt }, // don't care about 3rd arg + // { "", "", std::optional{false} }, // only when 3rd arg is FALSE + }; + + for (const auto &m : mappings) { + const auto &from_file = std::get<0>(m); + const auto &to_file = std::get<1>(m); + const auto &third_req = std::get<2>(m); + + ReplaceBindingRule cfg(from_file, to_file, third_req); + ReplaceBindingRule::Optimize(in, plan, cfg); + } + // No type changes; ResolveOperatorTypes() not required. + } +}; + + +// ----------------------------- +// Dynamic optimizer pass +// ----------------------------- +struct CactusDynamicOpt : public OptimizerExtension { + CactusDynamicOpt() { optimize_function = Run; } + + static void Run(OptimizerExtensionInput &in, unique_ptr &plan) { + if (!plan) return; + + // std::string pretty = plan->ToString(); + // std::cout << "\n[CactusDynamicOpt] Initial Plan:\n" << pretty << "\n"; + + // 1) figure out which schema this query uses + string query_schema = GetQuerySchema(*plan); + // std::cout << "[CactusDynamicOpt] query_schema = '" << query_schema << "'\n"; + + // 2) pull metrics from the in-code catalog + const auto &mc = MetricCatalog::Get(); + double join_ratio = mc.GetMetric(query_schema, "join_ratio", 0.0); + double sparsity = mc.GetMetric(query_schema, "sparsity", 1.0); + + const double HIGH_JOIN_RATIO = 1e-5; + const double SPARSE_THRESH = 0.30; + + // std::cout << "[CactusDynamicOpt] join_ratio=" << join_ratio + // << " sparsity=" << sparsity << "\n"; + + // 3) your existing setting to choose profiles + Value setting; + in.context.TryGetCurrentSetting("enable_extended_optimizer", setting); + std::string opt_setting = setting.ToString(); + + if (opt_setting.find("1") != std::string::npos) { + if (join_ratio > HIGH_JOIN_RATIO) { + MLDecompositionPushdownRewriteAction::apply(in, plan); + } + } + + if (opt_setting.find("2") != std::string::npos) { + if (join_ratio > HIGH_JOIN_RATIO && sparsity < SPARSE_THRESH) { + MLDecompositionPushdownRewriteAction::apply(in, plan); + MatMulDense2SparseRewriteAction::apply(in, plan); + } + } + + if (opt_setting.find("3") != std::string::npos) { + MultiLayerUDF2TorchNNRewriteAction::apply(in, plan); + } + + // pretty = plan->ToString(); + // std::cout << "\n[CactusDynamicOpt] Final Plan:\n" << pretty << "\n"; + } + +}; + +struct CactusDynamicOpt1 : public OptimizerExtension { + CactusDynamicOpt1() { optimize_function = Run; } + + static void Run(OptimizerExtensionInput &in, unique_ptr &plan) { + // DumpPlanBase64(in.context, *plan, "./plans/initial_plan.duckdbplan"); + std::string pretty = plan->ToString(); + std::cout << "\nInitial Plan:\n" << pretty << "\n"; + MatMulDense2SparseRewriteAction::apply(in, plan); + // DumpPlanBase64(in.context, *plan, "./plans/final_plan.duckdbplan"); + pretty = plan->ToString(); + std::cout << "\nFinal Plan:\n" << pretty << "\n"; + } +}; + +// ----------------------------- +// Registration hook +// ----------------------------- +void RegisterCactusOptimizers(DatabaseInstance &db) { + auto &cfg = DBConfig::GetConfig(db); + + // Order matters if passes depend on one another. + // Typically do argument/string rewrites early, then other expression rewrites. + // cfg.optimizer_extensions.push_back(CactusReplaceBindingOpt{}); + // cfg.optimizer_extensions.push_back(CactusAdd5Opt{}); + cfg.optimizer_extensions.push_back(CactusDynamicOpt{}); + // cfg.optimizer_extensions.push_back(CactusDynamicOpt1{}); +} + +} // namespace duckdb diff --git a/src/optimization/replace_binding_rule.cpp b/src/optimization/replace_binding_rule.cpp new file mode 100644 index 0000000..1d9ad8f --- /dev/null +++ b/src/optimization/replace_binding_rule.cpp @@ -0,0 +1,111 @@ +#include "optimization/replace_binding_rule.hpp" + +#include "duckdb/common/string_util.hpp" +#include "duckdb/planner/expression/bound_constant_expression.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/planner/expression_iterator.hpp" +#include "duckdb/planner/logical_operator.hpp" +#include "duckdb/planner/operator/logical_projection.hpp" +#include "duckdb/planner/operator/logical_filter.hpp" +#include "duckdb/planner/operator/logical_aggregate.hpp" +#include "duckdb/planner/operator/logical_comparison_join.hpp" +#include "duckdb/planner/operator/logical_join.hpp" + +#include "duckdb/planner/operator/logical_filter.hpp" +#include "duckdb/planner/operator/logical_projection.hpp" +#include "duckdb/planner/operator/logical_aggregate.hpp" +#include "duckdb/planner/operator/logical_comparison_join.hpp" +#include "duckdb/common/enums/expression_type.hpp" +#include "ml_functions/one_hot_encoder.hpp" +// #include "bind_utils/one_hot_encoder_bind_util.hpp" +// no need for if we use a template helper + + +using namespace duckdb; + +duckdb::unique_ptr +ReplaceBindingRule::RewriteExpr(duckdb::unique_ptr expr, const ReplaceBindingRule &cfg) { + if (!expr) return expr; + + std::cout << "Examining expr: " << ExpressionTypeToString(expr->type) << " " << ExpressionClassToString(expr->expression_class) << "\n"; + + // Recurse first + ExpressionIterator::EnumerateChildren(*expr, [&](duckdb::unique_ptr &child) { + child = RewriteExpr(std::move(child), cfg); + }); + + // std::cout << expr.ToString() << "\n\n"; + + // Match function node + if (expr->expression_class == ExpressionClass::BOUND_FUNCTION) { + std::cout << "\nBound function found\n"; + auto &f = expr->Cast(); + std::cout << f.function.name << " " << cfg.kFuncName << "\n"; + if (f.function.name == cfg.kFuncName) { + auto &bd = f.bind_info->Cast(); + + //change the argument from the path to another path + f.children[1] = make_uniq(Value(cfg.to_path)); + + //change the bind data to reflect data from new path + std::map str_map; + std::map int_map; + const bool is_string = bd.is_string; + + one_hot_encoder_bind_constructor(cfg.to_path, is_string, str_map, int_map); + f.bind_info = make_uniq(str_map, int_map, is_string); + } + } + return expr; +} + +// ---- Polyfill: enumerate/mutate expressions for common operator types ---- +// Visit & mutate expressions held by common logical operators in your query +template +static void ForEachExpression(LogicalOperator &op, Fn &&fn) { + switch (op.type) { + case LogicalOperatorType::LOGICAL_PROJECTION: { + auto &proj = op.Cast(); + for (auto &e : proj.expressions) fn(e); + break; + } + case LogicalOperatorType::LOGICAL_FILTER: { + auto &fil = op.Cast(); + for (auto &e : fil.expressions) fn(e); + break; + } + case LogicalOperatorType::LOGICAL_COMPARISON_JOIN: { + auto &cj = op.Cast(); + for (auto &cond : cj.conditions) { + fn(cond.left); + fn(cond.right); + } + break; + } + default: + // Add more operator cases as needed (WINDOW/ORDER/etc) for other queries + break; + } +} + + +void ReplaceBindingRule::RewriteOp(LogicalOperator &op, const ReplaceBindingRule &cfg) { + // Mutate expressions of this operator + ForEachExpression(op, [&](duckdb::unique_ptr &e) { + e = RewriteExpr(std::move(e), cfg); + }); + + // Recurse into children + for (auto &child : op.children) { + RewriteOp(*child, cfg); + } +} + + +void ReplaceBindingRule::Optimize(OptimizerExtensionInput &input, + duckdb::unique_ptr &plan, + const ReplaceBindingRule &cfg) { + if (!plan) return; + RewriteOp(*plan, cfg); + // plan->ResolveOperatorTypes(); // not needed +} diff --git a/test/sql/argmax.test b/test/sql/argmax.test new file mode 100644 index 0000000..7f80018 --- /dev/null +++ b/test/sql/argmax.test @@ -0,0 +1,54 @@ +# name: test/sql/ml_argmax.test +# description: test ml_argmax scalar function (LIST -> INTEGER index) +# group: [sql] + +require cactusdb + +statement ok +LOAD cactusdb; + +# ------------------------------------------------------------------ +# Inputs: several rows of DOUBLE[] with varying patterns +# ------------------------------------------------------------------ +statement ok +CREATE TEMP TABLE inputs(a DOUBLE[]); + +# [1,2,3] -> 2 +# [10,5,10] -> 0 (tie: first max) +# [-1,-5,-3] -> 0 (max is -1 at index 0) +# [7,7,7] -> 0 (all equal: first index) +statement ok +INSERT INTO inputs VALUES + ([1,2,3]), + ([10,5,10]), + ([-1,-5,-3]), + ([7,7,7]); + +# ------------------------------------------------------------------ +# Basic ml_argmax on valid rows +# ------------------------------------------------------------------ +query I +SELECT ml_argmax(a) AS idx +FROM inputs; +---- +2 +0 +0 +0 + +# ------------------------------------------------------------------ +# NULL row handling: result should be NULL +# ------------------------------------------------------------------ +statement ok +INSERT INTO inputs VALUES (NULL); + +query I +SELECT ml_argmax(a) AS idx +FROM inputs +ORDER BY 1 NULLS LAST; +---- +0 +0 +0 +2 +NULL \ No newline at end of file diff --git a/test/sql/cactusdb.test b/test/sql/cactusdb.test deleted file mode 100644 index 36cab74..0000000 --- a/test/sql/cactusdb.test +++ /dev/null @@ -1,31 +0,0 @@ -# name: test/sql/cactusdb.test -# description: test cactusdb extension -# group: [sql] - -require cactusdb - -#Create table -statement ok -CREATE TABLE array_data (id INTEGER, features FLOAT[8]); - -#Insert data -statement ok -INSERT INTO array_data VALUES - (1, [1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1]), - (2, [2.2, 2.2, 2.2, 2.2, 2.2, 2.2, 2.2, 2.2]), - (3, [3.3, 3.3, 3.3, 3.3, 3.3, 3.3, 3.3, 3.3]), - (4, [4.4, 4.4, 4.4, 4.4, 4.4, 4.4, 4.4, 4.4]), - (5, [5.5, 5.5, 5.5, 5.5, 5.5, 5.5, 5.5, 5.5]), - (6, [6.6, 6.6, 6.6, 6.6, 6.6, 6.6, 6.6, 6.6]); - -#Test query -statement ok -SELECT id, features, FROM array_data ORDER BY id - -#Test query -statement ok -SELECT id, array_inner_product(features, array_value(1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT)), FROM array_data ORDER BY id - -#Test query -statement ok -SELECT id, matmul(features, array_value(1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT), 8, 2), FROM array_data ORDER BY id diff --git a/test/sql/extension_test.test b/test/sql/extension_test.test new file mode 100644 index 0000000..17890b1 --- /dev/null +++ b/test/sql/extension_test.test @@ -0,0 +1,24 @@ +# name: test/sql/table_function/duckdb_settings_extension.test +# description: Test duckdb_settings function with extensions +# group: [table_function] + +statement ok +SET default_null_order='nulls_first'; + +require parquet + +statement ok +SELECT * FROM duckdb_settings(); + +query II +SELECT name, value FROM duckdb_settings() WHERE name='default_null_order'; +---- +default_null_order NULLS_FIRST + +statement ok +SET default_null_order='nulls_last' + +query II +SELECT name, value FROM duckdb_settings() WHERE name='default_null_order'; +---- +default_null_order NULLS_LAST \ No newline at end of file diff --git a/test/sql/matadd.test b/test/sql/matadd.test new file mode 100644 index 0000000..3be2f5e --- /dev/null +++ b/test/sql/matadd.test @@ -0,0 +1,62 @@ +# name: test/sql/mat_add.test +# description: test mat_add scalar function (broadcast 1xN biases from CSV) +# group: [sql] + +require cactusdb + +statement ok +LOAD cactusdb; + +# ------------------------------------------------------------------ +# Create biases.csv: a single 1x3 row +# [0.5, -1.0, 2.5] +# ------------------------------------------------------------------ +statement ok +CREATE TEMP TABLE bias_row(c1 DOUBLE, c2 DOUBLE, c3 DOUBLE); + +statement ok +INSERT INTO bias_row VALUES (0.5, -1.0, 2.5); + +statement ok +COPY bias_row TO 'biases.csv' (FORMAT CSV, HEADER FALSE); + +# ------------------------------------------------------------------ +# Inputs: three rows of equal-length vectors (length = 3) +# ------------------------------------------------------------------ +statement ok +CREATE TEMP TABLE inputs(a DOUBLE[]); + +statement ok +INSERT INTO inputs VALUES ([1,2,3]), ([4,5,6]), ([7,8,9]); + +# ------------------------------------------------------------------ +# Basic broadcast add: y = a + biases +# Expect: +# [1,2,3] + [0.5,-1.0,2.5] = [1.5, 1.0, 5.5] +# [4,5,6] + [0.5,-1.0,2.5] = [4.5, 4.0, 8.5] +# [7,8,9] + [0.5,-1.0,2.5] = [7.5, 7.0, 11.5] +# ------------------------------------------------------------------ +query I +SELECT mat_add(a, 'biases.csv') AS y +FROM inputs +ORDER BY 1; +---- +[1.5, 1.0, 5.5] +[4.5, 4.0, 8.5] +[7.5, 7.0, 11.5] + +# ------------------------------------------------------------------ +# NULL row handling: result should be NULL for that row +# ------------------------------------------------------------------ +statement ok +INSERT INTO inputs VALUES (NULL); + +query I +SELECT mat_add(a, 'biases.csv') AS y +FROM inputs +ORDER BY 1 NULLS LAST; +---- +[1.5, 1.0, 5.5] +[4.5, 4.0, 8.5] +[7.5, 7.0, 11.5] +NULL diff --git a/test/sql/matmul.test b/test/sql/matmul.test new file mode 100644 index 0000000..5d6bcc6 --- /dev/null +++ b/test/sql/matmul.test @@ -0,0 +1,56 @@ +# name: test/sql/cactusdb.test +# description: test cactusdb extension +# group: [sql] + +require cactusdb + +statement ok +LOAD cactusdb; + +# ------------------------------------------------------------------ +# Create weights file_1.csv: 3x2 matrix +# [[0.1, 0.2], +# [0.3, 0.4], +# [0.5, 0.6]] +# ------------------------------------------------------------------ +statement ok +CREATE TEMP TABLE w1(c1 DOUBLE, c2 DOUBLE); + +statement ok +INSERT INTO w1 VALUES (0.1,0.2),(0.3,0.4),(0.5,0.6); + +statement ok +COPY w1 TO 'file_1.csv' (FORMAT CSV, HEADER FALSE); + +# ------------------------------------------------------------------ +# Create weights file_2.csv: 2x2 matrix (scale by 10) +# [[10.0, 0.0], +# [0.0, 10.0]] +# ------------------------------------------------------------------ +statement ok +CREATE TEMP TABLE w2(c1 DOUBLE, c2 DOUBLE); + +statement ok +INSERT INTO w2 VALUES (10.0,0.0),(0.0,10.0); + +statement ok +COPY w2 TO 'file_2.csv' (FORMAT CSV, HEADER FALSE); + +# ------------------------------------------------------------------ +# Inputs: three rows of vectors +# ------------------------------------------------------------------ +statement ok +CREATE TABLE inputs(a DOUBLE[]); + +statement ok +INSERT INTO inputs VALUES ([1,2,3]), ([4,5,6]), ([7,8,9]); + + +query I +SELECT mat_mul(mat_mul(a, 'file_1.csv', 'dense'), 'file_2.csv', 'sparse') AS y +FROM inputs +ORDER BY 1; +---- +[22.0, 28.0] +[49.0, 63.99999999999999] +[76.0, 100.0] diff --git a/test/sql/min_max_scaler.test b/test/sql/min_max_scaler.test new file mode 100644 index 0000000..78422ec --- /dev/null +++ b/test/sql/min_max_scaler.test @@ -0,0 +1,64 @@ +# name: test/sql/minmaxscaler.test +# description: test min_max_scaler (file + constants) in cactusdb extension +# group: [sql] + +require cactusdb + +statement ok +LOAD cactusdb; + +# ------------------------------------------------------------------ +# Inputs: three rows of REAL vectors +# ------------------------------------------------------------------ +statement ok +CREATE TABLE inputs(a REAL[]); + +statement ok +INSERT INTO inputs VALUES ([1,5,9]), ([4,5,6]), ([7,8,9]); + +# ================================================================== +# A) FILE-BASED STATS (2 x k): mins on row 1, maxs on row 2 +# mins = [0,0,0], maxs = [10,10,10] +# Expected scaling: +# [1,5,9] -> [0.1,0.5,0.9] +# [4,5,6] -> [0.4,0.5,0.6] +# [7,8,9] -> [0.7,0.8,0.9] +# ================================================================== + +statement ok +CREATE TEMP TABLE s1(c1 REAL, c2 REAL, c3 REAL); + +statement ok +INSERT INTO s1 VALUES (0.0,0.0,0.0), (10.0,10.0,10.0); + +statement ok +COPY s1 TO 'scaler_2xk.csv' (FORMAT CSV, HEADER FALSE, DELIMITER ' '); + +query I +SELECT min_max_scaler(a, 'scaler_2xk.csv') AS y +FROM inputs +ORDER BY 1; +---- +[0.1, 0.5, 0.9] +[0.4, 0.5, 0.6] +[0.7, 0.8, 0.9] + +# ================================================================== +# F) LIST PARAMS OVERLOAD: +# (a DOUBLE[], min DOUBLE[], max DOUBLE[], cols INTEGER) +# mins/maxs = [0,0,0]/[10,10,10] -> same expected output as A/B/C +# ================================================================== + +query I +SELECT min_max_scaler( + a, -- cast input to DOUBLE[] + [0.0, 0.0, 0.0]::DOUBLE[], -- min list (DOUBLE[]) + [10.0, 10.0, 10.0]::DOUBLE[], -- max list (DOUBLE[]) + 3 -- number of columns + ) AS y +FROM inputs +ORDER BY 1; +---- +[0.1, 0.5, 0.9] +[0.4, 0.5, 0.6] +[0.7, 0.8, 0.9] diff --git a/test/sql/nn_test.test b/test/sql/nn_test.test new file mode 100644 index 0000000..37e87b8 --- /dev/null +++ b/test/sql/nn_test.test @@ -0,0 +1,114 @@ +# name: test/sql/cactusdb_nn_project.test +# description: compose mat_mul, mat_add, cdb_softmax, ml_argmax into a simple 1-layer NN in a Project node +# group: [sql] + +require cactusdb + +statement ok +LOAD cactusdb; + +# ------------------------------------------------------------------ +# Create a 2x2 weight matrix W for a simple classifier: +# +# logits = [x0 - x1, -x0 + x1] +# +# This chooses: +# class 0 if x0 > x1 +# class 1 if x1 > x0 +# class 0 on ties (since both logits equal) +# +# W = [[ 1, -1 ], +# [ -1, 1 ]] +# so that [x0, x1] * W = [x0 - x1, -x0 + x1] +# ------------------------------------------------------------------ +statement ok +CREATE TEMP TABLE w_nn(c1 DOUBLE, c2 DOUBLE); + +statement ok +INSERT INTO w_nn VALUES + ( 1.0, -1.0), + (-1.0, 1.0); + +statement ok +COPY w_nn TO 'w_nn.csv' (FORMAT CSV, HEADER FALSE); + +# ------------------------------------------------------------------ +# Create zero biases b = [0, 0] so mat_add is exercised but +# does not change logits: +# z = mat_add(mat_mul(a, W), b) +# ------------------------------------------------------------------ +statement ok +CREATE TEMP TABLE b_zero(c1 DOUBLE, c2 DOUBLE); + +statement ok +INSERT INTO b_zero VALUES (0.0, 0.0); + +statement ok +COPY b_zero TO 'b_zero.csv' (FORMAT CSV, HEADER FALSE); + +# ------------------------------------------------------------------ +# Inputs: +# id a +# 1 [2,1] -> logits [1,-1] -> softmax ~[0.880797, 0.119203] -> class 0 +# 2 [1,3] -> logits [-2,2] -> softmax ~[0.017986, 0.982014] -> class 1 +# 3 [5,5] -> logits [0,0] -> softmax [0.5, 0.5] -> class 0 (tie) +# ------------------------------------------------------------------ +statement ok +CREATE TEMP TABLE inputs(id INTEGER, a DOUBLE[]); + +statement ok +INSERT INTO inputs VALUES + (1, [2.0, 1.0]), + (2, [1.0, 3.0]), + (3, [5.0, 5.0]); + +# ------------------------------------------------------------------ +# 1) Check the softmax probabilities of the "NN": +# +# probs = softmax( mat_add( mat_mul(a, W), b ) ) +# +# Entire NN is a single project expression: +# cdb_softmax(mat_add(mat_mul(a, 'w_nn.csv','dense'), 'b_zero.csv')) +# ------------------------------------------------------------------ +query I +SELECT list_transform( + cdb_softmax( + mat_add( + mat_mul(a, 'w_nn.csv', 'dense'), + 'b_zero.csv' + ) + ), + x -> round(x::DOUBLE, 6) + ) AS probs +FROM inputs +ORDER BY id; +---- +[0.880797, 0.119203] +[0.017986, 0.982014] +[0.5, 0.5] + +# ------------------------------------------------------------------ +# 2) Full classification pipeline: +# +# pred = argmax( softmax( mat_add( mat_mul(a, W), b ) ) ) +# +# This matches: +# [2,1] -> class 0 +# [1,3] -> class 1 +# [5,5] -> class 0 +# ------------------------------------------------------------------ +query I +SELECT ml_argmax( + cdb_softmax( + mat_add( + mat_mul(a, 'w_nn.csv', 'dense'), + 'b_zero.csv' + ) + ) + ) AS cls +FROM inputs +ORDER BY id; +---- +0 +1 +0 diff --git a/test/sql/one_hot_encoder.test b/test/sql/one_hot_encoder.test new file mode 100644 index 0000000..7b5ae99 --- /dev/null +++ b/test/sql/one_hot_encoder.test @@ -0,0 +1,95 @@ +# name: test/sql/one_hot_encoder.test +# description: tests for one_hot_encode UDF (string & integer categories) +# group: [sql] + +require cactusdb + +statement ok +LOAD cactusdb; + +# =============================================================== +# STRING categories mapping CSV: red->0, green->1, blue->2 +# We emulate a header by adding a first dummy row and writing with HEADER FALSE. +# =============================================================== + +statement ok +CREATE TEMP TABLE str_map(value VARCHAR, position INTEGER); + +statement ok +INSERT INTO str_map VALUES + ('value', NULL), -- dummy header row to be skipped by Bind (getline once) + ('red', 0), + ('green', 1), + ('blue', 2); + +statement ok +COPY str_map TO 'str_map.csv' (FORMAT CSV, HEADER FALSE); + +# Inputs (includes an unknown 'yellow') +statement ok +CREATE TEMP TABLE inputs_str(x VARCHAR); + +statement ok +INSERT INTO inputs_str VALUES ('red'), ('green'), ('blue'), ('yellow'); + +# Order by x -> blue, green, red, yellow +# Expected one-hot (k=3): +# blue -> [0,0,1] +# green -> [0,1,0] +# red -> [1,0,0] +# yellow -> [0,0,0] (unknown) + +query I +SELECT one_hot_encode(x, 'str_map.csv', TRUE) AS oh +FROM inputs_str +ORDER BY x; +---- +[0.0, 0.0, 1.0] +[0.0, 1.0, 0.0] +[1.0, 0.0, 0.0] +[0.0, 0.0, 0.0] + +# =============================================================== +# INTEGER categories mapping CSV: 10->0, 20->1, 30->2, 50->3 +# Again, add a dummy header row and write with HEADER FALSE. +# =============================================================== + +statement ok +CREATE TEMP TABLE int_map(value BIGINT, position INTEGER); + +statement ok +INSERT INTO int_map VALUES + (NULL, NULL), -- dummy header + (10, 0), + (20, 1), + (30, 2), + (50, 3); + +statement ok +COPY int_map TO 'int_map.csv' (FORMAT CSV, HEADER FALSE); + +statement ok +CREATE TEMP TABLE inputs_int(v INTEGER); + +statement ok +INSERT INTO inputs_int VALUES (10), (20), (30), (40), (50); + +# Your UDF takes VARCHAR for the first arg, so CAST(v AS VARCHAR) +# Order by v: 10,20,30,40,50 +# Expected (k=4): +# 10 -> [1,0,0,0] +# 20 -> [0,1,0,0] +# 30 -> [0,0,1,0] +# 40 -> [0,0,0,0] (unknown) +# 50 -> [0,0,0,1] + +query I +SELECT one_hot_encode(CAST(v AS VARCHAR), 'int_map.csv', FALSE) AS oh +FROM inputs_int +ORDER BY v; +---- +[1.0, 0.0, 0.0, 0.0] +[0.0, 1.0, 0.0, 0.0] +[0.0, 0.0, 1.0, 0.0] +[0.0, 0.0, 0.0, 0.0] +[0.0, 0.0, 0.0, 1.0] diff --git a/test/sql/sigmoid.test b/test/sql/sigmoid.test new file mode 100644 index 0000000..ffaf4ba --- /dev/null +++ b/test/sql/sigmoid.test @@ -0,0 +1,53 @@ +# name: test/sql/sigmoid.test +# description: test sigmoid scalar function (LIST -> LIST) +# group: [sql] + +require cactusdb + +statement ok +LOAD cactusdb; + +# ------------------------------------------------------------------ +# Inputs +# ------------------------------------------------------------------ +statement ok +CREATE TEMP TABLE inputs(a DOUBLE[]); + +# Known cases: +# [0,0] -> [0.5, 0.5] +# [-1,1] -> [~0.268941, ~0.731059] +# [-5,0,5] -> [~0.006693, 0.5, ~0.993307] +# [7,7,7] -> [sigmoid(7)] * 3 +statement ok +INSERT INTO inputs VALUES + ([0,0]), + ([-1,1]), + ([-5,0,5]), + ([7,7,7]); + +# Round for deterministic output +query I +SELECT list_transform(sigmoid(a), x -> round(x::DOUBLE, 6)) AS y +FROM inputs; +---- +[0.5, 0.5] +[0.268941, 0.731059] +[0.006693, 0.5, 0.993307] +[0.999089, 0.999089, 0.999089] + +# ------------------------------------------------------------------ +# NULL row -> NULL +# ------------------------------------------------------------------ +statement ok +INSERT INTO inputs VALUES (NULL); + +query I +SELECT sigmoid(a) AS y +FROM inputs +ORDER BY 1 NULLS LAST; +---- +[0.0066928509242848554, 0.5, 0.9933071490757153] +[0.2689414213699951, 0.7310585786300049] +[0.5, 0.5] +[0.9990889488055994, 0.9990889488055994, 0.9990889488055994] +NULL diff --git a/test/sql/softmax.test b/test/sql/softmax.test new file mode 100644 index 0000000..c82dcc5 --- /dev/null +++ b/test/sql/softmax.test @@ -0,0 +1,57 @@ +# name: test/sql/softmax.test +# description: test cdb_softmax scalar function (LIST -> LIST) +# group: [sql] + +require cactusdb + +statement ok +LOAD cactusdb; + +# ------------------------------------------------------------------ +# Inputs +# ------------------------------------------------------------------ +statement ok +CREATE TEMP TABLE inputs(a DOUBLE[]); + +# Exact or easy-to-check cases: +# [0,0] -> [0.5, 0.5] +# [0,1] -> [~0.26894, ~0.73106] +# [1,2,3] -> [~0.09003, ~0.24473, ~0.66524] +# [-5, 0, 5] -> [~0.000045, ~0.006693, ~0.993262] (well-known sigmoid-ish spread) +# [7,7,7] -> [1/3, 1/3, 1/3] +statement ok +INSERT INTO inputs VALUES + ([0,0]), + ([0,1]), + ([1,2,3]), + ([-5,0,5]), + ([7,7,7]); + +# Round results to 6 decimals for deterministic comparison +query I +SELECT list_transform(cdb_softmax(a), x -> round(x::DOUBLE, 6)) AS y +FROM inputs; +---- +[0.5, 0.5] +[0.268941, 0.731059] +[0.090031, 0.244728, 0.665241] +[4.5e-05, 0.006693, 0.993262] +[0.333333, 0.333333, 0.333333] + +# ------------------------------------------------------------------ +# NULL row -> NULL +# ------------------------------------------------------------------ +statement ok +INSERT INTO inputs VALUES (NULL); + +query I +SELECT cdb_softmax(a) AS y +FROM inputs +ORDER BY 1 NULLS LAST; +---- +[4.509404123635488e-05, 0.006692549116589288, 0.9932623568421745] +[0.09003057317038048, 0.24472847105479764, 0.6652409557748219] +[0.2689414213699951, 0.7310585786300049] +[0.3333333333333333, 0.3333333333333333, 0.3333333333333333] +[0.5, 0.5] +NULL diff --git a/test/sql/torch.test b/test/sql/torch.test new file mode 100644 index 0000000..9244715 --- /dev/null +++ b/test/sql/torch.test @@ -0,0 +1,79 @@ +# name: test/sql/torchdnn_identity.test +# description: test torchdnn scalar function with dummy identity models +# group: [sql] + +require cactusdb + +statement ok +LOAD cactusdb; + +########################################################### +# Setup: simple input table of 3D vectors +########################################################### + +statement ok +CREATE TEMP TABLE torchdnn_input(v DOUBLE[]); + +statement ok +INSERT INTO torchdnn_input VALUES + ([1.0, 2.0, 3.0]), + ([3.0, 2.0, 1.0]); + +########################################################### +# 1) Vector output variant: torchdnn(v, 'identity3') +# identity3 implements y = x (3D identity layer) +########################################################### + +query I +SELECT torchdnn(v, 'identity3') AS y +FROM torchdnn_input +ORDER BY 1; +---- +[1.0, 2.0, 3.0] +[3.0, 2.0, 1.0] + +########################################################### +# 2) Argmax variant: torchdnn(v, 1, 'identity3_argmax') +# +# identity3_argmax implements y = x + argmax on output: +# [1,2,3] -> 2 +# [3,2,1] -> 0 +# We order results by idx so the output is deterministic. +########################################################### + +query I +SELECT torchdnn(v, 1, 'identity3_argmax') AS idx +FROM torchdnn_input +ORDER BY 1; +---- +0 +2 + +########################################################### +# 3) NULL row handling for vector output +########################################################### + +statement ok +INSERT INTO torchdnn_input VALUES (NULL); + +query I +SELECT torchdnn(v, 'identity3') AS y +FROM torchdnn_input +ORDER BY 1 NULLS LAST; +---- +[1.0, 2.0, 3.0] +[3.0, 2.0, 1.0] +NULL + +########################################################### +# 4) NULL row handling for argmax output +########################################################### + +query I +SELECT torchdnn(v, 1, 'identity3_argmax') AS idx +FROM torchdnn_input +ORDER BY 1 NULLS LAST; +---- +0 +2 +NULL