diff --git a/.claude/CLAUDE.md b/.claude/CLAUDE.md index 8d0f2c036e1..d89c979596f 100644 --- a/.claude/CLAUDE.md +++ b/.claude/CLAUDE.md @@ -116,7 +116,8 @@ or lightweight dependencies, use `add_executable` / `add_test`. ## Formatting ```bash -make format # format all changed files +pre-commit run --files # format/lint specific files +pre-commit run --all-files # sweep the whole repo ``` ## Coding Style diff --git a/CMake/resolve_dependency_modules/README.md b/CMake/resolve_dependency_modules/README.md index 0aedd1794d7..21fe605a356 100644 --- a/CMake/resolve_dependency_modules/README.md +++ b/CMake/resolve_dependency_modules/README.md @@ -32,11 +32,11 @@ by Velox. See details on bundling below. | fmt | 11.2.0 | Yes | Used API must be fmt 9 compatible | | simdjson | 4.1.0 | Yes || | faiss | 1.11.0 | Yes || -| folly | v2026.01.05.00 | Yes || -| fizz | v2026.01.05.00 | No || -| wangle | v2026.01.05.00 | No || -| mvfst | v2026.01.05.00 | No || -| fbthrift | v2026.01.05.00 | No || +| folly | v2026.07.13.00 | Yes || +| fizz | v2026.07.13.00 | No || +| wangle | v2026.07.13.00 | No || +| mvfst | v2026.07.13.00 | No || +| fbthrift | v2026.07.13.00 | No || | libstemmer | 2.2.0 | Yes || | DuckDB (testing) | 0.8.1 | Yes || | arrow | 15.0.0 | Yes || @@ -44,6 +44,8 @@ by Velox. See details on bundling below. | s2geometry | 0.12.0 | Yes || | fast_float | v8.0.2 | Yes || | xxhash | default | No || +| flatbuffers | 24.3.25 | Yes | Only with `VELOX_ENABLE_NIMBLE=ON`. `flatc` is required, not just the runtime. Held at the version cuDF expects | +| openzl | 6b48fa48 | Yes | Only with `VELOX_ENABLE_NIMBLE=ON`. Pinned to a commit; no suitable release tag | # Bundled Dependency Management This module provides a dependency management system that allows us to automatically fetch and build dependencies from source if needed. diff --git a/CMake/resolve_dependency_modules/fbthrift/compactv1-protocol-refiller.patch b/CMake/resolve_dependency_modules/fbthrift/compactv1-protocol-refiller.patch deleted file mode 100644 index 315590b44d3..00000000000 --- a/CMake/resolve_dependency_modules/fbthrift/compactv1-protocol-refiller.patch +++ /dev/null @@ -1,23 +0,0 @@ ---- a/thrift/lib/cpp2/protocol/ProtocolReaderWithRefill.h -+++ b/thrift/lib/cpp2/protocol/ProtocolReaderWithRefill.h -@@ -172,6 +172,20 @@ - protocol_.skipBytes(bytes); - } - -+ protected: -+ /** -+ * Allow derived refill readers to use this helper where they cannot access -+ * CompactProtocolReader::in_ directly (due to not having a friend -+ * declaration). -+ * Allows users to implement refill readers for CompactV1ProtocolReader if -+ * necessary, without supporting it in the Thrift codebase or having to patch -+ * in a friend declaration. -+ */ -+ template -+ T readLEFromBuffer() { -+ return protocol_.in_.template readLE(); -+ } -+ - private: - /** - * Make sure a varint can be read from the current buffer after idx bytes. diff --git a/CMake/resolve_dependency_modules/flatbuffers.cmake b/CMake/resolve_dependency_modules/flatbuffers.cmake new file mode 100644 index 00000000000..460d9660af4 --- /dev/null +++ b/CMake/resolve_dependency_modules/flatbuffers.cmake @@ -0,0 +1,71 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +include_guard(GLOBAL) + +# Held at the version cuDF expects; do not bump on its own. cuDF resolves +# FlatBuffers through rapids_cpm_find(flatbuffers 24.3.25), which calls +# find_package() before falling back to its own download, and FlatBuffers' +# config-version file reports every newer version as compatible. Any later +# version visible to that find_package() is therefore used in place of 24.3.25, +# and cuDF's checked-in generated headers static_assert on major 24 / minor 3. +# Nimble itself only requires 22.9.4 or later, for build_flatbuffers(). +set(VELOX_FLATBUFFERS_VERSION 24.3.25) +set( + VELOX_FLATBUFFERS_BUILD_SHA256_CHECKSUM + 4157c5cacdb59737c5d627e47ac26b140e9ee28b1102f812b36068aab728c1ed +) +string( + CONCAT + VELOX_FLATBUFFERS_SOURCE_URL + "https://github.com/google/flatbuffers/archive/refs/tags/" + "v${VELOX_FLATBUFFERS_VERSION}.tar.gz" +) + +velox_resolve_dependency_url(FLATBUFFERS) + +message(STATUS "Building FlatBuffers from source") + +# The flatc code generator is required, not just the runtime library: Nimble +# generates C++ headers from .fbs schemas at build time via build_flatbuffers(). +# This version exports it as the plain `flatc` target, which build_flatbuffers() +# falls back to on its own, so there is no need to set +# FLATBUFFERS_FLATC_EXECUTABLE here. FlatBuffers does set that variable itself, +# but only in its own directory scope, so it never reaches Velox. +set(FLATBUFFERS_BUILD_FLATC ON) +set(FLATBUFFERS_BUILD_FLATLIB ON) +set(FLATBUFFERS_BUILD_SHAREDLIB OFF) +set(FLATBUFFERS_BUILD_FLATHASH OFF) +set(FLATBUFFERS_BUILD_TESTS OFF) +set(FLATBUFFERS_INSTALL OFF) + +FetchContent_Declare( + flatbuffers + URL ${VELOX_FLATBUFFERS_SOURCE_URL} + URL_HASH ${VELOX_FLATBUFFERS_BUILD_SHA256_CHECKSUM} + OVERRIDE_FIND_PACKAGE + SYSTEM + EXCLUDE_FROM_ALL +) + +FetchContent_MakeAvailable(flatbuffers) + +# Consumers written against a system FlatBuffers use FLATBUFFERS_INCLUDE_DIR, +# which only FindFlatBuffers.cmake defines. Export it for the bundled build too +# so the same target_include_directories() call works either way. +set( + FLATBUFFERS_INCLUDE_DIR + "${flatbuffers_SOURCE_DIR}/include" + CACHE INTERNAL + "FlatBuffers include directory" +) diff --git a/CMake/resolve_dependency_modules/folly/CMakeLists.txt b/CMake/resolve_dependency_modules/folly/CMakeLists.txt index 8139c11a57e..0eef981ae54 100644 --- a/CMake/resolve_dependency_modules/folly/CMakeLists.txt +++ b/CMake/resolve_dependency_modules/folly/CMakeLists.txt @@ -17,10 +17,10 @@ cmake_minimum_required(VERSION 3.28) velox_set_source(FastFloat) velox_resolve_dependency(FastFloat CONFIG REQUIRED) -set(VELOX_FOLLY_BUILD_VERSION v2026.01.05.00) +set(VELOX_FOLLY_BUILD_VERSION v2026.07.13.00) set( VELOX_FOLLY_BUILD_SHA256_CHECKSUM - 8b41494b664fcde3d02f652d6a27381ec5e938ef58cf986d8ce4958d424af101 + b2d4bc0029d466328e807efa673ad01c1086c385ffa800f2055543c3b9640a72 ) set( VELOX_FOLLY_SOURCE_URL @@ -57,10 +57,6 @@ set(FOLLY_HAVE_INT128_T ON) FetchContent_MakeAvailable(folly) -# Folly::folly is not valid for FC but we want to match FindFolly -add_library(Folly::folly ALIAS folly) -add_library(Folly::follybenchmark ALIAS follybenchmark) - if(gflags_SOURCE STREQUAL "BUNDLED") add_dependencies(folly glog::glog gflags::gflags fmt::fmt) endif() diff --git a/CMake/resolve_dependency_modules/folly/folly-no-export.patch b/CMake/resolve_dependency_modules/folly/folly-no-export.patch index 375eaa16c79..2eaed5b4d72 100644 --- a/CMake/resolve_dependency_modules/folly/folly-no-export.patch +++ b/CMake/resolve_dependency_modules/folly/folly-no-export.patch @@ -13,7 +13,7 @@ # limitations under the License. --- a/CMakeLists.txt +++ b/CMakeLists.txt -@@ -500,14 +500,14 @@ target_link_libraries(folly_test_util +@@ -296,15 +296,15 @@ target_link_libraries(folly_test_util apply_folly_compile_options_to_target(folly_test_util) list(APPEND FOLLY_INSTALL_TARGETS folly_test_util) @@ -22,6 +22,7 @@ - RUNTIME DESTINATION bin - LIBRARY DESTINATION ${LIB_INSTALL_DIR} - ARCHIVE DESTINATION ${LIB_INSTALL_DIR}) +- -auto_install_files(folly ${FOLLY_DIR} - ${hfiles} -) @@ -30,13 +31,14 @@ +# RUNTIME DESTINATION bin +# LIBRARY DESTINATION ${LIB_INSTALL_DIR} +# ARCHIVE DESTINATION ${LIB_INSTALL_DIR}) ++# +#auto_install_files(folly ${FOLLY_DIR} +# ${hfiles} +#) install( FILES ${CMAKE_CURRENT_BINARY_DIR}/folly/folly-config.h DESTINATION ${INCLUDE_INSTALL_DIR}/folly -@@ -531,13 +531,13 @@ install( +@@ -328,13 +328,13 @@ install( DESTINATION ${CMAKE_INSTALL_DIR} COMPONENT dev ) diff --git a/CMake/resolve_dependency_modules/openzl.cmake b/CMake/resolve_dependency_modules/openzl.cmake new file mode 100644 index 00000000000..b839df38f99 --- /dev/null +++ b/CMake/resolve_dependency_modules/openzl.cmake @@ -0,0 +1,83 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +include_guard(GLOBAL) + +# Pinned to a raw commit rather than to the v0.2.0 tag, which predates both the +# cross-platform zstd/xxhash dependency handling OpenZL needs to configure +# cleanly here and the current descriptor API. The revision tracks what fbcode +# builds Nimble against, so that velox/dwio/nimble compiles identically in both +# trees; a pin older than fbcode's OpenZL breaks the internal build. OpenZL +# ships from fbcode to GitHub, so there is no tag to follow. Keep in sync with +# OPENZL_VERSION in scripts/setup-versions.sh, which installs the same revision +# for builds that resolve OpenZL as a system package. +set(VELOX_OPENZL_VERSION 7340a712cce1b8331bec3467600dba99a562e052) +set( + VELOX_OPENZL_BUILD_SHA256_CHECKSUM + ace6a975c3bb28da0f211c86b3f172b427acec4a2044ab637e66649cb022a953 +) +string( + CONCAT + VELOX_OPENZL_SOURCE_URL + "https://github.com/facebook/openzl/archive/" + "${VELOX_OPENZL_VERSION}.tar.gz" +) + +velox_resolve_dependency_url(OPENZL) + +message(STATUS "Building OpenZL from source") + +# Only the core `openzl` and `openzl_cpp` libraries are consumed. Everything +# else OpenZL can build pulls in dependencies Velox does not otherwise need. +set(OPENZL_BUILD_CLI OFF) +set(OPENZL_BUILD_EXAMPLES OFF) +set(OPENZL_BUILD_TOOLS OFF) +set(OPENZL_BUILD_CUSTOM_PARSERS OFF) +set(OPENZL_BUILD_TESTS OFF) +set(OPENZL_BUILD_BENCHMARKS OFF) +set(OPENZL_BUILD_PYTHON_EXT OFF) +set(OPENZL_INSTALL OFF) + +# OpenZL hard-codes CMAKE_CXX_STANDARD to 17, but Velox builds with C++20. Its +# C++ bindings select between std:: types and internal polyfills (poly::span, +# poly::source_location, poly::string_view, ...) at compile time via +# feature-test macros in openzl/cpp/include/openzl/cpp/detail/Portability.hpp. +# If openzl_cpp is compiled as C++17 while Velox consumes its headers as C++20, +# the same poly:: aliases resolve to different underlying types on either side +# of the ABI, producing mismatched mangled symbols and undefined-reference link +# errors. The patch makes OpenZL honor an externally provided standard, so this +# build inherits Velox's C++20 and scripts/setup-common.sh can install a +# matching copy. The tarball is not a git repository, hence the `git init`. +FetchContent_Declare( + openzl + URL ${VELOX_OPENZL_SOURCE_URL} + URL_HASH ${VELOX_OPENZL_BUILD_SHA256_CHECKSUM} + PATCH_COMMAND git init -q && git apply ${CMAKE_CURRENT_LIST_DIR}/openzl/openzl-cxx-standard.patch + OVERRIDE_FIND_PACKAGE + SYSTEM + EXCLUDE_FROM_ALL +) + +# OpenZL's C++ bindings brace initialize several descriptor structs without +# naming every member. TREAT_WARNINGS_AS_ERRORS puts -Werror into the global +# CMAKE_CXX_FLAGS, which FetchContent subprojects inherit, so those become build +# failures. Velox holds its own code to that standard, not its third party +# dependencies. Relax just this warning while OpenZL is configured and restore +# the flags afterwards, as duckdb and geos already do. Only the C++ bindings are +# affected; -Werror is never added to CMAKE_C_FLAGS. +set(PREVIOUS_CMAKE_CXX_FLAGS ${CMAKE_CXX_FLAGS}) +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-missing-field-initializers") + +FetchContent_MakeAvailable(openzl) + +set(CMAKE_CXX_FLAGS ${PREVIOUS_CMAKE_CXX_FLAGS}) diff --git a/CMake/resolve_dependency_modules/openzl/openzl-cxx-standard.patch b/CMake/resolve_dependency_modules/openzl/openzl-cxx-standard.patch new file mode 100644 index 00000000000..7fa5b71aca5 --- /dev/null +++ b/CMake/resolve_dependency_modules/openzl/openzl-cxx-standard.patch @@ -0,0 +1,17 @@ +--- a/CMakeLists.txt ++++ b/CMakeLists.txt +@@ -67,7 +67,13 @@ + # ---- Compiler Settings ---- + set(CMAKE_C_STANDARD 11) + set(CMAKE_C_STANDARD_REQUIRED ON) +-set(CMAKE_CXX_STANDARD 17) ++# Honor an externally provided C++ standard so the C++ bindings compile at the ++# consumer's standard (Velox is C++20). OpenZL's poly:: types select between ++# std:: and internal polyfills by standard, so the bindings must be built at ++# the same standard the consumer uses to avoid mismatched mangled symbols. ++if(NOT DEFINED CMAKE_CXX_STANDARD) ++ set(CMAKE_CXX_STANDARD 17) ++endif() + set(CMAKE_CXX_STANDARD_REQUIRED ON) + + # ---- Platform Checks ---- diff --git a/CMake/third-party/BuildFlatBuffers.cmake b/CMake/third-party/BuildFlatBuffers.cmake new file mode 100644 index 00000000000..052a0f41d73 --- /dev/null +++ b/CMake/third-party/BuildFlatBuffers.cmake @@ -0,0 +1,474 @@ +# Copyright 2015 Google Inc. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# General function to create FlatBuffer build rules for the given list of +# schemas. +# +# flatbuffers_schemas: A list of flatbuffer schema files to process. +# +# schema_include_dirs: A list of schema file include directories, which will be +# passed to flatc via the -I parameter. +# +# custom_target_name: The generated files will be added as dependencies for a +# new custom target with this name. You should add that target as a dependency +# for your main target to ensure these files are built. You can also retrieve +# various properties from this target, such as GENERATED_INCLUDES_DIR, +# BINARY_SCHEMAS_DIR, and COPY_TEXT_SCHEMAS_DIR. +# +# additional_dependencies: A list of additional dependencies that you'd like all +# generated files to depend on. Pass in a blank string if you have none. +# +# generated_includes_dir: Where to generate the C++ header files for these +# schemas. The generated includes directory will automatically be added to +# CMake's include directories, and will be where generated header files are +# placed. This parameter is optional; pass in empty string if you don't want to +# generate include files for these schemas. +# +# binary_schemas_dir: If you specify an optional binary schema directory, binary +# schemas will be generated for these schemas as well, and placed into the given +# directory. +# +# copy_text_schemas_dir: If you want all text schemas (including schemas from +# all schema include directories) copied into a directory (for example, if you +# need them within your project to build JSON files), you can specify that +# folder here. All text schemas will be copied to that folder. +# +# IMPORTANT: Make sure you quote all list arguments you pass to this function! +# Otherwise CMake will only pass in the first element. Example: +# build_flatbuffers("${fb_files}" "${include_dirs}" target_name ...) +function( + build_flatbuffers + flatbuffers_schemas + schema_include_dirs + custom_target_name + additional_dependencies + generated_includes_dir + binary_schemas_dir + copy_text_schemas_dir +) + # Test if including from FindFlatBuffers + if(FLATBUFFERS_FLATC_EXECUTABLE) + set(FLATC_TARGET "") + set(FLATC ${FLATBUFFERS_FLATC_EXECUTABLE}) + elseif(TARGET flatbuffers::flatc) + set(FLATC_TARGET flatbuffers::flatc) + set(FLATC flatbuffers::flatc) + else() + set(FLATC_TARGET flatc) + set(FLATC flatc) + endif() + set(FLATC_SCHEMA_ARGS --gen-mutable) + if(FLATBUFFERS_FLATC_SCHEMA_EXTRA_ARGS) + set( + FLATC_SCHEMA_ARGS + ${FLATBUFFERS_FLATC_SCHEMA_EXTRA_ARGS} + ${FLATC_SCHEMA_ARGS} + ) + endif() + + set(working_dir "${CMAKE_CURRENT_SOURCE_DIR}") + + set(schema_glob "*.fbs") + # Generate the include files parameters. + set(include_params "") + set(all_generated_files "") + foreach(include_dir ${schema_include_dirs}) + set(include_params -I ${include_dir} ${include_params}) + if(NOT ${copy_text_schemas_dir} STREQUAL "") + # Copy text schemas from dependent folders. + file(GLOB_RECURSE dependent_schemas ${include_dir}/${schema_glob}) + foreach(dependent_schema ${dependent_schemas}) + file(COPY ${dependent_schema} DESTINATION ${copy_text_schemas_dir}) + endforeach() + endif() + endforeach() + + foreach(schema ${flatbuffers_schemas}) + get_filename_component(filename ${schema} NAME_WE) + # For each schema, do the things we requested. + if(NOT ${generated_includes_dir} STREQUAL "") + set(generated_include ${generated_includes_dir}/${filename}Generated.h) + add_custom_command( + OUTPUT ${generated_include} + # Nimble code expects an upper case suffix to the generated file. + COMMAND + ${FLATC} "--filename-suffix" "Generated" ${FLATC_SCHEMA_ARGS} -o + ${generated_includes_dir} ${include_params} -c ${schema} + DEPENDS ${FLATC_TARGET} ${schema} ${additional_dependencies} + WORKING_DIRECTORY "${working_dir}" + ) + list(APPEND all_generated_files ${generated_include}) + endif() + + if(NOT ${binary_schemas_dir} STREQUAL "") + set(binary_schema ${binary_schemas_dir}/${filename}.bfbs) + add_custom_command( + OUTPUT ${binary_schema} + COMMAND + ${FLATC} -b --schema -o ${binary_schemas_dir} ${include_params} + ${schema} + DEPENDS ${FLATC_TARGET} ${schema} ${additional_dependencies} + WORKING_DIRECTORY "${working_dir}" + ) + list(APPEND all_generated_files ${binary_schema}) + endif() + + if(NOT ${copy_text_schemas_dir} STREQUAL "") + file(COPY ${schema} DESTINATION ${copy_text_schemas_dir}) + endif() + endforeach() + + # Create a custom target that depends on all the generated files. This is the + # target that you can depend on to trigger all these to be built. + add_custom_target( + ${custom_target_name} + DEPENDS ${all_generated_files} ${additional_dependencies} + ) + + # Register the include directory we are using. + if(NOT ${generated_includes_dir} STREQUAL "") + include_directories(${generated_includes_dir}) + set_property( + TARGET ${custom_target_name} + PROPERTY GENERATED_INCLUDES_DIR ${generated_includes_dir} + ) + endif() + + # Register the binary schemas dir we are using. + if(NOT ${binary_schemas_dir} STREQUAL "") + set_property( + TARGET ${custom_target_name} + PROPERTY BINARY_SCHEMAS_DIR ${binary_schemas_dir} + ) + endif() + + # Register the text schema copy dir we are using. + if(NOT ${copy_text_schemas_dir} STREQUAL "") + set_property( + TARGET ${custom_target_name} + PROPERTY COPY_TEXT_SCHEMAS_DIR ${copy_text_schemas_dir} + ) + endif() +endfunction() + +# Creates a target that can be linked against that generates flatbuffer headers. +# +# This function takes a target name and a list of schemas. You can also specify +# other flagc flags using the FLAGS option to change the behavior of the flatc +# tool. +# +# When the target_link_libraries is done within a different directory than +# flatbuffers_generate_headers is called, then the target should also be +# dependent the custom generation target called GENERATE_. +# +# Arguments: TARGET: The name of the target to generate. SCHEMAS: The list of +# schema files to generate code for. BINARY_SCHEMAS_DIR: Optional. The directory +# in which to generate binary schemas. Binary schemas will only be generated if +# a path is provided. INCLUDE: Optional. Search for includes in the specified +# paths. (Use this instead of "-I " and the FLAGS option so that CMake is +# aware of the directories that need to be searched). INCLUDE_PREFIX: Optional. +# The directory in which to place the generated files. Use this instead of the +# --include-prefix option. FLAGS: Optional. A list of any additional flags that +# you would like to pass to flatc. +# +# Example: +# +# flatbuffers_generate_headers( TARGET my_generated_headers_target +# INCLUDE_PREFIX ${MY_INCLUDE_PREFIX}" SCHEMAS ${MY_SCHEMA_FILES} +# BINARY_SCHEMAS_DIR "${MY_BINARY_SCHEMA_DIRECTORY}" FLAGS --gen-object-api) +# +# target_link_libraries(MyExecutableTarget PRIVATE my_generated_headers_target ) +# +# Optional (only needed within different directory): add_dependencies(app +# GENERATE_my_generated_headers_target) +function(flatbuffers_generate_headers) + # Parse function arguments. + set(options) + set(one_value_args "TARGET" "INCLUDE_PREFIX" "BINARY_SCHEMAS_DIR") + set(multi_value_args "SCHEMAS" "INCLUDE" "FLAGS") + cmake_parse_arguments( + PARSE_ARGV + 0 + FLATBUFFERS_GENERATE_HEADERS + "${options}" + "${one_value_args}" + "${multi_value_args}" + ) + + # Test if including from FindFlatBuffers + if(FLATBUFFERS_FLATC_EXECUTABLE) + set(FLATC_TARGET "") + set(FLATC ${FLATBUFFERS_FLATC_EXECUTABLE}) + elseif(TARGET flatbuffers::flatc) + set(FLATC_TARGET flatbuffers::flatc) + set(FLATC flatbuffers::flatc) + else() + set(FLATC_TARGET flatc) + set(FLATC flatc) + endif() + + set(working_dir "${CMAKE_CURRENT_SOURCE_DIR}") + + # Generate the include files parameters. + set(include_params "") + foreach(include_dir ${FLATBUFFERS_GENERATE_HEADERS_INCLUDE}) + set(include_params -I ${include_dir} ${include_params}) + endforeach() + + # Create a directory to place the generated code. + set( + generated_target_dir + "${CMAKE_CURRENT_BINARY_DIR}/${FLATBUFFERS_GENERATE_HEADERS_TARGET}" + ) + set(generated_include_dir "${generated_target_dir}") + if(NOT ${FLATBUFFERS_GENERATE_HEADERS_INCLUDE_PREFIX} STREQUAL "") + set( + generated_include_dir + "${generated_include_dir}/${FLATBUFFERS_GENERATE_HEADERS_INCLUDE_PREFIX}" + ) + list( + APPEND + FLATBUFFERS_GENERATE_HEADERS_FLAGS + "--include-prefix" + ${FLATBUFFERS_GENERATE_HEADERS_INCLUDE_PREFIX} + ) + endif() + + set(generated_custom_commands) + + # Create rules to generate the code for each schema. + foreach(schema ${FLATBUFFERS_GENERATE_HEADERS_SCHEMAS}) + get_filename_component(filename ${schema} NAME_WE) + set(generated_include "${generated_include_dir}/${filename}_generated.h") + + # Generate files for grpc if needed + set(generated_source_file) + if("${FLATBUFFERS_GENERATE_HEADERS_FLAGS}" MATCHES "--grpc") + # Check if schema file contain a rpc_service definition + file(STRINGS ${schema} has_grpc REGEX "rpc_service") + if(has_grpc) + list( + APPEND + generated_include + "${generated_include_dir}/${filename}.grpc.fb.h" + ) + set( + generated_source_file + "${generated_include_dir}/${filename}.grpc.fb.cc" + ) + endif() + endif() + + add_custom_command( + OUTPUT ${generated_include} ${generated_source_file} + COMMAND + ${FLATC} ${FLATC_ARGS} -o ${generated_include_dir} ${include_params} -c + ${schema} ${FLATBUFFERS_GENERATE_HEADERS_FLAGS} + DEPENDS ${FLATC_TARGET} ${schema} + WORKING_DIRECTORY "${working_dir}" + COMMENT "Building ${schema} flatbuffers..." + ) + list(APPEND all_generated_header_files ${generated_include}) + list(APPEND all_generated_source_files ${generated_source_file}) + list( + APPEND + generated_custom_commands + "${generated_include}" + "${generated_source_file}" + ) + + # Geneate the binary flatbuffers schemas if instructed to. + if(NOT ${FLATBUFFERS_GENERATE_HEADERS_BINARY_SCHEMAS_DIR} STREQUAL "") + set( + binary_schema + "${FLATBUFFERS_GENERATE_HEADERS_BINARY_SCHEMAS_DIR}/${filename}.bfbs" + ) + add_custom_command( + OUTPUT ${binary_schema} + COMMAND + ${FLATC} -b --schema -o + ${FLATBUFFERS_GENERATE_HEADERS_BINARY_SCHEMAS_DIR} ${include_params} + ${schema} + DEPENDS ${FLATC_TARGET} ${schema} + WORKING_DIRECTORY "${working_dir}" + ) + list(APPEND generated_custom_commands "${binary_schema}") + list(APPEND all_generated_binary_files ${binary_schema}) + endif() + endforeach() + + # Create an additional target as add_custom_command scope is only within same + # directory (CMakeFile.txt) + set(generate_target GENERATE_${FLATBUFFERS_GENERATE_HEADERS_TARGET}) + add_custom_target( + ${generate_target} + ALL + DEPENDS ${generated_custom_commands} + COMMENT + "Generating flatbuffer target ${FLATBUFFERS_GENERATE_HEADERS_TARGET}" + ) + + # Set up interface library + add_library(${FLATBUFFERS_GENERATE_HEADERS_TARGET} INTERFACE) + target_sources( + ${FLATBUFFERS_GENERATE_HEADERS_TARGET} + INTERFACE + ${all_generated_header_files} + ${all_generated_binary_files} + ${all_generated_source_files} + ${FLATBUFFERS_GENERATE_HEADERS_SCHEMAS} + ) + add_dependencies( + ${FLATBUFFERS_GENERATE_HEADERS_TARGET} + ${FLATC} + ${FLATBUFFERS_GENERATE_HEADERS_SCHEMAS} + ) + target_include_directories( + ${FLATBUFFERS_GENERATE_HEADERS_TARGET} + INTERFACE ${generated_target_dir} + ) + + # Organize file layout for IDEs. + source_group( + TREE "${generated_target_dir}" + PREFIX "Flatbuffers/Generated/Headers Files" + FILES ${all_generated_header_files} + ) + source_group( + TREE "${generated_target_dir}" + PREFIX "Flatbuffers/Generated/Source Files" + FILES ${all_generated_source_files} + ) + source_group( + TREE ${working_dir} + PREFIX "Flatbuffers/Schemas" + FILES ${FLATBUFFERS_GENERATE_HEADERS_SCHEMAS} + ) + if(NOT ${FLATBUFFERS_GENERATE_HEADERS_BINARY_SCHEMAS_DIR} STREQUAL "") + source_group( + TREE "${FLATBUFFERS_GENERATE_HEADERS_BINARY_SCHEMAS_DIR}" + PREFIX "Flatbuffers/Generated/Binary Schemas" + FILES ${all_generated_binary_files} + ) + endif() +endfunction() + +# Creates a target that can be linked against that generates flatbuffer binaries +# from json files. +# +# This function takes a target name and a list of schemas and Json files. You +# can also specify other flagc flags and options to change the behavior of the +# flatc compiler. +# +# Adding this target to your executable ensurses that the flatbuffer binaries +# are compiled before your executable is run. +# +# Arguments: TARGET: The name of the target to generate. JSON_FILES: The list of +# json files to compile to flatbuffers binaries. SCHEMA: The flatbuffers schema +# of the Json files to be compiled. INCLUDE: Optional. Search for includes in +# the specified paths. (Use this instead of "-I " and the FLAGS option so +# that CMake is aware of the directories that need to be searched). OUTPUT_DIR: +# The directly where the generated flatbuffers binaries should be placed. FLAGS: +# Optional. A list of any additional flags that you would like to pass to flatc. +# +# Example: +# +# flatbuffers_generate_binary_files( TARGET my_binary_data SCHEMA +# "${MY_SCHEMA_DIR}/my_example_schema.fbs" JSON_FILES ${MY_JSON_FILES} +# OUTPUT_DIR "${MY_BINARY_DATA_DIRECTORY}" FLAGS --strict-json) +# +# target_link_libraries(MyExecutableTarget PRIVATE my_binary_data ) +function(flatbuffers_generate_binary_files) + # Parse function arguments. + set(options) + set(one_value_args "TARGET" "SCHEMA" "OUTPUT_DIR") + set(multi_value_args "JSON_FILES" "INCLUDE" "FLAGS") + cmake_parse_arguments( + PARSE_ARGV + 0 + FLATBUFFERS_GENERATE_BINARY_FILES + "${options}" + "${one_value_args}" + "${multi_value_args}" + ) + + # Test if including from FindFlatBuffers + if(FLATBUFFERS_FLATC_EXECUTABLE) + set(FLATC_TARGET "") + set(FLATC ${FLATBUFFERS_FLATC_EXECUTABLE}) + elseif(TARGET flatbuffers::flatc) + set(FLATC_TARGET flatbuffers::flatc) + set(FLATC flatbuffers::flatc) + else() + set(FLATC_TARGET flatc) + set(FLATC flatc) + endif() + + set(working_dir "${CMAKE_CURRENT_SOURCE_DIR}") + + # Generate the include files parameters. + set(include_params "") + foreach(include_dir ${FLATBUFFERS_GENERATE_BINARY_FILES_INCLUDE}) + set(include_params -I ${include_dir} ${include_params}) + endforeach() + + # Create rules to generate the flatbuffers binary for each json file. + foreach(json_file ${FLATBUFFERS_GENERATE_BINARY_FILES_JSON_FILES}) + get_filename_component(filename ${json_file} NAME_WE) + set( + generated_binary_file + "${FLATBUFFERS_GENERATE_BINARY_FILES_OUTPUT_DIR}/${filename}.bin" + ) + add_custom_command( + OUTPUT ${generated_binary_file} + COMMAND + ${FLATC} ${FLATC_ARGS} -o + ${FLATBUFFERS_GENERATE_BINARY_FILES_OUTPUT_DIR} ${include_params} -b + ${FLATBUFFERS_GENERATE_BINARY_FILES_SCHEMA} ${json_file} + ${FLATBUFFERS_GENERATE_BINARY_FILES_FLAGS} + DEPENDS ${FLATC_TARGET} ${json_file} + WORKING_DIRECTORY "${working_dir}" + COMMENT "Building ${json_file} binary flatbuffers..." + ) + list(APPEND all_generated_binary_files ${generated_binary_file}) + endforeach() + + # Set up interface library + add_library(${FLATBUFFERS_GENERATE_BINARY_FILES_TARGET} INTERFACE) + target_sources( + ${FLATBUFFERS_GENERATE_BINARY_FILES_TARGET} + INTERFACE + ${all_generated_binary_files} + ${FLATBUFFERS_GENERATE_BINARY_FILES_JSON_FILES} + ${FLATBUFFERS_GENERATE_BINARY_FILES_SCHEMA} + ) + add_dependencies(${FLATBUFFERS_GENERATE_BINARY_FILES_TARGET} ${FLATC}) + + # Organize file layout for IDEs. + source_group( + TREE ${working_dir} + PREFIX "Flatbuffers/JSON Files" + FILES ${FLATBUFFERS_GENERATE_BINARY_FILES_JSON_FILES} + ) + source_group( + TREE ${working_dir} + PREFIX "Flatbuffers/Schemas" + FILES ${FLATBUFFERS_GENERATE_BINARY_FILES_SCHEMA} + ) + source_group( + TREE ${FLATBUFFERS_GENERATE_BINARY_FILES_OUTPUT_DIR} + PREFIX "Flatbuffers/Generated/Binary Files" + FILES ${all_generated_binary_files} + ) +endfunction() diff --git a/CMake/third-party/FBThriftCppLibrary.cmake b/CMake/third-party/FBThriftCppLibrary.cmake index 955f8a1e2d9..51da53c0e33 100644 --- a/CMake/third-party/FBThriftCppLibrary.cmake +++ b/CMake/third-party/FBThriftCppLibrary.cmake @@ -131,7 +131,17 @@ function(add_fbthrift_cpp_library LIB_NAME THRIFT_FILE) "${FBTHRIFT_COMPILER}" ) - add_library("${LIB_NAME}" STATIC ${generated_sources}) + # Now emit the library rule to compile the sources + if (BUILD_SHARED_LIBS) + set(LIB_TYPE SHARED) + else () + set(LIB_TYPE STATIC) + endif () + + add_library( + "${LIB_NAME}" ${LIB_TYPE} + ${generated_sources} + ) target_include_directories( "${LIB_NAME}" diff --git a/CMakeLists.txt b/CMakeLists.txt index c4564d54e5b..bd5b8f7d9d9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -66,6 +66,7 @@ list( include(ResolveDependency) include(VeloxUtils) include(CMakeDependentOption) +include(CheckIPOSupported) velox_set_with_default(VELOX_DEPENDENCY_SOURCE_DEFAULT VELOX_DEPENDENCY_SOURCE AUTO) message(STATUS "Build type: ${CMAKE_BUILD_TYPE}") @@ -91,6 +92,7 @@ option(VELOX_MONO_LIBRARY "Build single unified library." ON) option(ENABLE_ALL_WARNINGS "Enable -Wall and -Wextra compiler warnings." ON) option(VELOX_BUILD_SHARED "Build Velox as shared libraries." OFF) option(VELOX_BUILD_CMAKE_PACKAGE "Build CMake package for Velox." OFF) +option(VELOX_ENABLE_LTO "Enable link-time optimization (IPO/LTO)." OFF) option(VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST "Disable Wave branch kernel test." OFF) # While it's possible to build both in one go we currently want to build either # static or shared. @@ -116,6 +118,18 @@ if(VELOX_BUILD_SHARED) ) endif() +if(VELOX_ENABLE_LTO) + check_ipo_supported(RESULT VELOX_IPO_SUPPORTED OUTPUT VELOX_IPO_ERROR LANGUAGES CXX) + if(NOT VELOX_IPO_SUPPORTED) + message( + FATAL_ERROR + "VELOX_ENABLE_LTO requested, but IPO/LTO is not supported: ${VELOX_IPO_ERROR}" + ) + endif() + set(CMAKE_INTERPROCEDURAL_OPTIMIZATION ON) + message(STATUS "VELOX_ENABLE_LTO enabled") +endif() + # option() always creates a BOOL variable so we have to use a normal cache # variable with STRING type for this option. # @@ -155,9 +169,15 @@ option(VELOX_ENABLE_GCS "Build GCS Connector" OFF) option(VELOX_ENABLE_ABFS "Build Abfs Connector" OFF) option(VELOX_ENABLE_HDFS "Build Hdfs Connector" OFF) option(VELOX_ENABLE_PARQUET "Enable Parquet support" ON) +option(VELOX_ENABLE_NIMBLE "Enable Nimble support" OFF) option(VELOX_ENABLE_ARROW "Enable Arrow support" OFF) option(VELOX_ENABLE_GEO "Enable Geospatial support" ON) option(VELOX_ENABLE_REMOTE_FUNCTIONS "Enable remote function support" OFF) +option( + VELOX_ENABLE_LOCAL_RUNNER_SERVICE + "Enable LocalRunnerService and VeloxQueryRunner for expression fuzzer regression testing" + OFF +) option(VELOX_ENABLE_CCACHE "Use ccache if installed." ON) option(VELOX_ENABLE_COMPRESSION_LZ4 "Enable Lz4 compression support." OFF) @@ -660,8 +680,13 @@ if(${VELOX_BUILD_PYTHON_PACKAGE}) velox_resolve_dependency(pybind11 2.10.0) endif() -# DWIO (ORC/DWRF) depends on protobuf. -if(${VELOX_BUILD_MINIMAL_WITH_DWIO} OR ${VELOX_ENABLE_HIVE_CONNECTOR} OR VELOX_ENABLE_GCS) +# DWIO (ORC/DWRF) depends on protobuf. Nimble's serializer does too. +if( + ${VELOX_BUILD_MINIMAL_WITH_DWIO} + OR ${VELOX_ENABLE_HIVE_CONNECTOR} + OR VELOX_ENABLE_GCS + OR VELOX_ENABLE_NIMBLE +) # Locate or build protobuf. velox_set_source(Protobuf) velox_resolve_dependency(Protobuf 3.21.7 REQUIRED) @@ -691,9 +716,14 @@ if(${VELOX_BUILD_TESTING}) velox_resolve_dependency(gRPC) endif() -if(VELOX_ENABLE_REMOTE_FUNCTIONS OR VELOX_ENABLE_PARQUET) - # TODO: Move this to use resolve_dependency(). For some reason, FBThrift - # requires clients to explicitly install fizz and wangle. +# FBThrift is required by remote functions and LocalRunnerService. +if(VELOX_ENABLE_REMOTE_FUNCTIONS OR VELOX_ENABLE_PARQUET OR VELOX_ENABLE_LOCAL_RUNNER_SERVICE) + set(VELOX_NEEDS_FBTHRIFT ON) +endif() + +# TODO: Move this to use resolve_dependency(). For some reason, FBThrift +# requires clients to explicitly install fizz and wangle. +if(VELOX_NEEDS_FBTHRIFT) find_package(fizz CONFIG REQUIRED) find_package(wangle CONFIG REQUIRED) find_package(FBThrift CONFIG REQUIRED) @@ -856,4 +886,41 @@ if(VELOX_ENABLE_GEO) velox_resolve_dependency(s2geometry) endif() +# Resolved last, after every other dependency. OpenZL vendors its own copy of +# zstd, whose CMake defines an `uninstall` target. GEOS defines one too, but +# without the `if(NOT TARGET uninstall)` guard that zstd has, so GEOS fails +# outright if anything claimed the name first. Ordering Nimble after GEOS lets +# GEOS win the race and zstd skip, which is the only arrangement in which a +# BUNDLED build of both succeeds. Nothing between here and add_subdirectory() +# needs these targets. +if(VELOX_ENABLE_NIMBLE) + # Nimble is the only consumer of both, so they are resolved here rather than + # unconditionally. FlatBuffers supplies both the runtime library and the flatc + # code generator that turns Nimble's .fbs schemas into C++ headers. Nimble's + # third dependency, FSST, is vendored under velox/external/fsst instead. + velox_set_source(flatbuffers) + velox_resolve_dependency(flatbuffers) + # build_flatbuffers() only ships with FlatBuffers 22.9.4 and later, and a + # bundled build does not put the package's own modules on CMAKE_MODULE_PATH at + # all, so CMake/third-party/BuildFlatBuffers.cmake is used unconditionally to + # keep a single code path. It resolves flatc from FLATBUFFERS_FLATC_EXECUTABLE, + # the flatbuffers::flatc target or the flatc target, whichever the resolved + # package provides. + include(BuildFlatBuffers) + + velox_set_source(openzl) + velox_resolve_dependency(openzl) + # A bundled build exposes the bare target names while find_package() exposes + # namespaced ones. Aliasing between them is not safe: a find_package() that + # reports the package as not found still leaves its imported OpenZL:: targets + # behind, so `if(NOT TARGET OpenZL::openzl_cpp)` would skip the alias and + # consumers would link the rejected system package. Select on the resolution + # mode instead, which is unambiguous. + if(openzl_SOURCE STREQUAL "SYSTEM") + set(VELOX_OPENZL_LIBRARIES OpenZL::openzl OpenZL::openzl_cpp) + else() + set(VELOX_OPENZL_LIBRARIES openzl openzl_cpp) + endif() +endif() + add_subdirectory(velox) diff --git a/CODING_STYLE.md b/CODING_STYLE.md index d5432a9216b..cd88b71346c 100644 --- a/CODING_STYLE.md +++ b/CODING_STYLE.md @@ -555,6 +555,41 @@ using ContinuePromise = VeloxPromise; needs access to private members, redesign the API or test through public methods instead. +### Breaking API changes and `VELOX_ENABLE_BACKWARD_COMPATIBILITY` + +Velox is synced into Meta's internal repository. Prestissimo's copy there is +read-only — its source of truth is the Presto GitHub repo — so it cannot be +updated in the same change as a breaking Velox API change. To let it keep +compiling, Velox can carry the old signature behind the +`VELOX_ENABLE_BACKWARD_COMPATIBILITY` macro. Only Prestissimo's internal build +defines it; CMake and every open-source build see the new API alone. + +* **Additive only.** The old and new signatures must coexist in one binary. + Never select between them with `#ifdef`/`#else` — a change that alters a + return type or a virtual's signature gives different translation units + different vtable layouts, which is an ODR violation rather than a + compatibility shim. A change that cannot be made additive cannot use the + macro. +* **Define the legacy form inline in the header,** delegating to the new one. + Prestissimo's build compiles the header, so an out-of-line definition would + not reach it. +* **It is temporary.** Name the replacement on the legacy declaration and + delete the guarded block once Prestissimo has migrated. The macro is a + bridge, not a deprecation mechanism — do not reach for it to spare callers + you can update yourself, and never for callers that live in this repository. + +```cpp + // Current API. + ExchangeNode(const PlanNodeId& id, RowTypePtr type, std::string serdeKind); + +#ifdef VELOX_ENABLE_BACKWARD_COMPATIBILITY + /// Legacy constructor. Prefer the std::string overload above. Removed once + /// all callers have migrated. + ExchangeNode(const PlanNodeId& id, RowTypePtr type, VectorSerde::Kind kind) + : ExchangeNode(id, std::move(type), VectorSerde::kindName(kind)) {} +#endif // VELOX_ENABLE_BACKWARD_COMPATIBILITY +``` + ## Tests * **Each test file should have one test suite with a matching name.** E.g., diff --git a/README.md b/README.md index ce3013a9629..42fe7dd9c99 100644 --- a/README.md +++ b/README.md @@ -73,9 +73,9 @@ of available functions [can be found here.](https://facebookincubator.github.io/ Recent blog posts ([all posts](https://velox-lib.io/blog)): +- [Native Delta Statistics with Velox Task Barriers](https://velox-lib.io/blog/native-delta-statistics) (2026-08-25) - [Build Once, Probe Many: Hash Table Caching in Velox](https://velox-lib.io/blog/hash-table-caching) (2026-08-03) - [War of the Allocators](https://velox-lib.io/blog/war-of-the-allocators) (2026-07-27) -- [Making OpenZL Available in Nimble OSS](https://velox-lib.io/blog/openzl-in-nimble-oss) (2026-07-05) ## Community diff --git a/scripts/checks/license-header.py b/scripts/checks/license-header.py index fb8f8ac1b6b..891ac928037 100755 --- a/scripts/checks/license-header.py +++ b/scripts/checks/license-header.py @@ -119,6 +119,50 @@ def wrapper_hash(header, args): } ) +# Copyright lines that also count as an existing license header. The canonical +# line lives in the header file passed via --header and is the one written into +# files that have none. These variants are only recognised, never inserted, so +# that code imported with an equally valid Meta copyright line is not given a +# second header on top of the one it already carries. +ALTERNATE_COPYRIGHT_LINES = [ + " Copyright (c) Meta Platforms, Inc. and affiliates.", +] + + +def header_variants(header_text): + """Canonical header first, then the same header with each accepted + copyright line substituted for the first.""" + return [header_text] + [ + [line] + header_text[1:] for line in ALTERNATE_COPYRIGHT_LINES + ] + + +def find_header(content, header_comment, args): + """Locate an existing license header, exactly or fuzzily. Returns the + (start, end) span of the match, or None.""" + found = content.find(header_comment, 0, len(header_comment) + args.extra) + if found >= 0: + return found, found + len(header_comment) + + # Look for a fuzzy match in the first 60 chars. + found = regex.search( + "(?be)(%s){e<=%d}" % (regex.escape(header_comment[0:60]), 6), + content[0 : 80 + args.extra], + ) + if not found: + return None + + # If the first 80 chars match, try harder for the rest of the header. + fuzzy = regex.compile( + "(?be)(%s){e<=%d}" % (regex.escape(header_comment), args.editdist) + ) + found = fuzzy.search(content[0 : len(header_comment) + args.extra], found.start()) + if not found: + return None + + return found.start(), found.end() + + file_pattern = regex.compile( "|".join( [ @@ -169,6 +213,7 @@ def main(): log_to = sys.stdout header_text = file_lines(args.header) + variants = header_variants(header_text) if len(args.files) == 1 and args.files[0] == "-": files = [file.strip() for file in sys.stdin.readlines()] @@ -192,36 +237,20 @@ def main(): start = 0 end = 0 - # Look for an exact substr match - # - found = content.find(header_comment, 0, len(header_comment) + args.extra) - if found >= 0: + # A file is licensed if it carries the canonical header or any accepted + # variant of it. + span = None + for variant in variants: + span = find_header(content, wrap.wrapper(variant, args), args) + if span: + break + + if span: if not args.remove: message(log_to, "OK : " + filepath) continue - start = found - end = found + len(header_comment) - else: - # Look for a fuzzy match in the first 60 chars - # - found = regex.search( - "(?be)(%s){e<=%d}" % (regex.escape(header_comment[0:60]), 6), - content[0 : 80 + args.extra], - ) - if found: - fuzzy = regex.compile( - "(?be)(%s){e<=%d}" % (regex.escape(header_comment), args.editdist) - ) - - # If the first 80 chars match - try harder for the rest of the header - # - found = fuzzy.search( - content[0 : len(header_comment) + args.extra], found.start() - ) - if found: - start = found.start() - end = found.end() + start, end = span if args.remove: if start == 0 and end == 0: diff --git a/scripts/checks/run-clang-tidy.py b/scripts/checks/run-clang-tidy.py index f66c3ec06e7..9eeea38fa06 100755 --- a/scripts/checks/run-clang-tidy.py +++ b/scripts/checks/run-clang-tidy.py @@ -30,7 +30,7 @@ def __setitem__(self, key, value): self[key].append(value) -def git_changed_lines(commit): +def git_changed_lines(commit, exclude_dirs=None): file = "" changed_lines = Multimap() @@ -49,11 +49,16 @@ def git_changed_lines(commit): # as clang-tidy doesn't support CUDA compiler flags and CUDA # headers. Exclude *-inl.h files: they are designed to be # included from their corresponding header and cannot be - # compiled as standalone translation units. + # compiled as standalone translation units. Nimble is not enabled + # in the adapters build that produces the compilation database. if ( "cudf/" not in matched_file and "wave/" not in matched_file and "ucx-exchange/" not in matched_file + and "velox/dwio/nimble/" not in matched_file + and not any( + exclude_dir in matched_file for exclude_dir in (exclude_dirs or []) + ) and not matched_file.endswith("-inl.h") ): file = matched_file @@ -80,6 +85,11 @@ def check_output(output): def tidy(args): files = util.input_files(args.files) + files = [ + file + for file in files + if not any(exclude_dir in file for exclude_dir in args.exclude_dirs) + ] files = [file for file in files if re.match(r".*(\.cpp|\.h|\.hpp)$", file)] # Exclude files in cudf, wave, and torchwave directories @@ -96,7 +106,7 @@ def tidy(args): files = [file for file in files if not file.endswith("-inl.h")] in_gha = os.environ.get("GITHUB_ACTIONS") is not None - changed_lines = git_changed_lines(args.commit) + changed_lines = git_changed_lines(args.commit, args.exclude_dirs) line_filter = json.dumps( [{"name": key, "lines": value} for key, value in changed_lines.items()] @@ -162,6 +172,13 @@ def parse_args(): parser.add_argument("--fix") parser.add_argument("-p", help="Path containing 'compile_commands.json'") + parser.add_argument( + "--exclude-dir", + action="append", + default=[], + dest="exclude_dirs", + help="Additional path substring to exclude from clang-tidy.", + ) parser.add_argument("files", metavar="FILES", nargs="+", help="files to process") return parser.parse_args() diff --git a/scripts/docker/centos-multi.dockerfile b/scripts/docker/centos-multi.dockerfile index b516c28c561..f168b38f1c0 100644 --- a/scripts/docker/centos-multi.dockerfile +++ b/scripts/docker/centos-multi.dockerfile @@ -44,7 +44,7 @@ COPY scripts/setup-common.sh / COPY scripts/setup-centos9.sh / COPY CMake/resolve_dependency_modules/arrow/arrow-testing-boost.patch / COPY CMake/resolve_dependency_modules/arrow/cmake-compatibility.patch / -COPY CMake/resolve_dependency_modules/fbthrift/compactv1-protocol-refiller.patch / +COPY CMake/resolve_dependency_modules/openzl/openzl-cxx-standard.patch / ARG VELOX_BUILD_SHARED=ON # Building libvelox.so requires folly and gflags to be built shared as well for now @@ -70,7 +70,7 @@ ENV UV_TOOL_BIN_DIR=/usr/local/bin \ # https://github.com/apache/arrow/pull/45424 ENV CMAKE_POLICY_VERSION_MINIMUM="3.5" \ VELOX_ARROW_CMAKE_PATCH="/arrow-testing-boost.patch /cmake-compatibility.patch" \ - VELOX_FBTHRIFT_CMAKE_PATCH="/compactv1-protocol-refiller.patch" + VELOX_OPENZL_CMAKE_PATCH="/openzl-cxx-standard.patch" # Ensure libraries installed to INSTALL_PREFIX are found at runtime (e.g. # thrift1 needs libgflags.so.2.2 when folly links gflags statically but diff --git a/scripts/docker/fedora.dockerfile b/scripts/docker/fedora.dockerfile index 0e1428fe91b..6d3f55709b9 100644 --- a/scripts/docker/fedora.dockerfile +++ b/scripts/docker/fedora.dockerfile @@ -25,7 +25,7 @@ COPY scripts/setup-centos9.sh / COPY scripts/setup-fedora.sh / COPY CMake/resolve_dependency_modules/arrow/cmake-compatibility.patch / COPY CMake/resolve_dependency_modules/arrow/arrow-testing-boost.patch / -COPY CMake/resolve_dependency_modules/fbthrift/compactv1-protocol-refiller.patch / +COPY CMake/resolve_dependency_modules/openzl/openzl-cxx-standard.patch / ARG VELOX_BUILD_SHARED=ON # Building libvelox.so requires folly and gflags to be built shared as well for now @@ -42,7 +42,7 @@ ENV UV_TOOL_BIN_DIR=/usr/local/bin \ # CMake 4.0 removed support for cmake minimums of <=3.5 and will fail builds, this overrides it ENV CMAKE_POLICY_VERSION_MINIMUM="3.5" \ VELOX_ARROW_CMAKE_PATCH="/cmake-compatibility.patch /arrow-testing-boost.patch" \ - VELOX_FBTHRIFT_CMAKE_PATCH="/compactv1-protocol-refiller.patch" + VELOX_OPENZL_CMAKE_PATCH="/openzl-cxx-standard.patch" # Some CMake configs contain the hard coded prefix '/deps', we need to replace that with # the future location to avoid build errors in the base-image diff --git a/scripts/docker/ubuntu-22.04-cpp.dockerfile b/scripts/docker/ubuntu-22.04-cpp.dockerfile index 20f956cade3..145afe992cb 100644 --- a/scripts/docker/ubuntu-22.04-cpp.dockerfile +++ b/scripts/docker/ubuntu-22.04-cpp.dockerfile @@ -11,14 +11,22 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# tzdata is pinned to a known-good version so docker rebuilds (any time -# scripts/docker/*.dockerfile or scripts/setup-*.sh changes) don't +# tzdata is pinned to an exact version so docker rebuilds (any time +# scripts/docker/*.dockerfile or scripts/setup-*.sh changes) do not # silently bump tzdata in the image. See issue #17522 for the bug class -# this prevents — version mismatch between OS tzdata and consumers' -# bundled tzdb code can produce silent 1-hour offsets in TIMESTAMP -# WITH TIME ZONE values. To bump intentionally: change the default and -# rebuild locally to confirm before merging. -ARG UBUNTU_TZDATA_VERSION=2026a-0ubuntu0.22.04.1 +# this prevents: a version mismatch between OS tzdata and consumers' +# bundled tzdb code can produce silent one hour offsets in TIMESTAMP +# WITH TIME ZONE values. +# +# Ubuntu keeps only the latest build of tzdata in the jammy pockets, so an +# exact pin is purged once Canonical publishes a newer build, which then +# fails the image bake with "Version ... not found" (issue #18440). Track +# the build currently published for 22.04. The Presto source of truth +# fuzzers, the tz sensitive consumer from #17522, run on the centos9 and +# presto-java images and are governed by CENTOS_TZDATA_VERSION, not this +# pin, so this version can move independently. To bump: set it to the +# current 22.04 build and rebuild to confirm before merging. +ARG UBUNTU_TZDATA_VERSION=2026c-0ubuntu0.22.04.1 ARG base=ubuntu:22.04 FROM ${base} @@ -38,10 +46,10 @@ RUN apt update && \ COPY scripts /velox/scripts/ COPY CMake/resolve_dependency_modules/arrow/cmake-compatibility.patch / COPY CMake/resolve_dependency_modules/arrow/arrow-testing-boost.patch / -COPY CMake/resolve_dependency_modules/fbthrift/compactv1-protocol-refiller.patch / +COPY CMake/resolve_dependency_modules/openzl/openzl-cxx-standard.patch / ENV VELOX_ARROW_CMAKE_PATCH="/cmake-compatibility.patch /arrow-testing-boost.patch" \ - VELOX_FBTHRIFT_CMAKE_PATCH="/compactv1-protocol-refiller.patch" \ + VELOX_OPENZL_CMAKE_PATCH="/openzl-cxx-standard.patch" \ UV_TOOL_BIN_DIR=/usr/local/bin \ UV_INSTALL_DIR=/usr/local/bin diff --git a/scripts/setup-centos9.sh b/scripts/setup-centos9.sh index 538fdf48156..c218e50dacd 100755 --- a/scripts/setup-centos9.sh +++ b/scripts/setup-centos9.sh @@ -112,6 +112,8 @@ function install_velox_deps { run_and_time install_protobuf run_and_time install_fmt run_and_time install_fast_float + run_and_time install_flatbuffers + run_and_time install_openzl run_and_time install_folly run_and_time install_fizz run_and_time install_wangle diff --git a/scripts/setup-common.sh b/scripts/setup-common.sh index bf0af90fde7..10bbd990775 100755 --- a/scripts/setup-common.sh +++ b/scripts/setup-common.sh @@ -21,7 +21,7 @@ source "$SCRIPT_DIR"/setup-versions.sh VELOX_BUILD_SHARED=${VELOX_BUILD_SHARED:-"OFF"} #Build folly and gflags shared for use in libvelox.so. VELOX_ARROW_CMAKE_PATCH=${VELOX_ARROW_CMAKE_PATCH:-""} # avoid error due to +u -VELOX_FBTHRIFT_CMAKE_PATCH=${VELOX_FBTHRIFT_CMAKE_PATCH:-""} +VELOX_OPENZL_CMAKE_PATCH=${VELOX_OPENZL_CMAKE_PATCH:-""} CMAKE_BUILD_TYPE="${BUILD_TYPE:-Release}" DEPENDENCY_DIR=${DEPENDENCY_DIR:-$(pwd)} BUILD_GEOS="${BUILD_GEOS:-true}" @@ -49,7 +49,7 @@ function install_fmt { function install_folly { wget_and_untar https://github.com/facebook/folly/archive/refs/tags/"${FB_OS_VERSION}".tar.gz folly - local FOLLY_FLAGS=(-DBUILD_SHARED_LIBS="$VELOX_BUILD_SHARED" -DBUILD_TESTS=OFF -DFOLLY_HAVE_INT128_T=ON) + local FOLLY_FLAGS=(-DBUILD_SHARED_LIBS="$VELOX_BUILD_SHARED" -DBUILD_TESTS=OFF -DFOLLY_HAVE_INT128_T=ON -DCMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE}") # When folly is static, use static gflags to avoid dual gflags flag # registration when .so plugins are dlopen'd (both the binary and plugin # would register the same flags in a shared gflags registry). @@ -69,6 +69,47 @@ function install_fast_float { cmake_install_dir fast_float -DBUILD_TESTS=OFF } +# Only required for VELOX_ENABLE_NIMBLE=ON. Both the runtime library and the +# flatc code generator are needed, since Nimble generates C++ headers from .fbs +# schemas at build time. +function install_flatbuffers { + wget_and_untar https://github.com/google/flatbuffers/archive/refs/tags/v"${FLATBUFFERS_VERSION}".tar.gz flatbuffers + cmake_install_dir flatbuffers -DFLATBUFFERS_BUILD_TESTS=OFF -DFLATBUFFERS_BUILD_FLATC=ON -DFLATBUFFERS_BUILD_SHAREDLIB=OFF +} + +# Only required for VELOX_ENABLE_NIMBLE=ON. Only the core library and its C++ +# bindings are consumed; everything else OpenZL can build pulls in dependencies +# Velox does not otherwise need. +function install_openzl { + wget_and_untar https://github.com/facebook/openzl/archive/"${OPENZL_VERSION}".tar.gz openzl + ( + # OpenZL hard-codes C++17, which would leave openzl_cpp ABI-incompatible + # with C++20 Velox. Apply the same patch the BUNDLED CMake resolver uses so + # both resolution modes produce a C++20 library. + if [ -z "$VELOX_OPENZL_CMAKE_PATCH" ]; then + # A different path is needed when building the Dockerfile. + ABSOLUTE_SCRIPTDIR=$(realpath "$SCRIPT_DIR") + VELOX_OPENZL_CMAKE_PATCH="$ABSOLUTE_SCRIPTDIR/../CMake/resolve_dependency_modules/openzl/openzl-cxx-standard.patch" + fi + + cd "$DEPENDENCY_DIR"/openzl || exit 1 + if command -v patch >/dev/null 2>&1; then + patch -p1 -i "$VELOX_OPENZL_CMAKE_PATCH" || exit 1 + else + git apply "$VELOX_OPENZL_CMAKE_PATCH" || exit 1 + fi + ) || exit 1 + cmake_install_dir openzl \ + -DCMAKE_CXX_STANDARD=20 \ + -DOPENZL_BUILD_CLI=OFF \ + -DOPENZL_BUILD_EXAMPLES=OFF \ + -DOPENZL_BUILD_TOOLS=OFF \ + -DOPENZL_BUILD_CUSTOM_PARSERS=OFF \ + -DOPENZL_BUILD_TESTS=OFF \ + -DOPENZL_BUILD_BENCHMARKS=OFF \ + -DOPENZL_BUILD_PYTHON_EXT=OFF +} + function install_wangle { wget_and_untar https://github.com/facebook/wangle/archive/refs/tags/"${FB_OS_VERSION}".tar.gz wangle cmake_install_dir wangle/wangle -DBUILD_TESTS=OFF @@ -81,31 +122,7 @@ function install_mvfst { function install_fbthrift { wget_and_untar https://github.com/facebook/fbthrift/archive/refs/tags/"${FB_OS_VERSION}".tar.gz fbthrift - - # This patch is integrated into the latest FBOS version of folly and can be removed on upgrade. - if [ -z "${VELOX_FBTHRIFT_CMAKE_PATCH}" ]; then - # We need to set a different path when building the Dockerfile. - ABSOLUTE_SCRIPTDIR=$(realpath "${SCRIPT_DIR}") - - VELOX_FBTHRIFT_CMAKE_PATCH="${ABSOLUTE_SCRIPTDIR}/../CMake/resolve_dependency_modules/fbthrift/compactv1-protocol-refiller.patch" - fi - ( - cd "$DEPENDENCY_DIR"/fbthrift || exit 1 - # Skip applying the patch if it is already applied. - git apply --reverse --check "${VELOX_FBTHRIFT_CMAKE_PATCH}" 2>/dev/null || - git apply "${VELOX_FBTHRIFT_CMAKE_PATCH}" || exit 1 - ) - - # Apple Clang's libc++ no longer defines _LIBCPP_HAS_NO_ASAN (renamed to - # _LIBCPP_INSTRUMENTED_WITH_ASAN), so folly's UninitializedMemoryHacks.h - # causing undefined symbol. This is fixed in the latest FBOS versions and - # can be removed on FBOS upgrade. - local FBTHRIFT_EXTRA_CXXFLAGS="" - if [[ "$(uname)" == "Darwin" ]]; then - FBTHRIFT_EXTRA_CXXFLAGS=" -D_LIBCPP_HAS_NO_ASAN" - fi - EXTRA_PKG_CXXFLAGS="$FBTHRIFT_EXTRA_CXXFLAGS" \ - cmake_install_dir fbthrift -Denable_tests=OFF -DBUILD_TESTS=OFF -DBUILD_SHARED_LIBS=OFF + cmake_install_dir fbthrift -Denable_tests=OFF -DBUILD_TESTS=OFF -DBUILD_SHARED_LIBS=OFF } function install_duckdb { @@ -429,7 +446,11 @@ function install_azure_storage_sdk_cpp { function install_hdfs_deps { # Dependencies for Hadoop testing - wget_and_untar https://dlcdn.apache.org/hadoop/common/hadoop-"${HADOOP_VERSION}"/hadoop-"${HADOOP_VERSION}".tar.gz hadoop + local arch + arch=$(uname -m) + local hadoop_tarball="hadoop-${HADOOP_VERSION}.tar.gz" + [[ ${arch} == "aarch64" ]] && hadoop_tarball="hadoop-${HADOOP_VERSION}-aarch64.tar.gz" + wget_and_untar "https://dlcdn.apache.org/hadoop/common/hadoop-${HADOOP_VERSION}/${hadoop_tarball}" hadoop cp -a "${DEPENDENCY_DIR}"/hadoop "$INSTALL_PREFIX" wget "${WGET_OPTS[@]}" -P "$INSTALL_PREFIX"/hadoop/share/hadoop/common/lib/ https://repo1.maven.org/maven2/junit/junit/4.11/junit-4.11.jar # Needed for HADOOP 3.3.6 minicluster. Can remove after updating to 3.4.2. @@ -437,17 +458,22 @@ function install_hdfs_deps { } function install_uv { + # Default the uv tool/install dirs to INSTALL_PREFIX only when it is writable. + # On bare CI runners INSTALL_PREFIX (/usr/local) is root-owned, so a non-sudo + # `uv tool install` cannot symlink there and fails with "Permission denied"; + # leaving these unset lets uv use its writable default (~/.local/bin). + # Container images set UV_TOOL_BIN_DIR via ENV, which is preserved here. + if [[ -w "$INSTALL_PREFIX/bin" ]]; then + export UV_TOOL_BIN_DIR="${UV_TOOL_BIN_DIR:-$INSTALL_PREFIX/bin}" + export UV_INSTALL_DIR="${UV_INSTALL_DIR:-$UV_TOOL_BIN_DIR}" + fi if command -v uv >/dev/null 2>&1; then echo "uv is already installed." else echo "Installing uv..." - - export UV_TOOL_BIN_DIR="${UV_TOOL_BIN_DIR:-$INSTALL_PREFIX/bin}" - export UV_INSTALL_DIR=${UV_INSTALL_DIR:-"$UV_TOOL_BIN_DIR"} - curl -LsSf https://astral.sh/uv/install.sh | sh - uv tool update-shell fi + uv tool update-shell } function uv_install { diff --git a/scripts/setup-fedora.sh b/scripts/setup-fedora.sh index d6b31b44cd6..ca9b654a9b1 100755 --- a/scripts/setup-fedora.sh +++ b/scripts/setup-fedora.sh @@ -71,6 +71,8 @@ function install_velox_deps { run_and_time install_gcs_sdk_cpp #grpc, abseil, protobuf run_and_time install_fmt run_and_time install_fast_float + run_and_time install_flatbuffers + run_and_time install_openzl run_and_time install_folly run_and_time install_fizz run_and_time install_wangle diff --git a/scripts/setup-helper-functions.sh b/scripts/setup-helper-functions.sh index ac442c5aa11..4a64b9d623c 100755 --- a/scripts/setup-helper-functions.sh +++ b/scripts/setup-helper-functions.sh @@ -282,6 +282,7 @@ function cmake_install { -DCMAKE_POSITION_INDEPENDENT_CODE=ON \ "${INSTALL_PREFIX+-DCMAKE_PREFIX_PATH=}${INSTALL_PREFIX-}" \ "${INSTALL_PREFIX+-DCMAKE_INSTALL_PREFIX=}${INSTALL_PREFIX-}" \ + -DCMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE:-Release}" \ -DCMAKE_CXX_FLAGS="$COMPILER_FLAGS" \ -DCMAKE_C_COMPILER_LAUNCHER=${CCACHE} \ -DCMAKE_CXX_COMPILER_LAUNCHER=${CCACHE} \ diff --git a/scripts/setup-macos.sh b/scripts/setup-macos.sh index e8210b30577..b8a8e8d56ff 100755 --- a/scripts/setup-macos.sh +++ b/scripts/setup-macos.sh @@ -188,6 +188,8 @@ function install_velox_deps { run_and_time install_protobuf run_and_time install_fmt run_and_time install_fast_float + run_and_time install_flatbuffers + run_and_time install_openzl run_and_time install_folly run_and_time install_fizz run_and_time install_wangle diff --git a/scripts/setup-manylinux.sh b/scripts/setup-manylinux.sh index ca3fc80a414..42d1d0a0314 100755 --- a/scripts/setup-manylinux.sh +++ b/scripts/setup-manylinux.sh @@ -117,6 +117,8 @@ function install_velox_deps { run_and_time install_protobuf run_and_time install_fmt run_and_time install_fast_float + run_and_time install_flatbuffers + run_and_time install_openzl run_and_time install_folly run_and_time install_fizz run_and_time install_wangle diff --git a/scripts/setup-ubuntu.sh b/scripts/setup-ubuntu.sh index b2505ae7ff7..5ba7f127710 100755 --- a/scripts/setup-ubuntu.sh +++ b/scripts/setup-ubuntu.sh @@ -240,6 +240,8 @@ function install_velox_deps { run_and_time install_grpc run_and_time install_boost run_and_time install_fast_float + run_and_time install_flatbuffers + run_and_time install_openzl run_and_time install_folly run_and_time install_fizz run_and_time install_wangle diff --git a/scripts/setup-versions.sh b/scripts/setup-versions.sh index bcabe8dc50f..a2b35ceb9dd 100755 --- a/scripts/setup-versions.sh +++ b/scripts/setup-versions.sh @@ -25,7 +25,7 @@ # /build/fbcode_builder/CMake/FBThriftCppLibrary.cmake # The new FB_OS version of fbthrift might require changes such that thrift # files are generated properly on all platforms. -FB_OS_VERSION="v2026.01.05.00" +FB_OS_VERSION="v2026.07.13.00" FMT_VERSION="11.2.0" BOOST_VERSION="boost-1.84.0" ARROW_VERSION="18.0.0" @@ -48,6 +48,18 @@ S2GEOMETRY_VERSION="0.12.0" FAISS_VERSION="1.11.0" FAST_FLOAT_VERSION="v8.0.2" CCACHE_VERSION="4.11.3" +# Only needed for VELOX_ENABLE_NIMBLE=ON. FSST is vendored under +# velox/external/fsst, so it has no pin here. OpenZL is pinned to the revision +# fbcode builds Nimble against rather than to v0.2.0, which predates the +# cross-platform zstd handling it needs to configure here; keep it in sync with +# VELOX_OPENZL_VERSION in CMake/resolve_dependency_modules/openzl.cmake, or a +# system install and a bundled build resolve different descriptor APIs. +# FlatBuffers is held at the version cuDF expects, since a system install of +# anything newer is picked up by cuDF's own find_package() and breaks its +# build; see +# CMake/resolve_dependency_modules/flatbuffers.cmake. Keep the two in sync. +FLATBUFFERS_VERSION="24.3.25" +OPENZL_VERSION="7340a712cce1b8331bec3467600dba99a562e052" # Adapter related versions. ABSEIL_VERSION="20240116.2" diff --git a/velox/CMakeLists.txt b/velox/CMakeLists.txt index 3efa70d3f65..f77865a423c 100644 --- a/velox/CMakeLists.txt +++ b/velox/CMakeLists.txt @@ -26,6 +26,11 @@ add_subdirectory(external/date) add_subdirectory(external/tzdb) add_subdirectory(external/md5) add_subdirectory(external/hdfs) + +if(VELOX_ENABLE_NIMBLE) + # Nimble's string encoding is the only consumer of FSST. + add_subdirectory(external/fsst) +endif() # # examples depend on expression diff --git a/velox/benchmarks/basic/CastBenchmark.cpp b/velox/benchmarks/basic/CastBenchmark.cpp index 42f2329cab9..53c6eaefb06 100644 --- a/velox/benchmarks/basic/CastBenchmark.cpp +++ b/velox/benchmarks/basic/CastBenchmark.cpp @@ -99,6 +99,11 @@ int main(int argc, char** argv) { [](auto row) { return fmt::format("2024-05-{:02d}", 1 + row % 30); }); auto invalidDateStrings = vectorMaker.flatVector( vectorSize, [](auto row) { return fmt::format("2024-05...{}", row); }); + auto validTimeStrings = + vectorMaker.flatVector(vectorSize, [](auto row) { + return fmt::format( + "{:02d}:{:02d}:{:02d}", row % 24, row % 60, row % 60); + }); auto timeInput = vectorMaker.flatVector( vectorSize, [](auto j) { return j % 86'400'000; }, nullptr, TIME()); @@ -155,6 +160,12 @@ int main(int argc, char** argv) { vectorMaker.rowVector({"timestamp"}, {timestampInput})) .addExpression("cast", "cast (timestamp as varchar)"); + benchmarkBuilder + .addBenchmarkSet( + "cast_varchar_as_time", + vectorMaker.rowVector({"valid_time"}, {validTimeStrings})) + .addExpression("cast_valid", "cast (valid_time as time)"); + benchmarkBuilder .addBenchmarkSet( "cast_varchar_as_double", diff --git a/velox/common/Casts.h b/velox/common/Casts.h index 8687d56c149..c02b5a98f38 100644 --- a/velox/common/Casts.h +++ b/velox/common/Casts.h @@ -76,6 +76,14 @@ To* checkedPointerCast(From* input) { return casted; } +/// Checks that pointer is not null and returns it for expression contexts such +/// as constructor initializer lists. +template +T* checkedNotNull(T* pointer) { + VELOX_CHECK_NOT_NULL(pointer); + return pointer; +} + template std::unique_ptr staticUniquePointerCast(std::unique_ptr input) { VELOX_CHECK_NOT_NULL(input.get()); diff --git a/velox/common/base/CMakeLists.txt b/velox/common/base/CMakeLists.txt index 5080424f40f..3d78b4a8184 100644 --- a/velox/common/base/CMakeLists.txt +++ b/velox/common/base/CMakeLists.txt @@ -32,6 +32,7 @@ velox_add_library( AdmissionController.cpp BitUtil.cpp BloomFilter.cpp + ConcurrentRuntimeStatWriter.cpp Counters.cpp Fs.cpp PeriodicStatsReporter.cpp @@ -53,6 +54,7 @@ velox_add_library( ClassName.h CoalesceIo.h ConcurrentCounter.h + ConcurrentRuntimeStatWriter.h CountBits.h Counters.h Crc.h diff --git a/velox/common/base/ConcurrentRuntimeStatWriter.cpp b/velox/common/base/ConcurrentRuntimeStatWriter.cpp new file mode 100644 index 00000000000..ee099eab5b8 --- /dev/null +++ b/velox/common/base/ConcurrentRuntimeStatWriter.cpp @@ -0,0 +1,46 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/common/base/ConcurrentRuntimeStatWriter.h" + +#include "velox/common/base/Exceptions.h" + +namespace facebook::velox { + +void ConcurrentRuntimeStatWriter::addRuntimeStat( + std::string_view name, + const RuntimeCounter& value) { + auto lockedStats = runtimeStats_.wlock(); + auto [it, unused] = lockedStats->try_emplace(std::string(name), value.unit); + it->second.merge(value); +} + +void ConcurrentRuntimeStatWriter::setRuntimeStat( + std::string_view name, + const RuntimeMetric& metric) { + runtimeStats_.wlock()->insert_or_assign(std::string(name), metric); +} + +std::unordered_map +ConcurrentRuntimeStatWriter::runtimeStats() const { + return *runtimeStats_.rlock(); +} + +void ConcurrentRuntimeStatWriter::clear() { + runtimeStats_.wlock()->clear(); +} + +} // namespace facebook::velox diff --git a/velox/common/base/ConcurrentRuntimeStatWriter.h b/velox/common/base/ConcurrentRuntimeStatWriter.h new file mode 100644 index 00000000000..4b7e3d57d73 --- /dev/null +++ b/velox/common/base/ConcurrentRuntimeStatWriter.h @@ -0,0 +1,54 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include + +#include "velox/common/base/RuntimeMetrics.h" + +namespace facebook::velox { + +/// Accumulates runtime metrics by name behind a lock, and exposes a snapshot. +/// Unlike writers that forward each sample to an owning operator, this one owns +/// the values, so any number of threads may record into it concurrently. +class ConcurrentRuntimeStatWriter : public BaseRuntimeStatWriter { + public: + /// Merges 'value' under 'name'. All samples under a name must share the same + /// unit; a mismatched unit throws. + void addRuntimeStat(std::string_view name, const RuntimeCounter& value) + override; + + /// Replaces any existing metric under 'name'. + void setRuntimeStat(std::string_view name, const RuntimeMetric& metric) + override; + + /// Returns a snapshot of the accumulated metrics. + std::unordered_map runtimeStats() const; + + /// Drops all accumulated metrics. + void clear(); + + private: + // All samples under a name share one unit; addRuntimeStat enforces it. + folly::Synchronized> + runtimeStats_; +}; + +} // namespace facebook::velox diff --git a/velox/common/base/Macros.h b/velox/common/base/Macros.h index 8637e9eee75..98573ff5c3b 100644 --- a/velox/common/base/Macros.h +++ b/velox/common/base/Macros.h @@ -43,25 +43,6 @@ #define VELOX_UNSUPPRESS_DEPRECATED_WARNING _Pragma("GCC diagnostic pop") #endif -// Disable missing-field-initializers for Clang and GCC -#ifdef __clang__ -#if defined(__has_warning) && \ - __has_warning("-Wmissing-designated-field-initializers") -#define VELOX_SUPPRESS_MISSING_DESIGNATED_FIELD_INITIALIZERS_WARNING \ - _Pragma("clang diagnostic push"); \ - _Pragma( \ - "clang diagnostic ignored \"-Wmissing-designated-field-initializers\"") -#define VELOX_UNSUPPRESS_MISSING_DESIGNATED_FIELD_INITIALIZERS_WARNING \ - _Pragma("clang diagnostic pop"); -#else -#define VELOX_SUPPRESS_MISSING_DESIGNATED_FIELD_INITIALIZERS_WARNING -#define VELOX_UNSUPPRESS_MISSING_DESIGNATED_FIELD_INITIALIZERS_WARNING -#endif -#else -#define VELOX_SUPPRESS_MISSING_DESIGNATED_FIELD_INITIALIZERS_WARNING -#define VELOX_UNSUPPRESS_MISSING_DESIGNATED_FIELD_INITIALIZERS_WARNING -#endif - #define VELOX_CONCAT(x, y) x##y // Need this extra layer to expand __COUNTER__. #define VELOX_VARNAME_IMPL(x, y) VELOX_CONCAT(x, y) @@ -76,3 +57,28 @@ #else #define VELOX_CONSTEXPR_SINGLETON static constexpr #endif + +// Marks a function as one that GPU execution is expected to be able to call. +// +// Under nvcc this expands to `__host__ __device__`, so the function is +// compiled for the device as well as the host. Under any host-only compiler it +// expands to nothing, leaving CPU builds unchanged. +// +// Two things it deliberately is not: +// +// - It is not an inlining hint. A non-template function defined in a header +// still needs its own `inline`; write `VELOX_GPU_COMPATIBLE inline void +// f()`. Templates and in-class member functions are already implicitly +// inline and need nothing extra. +// - It is not a guarantee. nvcc reports a call from an annotated function +// into host-only code as a *warning*, then emits a kernel that silently +// computes the wrong answer. Anything carrying this macro therefore has to +// stay within what device code can reach: no exceptions, no allocation, no +// runtime indexing of a namespace- or class-scope constexpr table, and no +// compiler builtin lacking a device implementation. Build GPU targets with +// `--diag-error=20011` so the compiler enforces that rather than a reader. +#ifdef __CUDACC__ +#define VELOX_GPU_COMPATIBLE __host__ __device__ +#else +#define VELOX_GPU_COMPATIBLE +#endif diff --git a/velox/common/base/RuntimeMetrics.cpp b/velox/common/base/RuntimeMetrics.cpp index 06e1f8e1542..bdda1c3cbf8 100644 --- a/velox/common/base/RuntimeMetrics.cpp +++ b/velox/common/base/RuntimeMetrics.cpp @@ -34,6 +34,11 @@ void RuntimeMetric::aggregate() { min = max = sum; } +void RuntimeMetric::merge(const RuntimeCounter& value) { + VELOX_CHECK_EQ(unit, value.unit, "Unit mismatch for runtime stat"); + addValue(value.value); +} + void RuntimeMetric::merge(const RuntimeMetric& other) #if defined(__has_feature) #if __has_feature(__address_sanitizer__) diff --git a/velox/common/base/RuntimeMetrics.h b/velox/common/base/RuntimeMetrics.h index f9c910e3adc..14838c1f1b8 100644 --- a/velox/common/base/RuntimeMetrics.h +++ b/velox/common/base/RuntimeMetrics.h @@ -18,7 +18,9 @@ #include #include +#include #include +#include #include namespace facebook::velox { @@ -72,6 +74,9 @@ struct RuntimeMetric { void printMetric(std::ostream& stream) const; + /// Merges 'value' into this metric. Both must have the same unit. + void merge(const RuntimeCounter& value); + void merge(const RuntimeMetric& other); std::string toString() const; @@ -92,8 +97,27 @@ class BaseRuntimeStatWriter { virtual void setRuntimeStat( std::string_view /* name */, const RuntimeMetric& /* metric */) {} + + /// Adds a wall or cpu duration sample under 'name', tagged as nanoseconds. + void addTiming(std::string_view name, std::chrono::nanoseconds duration) { + addRuntimeStat( + name, RuntimeCounter(duration.count(), RuntimeCounter::Unit::kNanos)); + } + + /// Adds a unitless count sample under 'name'. + void addCount(std::string_view name, int64_t value) { + addRuntimeStat(name, RuntimeCounter(value)); + } + + /// Adds a size sample under 'name', tagged as bytes. + void addBytes(std::string_view name, int64_t bytes) { + addRuntimeStat(name, RuntimeCounter(bytes, RuntimeCounter::Unit::kBytes)); + } }; +/// Discards every metric written to it. +class NoopRuntimeStatWriter : public BaseRuntimeStatWriter {}; + /// Setting a concrete runtime stats writer on the thread will ensure that any /// code can add runtime counters to the current Operator running on that /// thread. diff --git a/velox/common/base/TrackedExecutor.cpp b/velox/common/base/TrackedExecutor.cpp index bd34c9e3031..680a091a0ba 100644 --- a/velox/common/base/TrackedExecutor.cpp +++ b/velox/common/base/TrackedExecutor.cpp @@ -51,17 +51,11 @@ TrackedExecutor::Func TrackedExecutor::wrapFunc(Func func) { }; } -void TrackedExecutor::reportTo( - BaseRuntimeStatWriter& writer, - std::string_view prefix) const { +void TrackedExecutor::reportTo(BaseRuntimeStatWriter& writer) const { + writer.setRuntimeStat(kExecutorWaitNanos, metrics_->waitTime); writer.setRuntimeStat( - fmt::format("{}-{}", prefix, kExecutorWaitNanos), metrics_->waitTime); - writer.setRuntimeStat( - fmt::format("{}-{}", prefix, kExecutorExecutionWallNanos), - metrics_->executionWallTime); - writer.setRuntimeStat( - fmt::format("{}-{}", prefix, kExecutorExecutionCpuNanos), - metrics_->executionCpuTime); + kExecutorExecutionWallNanos, metrics_->executionWallTime); + writer.setRuntimeStat(kExecutorExecutionCpuNanos, metrics_->executionCpuTime); } } // namespace facebook::velox diff --git a/velox/common/base/TrackedExecutor.h b/velox/common/base/TrackedExecutor.h index e1ba7bf2526..d1663a1f5ab 100644 --- a/velox/common/base/TrackedExecutor.h +++ b/velox/common/base/TrackedExecutor.h @@ -67,8 +67,8 @@ class TrackedExecutor final : public folly::Executor { return executor_->getNumPriorities(); } - /// Reports accumulated metrics to 'writer', naming each 'prefix-'. - void reportTo(BaseRuntimeStatWriter& writer, std::string_view prefix) const; + /// Reports accumulated metrics to 'writer' under the kExecutor* names. + void reportTo(BaseRuntimeStatWriter& writer) const; private: // Instruments 'func' to record its enqueue-wait, wall, and cpu time into diff --git a/velox/common/base/tests/CMakeLists.txt b/velox/common/base/tests/CMakeLists.txt index 6d92940b52c..82d0dc03065 100644 --- a/velox/common/base/tests/CMakeLists.txt +++ b/velox/common/base/tests/CMakeLists.txt @@ -20,6 +20,7 @@ add_executable( BloomFilterTest.cpp CoalesceIoTest.cpp ConcurrentCounterTest.cpp + ConcurrentRuntimeStatWriterTest.cpp ExceptionTest.cpp FsTest.cpp IndexedPriorityQueueTest.cpp diff --git a/velox/common/base/tests/ConcurrentRuntimeStatWriterTest.cpp b/velox/common/base/tests/ConcurrentRuntimeStatWriterTest.cpp new file mode 100644 index 00000000000..c9094d12732 --- /dev/null +++ b/velox/common/base/tests/ConcurrentRuntimeStatWriterTest.cpp @@ -0,0 +1,153 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/common/base/ConcurrentRuntimeStatWriter.h" + +#include +#include +#include + +#include +#include + +#include "velox/common/base/Exceptions.h" +#include "velox/common/base/tests/GTestUtils.h" + +namespace facebook::velox { + +class ConcurrentRuntimeStatWriterTest : public testing::Test { + protected: + ConcurrentRuntimeStatWriter writer_; +}; + +TEST_F(ConcurrentRuntimeStatWriterTest, addAccumulates) { + writer_.addRuntimeStat( + "wall", RuntimeCounter(10, RuntimeCounter::Unit::kNanos)); + writer_.addRuntimeStat( + "wall", RuntimeCounter(30, RuntimeCounter::Unit::kNanos)); + + const auto stats = writer_.runtimeStats(); + const auto& wall = stats.at("wall"); + EXPECT_EQ(wall.count, 2); + EXPECT_EQ(wall.sum, 40); + EXPECT_EQ(wall.min, 10); + EXPECT_EQ(wall.max, 30); + EXPECT_EQ(wall.unit, RuntimeCounter::Unit::kNanos); +} + +TEST_F(ConcurrentRuntimeStatWriterTest, setReplaces) { + RuntimeMetric preset(RuntimeCounter::Unit::kBytes); + preset.addValue(5); + preset.addValue(15); + writer_.setRuntimeStat("bytes", preset); + + RuntimeMetric replacement(RuntimeCounter::Unit::kBytes); + replacement.addValue(100); + writer_.setRuntimeStat("bytes", replacement); + + // The second metric replaces the first rather than merging into it. + const auto stats = writer_.runtimeStats(); + const auto& bytes = stats.at("bytes"); + EXPECT_EQ(bytes.count, 1); + EXPECT_EQ(bytes.sum, 100); +} + +TEST_F(ConcurrentRuntimeStatWriterTest, unitMismatchThrows) { + writer_.addRuntimeStat("m", RuntimeCounter(10, RuntimeCounter::Unit::kNanos)); + VELOX_ASSERT_THROW( + writer_.addRuntimeStat( + "m", RuntimeCounter(20, RuntimeCounter::Unit::kBytes)), + "Unit mismatch for runtime stat"); + + // The rejected sample leaves the metric untouched. + const auto stats = writer_.runtimeStats(); + const auto& metric = stats.at("m"); + EXPECT_EQ(metric.count, 1); + EXPECT_EQ(metric.sum, 10); + EXPECT_EQ(metric.unit, RuntimeCounter::Unit::kNanos); +} + +TEST_F(ConcurrentRuntimeStatWriterTest, unitMismatchThrowsAfterSet) { + // setRuntimeStat replaces without checking the unit, so the add path is what + // catches a mismatch on a name seeded by a set. + RuntimeMetric seeded(RuntimeCounter::Unit::kBytes); + seeded.addValue(10); + writer_.setRuntimeStat("m", seeded); + + VELOX_ASSERT_THROW( + writer_.addRuntimeStat( + "m", RuntimeCounter(20, RuntimeCounter::Unit::kNanos)), + "Unit mismatch for runtime stat"); +} + +TEST_F(ConcurrentRuntimeStatWriterTest, unitHelpersApplyTheirUnit) { + writer_.addCount("splits", 3); + writer_.addBytes("read", 128); + + const auto stats = writer_.runtimeStats(); + EXPECT_EQ(stats.at("splits").sum, 3); + EXPECT_EQ(stats.at("splits").unit, RuntimeCounter::Unit::kNone); + EXPECT_EQ(stats.at("read").sum, 128); + EXPECT_EQ(stats.at("read").unit, RuntimeCounter::Unit::kBytes); +} + +TEST_F(ConcurrentRuntimeStatWriterTest, clearDropsEverything) { + writer_.addRuntimeStat("a", RuntimeCounter(1)); + writer_.addRuntimeStat("b", RuntimeCounter(2)); + writer_.clear(); + EXPECT_THAT(writer_.runtimeStats(), testing::IsEmpty()); + + // A cleared name may take a new unit, since nothing survives the clear. + writer_.addRuntimeStat("a", RuntimeCounter(5, RuntimeCounter::Unit::kBytes)); + const auto stats = writer_.runtimeStats(); + EXPECT_EQ(stats.at("a").sum, 5); + EXPECT_EQ(stats.at("a").unit, RuntimeCounter::Unit::kBytes); +} + +TEST_F(ConcurrentRuntimeStatWriterTest, concurrentAddIsLossless) { + constexpr int32_t kNumThreads{8}; + constexpr int32_t kSamplesPerThread{2'000}; + + // Every thread adds to one shared name and to a name only it uses, so the + // test covers both contended and uncontended keys. + std::vector threads; + threads.reserve(kNumThreads); + for (int32_t thread = 0; thread < kNumThreads; ++thread) { + threads.emplace_back([&, thread] { + const auto ownName = fmt::format("perThread{}", thread); + for (int32_t sample = 0; sample < kSamplesPerThread; ++sample) { + writer_.addCount("shared", 1); + writer_.addCount(ownName, 1); + } + }); + } + for (auto& thread : threads) { + thread.join(); + } + + const auto stats = writer_.runtimeStats(); + const auto& shared = stats.at("shared"); + EXPECT_EQ(shared.count, kNumThreads * kSamplesPerThread); + EXPECT_EQ(shared.sum, kNumThreads * kSamplesPerThread); + + for (int32_t thread = 0; thread < kNumThreads; ++thread) { + const auto& perThread = stats.at(fmt::format("perThread{}", thread)); + EXPECT_EQ(perThread.count, kSamplesPerThread); + EXPECT_EQ(perThread.sum, kSamplesPerThread); + } +} + +} // namespace facebook::velox diff --git a/velox/common/base/tests/RuntimeMetricsTest.cpp b/velox/common/base/tests/RuntimeMetricsTest.cpp index b4a11ff6859..150ad67cdc6 100644 --- a/velox/common/base/tests/RuntimeMetricsTest.cpp +++ b/velox/common/base/tests/RuntimeMetricsTest.cpp @@ -15,7 +15,11 @@ */ #include "velox/common/base/RuntimeMetrics.h" +#include #include +#include "velox/common/base/ConcurrentRuntimeStatWriter.h" +#include "velox/common/base/VeloxException.h" +#include "velox/common/base/tests/GTestUtils.h" namespace facebook::velox { @@ -99,27 +103,24 @@ TEST_F(RuntimeMetricsTest, saturateCast) { EXPECT_EQ(rm.max, maxInt64); } -class SetThreadLocalRuntimeStatTest : public testing::Test { - protected: - class RuntimeStatCollector : public BaseRuntimeStatWriter { - public: - void setRuntimeStat(std::string_view name, const RuntimeMetric& metric) - override { - stats_.insert_or_assign(std::string(name), metric); - } - - const RuntimeMetric* getMetric(const std::string& name) const { - auto it = stats_.find(name); - return it != stats_.end() ? &it->second : nullptr; - } - - private: - std::unordered_map stats_; - }; -}; +TEST_F(RuntimeMetricsTest, mergeCounter) { + RuntimeMetric rm(RuntimeCounter::Unit::kBytes); + + rm.merge(RuntimeCounter(10, RuntimeCounter::Unit::kBytes)); + testMetric(rm, 10, 1, 10, 10); + + rm.merge(RuntimeCounter(30, RuntimeCounter::Unit::kBytes)); + testMetric(rm, 40, 2, 10, 30); + + VELOX_ASSERT_THROW( + rm.merge(RuntimeCounter(1, RuntimeCounter::Unit::kNanos)), + "Unit mismatch for runtime stat"); +} + +class SetThreadLocalRuntimeStatTest : public testing::Test {}; TEST_F(SetThreadLocalRuntimeStatTest, singleMetric) { - RuntimeStatCollector collector; + ConcurrentRuntimeStatWriter collector; RuntimeStatWriterScopeGuard guard(&collector); RuntimeMetric metric(RuntimeCounter::Unit::kNone); @@ -129,16 +130,17 @@ TEST_F(SetThreadLocalRuntimeStatTest, singleMetric) { setThreadLocalRuntimeStat("test.metric", metric); - const auto* result = collector.getMetric("test.metric"); - ASSERT_NE(result, nullptr); - EXPECT_EQ(result->count, 3); - EXPECT_EQ(result->sum, 60); - EXPECT_EQ(result->min, 10); - EXPECT_EQ(result->max, 30); + const auto stats = collector.runtimeStats(); + ASSERT_EQ(stats.count("test.metric"), 1); + const auto& result = stats.at("test.metric"); + EXPECT_EQ(result.count, 3); + EXPECT_EQ(result.sum, 60); + EXPECT_EQ(result.min, 10); + EXPECT_EQ(result.max, 30); } TEST_F(SetThreadLocalRuntimeStatTest, existingMetric) { - RuntimeStatCollector collector; + ConcurrentRuntimeStatWriter collector; RuntimeStatWriterScopeGuard guard(&collector); RuntimeMetric first(RuntimeCounter::Unit::kNone); @@ -150,24 +152,25 @@ TEST_F(SetThreadLocalRuntimeStatTest, existingMetric) { second.addValue(15); setThreadLocalRuntimeStat("test.metric", second); - const auto* result = collector.getMetric("test.metric"); - ASSERT_NE(result, nullptr); - EXPECT_EQ(result->count, 2); - EXPECT_EQ(result->sum, 20); - EXPECT_EQ(result->min, 5); - EXPECT_EQ(result->max, 15); + const auto stats = collector.runtimeStats(); + ASSERT_EQ(stats.count("test.metric"), 1); + const auto& result = stats.at("test.metric"); + EXPECT_EQ(result.count, 2); + EXPECT_EQ(result.sum, 20); + EXPECT_EQ(result.min, 5); + EXPECT_EQ(result.max, 15); } TEST_F(SetThreadLocalRuntimeStatTest, emptyMetric) { - RuntimeStatCollector collector; + ConcurrentRuntimeStatWriter collector; RuntimeStatWriterScopeGuard guard(&collector); RuntimeMetric empty(RuntimeCounter::Unit::kNone); setThreadLocalRuntimeStat("test.empty", empty); - const auto* result = collector.getMetric("test.empty"); - ASSERT_NE(result, nullptr); - EXPECT_EQ(result->count, 0); + const auto stats = collector.runtimeStats(); + ASSERT_EQ(stats.count("test.empty"), 1); + EXPECT_EQ(stats.at("test.empty").count, 0); } TEST_F(SetThreadLocalRuntimeStatTest, noWriter) { @@ -176,4 +179,19 @@ TEST_F(SetThreadLocalRuntimeStatTest, noWriter) { setThreadLocalRuntimeStat("test.nowriter", metric); } +TEST_F(SetThreadLocalRuntimeStatTest, noopWriter) { + ConcurrentRuntimeStatWriter collector; + RuntimeStatWriterScopeGuard collecting(&collector); + + { + NoopRuntimeStatWriter noop; + RuntimeStatWriterScopeGuard discarding(&noop); + addThreadLocalRuntimeStat("test.noop", RuntimeCounter(1)); + } + EXPECT_THAT(collector.runtimeStats(), testing::IsEmpty()); + + addThreadLocalRuntimeStat("test.after", RuntimeCounter(1)); + EXPECT_EQ(collector.runtimeStats().count("test.after"), 1); +} + } // namespace facebook::velox diff --git a/velox/common/base/tests/SimdUtilTest.cpp b/velox/common/base/tests/SimdUtilTest.cpp index 5573df46028..3e2e6126cf0 100644 --- a/velox/common/base/tests/SimdUtilTest.cpp +++ b/velox/common/base/tests/SimdUtilTest.cpp @@ -393,7 +393,7 @@ TEST_F(SimdUtilTest, crc32) { EXPECT_EQ(checksum, 121285919); } -TEST_F(SimdUtilTest, Batch64_assign) { +TEST_F(SimdUtilTest, batch64Assign) { auto b = simd::Batch64::from({0, 1}); EXPECT_EQ(b.data[0], 0); EXPECT_EQ(b.data[1], 1); @@ -402,7 +402,7 @@ TEST_F(SimdUtilTest, Batch64_assign) { EXPECT_EQ(b.data[1], 1); } -TEST_F(SimdUtilTest, Batch64_arithmetics) { +TEST_F(SimdUtilTest, batch64Arithmetics) { auto b = simd::Batch64::from({0, 1}); auto bb = b + 42; EXPECT_EQ(bb.data[0], 42); @@ -412,7 +412,7 @@ TEST_F(SimdUtilTest, Batch64_arithmetics) { EXPECT_EQ(bb.data[1], 0); } -TEST_F(SimdUtilTest, Batch64_memory) { +TEST_F(SimdUtilTest, batch64Memory) { int32_t data[] = {0, 1}; auto b = simd::Batch64::load_unaligned(data); EXPECT_EQ(b.data[0], 0); diff --git a/velox/common/base/tests/StatsReporterTest.cpp b/velox/common/base/tests/StatsReporterTest.cpp index daddca74b25..3b09f444d65 100644 --- a/velox/common/base/tests/StatsReporterTest.cpp +++ b/velox/common/base/tests/StatsReporterTest.cpp @@ -791,7 +791,7 @@ class DynamicQuantilePatternTest : public StatsReporterTest, public testing::WithParamInterface {}; -TEST_P(DynamicQuantilePatternTest, PatternScenarios) { +TEST_P(DynamicQuantilePatternTest, patternScenarios) { const auto& testCase = GetParam(); testCase.testFunc(this); } diff --git a/velox/common/base/tests/TrackedExecutorTest.cpp b/velox/common/base/tests/TrackedExecutorTest.cpp index db1973c0ece..dbb656e0fe2 100644 --- a/velox/common/base/tests/TrackedExecutorTest.cpp +++ b/velox/common/base/tests/TrackedExecutorTest.cpp @@ -15,6 +15,7 @@ */ #include "velox/common/base/TrackedExecutor.h" +#include "velox/common/base/ConcurrentRuntimeStatWriter.h" #include #include @@ -22,7 +23,6 @@ #include #include -#include #include #include @@ -31,33 +31,9 @@ namespace facebook::velox { namespace { -// Captures the metrics that TrackedExecutor::reportTo writes so the test can -// inspect them by name. -class MapStatWriter : public BaseRuntimeStatWriter { - public: - void setRuntimeStat(std::string_view name, const RuntimeMetric& metric) - override { - metrics_.insert_or_assign(std::string{name}, metric); - } - - void addRuntimeStat(std::string_view name, const RuntimeCounter& value) - override { - auto [it, inserted] = - metrics_.try_emplace(std::string{name}, RuntimeMetric(value.unit)); - it->second.addValue(value.value); - } - - const std::map& metrics() const { - return metrics_; - } - - private: - std::map metrics_; -}; - class TrackedExecutorTest : public testing::Test {}; -TEST_F(TrackedExecutorTest, reportsPerCallbackMetricsUnderPrefix) { +TEST_F(TrackedExecutorTest, reportsOneSamplePerCallback) { // Run callbacks inline so the per-metric counts are deterministic. TrackedExecutor tracked{ folly::getKeepAliveToken(folly::InlineExecutor::instance())}; @@ -74,20 +50,20 @@ TEST_F(TrackedExecutorTest, reportsPerCallbackMetricsUnderPrefix) { }); } - MapStatWriter writer; - tracked.reportTo(writer, "myOp"); - const auto& metrics = writer.metrics(); + ConcurrentRuntimeStatWriter writer; + tracked.reportTo(writer); + const auto metrics = writer.runtimeStats(); ASSERT_THAT( metrics, testing::UnorderedElementsAre( - testing::Key("myOp-executorWaitNanos"), - testing::Key("myOp-executorExecutionWallNanos"), - testing::Key("myOp-executorExecutionCpuNanos"))); + testing::Key("executorWaitNanos"), + testing::Key("executorExecutionWallNanos"), + testing::Key("executorExecutionCpuNanos"))); - const auto& wait = metrics.at("myOp-executorWaitNanos"); - const auto& wall = metrics.at("myOp-executorExecutionWallNanos"); - const auto& cpu = metrics.at("myOp-executorExecutionCpuNanos"); + const auto& wait = metrics.at("executorWaitNanos"); + const auto& wall = metrics.at("executorExecutionWallNanos"); + const auto& cpu = metrics.at("executorExecutionCpuNanos"); // Every scheduled callback contributes one sample to each metric. EXPECT_EQ(wait.count, kNumTasks); @@ -117,13 +93,13 @@ TEST_F(TrackedExecutorTest, keepsMetricCountsAlignedWhenCallbackThrows) { }), std::runtime_error); - MapStatWriter writer; - tracked.reportTo(writer, "op"); - const auto& metrics = writer.metrics(); + ConcurrentRuntimeStatWriter writer; + tracked.reportTo(writer); + const auto metrics = writer.runtimeStats(); - EXPECT_EQ(metrics.at("op-executorWaitNanos").count, 1); - EXPECT_EQ(metrics.at("op-executorExecutionWallNanos").count, 1); - EXPECT_EQ(metrics.at("op-executorExecutionCpuNanos").count, 1); + EXPECT_EQ(metrics.at("executorWaitNanos").count, 1); + EXPECT_EQ(metrics.at("executorExecutionWallNanos").count, 1); + EXPECT_EQ(metrics.at("executorExecutionCpuNanos").count, 1); } } // namespace diff --git a/velox/common/caching/AsyncDataCache.cpp b/velox/common/caching/AsyncDataCache.cpp index 0549f759b0f..63c8ebafdf4 100644 --- a/velox/common/caching/AsyncDataCache.cpp +++ b/velox/common/caching/AsyncDataCache.cpp @@ -922,6 +922,8 @@ void AsyncDataCache::shutdown() { } void CacheShard::shutdown() { + std::lock_guard l(mutex_); + entryMap_.clear(); entries_.clear(); freeEntries_.clear(); } diff --git a/velox/common/caching/CMakeLists.txt b/velox/common/caching/CMakeLists.txt index 2829d2c798b..c3031571b47 100644 --- a/velox/common/caching/CMakeLists.txt +++ b/velox/common/caching/CMakeLists.txt @@ -18,6 +18,7 @@ velox_add_library( CacheTTLController.cpp FileHandle.cpp FileIds.cpp + FileProperties.cpp ScanTracker.cpp SsdCache.cpp SsdFile.cpp diff --git a/velox/common/caching/FileHandle.cpp b/velox/common/caching/FileHandle.cpp index 20adf910f21..5ac21907fd0 100644 --- a/velox/common/caching/FileHandle.cpp +++ b/velox/common/caching/FileHandle.cpp @@ -59,6 +59,7 @@ std::unique_ptr FileHandleGenerator::operator()( options.readRangeHint = properties->readRangeHint; options.extraFileInfo = properties->extraFileInfo; options.fileReadOps = properties->fileReadOps; + options.ioStatistics = properties->ioStatistics; } const auto& filename = key.filename; fileHandle->file = filesystems::getFileSystem(filename, properties_) diff --git a/velox/common/caching/FileProperties.cpp b/velox/common/caching/FileProperties.cpp new file mode 100644 index 00000000000..01e8dbef4f2 --- /dev/null +++ b/velox/common/caching/FileProperties.cpp @@ -0,0 +1,73 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/common/caching/FileProperties.h" + +namespace facebook::velox { + +folly::dynamic FileProperties::serialize() const { + folly::dynamic obj = folly::dynamic::object; + obj["fileSize"] = + fileSize.has_value() ? folly::dynamic(fileSize.value()) : nullptr; + obj["modificationTime"] = modificationTime.has_value() + ? folly::dynamic(modificationTime.value()) + : nullptr; + obj["readRangeHint"] = readRangeHint.has_value() + ? folly::dynamic(readRangeHint.value()) + : nullptr; + obj["extraFileInfo"] = + extraFileInfo == nullptr ? nullptr : folly::dynamic(*extraFileInfo); + + folly::dynamic fileReadOpsObj = folly::dynamic::object; + for (const auto& [key, value] : fileReadOps) { + fileReadOpsObj[key] = value; + } + obj["fileReadOps"] = fileReadOpsObj; + + return obj; +} + +namespace { +std::optional optionalInt(const folly::dynamic& obj, const char* key) { + const auto& value = obj.getDefault(key, nullptr); + return value.isNull() ? std::nullopt : std::optional(value.asInt()); +} +} // namespace + +// static +FileProperties FileProperties::create(const folly::dynamic& obj) { + FileProperties properties; + properties.fileSize = optionalInt(obj, "fileSize"); + properties.modificationTime = optionalInt(obj, "modificationTime"); + properties.readRangeHint = optionalInt(obj, "readRangeHint"); + + const auto& extraFileInfoObj = obj.getDefault("extraFileInfo", nullptr); + if (!extraFileInfoObj.isNull()) { + properties.extraFileInfo = + std::make_shared(extraFileInfoObj.asString()); + } + + const auto& fileReadOpsObj = obj.getDefault("fileReadOps", nullptr); + if (!fileReadOpsObj.isNull()) { + for (const auto& [key, value] : fileReadOpsObj.items()) { + properties.fileReadOps[key.asString()] = value.asString(); + } + } + + return properties; +} + +} // namespace facebook::velox diff --git a/velox/common/caching/FileProperties.h b/velox/common/caching/FileProperties.h index 9827a7e71f9..866ea7e35c1 100644 --- a/velox/common/caching/FileProperties.h +++ b/velox/common/caching/FileProperties.h @@ -17,18 +17,34 @@ #pragma once #include +#include #include #include +#include namespace facebook::velox { +namespace io { +class IoStatistics; +} + struct FileProperties { std::optional fileSize; std::optional modificationTime; std::optional readRangeHint{std::nullopt}; std::shared_ptr extraFileInfo{nullptr}; folly::F14FastMap fileReadOps{}; + + /// Non-owning pointer to the statistics of the data source that opened the + /// file. Set locally by the reader at open time, so 'serialize()' skips it. + io::IoStatistics* ioStatistics{nullptr}; + + folly::dynamic serialize() const; + + /// Reads what 'serialize()' wrote. Absent keys fall back to the member + /// default. + static FileProperties create(const folly::dynamic& obj); }; } // namespace facebook::velox diff --git a/velox/common/file/FileSystems.h b/velox/common/file/FileSystems.h index 70e6f439cfc..4202e7efaf7 100644 --- a/velox/common/file/FileSystems.h +++ b/velox/common/file/FileSystems.h @@ -28,6 +28,9 @@ namespace facebook::velox { namespace config { class ConfigBase; } +namespace io { +class IoStatistics; +} class IoStats; class ReadFile; class WriteFile; @@ -78,6 +81,12 @@ struct FileOptions { IoStats* stats{nullptr}; + /// Per-operation counters, keyed by operation name, unlike 'stats' above + /// which has no operation dimension. Non-owning: a file system may retain + /// this and record into it on every read, so it must outlive any file opened + /// with these options. + io::IoStatistics* ioStatistics{nullptr}; + /// A raw string that client can encode as anything they want to describe the /// file. For example, extraFileInfo can contain serialized file descriptors /// or other specific backend filesystem metadata can be used during for a diff --git a/velox/common/file/IoUringReader.cpp b/velox/common/file/IoUringReader.cpp index 17923f91aed..209f7c19164 100644 --- a/velox/common/file/IoUringReader.cpp +++ b/velox/common/file/IoUringReader.cpp @@ -275,7 +275,17 @@ IoUringReader::Stats getIoUringReaderStats(uint64_t& numReaders) { #if FOLLY_HAS_LIBURING bool IoUringReader::available() { - return folly::IoUringBackend::isAvailable(); + // folly::IoUringBackend::isAvailable() reports unavailability by returning + // false, except when ring setup fails with ENOMEM: that path throws + // std::runtime_error instead. ENOMEM here usually means the host ran out of + // RLIMIT_MEMLOCK rather than out of memory, which is a transient, host-wide + // condition. Report it as unavailable so callers can fall back. + try { + return folly::IoUringBackend::isAvailable(); + } catch (const std::exception& e) { + LOG(WARNING) << "io_uring availability probe failed: " << e.what(); + return false; + } } #else diff --git a/velox/common/file/tests/FileUtilsTest.cpp b/velox/common/file/tests/FileUtilsTest.cpp index eaa2b72cb94..1e9a01964bc 100644 --- a/velox/common/file/tests/FileUtilsTest.cpp +++ b/velox/common/file/tests/FileUtilsTest.cpp @@ -116,7 +116,7 @@ auto getReader( } // namespace -TEST(CoalesceSegmentsTest, EmptyCase) { +TEST(CoalesceSegmentsTest, emptyCase) { const Regions r = {}; MockShouldCoalesce shouldCoalesce; @@ -129,7 +129,7 @@ TEST(CoalesceSegmentsTest, EmptyCase) { EXPECT_EQ(resultRegions, expected); } -TEST(CoalesceSegmentsTest, MergeAll) { +TEST(CoalesceSegmentsTest, mergeAll) { const Regions r = {{0, 4}, {4, 4}, {8, 1}, {9, 2}, {11, 3}}; MockShouldCoalesce shouldCoalesce; @@ -146,7 +146,7 @@ TEST(CoalesceSegmentsTest, MergeAll) { EXPECT_EQ(resultRegions, expected); } -TEST(CoalesceSegmentsTest, MergeNone) { +TEST(CoalesceSegmentsTest, mergeNone) { const Regions r = {{0, 4}, {4, 4}, {8, 1}, {9, 2}, {11, 3}}; MockShouldCoalesce shouldCoalesce; @@ -164,7 +164,7 @@ TEST(CoalesceSegmentsTest, MergeNone) { EXPECT_EQ(resultRegions, expected); } -TEST(CoalesceSegmentsTest, MergeOdd) { +TEST(CoalesceSegmentsTest, mergeOdd) { const Regions r = {{0, 4}, {4, 4}, {8, 1}, {9, 2}, {11, 3}}; auto isOdd = [](size_t i) { return i % 2 == 1; }; @@ -184,7 +184,7 @@ TEST(CoalesceSegmentsTest, MergeOdd) { EXPECT_EQ(resultRegions, expected); } -TEST(CoalesceSegmentsTest, MergeEven) { +TEST(CoalesceSegmentsTest, mergeEven) { const Regions r = {{{0, 4}, {4, 4}, {8, 1}, {9, 2}, {11, 3}}}; auto isEven = [](size_t i) { return i % 2 == 0; }; @@ -203,7 +203,7 @@ TEST(CoalesceSegmentsTest, MergeEven) { EXPECT_EQ(resultRegions, expected); } -TEST(CoalesceIfDistanceLETest, MultipleCases) { +TEST(CoalesceIfDistanceLETest, multipleCases) { EXPECT_TRUE(willCoalesceIfDistanceLE(0, {0, 1}, {1, 1}, 0)); EXPECT_FALSE(willCoalesceIfDistanceLE(0, {0, 1}, {2, 1}, 0)); @@ -218,7 +218,7 @@ TEST(CoalesceIfDistanceLETest, MultipleCases) { EXPECT_TRUE(willCoalesceIfDistanceLE(0, {0, 0}, {0, 1}, 0)); } -TEST(CoalesceIfDistanceLETest, MultipleSegments) { +TEST(CoalesceIfDistanceLETest, multipleSegments) { uint64_t coalescedBytes = 0; auto willCoalesce = CoalesceIfDistanceLE(10, &coalescedBytes); EXPECT_TRUE(willCoalesce({0, 1}, {1, 1})); // 0 @@ -229,11 +229,11 @@ TEST(CoalesceIfDistanceLETest, MultipleSegments) { EXPECT_EQ(coalescedBytes, 19); } -TEST(CoalesceIfDistanceLETest, SupportsNullArgument) { +TEST(CoalesceIfDistanceLETest, supportsNullArgument) { EXPECT_NO_THROW(CoalesceIfDistanceLE(10, nullptr)({0, 10}, {20, 5})); // 10 } -TEST(CoalesceIfDistanceLETest, SegmentsMustBeSorted) { +TEST(CoalesceIfDistanceLETest, segmentsMustBeSorted) { EXPECT_THROW( willCoalesceIfDistanceLE(0, {1, 1}, {0, 1}, 0), ::facebook::velox::VeloxRuntimeError); @@ -248,7 +248,7 @@ TEST(CoalesceIfDistanceLETest, SegmentsMustBeSorted) { ::facebook::velox::VeloxRuntimeError); } -TEST(CoalesceIfDistanceLETest, SegmentsCantOverlap) { +TEST(CoalesceIfDistanceLETest, segmentsCantOverlap) { EXPECT_THROW( willCoalesceIfDistanceLE(0, {0, 1}, {0, 1}, 0), ::facebook::velox::VeloxRuntimeError); @@ -271,7 +271,7 @@ TEST(CoalesceIfDistanceLETest, SegmentsCantOverlap) { class ReadToIOBufsTest : public testing::TestWithParam {}; -TEST_P(ReadToIOBufsTest, CanRead) { +TEST_P(ReadToIOBufsTest, canRead) { Regions r = {{0, 1}, {5, 1}, {10, 6}, {16, 5}}; std::vector iobufs; iobufs.reserve(r.size()); diff --git a/velox/common/file/tests/LocalFileTest.cpp b/velox/common/file/tests/LocalFileTest.cpp index a9b21b002a8..9245f9b8c18 100644 --- a/velox/common/file/tests/LocalFileTest.cpp +++ b/velox/common/file/tests/LocalFileTest.cpp @@ -576,6 +576,10 @@ class LocalFileIoUringTest : public ::testing::Test { } void SetUp() override { + if (!IoUringReader::available()) { + GTEST_SKIP() << "io_uring is unavailable"; + } + ThreadLocalIoUringReader::testingClear(); IoUringReader::Options options; diff --git a/velox/common/fuzzer/ConstrainedGenerators.h b/velox/common/fuzzer/ConstrainedGenerators.h index 20d9565e38c..0138338b6e0 100644 --- a/velox/common/fuzzer/ConstrainedGenerators.h +++ b/velox/common/fuzzer/ConstrainedGenerators.h @@ -63,8 +63,7 @@ class RandomInputGenerator : public AbstractInputGenerator { return variant(randDate(rng_)); } if (type_->isTime()) { - VELOX_DCHECK(type_->equivalent(*TIME())); - return variant(randTime(rng_)); + return variant(randTime(rng_, type_)); } return variant(rand(rng_)); } @@ -583,9 +582,9 @@ class QDigestInputGenerator : public AbstractInputGenerator { auto makeDist = []() { if constexpr (std::is_integral_v) { - return std::uniform_int_distribution(0, 10000); + return std::uniform_int_distribution(-10'000, 10'000); } else { - return std::uniform_real_distribution(0.0, 10000.0); + return std::uniform_real_distribution(-10'000.0, 10'000.0); } }; diff --git a/velox/common/fuzzer/Utils.cpp b/velox/common/fuzzer/Utils.cpp index 2e5580396e5..ed77f6e3f54 100644 --- a/velox/common/fuzzer/Utils.cpp +++ b/velox/common/fuzzer/Utils.cpp @@ -119,8 +119,14 @@ int32_t randDate(FuzzerGenerator& rng) { return rand(rng, min, max); } -int32_t randTime(FuzzerGenerator& rng) { - return rand(rng, TIME()->getMin(), TIME()->getMax()); +int64_t randTime(FuzzerGenerator& rng, const TypePtr& type) { + VELOX_DCHECK(type->isTime()); + const bool isTimeMicroUtc = type->equivalent(*TIME_MICRO_UTC()); + const int64_t min = + isTimeMicroUtc ? TIME_MICRO_UTC()->getMin() : TIME()->getMin(); + const int64_t max = + isTimeMicroUtc ? TIME_MICRO_UTC()->getMax() : TIME()->getMax(); + return rand(rng, min, max); } /// Unicode character ranges. Ensure the vector indexes match the UTF8CharList diff --git a/velox/common/fuzzer/Utils.h b/velox/common/fuzzer/Utils.h index a9b5cdc2309..9fa49e6cefd 100644 --- a/velox/common/fuzzer/Utils.h +++ b/velox/common/fuzzer/Utils.h @@ -178,7 +178,7 @@ inline Timestamp rand(FuzzerGenerator& rng, DataSpec /*dataSpec*/) { int32_t randDate(FuzzerGenerator& rng); -int32_t randTime(FuzzerGenerator& rng); +int64_t randTime(FuzzerGenerator& rng, const TypePtr& type); /// Generate random timezone offset using biased distribution /// 25% probability: picks from frequently used offsets diff --git a/velox/common/hyperloglog/Murmur3Hash128.cpp b/velox/common/hyperloglog/Murmur3Hash128.cpp index f480b67a7b8..670ed8a5238 100644 --- a/velox/common/hyperloglog/Murmur3Hash128.cpp +++ b/velox/common/hyperloglog/Murmur3Hash128.cpp @@ -24,8 +24,10 @@ int64_t getLong(const void* data, int32_t offset) { return folly::loadUnaligned(static_cast(data) + offset); } -char getByte(const void* data, int32_t offset) { - return *(static_cast(data) + offset); +// Returns the byte zero extended. A signed type would sign extend a byte +// >= 0x80 across every higher bit of the 64-bit lane the tail shifts into. +uint8_t getByte(const void* data, int32_t offset) { + return *(static_cast(data) + offset); } // static diff --git a/velox/common/hyperloglog/tests/CMakeLists.txt b/velox/common/hyperloglog/tests/CMakeLists.txt index 03898321e46..36c3ada4a2c 100644 --- a/velox/common/hyperloglog/tests/CMakeLists.txt +++ b/velox/common/hyperloglog/tests/CMakeLists.txt @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -add_executable(velox_common_hyperloglog_test DenseHllTest.cpp SparseHllTest.cpp) +add_executable( + velox_common_hyperloglog_test + DenseHllTest.cpp + Murmur3Hash128Test.cpp + SparseHllTest.cpp +) add_test(NAME velox_common_hyperloglog_test COMMAND velox_common_hyperloglog_test) diff --git a/velox/common/hyperloglog/tests/Murmur3Hash128Test.cpp b/velox/common/hyperloglog/tests/Murmur3Hash128Test.cpp new file mode 100644 index 00000000000..4bb3284bfe4 --- /dev/null +++ b/velox/common/hyperloglog/tests/Murmur3Hash128Test.cpp @@ -0,0 +1,112 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "velox/common/hyperloglog/Murmur3Hash128.h" + +#include + +#include +#include +#include +#include + +using namespace facebook::velox::common::hll; + +namespace { + +int64_t hash64OfBytes(int64_t value, int64_t seed) { + char bytes[sizeof(int64_t)]; + std::memcpy(bytes, &value, sizeof(bytes)); + return Murmur3Hash128::hash64(bytes, sizeof(bytes), seed); +} + +} // namespace + +// hash64ForLong is hash64 specialized for an 8 byte input, so the two must +// return the same value for every input and seed. +TEST(Murmur3Hash128Test, hash64AgreesWithHash64ForLong) { + std::vector values = { + 0, + 1, + -1, + std::numeric_limits::min(), + std::numeric_limits::max(), + 0x80, + 0xff, + 0x7f7f7f7f7f7f7f7fLL, + static_cast(0x8080808080808080ULL), + 100012345678901LL, + }; + + std::mt19937_64 rng(12345); + for (int i = 0; i < 10000; ++i) { + values.push_back(static_cast(rng())); + } + + for (auto value : values) { + for (int64_t seed : {int64_t{0}, int64_t{1}, int64_t{0x5bd1e995}}) { + ASSERT_EQ( + hash64OfBytes(value, seed), + Murmur3Hash128::hash64ForLong(value, seed)) + << "value=" << value << " seed=" << seed; + } + } +} + +TEST(Murmur3Hash128Test, hash64ZeroExtendsHighTailBytes) { + const std::string highByteInput = "aaaaaa\xc3\xa9"; + ASSERT_EQ(highByteInput.size(), sizeof(int64_t)); + + int64_t asLong; + std::memcpy(&asLong, highByteInput.data(), sizeof(asLong)); + ASSERT_EQ(asLong, -6214015989967658655LL); + + EXPECT_EQ( + Murmur3Hash128::hash64( + highByteInput.data(), highByteInput.size(), /*seed=*/0), + Murmur3Hash128::hash64ForLong(asLong, /*seed=*/0)); + + const std::string asciiInput = "aaaaaaaa"; + int64_t asciiAsLong; + std::memcpy(&asciiAsLong, asciiInput.data(), sizeof(asciiAsLong)); + EXPECT_EQ( + Murmur3Hash128::hash64(asciiInput.data(), asciiInput.size(), /*seed=*/0), + Murmur3Hash128::hash64ForLong(asciiAsLong, /*seed=*/0)); +} + +// Inputs of 9 to 15 bytes reach tail positions an 8 byte input cannot, so +// hash64ForLong is unusable as the oracle. Expected values were computed from +// airlift/slice Murmur3Hash128.java#L152, the reference cited on hash64 in +// Murmur3Hash128.h, which masks each tail byte with & 0xFF. +TEST(Murmur3Hash128Test, hash64MatchesReferenceForNineToFifteenByteInputs) { + const std::vector> expected = { + {9, -3243918217540306855LL}, + {10, -7998242445834668767LL}, + {11, -4701723541745919412LL}, + {12, 1137560024168287992LL}, + {13, -2349531979471542545LL}, + {14, -2514671622296473736LL}, + {15, 5186136466562847178LL}, + }; + + for (const auto& [length, hash] : expected) { + std::string input(length - 2, 'a'); + input.append("\xc3\xa9"); + ASSERT_EQ(input.size(), length); + EXPECT_EQ( + Murmur3Hash128::hash64(input.data(), input.size(), /*seed=*/0), hash) + << "length=" << length; + } +} diff --git a/velox/common/io/IoStatistics.h b/velox/common/io/IoStatistics.h index 33cfd750a01..03bdb3fd36d 100644 --- a/velox/common/io/IoStatistics.h +++ b/velox/common/io/IoStatistics.h @@ -158,9 +158,12 @@ class IoStatistics { // or for an in-progress read-ahead to finish. IoCounter queryThreadIoLatencyUs_; - // Breakdown of queryThreadIoLatencyUs_ by I/O type: + // IO latency by type. Does not add up to queryThreadIoLatencyUs_ - a read + // issued by the prefetch thread is included in the metrics below but does not + // block the query thread and thus does not contribute to + // queryThreadIoLatencyUs_. - // Time spent waiting for remote storage reads (S3, HDFS, etc.) + // Time spent reading from remote storage (S3, HDFS, etc.) IoCounter storageReadLatencyUs_; // Time spent waiting for SSD cache reads diff --git a/velox/common/io/Options.h b/velox/common/io/Options.h index 568c581cdfc..853708db442 100644 --- a/velox/common/io/Options.h +++ b/velox/common/io/Options.h @@ -107,6 +107,13 @@ class ReaderOptions { return *this; } + /// Modifies whether a coalesced direct read packs its buffers into one + /// shared allocation. + ReaderOptions& setDirectBufferedInputSharedAllocation(bool enabled) { + directBufferedInputSharedAllocation_ = enabled; + return *this; + } + /// Modifies the maximum load coalesce distance. ReaderOptions& setMaxCoalesceDistance(int32_t distance) { maxCoalesceDistance_ = distance; @@ -142,6 +149,10 @@ class ReaderOptions { return loadQuantum_; } + bool directBufferedInputSharedAllocation() const { + return directBufferedInputSharedAllocation_; + } + int32_t maxCoalesceDistance() const { return maxCoalesceDistance_; } @@ -197,6 +208,7 @@ class ReaderOptions { uint64_t autoPreloadLength_{DEFAULT_AUTO_PRELOAD_SIZE}; PrefetchMode prefetchMode_{PrefetchMode::PREFETCH}; int32_t loadQuantum_{kDefaultLoadQuantum}; + bool directBufferedInputSharedAllocation_{false}; int32_t maxCoalesceDistance_{kDefaultCoalesceDistance}; int64_t maxCoalesceBytes_{kDefaultCoalesceBytes}; int32_t prefetchRowGroups_{kDefaultPrefetchRowGroups}; diff --git a/velox/common/memory/ArbitrationOperation.cpp b/velox/common/memory/ArbitrationOperation.cpp index b5c173c0dbc..211c589c4cb 100644 --- a/velox/common/memory/ArbitrationOperation.cpp +++ b/velox/common/memory/ArbitrationOperation.cpp @@ -20,13 +20,9 @@ #include "velox/common/base/Exceptions.h" #include "velox/common/base/RuntimeMetrics.h" #include "velox/common/memory/Memory.h" -#include "velox/common/testutil/TestValue.h" #include "velox/common/time/Timer.h" -using facebook::velox::common::testutil::TestValue; - namespace facebook::velox::memory { -using namespace facebook::velox::memory; ArbitrationOperation::ArbitrationOperation( ScopedArbitrationParticipant&& participant, diff --git a/velox/common/memory/ArbitrationOperation.h b/velox/common/memory/ArbitrationOperation.h index d9159c583c8..12eddabcee6 100644 --- a/velox/common/memory/ArbitrationOperation.h +++ b/velox/common/memory/ArbitrationOperation.h @@ -111,11 +111,39 @@ class ArbitrationOperation { globalArbitrationStartTimeNs_ = getCurrentTimeNano(); } - /// The execution stats of this arbitration operation after completion. + /// Holds the execution timing of this arbitration operation, valid only + /// after it finishes. The four fields are the gaps between consecutive + /// timestamps: + /// + /// createTimeNs_ created / enqueued + /// | + /// | localArbitrationWaitTimeNs + /// v + /// startTimeNs_ starts running (kRunning) + /// | + /// | localArbitrationExecTimeNs + /// v + /// globalArbitrationStartTimeNs_ global arbitration wait begins + /// | + /// | globalArbitrationWaitTimeNs + /// v + /// finishTimeNs_ finished (kFinished) + /// + /// executionTimeNs spans createTimeNs_..finishTimeNs_, i.e. the sum of the + /// three gaps. When no global arbitration happens, + /// globalArbitrationStartTimeNs_ is 0, localArbitrationExecTimeNs spans + /// startTimeNs_..finishTimeNs_, and globalArbitrationWaitTimeNs is 0. struct Stats { + /// Time queued on the participant before this operation starts running. uint64_t localArbitrationWaitTimeNs{0}; + + /// Time running local arbitration before any global arbitration wait. uint64_t localArbitrationExecTimeNs{0}; + + /// Time waiting for global arbitration; zero if none happened. uint64_t globalArbitrationWaitTimeNs{0}; + + /// Total time from creation to finish. uint64_t executionTimeNs{0}; }; @@ -128,19 +156,24 @@ class ArbitrationOperation { const uint64_t requestBytes_; const uint64_t timeoutNs_; - // The start time of this arbitration operation. + // Time when this operation was created and enqueued, before it starts + // running. See Stats for how the timestamps carve the timeline. const uint64_t createTimeNs_; const ScopedArbitrationParticipant participant_; State state_{State::kInit}; + // Time when this operation starts running (kRunning), after waiting its turn + // on the participant. uint64_t startTimeNs_{0}; + // Time when this operation finishes (kFinished). uint64_t finishTimeNs_{0}; uint64_t maxGrowBytes_{0}; uint64_t minGrowBytes_{0}; - // The time that starts global arbitration wait + // Time when this operation starts waiting for global arbitration. Zero if + // global arbitration does not happen. uint64_t globalArbitrationStartTimeNs_{}; friend class ArbitrationParticipant; diff --git a/velox/common/memory/ArbitrationParticipant.cpp b/velox/common/memory/ArbitrationParticipant.cpp index 7d74c98c731..aa2b38d3669 100644 --- a/velox/common/memory/ArbitrationParticipant.cpp +++ b/velox/common/memory/ArbitrationParticipant.cpp @@ -399,11 +399,10 @@ std::string ArbitrationParticipant::Stats::toString() const { } ScopedArbitrationParticipant::ScopedArbitrationParticipant( - std::shared_ptr ArbitrationParticipant, + std::shared_ptr participant, std::shared_ptr pool) - : ArbitrationParticipant_(std::move(ArbitrationParticipant)), - pool_(std::move(pool)) { - VELOX_CHECK_NOT_NULL(ArbitrationParticipant_); + : participant_(std::move(participant)), pool_(std::move(pool)) { + VELOX_CHECK_NOT_NULL(participant_); VELOX_CHECK_NOT_NULL(pool_); } diff --git a/velox/common/memory/ArbitrationParticipant.h b/velox/common/memory/ArbitrationParticipant.h index 530b04b9dda..8857007a0c5 100644 --- a/velox/common/memory/ArbitrationParticipant.h +++ b/velox/common/memory/ArbitrationParticipant.h @@ -371,27 +371,27 @@ class ArbitrationParticipant class ScopedArbitrationParticipant { public: ScopedArbitrationParticipant( - std::shared_ptr ArbitrationParticipant, + std::shared_ptr participant, std::shared_ptr pool); ArbitrationParticipant* operator->() const { - return ArbitrationParticipant_.get(); + return participant_.get(); } ArbitrationParticipant& operator*() const { - return *ArbitrationParticipant_; + return *participant_; } ArbitrationParticipant& operator()() const { - return *ArbitrationParticipant_; + return *participant_; } ArbitrationParticipant* get() const { - return ArbitrationParticipant_.get(); + return participant_.get(); } private: - std::shared_ptr ArbitrationParticipant_; + std::shared_ptr participant_; std::shared_ptr pool_; }; diff --git a/velox/common/memory/Memory.cpp b/velox/common/memory/Memory.cpp index 9df9199c79b..ef71a1808f5 100644 --- a/velox/common/memory/Memory.cpp +++ b/velox/common/memory/Memory.cpp @@ -190,7 +190,7 @@ MemoryManager::~MemoryManager() { if (checkUsageLeak_) { VELOX_FAIL(errMsg); } else { - LOG(ERROR) << errMsg; + VELOX_MEM_LOG(ERROR) << errMsg; } } } diff --git a/velox/common/memory/MemoryArbitrator.cpp b/velox/common/memory/MemoryArbitrator.cpp index 6856d301ef1..c1134aa497c 100644 --- a/velox/common/memory/MemoryArbitrator.cpp +++ b/velox/common/memory/MemoryArbitrator.cpp @@ -81,9 +81,10 @@ class NoopArbitrator : public MemoryArbitrator { explicit NoopArbitrator(const Config& config) : MemoryArbitrator(config) { VELOX_CHECK(config.kind.empty()); if (config_.capacity != kMaxMemory) { - LOG(WARNING) << "Query memory capacity[" - << succinctBytes(config_.capacity) << "] is set for " - << kind() << " arbitrator which has no capacity enforcement"; + VELOX_MEM_LOG(WARNING) + << "Query memory capacity[" << succinctBytes(config_.capacity) + << "] is set for " << kind() + << " arbitrator which has no capacity enforcement"; } } @@ -543,15 +544,17 @@ ScopedReclaimedBytesRecorder::~ScopedReclaimedBytesRecorder() { } const int64_t reservedBytesAfterReclaim = pool_->reservedBytes(); if (reservedBytesAfterReclaim > reservedBytesBeforeReclaim_) { - LOG(ERROR) << "Unexpected reserved bytes growth from " << pool_->name() - << ", root pool: " << pool_->root()->name() - << " after memory reclaim from " - << succinctBytes(reservedBytesBeforeReclaim_) << " to " - << succinctBytes(reservedBytesAfterReclaim) - << ", used: " << succinctBytes(pool_->usedBytes()) - << ", reservation: " << succinctBytes(pool_->reservedBytes()) - << ", root pool reservation: " - << succinctBytes(pool_->root()->reservedBytes()); + VELOX_MEM_LOG(ERROR) << "Unexpected reserved bytes growth from " + << pool_->name() + << ", root pool: " << pool_->root()->name() + << " after memory reclaim from " + << succinctBytes(reservedBytesBeforeReclaim_) << " to " + << succinctBytes(reservedBytesAfterReclaim) + << ", used: " << succinctBytes(pool_->usedBytes()) + << ", reservation: " + << succinctBytes(pool_->reservedBytes()) + << ", root pool reservation: " + << succinctBytes(pool_->root()->reservedBytes()); } *reclaimedBytes_ = reservedBytesBeforeReclaim_ - reservedBytesAfterReclaim; } diff --git a/velox/common/memory/MemoryPool.cpp b/velox/common/memory/MemoryPool.cpp index dec20a27eec..2d889b78d56 100644 --- a/velox/common/memory/MemoryPool.cpp +++ b/velox/common/memory/MemoryPool.cpp @@ -463,7 +463,10 @@ MemoryPoolImpl::MemoryPoolImpl( // The memory manager sets the capacity through grow() according to the // actually used memory arbitration policy. capacity_(parent_ != nullptr ? kMaxMemory : 0) { - VELOX_CHECK(options.threadSafe || isLeaf()); + VELOX_CHECK( + options.threadSafe || isLeaf(), + "Only a leaf memory pool can be non-thread-safe: {}", + name_); } MemoryPoolImpl::~MemoryPoolImpl() { @@ -561,11 +564,7 @@ void* MemoryPoolImpl::allocate( void MemoryPoolImpl::reportExternalAllocation(int64_t size) { VELOX_CHECK_GT(size, 0, "reportExternalAllocation requires positive size"); - if (FOLLY_UNLIKELY(kind_ != Kind::kLeaf)) { - VELOX_FAIL( - "Memory operation is only allowed on leaf memory pool: {}", toString()); - } - ++numExternalAllocs_; + CHECK_AND_INC_MEM_OP_STATS(this, ExternalAllocs); reserve(size); cumulativeExternalBytes_ += size; } @@ -712,11 +711,7 @@ bool MemoryPoolImpl::transferTo(MemoryPool* dest, void* buffer, uint64_t size) { void MemoryPoolImpl::reportExternalFree(int64_t size) { VELOX_CHECK_GT(size, 0, "reportExternalFree requires positive size"); - if (FOLLY_UNLIKELY(kind_ != Kind::kLeaf)) { - VELOX_FAIL( - "Memory operation is only allowed on leaf memory pool: {}", toString()); - } - ++numExternalFrees_; + CHECK_AND_INC_MEM_OP_STATS(this, ExternalFrees); release(size); } diff --git a/velox/common/memory/MemoryPool.h b/velox/common/memory/MemoryPool.h index 9715dda847b..23417357e0a 100644 --- a/velox/common/memory/MemoryPool.h +++ b/velox/common/memory/MemoryPool.h @@ -147,7 +147,7 @@ class MemoryPool : public std::enable_shared_from_this { /// If true, tracks the leaf memory pool usage in a thread-safe mode /// otherwise not. This only applies for leaf memory pool with memory usage - /// tracking enabled. We use non-thread safe tracking mode for single + /// tracking enabled. We use non-thread-safe tracking mode for single /// threaded use case. /// /// NOTE: user can turn on/off the thread-safe mode of each individual leaf @@ -218,8 +218,10 @@ class MemoryPool : public std::enable_shared_from_this { return trackUsage_; } - /// Returns true if this memory pools is thread safe which only applies for a - /// leaf memory pool with memory usage tracking enabled. + /// Returns true if this memory pool is thread-safe. Only a leaf memory pool + /// with memory usage tracking enabled can be non-thread-safe (see + /// Options::threadSafe). Aggregate memory pools, including the root, are + /// always thread-safe. virtual bool threadSafe() const { return threadSafe_; } @@ -241,6 +243,9 @@ class MemoryPool : public std::enable_shared_from_this { /// Invoked to create a named aggregate child memory pool. /// + /// Unlike a leaf child, an aggregate child is always thread-safe and has no + /// 'threadSafe' option, as its leaf descendants may update it concurrently. + /// /// NOTE: 'reclaimer' only applies if the aggregate memory pool has enabled /// memory usage tracking which inherits from its parent. virtual std::shared_ptr addAggregateChild( @@ -909,6 +914,10 @@ class MemoryPoolImpl : public MemoryPool { // if max capacity is exceeded or arbitration fails. void incrementReservationThreadSafe(MemoryPool* requestor, uint64_t size); + // Increments the reservation for a non-thread-safe leaf pool. Even though + // this leaf skips its own locking, it propagates the increment to the parent + // via the thread-safe path, because the aggregate ancestors are shared by + // other leaves and updated concurrently. FOLLY_ALWAYS_INLINE void incrementReservationNonThreadSafe( MemoryPool* requestor, uint64_t size) { diff --git a/velox/common/memory/MmapAllocator.cpp b/velox/common/memory/MmapAllocator.cpp index eb6eaabb5b9..f2b4d5f5d48 100644 --- a/velox/common/memory/MmapAllocator.cpp +++ b/velox/common/memory/MmapAllocator.cpp @@ -17,6 +17,11 @@ #include "velox/common/memory/MmapAllocator.h" #include +#include + +#include + +#include #include "velox/common/base/Counters.h" #include "velox/common/base/Portability.h" @@ -24,6 +29,23 @@ #include "velox/common/memory/Memory.h" namespace facebook::velox::memory { +uint64_t MmapAllocator::systemPageSize() { + static const uint64_t pageSize = [] { + const long value{::sysconf(_SC_PAGESIZE)}; + VELOX_CHECK_GT( + value, + 0, + "Failed to determine the system page size: {}", + folly::errnoStr(errno)); + return static_cast(value); + }(); + return pageSize; +} + +bool MmapAllocator::isPageSizeSupported() { + return systemPageSize() == AllocationTraits::kPageSize; +} + MmapAllocator::MmapAllocator(const Options& options) : MemoryAllocator(options.largestSizeClass), kind_(MemoryAllocator::Kind::kMmap), @@ -39,6 +61,24 @@ MmapAllocator::MmapAllocator(const Options& options) AllocationTraits::numPages( options.capacity - mallocReservedBytes_), 64 * sizeClassSizes_.back())) { + // MmapAllocator tracks memory at AllocationTraits::kPageSize granularity + // and calls madvise() on individual pages, which requires the address and + // length to line up with the OS's actual page size. NVIDIA's documented + // recommended default page size for Grace / Grace-Hopper systems is 64KB, + // not 4KB (see + // https://docs.nvidia.com/dccpu/grace-perf-tuning-guide/os-settings.html), + // and on such systems madvise() silently fails with EINVAL on sub-64KB + // regions instead of throwing, which corrupts this allocator's internal + // page-count accounting rather than surfacing a clear error. This isn't + // Velox-specific: jemalloc has the identical >4KB-page limitation (see + // https://github.com/arangodb/arangodb/issues/22177). Fail fast here + // instead of silently corrupting state. + VELOX_CHECK( + isPageSizeSupported(), + "MmapAllocator requires the system page size to match AllocationTraits::kPageSize ({} bytes); system page size is {} bytes. Use MemoryAllocator::Kind::kMalloc on this system instead.", + AllocationTraits::kPageSize, + systemPageSize()); + for (const auto& size : sizeClassSizes_) { sizeClasses_.push_back(std::make_unique(capacity_ / size, size)); } diff --git a/velox/common/memory/MmapAllocator.h b/velox/common/memory/MmapAllocator.h index b68c7f910f3..6e64687fbdb 100644 --- a/velox/common/memory/MmapAllocator.h +++ b/velox/common/memory/MmapAllocator.h @@ -55,6 +55,18 @@ class MmapAllocator : public MemoryAllocator { ~MmapAllocator(); + /// Returns the page size the OS reports, queried once on first call and + /// cached. Throws if the query fails. + static uint64_t systemPageSize(); + + /// Returns true if systemPageSize() matches AllocationTraits::kPageSize, the + /// granularity MmapAllocator assumes for its mmap/madvise calls (see the + /// constructor in MmapAllocator.cpp for why the two can differ). Construct a + /// MallocAllocator instead when it does not; the constructor enforces this + /// via VELOX_CHECK, so callers that want to select an allocator kind without + /// throwing should check this first. + static bool isPageSizeSupported(); + Kind kind() const override { return kind_; } diff --git a/velox/common/memory/SharedArbitrator.cpp b/velox/common/memory/SharedArbitrator.cpp index 43121a71177..267e9b42e27 100644 --- a/velox/common/memory/SharedArbitrator.cpp +++ b/velox/common/memory/SharedArbitrator.cpp @@ -355,7 +355,11 @@ void SharedArbitrator::setupGlobalArbitration( ",", globalArbitrationSpillCapacityLimits_); globalArbitrationController_ = std::make_unique([&]() { - folly::setThreadName("GlobalArbitrationController"); + // Thread names are truncated to 15 characters (Linux TASK_COMM_LEN - 1), + // so keep this name short enough to stay fully visible in top and pstack. + static constexpr std::string_view kThreadName{"ArbitrationMain"}; + static_assert(kThreadName.size() <= 15); + folly::setThreadName(kThreadName); globalArbitrationMain(); }); } @@ -1289,7 +1293,7 @@ uint64_t SharedArbitrator::reclaimUsedMemoryBySpill( } } if (victims.empty()) { - FB_LOG_EVERY_MS(WARNING, 1'000) + VELOX_MEM_LOG_EVERY_MS(WARNING, 1'000) << "No spill victim participant found with global arbitration target: " << succinctBytes(targetBytes); return 0; @@ -1408,8 +1412,8 @@ uint64_t SharedArbitrator::reclaim( freeCapacity(reclaimedBytes); if (reclaimedBytes == 0) { - FB_LOG_EVERY_MS(WARNING, 1'000) << fmt::format( - "Nothing reclaimed from memory pool {} with reclaim target {}, memory pool stats:\n{}\n{}", + VELOX_MEM_LOG_EVERY_MS(WARNING, 1'000) << fmt::format( + "Nothing reclaimed from memory pool {} with reclaim target {}, memory pool stats:\n{}\n{}", participant->name(), succinctBytes(targetBytes), participant->pool()->toString(), diff --git a/velox/common/memory/SharedArbitrator.h b/velox/common/memory/SharedArbitrator.h index b5e210034de..72ff75c38e3 100644 --- a/velox/common/memory/SharedArbitrator.h +++ b/velox/common/memory/SharedArbitrator.h @@ -387,19 +387,6 @@ class SharedArbitrator : public memory::MemoryArbitrator { // capacity limit by reclaiming used memory from the participant itself. bool ensureCapacity(ArbitrationOperation& op); - // Invoked to run local arbitration on the request memory pool. It first - // ensures the memory growth is within both memory pool and arbitrator - // capacity limits. This step might reclaim the used memory from the request - // memory pool itself. Then it tries to obtain free capacity from the - // arbitrator. At last, it tries to reclaim free memory from itself before it - // falls back to the global arbitration. The local arbitration run is - // protected by shared lock of 'arbitrationLock_' which can run in parallel - // for different query pools. The free memory reclamation is protected by - // arbitrator 'mutex_' which is an in-memory fast operation. The function - // returns false on failure. Otherwise, it needs to further check if - // 'needGlobalArbitration' is true or not. If true, needs to proceed with the - // global arbitration run. - // Invoked to initialize the global arbitration on arbitrator start-up. It // starts the background threads to used memory from running queries // on-demand. diff --git a/velox/common/memory/tests/MemoryAllocatorTest.cpp b/velox/common/memory/tests/MemoryAllocatorTest.cpp index 5b402ff9235..357c93a1c07 100644 --- a/velox/common/memory/tests/MemoryAllocatorTest.cpp +++ b/velox/common/memory/tests/MemoryAllocatorTest.cpp @@ -1606,7 +1606,7 @@ TEST_P(MemoryAllocatorTest, allocateZeroFilled) { ASSERT_TRUE(instance_->checkConsistency()); } -TEST_P(MemoryAllocatorTest, StlMemoryAllocator) { +TEST_P(MemoryAllocatorTest, stlMemoryAllocator) { { std::vector> data( 0, StlAllocator(*pool_)); diff --git a/velox/common/memory/tests/MemoryPoolTest.cpp b/velox/common/memory/tests/MemoryPoolTest.cpp index c3780552de8..3c408feee95 100644 --- a/velox/common/memory/tests/MemoryPoolTest.cpp +++ b/velox/common/memory/tests/MemoryPoolTest.cpp @@ -677,7 +677,7 @@ TEST_P(MemoryPoolTest, releasableMemory) { } } -TEST_P(MemoryPoolTest, ReallocTestSameSize) { +TEST_P(MemoryPoolTest, reallocTestSameSize) { auto manager = getMemoryManager(); auto root = manager->addRootPool(); @@ -699,7 +699,7 @@ TEST_P(MemoryPoolTest, ReallocTestSameSize) { ASSERT_EQ(2 * kChunkSize, pool->stats().peakBytes); } -TEST_P(MemoryPoolTest, ReallocTestHigher) { +TEST_P(MemoryPoolTest, reallocTestHigher) { auto manager = getMemoryManager(); auto root = manager->addRootPool(); @@ -720,7 +720,7 @@ TEST_P(MemoryPoolTest, ReallocTestHigher) { EXPECT_EQ(4 * kChunkSize, pool->stats().peakBytes); } -TEST_P(MemoryPoolTest, ReallocTestLower) { +TEST_P(MemoryPoolTest, reallocTestLower) { auto manager = getMemoryManager(); auto root = manager->addRootPool(); auto pool = root->addLeafChild("elastic_quota", isLeafThreadSafe_); @@ -905,7 +905,7 @@ TEST_P(MemoryPoolTest, memoryCapExceptions) { } } -TEST(MemoryPoolTest, GetAlignment) { +TEST(MemoryPoolTest, getAlignment) { { MemoryManager::Options options; options.allocatorCapacity = kMaxMemory; @@ -972,7 +972,7 @@ TEST(MemoryPoolTest, allocateAlignedTracksUsage) { EXPECT_GT(pool->stats().numFrees, statsBefore.numFrees); } -TEST_P(MemoryPoolTest, MemoryManagerGlobalCap) { +TEST_P(MemoryPoolTest, memoryManagerGlobalCap) { MemoryManager::Options options; options.allocatorCapacity = 32L * MB; options.arbitratorCapacity = 32L * MB; diff --git a/velox/common/memory/tests/MockSharedArbitratorTest.cpp b/velox/common/memory/tests/MockSharedArbitratorTest.cpp index 830aa1e699b..b7974d503eb 100644 --- a/velox/common/memory/tests/MockSharedArbitratorTest.cpp +++ b/velox/common/memory/tests/MockSharedArbitratorTest.cpp @@ -22,6 +22,7 @@ #include #include #include "folly/synchronization/EventCount.h" +#include "velox/common/base/ConcurrentRuntimeStatWriter.h" #include "velox/common/base/tests/GTestUtils.h" #include "velox/common/memory/MallocAllocator.h" #include "velox/common/memory/Memory.h" @@ -42,22 +43,6 @@ using namespace facebook::velox::exec; using namespace facebook::velox::exec::test; namespace facebook::velox::memory { -// Class to write runtime stats in the tests to the stats container. -class TestRuntimeStatWriter : public BaseRuntimeStatWriter { - public: - explicit TestRuntimeStatWriter( - std::unordered_map& stats) - : stats_{stats} {} - - void addRuntimeStat(std::string_view name, const RuntimeCounter& value) - override { - addOperatorRuntimeStats(name, value, stats_); - } - - private: - std::unordered_map& stats_; -}; - constexpr int64_t KB = 1024L; constexpr int64_t MB = 1024L * KB; @@ -1446,10 +1431,10 @@ DEBUG_ONLY_TEST_F(MockSharedArbitrationTest, localArbitrationsFromSameQuery) { std::atomic_int allocationCount{0}; auto runThread = std::thread([&]() { - std::unordered_map runtimeStats; - auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + ConcurrentRuntimeStatWriter statsWriter; + setThreadLocalRunTimeStatWriter(&statsWriter); runPool->allocate(memoryCapacity / 2); + auto runtimeStats = statsWriter.runtimeStats(); ASSERT_EQ( runtimeStats[std::string(SharedArbitrator::kMemoryArbitrationWallNanos)] .count, @@ -1471,10 +1456,10 @@ DEBUG_ONLY_TEST_F(MockSharedArbitrationTest, localArbitrationsFromSameQuery) { auto waitThread = std::thread([&]() { allocationWait.await([&]() { return !allocationWaitFlag.load(); }); - std::unordered_map runtimeStats; - auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + ConcurrentRuntimeStatWriter statsWriter; + setThreadLocalRunTimeStatWriter(&statsWriter); waitPool->allocate(memoryCapacity / 2 + MB); + auto runtimeStats = statsWriter.runtimeStats(); ASSERT_EQ( runtimeStats[std::string(SharedArbitrator::kMemoryArbitrationWallNanos)] .count, @@ -1548,11 +1533,11 @@ DEBUG_ONLY_TEST_F( std::atomic_int allocationCount{0}; auto taskThread1 = std::thread([&]() { - std::unordered_map runtimeStats; - auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + ConcurrentRuntimeStatWriter statsWriter; + setThreadLocalRunTimeStatWriter(&statsWriter); op1->allocate(MB); ASSERT_EQ(task1->capacity(), 8 * MB); + auto runtimeStats = statsWriter.runtimeStats(); ASSERT_EQ( runtimeStats[std::string(SharedArbitrator::kMemoryArbitrationWallNanos)] .count, @@ -1573,11 +1558,11 @@ DEBUG_ONLY_TEST_F( }); auto taskThread2 = std::thread([&]() { - std::unordered_map runtimeStats; - auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + ConcurrentRuntimeStatWriter statsWriter; + setThreadLocalRunTimeStatWriter(&statsWriter); op2->allocate(MB); ASSERT_EQ(task2->capacity(), 8 * MB); + auto runtimeStats = statsWriter.runtimeStats(); ASSERT_EQ( runtimeStats[std::string(SharedArbitrator::kMemoryArbitrationWallNanos)] .count, @@ -1801,9 +1786,8 @@ DEBUG_ONLY_TEST_F( .memoryPoolReserveCapacity = memoryPoolReservedCapacity}); auto globalArbitrationTriggerThread = std::thread([&]() { - std::unordered_map runtimeStats; - auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + ConcurrentRuntimeStatWriter statsWriter; + setThreadLocalRunTimeStatWriter(&statsWriter); std::vector> tasks; std::vector ops; @@ -1817,6 +1801,7 @@ DEBUG_ONLY_TEST_F( ops[i]->allocate(memoryPoolCapacity); } // We expect global arbitration has been triggered. + auto runtimeStats = statsWriter.runtimeStats(); ASSERT_GE( runtimeStats[std::string(SharedArbitrator::kMemoryArbitrationWallNanos)] .count, @@ -1921,9 +1906,8 @@ DEBUG_ONLY_TEST_F( }))); auto globalArbitrationTriggerThread = std::thread([&]() { - std::unordered_map runtimeStats; - auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + ConcurrentRuntimeStatWriter statsWriter; + setThreadLocalRunTimeStatWriter(&statsWriter); std::vector> tasks; std::vector ops; @@ -1937,6 +1921,7 @@ DEBUG_ONLY_TEST_F( ops[i]->allocate(memoryPoolCapacity); } // We expect global arbitration has been triggered. + auto runtimeStats = statsWriter.runtimeStats(); ASSERT_GE( runtimeStats[std::string(SharedArbitrator::kMemoryArbitrationWallNanos)] .count, @@ -1975,9 +1960,8 @@ DEBUG_ONLY_TEST_F( globalArbitrationStartWait.await( [&]() { return globalArbitrationStarted.load(); }); - std::unordered_map runtimeStats; - auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + ConcurrentRuntimeStatWriter statsWriter; + setThreadLocalRunTimeStatWriter(&statsWriter); localArbitrationOp->allocate(memoryPoolReservedCapacity); // Inject some delay for global arbitration. @@ -1987,6 +1971,7 @@ DEBUG_ONLY_TEST_F( globalArbitrationTriggerThread.join(); ASSERT_EQ(localArbitrationOp->capacity(), memoryPoolReservedCapacity); + auto runtimeStats = statsWriter.runtimeStats(); ASSERT_EQ( runtimeStats[std::string(SharedArbitrator::kGlobalArbitrationWaitCount)] .count, @@ -2043,13 +2028,13 @@ DEBUG_ONLY_TEST_F(MockSharedArbitrationTest, globalArbitrationAbortTimeRatio) { std::chrono::nanoseconds(pauseTimeNs)); }))); - std::unordered_map runtimeStats; - auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + ConcurrentRuntimeStatWriter statsWriter; + setThreadLocalRunTimeStatWriter(&statsWriter); const auto prevGlobalArbitrationRuns = arbitratorHelper.globalArbitrationRuns(); op1->allocate(memoryCapacity / 2); + auto runtimeStats = statsWriter.runtimeStats(); ASSERT_EQ( runtimeStats[std::string(SharedArbitrator::kMemoryArbitrationWallNanos)] .count, @@ -2109,11 +2094,11 @@ TEST_F(MockSharedArbitrationTest, globalArbitrationWithoutSpill) { abortOp->allocate(memoryCapacity / 2); ASSERT_EQ(triggerTask->capacity(), memoryCapacity / 2); - std::unordered_map runtimeStats; - auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + ConcurrentRuntimeStatWriter statsWriter; + setThreadLocalRunTimeStatWriter(&statsWriter); triggerOp->allocate(memoryCapacity / 2); + auto runtimeStats = statsWriter.runtimeStats(); ASSERT_EQ( runtimeStats[std::string(SharedArbitrator::kMemoryArbitrationWallNanos)] .count, @@ -2164,15 +2149,15 @@ TEST_F(MockSharedArbitrationTest, globalArbitrationSmallParticipantLargeGrow) { op1->allocate(kMemoryCapacity / 2); ASSERT_EQ(task0->capacity(), kMemoryCapacity / 2); - std::unordered_map runtimeStats; - auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + ConcurrentRuntimeStatWriter statsWriter; + setThreadLocalRunTimeStatWriter(&statsWriter); // task0 has 256MB + 256MB (attempt) = 512MB in top abort capacity limit // bucket, which shall be evaluated first, and hence killed by global // arbitration. VELOX_ASSERT_THROW(op0->allocate(kMemoryCapacity / 2), "aborted"); + auto runtimeStats = statsWriter.runtimeStats(); ASSERT_EQ( runtimeStats[std::string(SharedArbitrator::kMemoryArbitrationWallNanos)] .count, @@ -2392,10 +2377,10 @@ DEBUG_ONLY_TEST_F(MockSharedArbitrationTest, multipleGlobalRuns) { std::atomic_int allocations{0}; auto waitThread = std::thread([&]() { allocationWait.await([&]() { return !allocationWaitFlag.load(); }); - std::unordered_map runtimeStats; - auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + ConcurrentRuntimeStatWriter statsWriter; + setThreadLocalRunTimeStatWriter(&statsWriter); waitPool->allocate(memoryCapacity / 2); + auto runtimeStats = statsWriter.runtimeStats(); ASSERT_EQ( runtimeStats[std::string(SharedArbitrator::kMemoryArbitrationWallNanos)] .count, @@ -2420,10 +2405,10 @@ DEBUG_ONLY_TEST_F(MockSharedArbitrationTest, multipleGlobalRuns) { }); auto runThread = std::thread([&]() { - std::unordered_map runtimeStats; - auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + ConcurrentRuntimeStatWriter statsWriter; + setThreadLocalRunTimeStatWriter(&statsWriter); runPool->allocate(memoryCapacity / 2); + auto runtimeStats = statsWriter.runtimeStats(); ASSERT_EQ( runtimeStats[std::string(SharedArbitrator::kMemoryArbitrationWallNanos)] .count, diff --git a/velox/common/process/ProcessBase.cpp b/velox/common/process/ProcessBase.cpp index 0b9a4df2c64..2fe50f740ff 100644 --- a/velox/common/process/ProcessBase.cpp +++ b/velox/common/process/ProcessBase.cpp @@ -25,6 +25,7 @@ #include #include #include +#include constexpr const char* kProcSelfCmdline = "/proc/self/cmdline"; @@ -116,6 +117,16 @@ bool hasAvx2() { #endif } +bool hasSimd() { +#if XSIMD_WITH_AVX2 + return avx2CpuFlag && FLAGS_avx2; +#elif XSIMD_WITH_NEON64 + return true; +#else + return false; +#endif +} + bool hasBmi2() { #ifdef __BMI2__ return bmi2CpuFlag && FLAGS_bmi2; diff --git a/velox/common/process/ProcessBase.h b/velox/common/process/ProcessBase.h index 34edd6d1467..c04340a26cd 100644 --- a/velox/common/process/ProcessBase.h +++ b/velox/common/process/ProcessBase.h @@ -46,6 +46,12 @@ uint64_t threadCpuNanos(); /// by flag. bool hasAvx2(); +/// True if the platform has SIMD instructions suitable for the bulk reader +/// path (xsimd-based, no raw intrinsics). On x86 this requires AVX2 and +/// is gated by --avx2; on aarch64, NEON is always available and the bulk +/// path is unconditionally enabled. +bool hasSimd(); + /// True if the machine has Intel BMI2 instructions and these are not disabled /// by flag. bool hasBmi2(); diff --git a/velox/common/rpc/clients/MockRPCClient.cpp b/velox/common/rpc/clients/MockRPCClient.cpp index 2a266a2c686..cad6068347a 100644 --- a/velox/common/rpc/clients/MockRPCClient.cpp +++ b/velox/common/rpc/clients/MockRPCClient.cpp @@ -16,9 +16,12 @@ #include "velox/common/rpc/clients/MockRPCClient.h" +#include #include #include +#include "velox/common/base/Exceptions.h" + namespace facebook::velox::rpc { namespace { @@ -28,6 +31,47 @@ std::mt19937& threadLocalRng() { thread_local std::mt19937 rng{std::random_device{}()}; return rng; } + +// Builds an error response tagged with a typed cause, for the error-burst +// path. Unlike MockRPCClient::generateResponse(request, /*isError=*/true), +// which leaves errorKind at kNone, this sets errorKind so the congestion path +// can classify the failure. +RPCResponse makeErrorResponse(const RPCRequest& request, RPCErrorKind kind) { + // Every enumerator is listed explicitly (no default) so a newly added kind + // trips -Wswitch-enum instead of silently reusing another kind's label. + std::string_view label; + switch (kind) { + case RPCErrorKind::kRateLimited: + label = "rate-limit"; + break; + case RPCErrorKind::kTimeout: + label = "timeout"; + break; + case RPCErrorKind::kNullInput: + label = "null-input"; + break; + case RPCErrorKind::kBackendError: + label = "backend"; + break; + case RPCErrorKind::kEmptyResponse: + label = "empty-response"; + break; + case RPCErrorKind::kInvalidRequest: + label = "invalid-request"; + break; + case RPCErrorKind::kNone: + VELOX_UNREACHABLE( + "makeErrorResponse() requires a non-kNone error kind, " + "it is only reached on the error-burst path"); + } + return RPCResponse{ + .rowId = request.rowId, + .result = "", + .metadata = {}, + .error = + fmt::format("Simulated {} error for row {}", label, request.rowId), + .errorKind = kind}; +} } // namespace MockRPCClient::MockRPCClient( @@ -71,18 +115,63 @@ RPCResponse MockRPCClient::generateResponse( .error = std::nullopt}; } -folly::SemiFuture MockRPCClient::call(const RPCRequest& request) { - callCount_.fetch_add(1); +void MockRPCClient::setErrorBurst(const ErrorBurst& burst) { + // Install-once, before dispatch. burstErrorKind() reads errorBurst_ without + // a lock, so the struct must never be written while a request could be + // reading it. Rejecting a second install closes the resetCallCount() path: + // the counter check alone would let a re-arm race with requests still in + // flight from the previous burst. + VELOX_CHECK( + !burstInstalled_.load(std::memory_order_acquire), + "setErrorBurst() may only be called once per client"); + VELOX_CHECK_EQ( + callCount_.load(), + 0, + "setErrorBurst() must be called before the first request is dispatched"); + errorBurst_ = burst; + // Release: publishes the errorBurst_ write to any thread that subsequently + // observes burstInstalled_ as true through the acquire load below. + burstInstalled_.store(true, std::memory_order_release); +} - // Determine if this request should fail - std::uniform_real_distribution dist(0.0, 1.0); - bool shouldError = dist(threadLocalRng()) < errorRate_; +RPCErrorKind MockRPCClient::burstErrorKind(int64_t ordinal) const { + // Acquire: pairs with the release store in setErrorBurst(), so the + // errorBurst_ fields read below are guaranteed visible on this thread. + // errorBurst_ is never mutated once a burst is installed, so no further + // synchronization is needed. + if (!burstInstalled_.load(std::memory_order_acquire)) { + return RPCErrorKind::kNone; + } + if (errorBurst_.firstCall < errorBurst_.lastCall && + ordinal >= errorBurst_.firstCall && ordinal < errorBurst_.lastCall) { + return errorBurst_.errorKind; + } + return RPCErrorKind::kNone; +} + +folly::SemiFuture MockRPCClient::call(const RPCRequest& request) { + const int64_t ordinal = callCount_.fetch_add(1); + const RPCErrorKind burstKind = burstErrorKind(ordinal); + + // Draw the random error decision only on the non-burst path, so the RNG + // stream is consumed identically to callBatch(), which skips the draw for + // bursted requests. This keeps the two paths reproducible under a fixed seed. + bool shouldError = false; + if (burstKind == RPCErrorKind::kNone) { + std::uniform_real_distribution dist(0.0, 1.0); + shouldError = dist(threadLocalRng()) < errorRate_; + } // Use folly::via with the thread pool executor for safe async execution return folly::via( executor_.get(), - [this, request = request, shouldError, latency = latency_]() + [this, request = request, shouldError, burstKind, latency = latency_]() -> RPCResponse { + // Deterministic overload burst: fail fast with a typed cause, skipping + // the latency sleep (overload rejections come back immediately). + if (burstKind != RPCErrorKind::kNone) { + return makeErrorResponse(request, burstKind); + } // Simulate network latency /* sleep override */ std::this_thread::sleep_for(latency); // Generate and return the response @@ -95,12 +184,23 @@ folly::SemiFuture> MockRPCClient::callBatch( // Capture error rate for thread safety double errorRate = errorRate_; + // Reserve this batch's ordinals on the caller thread, as call() does, so the + // burst window covers a fixed set of requests even when several batches are + // in flight. Assigning them inside the lambda would tie the mapping to + // executor scheduling order. + const int64_t firstOrdinal = + callCount_.fetch_add(static_cast(requests.size())); + // Use folly::via with the thread pool executor for safe async execution return folly::via( executor_.get(), - [this, requests, errorRate, latency = latency_]() + [this, requests, errorRate, firstOrdinal, latency = latency_]() -> std::vector { - // Simulate network latency (single batch = single latency) + // Simulate network latency (single batch = single latency). Unlike + // call(), which returns a bursted request immediately, a batch pays + // this once up front even when some of its requests are in the burst + // window: the batch is one round trip, and a backend that rejects part + // of it still costs the caller that trip. /* sleep override */ std::this_thread::sleep_for(latency); std::vector responses; @@ -111,8 +211,14 @@ folly::SemiFuture> MockRPCClient::callBatch( thread_local std::mt19937 localRng{std::random_device{}()}; std::uniform_real_distribution dist(0.0, 1.0); - for (const auto& request : requests) { - callCount_.fetch_add(1); + for (size_t i = 0; i < requests.size(); ++i) { + const auto& request = requests[i]; + const RPCErrorKind burstKind = + burstErrorKind(firstOrdinal + static_cast(i)); + if (burstKind != RPCErrorKind::kNone) { + responses.push_back(makeErrorResponse(request, burstKind)); + continue; + } bool shouldError = dist(localRng) < errorRate; responses.push_back(generateResponse(request, shouldError)); } diff --git a/velox/common/rpc/clients/MockRPCClient.h b/velox/common/rpc/clients/MockRPCClient.h index 417196701cb..57f0fa49115 100644 --- a/velox/common/rpc/clients/MockRPCClient.h +++ b/velox/common/rpc/clients/MockRPCClient.h @@ -55,17 +55,46 @@ class MockRPCClient : public IRPCClient { return callCount_.load(); } - /// Resets the call counter. + /// Resets the call counter. An installed error burst stays installed and + /// cannot be replaced, so ordinals restart at 0 and the same burst window + /// applies again. Use a fresh client to test a different burst. void resetCallCount() { callCount_.store(0); } + /// Configures a deterministic error burst for congestion / AIMD tests. + /// Requests whose 0-based call ordinal falls in [firstCall, lastCall) are + /// failed with a response tagged `errorKind`, simulating a timed + /// backend-overload window (rate-limit / timeout). Ordinals are reserved on + /// the caller thread, one per call() request and a contiguous run per + /// callBatch() request, so the burst covers a fixed, timing-independent set + /// of requests. Disabled when firstCall >= lastCall (the default). + struct ErrorBurst { + int64_t firstCall{0}; + int64_t lastCall{0}; + RPCErrorKind errorKind{RPCErrorKind::kRateLimited}; + }; + + /// Installs the error burst. May be called at most once, and only before the + /// first request is dispatched; thereafter the burst is only read, never + /// mutated. Throws if called twice or after a request has been dispatched. + void setErrorBurst(const ErrorBurst& burst); + private: RPCResponse generateResponse(const RPCRequest& request, bool isError); + // Returns the burst error kind for a request at 0-based `ordinal`, or kNone + // when the ordinal is outside the configured burst window. + RPCErrorKind burstErrorKind(int64_t ordinal) const; + const std::chrono::milliseconds latency_; const double errorRate_; std::atomic callCount_{0}; + // Guards publication of errorBurst_ to the dispatch threads: written with + // release in setErrorBurst(), read with acquire in burstErrorKind(). + std::atomic burstInstalled_{false}; + // Configured error-burst window; read only once burstInstalled_ is true. + ErrorBurst errorBurst_{}; /// Shared executor (may be shared across clients for global throttling). std::shared_ptr executor_; diff --git a/velox/common/serialization/tests/TestRegistry.cpp b/velox/common/serialization/tests/TestRegistry.cpp index f2162367b63..57c3bc4971c 100644 --- a/velox/common/serialization/tests/TestRegistry.cpp +++ b/velox/common/serialization/tests/TestRegistry.cpp @@ -21,7 +21,7 @@ using namespace ::facebook::velox; namespace { -TEST(Registry, SmartPointerFactoryWithNoArgument) { +TEST(Registry, smartPointerFactoryWithNoArgument) { Registry()> registry; const size_t key = 0; @@ -38,7 +38,7 @@ TEST(Registry, SmartPointerFactoryWithNoArgument) { EXPECT_EQ(*registry.Create(key), value); } -TEST(Registry, ValueFactoryWithArguments) { +TEST(Registry, valueFactoryWithArguments) { Registry registry; const size_t key = 0; diff --git a/velox/common/testutil/tests/CastsTest.cpp b/velox/common/testutil/tests/CastsTest.cpp index b5a9cf60052..8908a26f4df 100644 --- a/velox/common/testutil/tests/CastsTest.cpp +++ b/velox/common/testutil/tests/CastsTest.cpp @@ -93,6 +93,16 @@ class CastsTest : public ::testing::Test { DerivedClass* derivedRawPtr_; }; +TEST_F(CastsTest, checkedNotNull) { + int value = 1; + EXPECT_EQ(&value, checkedNotNull(&value)); + + const int constValue = 2; + EXPECT_EQ(&constValue, checkedNotNull(&constValue)); + + VELOX_ASSERT_THROW(checkedNotNull(static_cast(nullptr)), ""); +} + // Tests for checkedPointerCast with shared_ptr TEST_F(CastsTest, checkedPointerCastSharedPtrSuccess) { // Cast derived to base (should always work) diff --git a/velox/common/time/Timer.cpp b/velox/common/time/Timer.cpp index 78416f03e67..132ff5815c3 100644 --- a/velox/common/time/Timer.cpp +++ b/velox/common/time/Timer.cpp @@ -16,13 +16,45 @@ #include "velox/common/time/Timer.h" +#include + #include "velox/common/testutil/ScopedTestTime.h" +#include "velox/common/time/CpuWallTimer.h" namespace facebook::velox { using namespace std::chrono; using common::testutil::ScopedTestTime; +ProcessCpuWallTimer::ProcessCpuWallTimer(CpuWallTiming& timing) + : wallTimeStart_{steady_clock::now()}, + cpuTimeStart_{processCpuNanos()}, + timing_{timing} { + ++timing_.count; +} + +ProcessCpuWallTimer::~ProcessCpuWallTimer() { + const auto cpuTimeEnd = processCpuNanos(); + if (cpuTimeStart_ != kUnavailableCpuTime && + cpuTimeEnd != kUnavailableCpuTime) { + timing_.cpuNanos += cpuTimeEnd - cpuTimeStart_; + } + timing_.wallNanos += + duration_cast(steady_clock::now() - wallTimeStart_).count(); +} + +uint64_t ProcessCpuWallTimer::processCpuNanos() noexcept { + rusage usage{}; + if (getrusage(RUSAGE_SELF, &usage) != 0) { + return kUnavailableCpuTime; + } + const auto toNanos = [](const timeval& value) { + return static_cast(value.tv_sec) * 1'000'000'000 + + static_cast(value.tv_usec) * 1'000; + }; + return toNanos(usage.ru_utime) + toNanos(usage.ru_stime); +} + #ifndef NDEBUG uint64_t getCurrentTimeSec() { diff --git a/velox/common/time/Timer.h b/velox/common/time/Timer.h index 2b7b53fcfb8..19027e9ac19 100644 --- a/velox/common/time/Timer.h +++ b/velox/common/time/Timer.h @@ -19,12 +19,15 @@ #include #include #include +#include #include #include "velox/common/process/ProcessBase.h" namespace facebook::velox { +struct CpuWallTiming; + /// Measures wall time (steady_clock) between construction and destruction, /// incrementing a user-supplied counter in microseconds. class MicrosecondWallTimer { @@ -99,6 +102,32 @@ class NanosecondCPUTimer { uint64_t start_; }; +/// Adds process-wide CPU and wall time to a CpuWallTiming. Unlike +/// CpuWallTimer, which measures CPU consumed by the calling thread, this timer +/// includes user and system CPU consumed by all process threads. CPU time can +/// exceed wall time when background thread pools run work concurrently. +class ProcessCpuWallTimer { + public: + explicit ProcessCpuWallTimer(CpuWallTiming& timing); + ~ProcessCpuWallTimer(); + + ProcessCpuWallTimer(const ProcessCpuWallTimer&) = delete; + ProcessCpuWallTimer& operator=(const ProcessCpuWallTimer&) = delete; + ProcessCpuWallTimer(ProcessCpuWallTimer&&) = delete; + ProcessCpuWallTimer& operator=(ProcessCpuWallTimer&&) = delete; + + private: + // Returns user and system CPU consumed across all process threads. + static uint64_t processCpuNanos() noexcept; + + static constexpr uint64_t kUnavailableCpuTime{ + std::numeric_limits::max()}; + + const std::chrono::steady_clock::time_point wallTimeStart_; + const uint64_t cpuTimeStart_; + CpuWallTiming& timing_; +}; + using MicrosecondTimer = MicrosecondWallTimer; using NanosecondTimer = NanosecondWallTimer; diff --git a/velox/common/time/tests/CpuWallTimerTest.cpp b/velox/common/time/tests/CpuWallTimerTest.cpp index 39bbf89b977..a4d199d05e8 100644 --- a/velox/common/time/tests/CpuWallTimerTest.cpp +++ b/velox/common/time/tests/CpuWallTimerTest.cpp @@ -14,11 +14,15 @@ * limitations under the License. */ +#include #include #include +#include +#include #include #include "velox/common/time/CpuWallTimer.h" +#include "velox/common/time/Timer.h" using namespace facebook::velox; @@ -104,6 +108,45 @@ TEST_F(CpuWallTimerTest, cpuWallTimer) { EXPECT_LT(cpuFirstTime, timing.cpuNanos); } +TEST_F(CpuWallTimerTest, processCpuWallTimerIncludesBackgroundThreads) { + constexpr uint32_t numWorkers{4}; + folly::CPUThreadPoolExecutor executor{numWorkers}; + std::latch allStarted{numWorkers}; + std::latch startWork{1}; + std::latch finished{numWorkers}; + std::atomic totalIterations{0}; + constexpr uint64_t iterationsPerWorker{2'000'000}; + for (uint32_t worker = 0; worker < numWorkers; ++worker) { + executor.add([&] { + allStarted.count_down(); + startWork.wait(); + for (uint64_t iteration = 0; iteration < iterationsPerWorker; + ++iteration) { + totalIterations.fetch_add(1, std::memory_order_relaxed); + } + finished.count_down(); + }); + } + allStarted.wait(); + + CpuWallTiming threadTiming; + CpuWallTiming processTiming; + { + ProcessCpuWallTimer processTimer{processTiming}; + CpuWallTimer threadTimer{threadTiming}; + startWork.count_down(); + finished.wait(); + } + + EXPECT_EQ( + totalIterations.load(std::memory_order_relaxed), + numWorkers * iterationsPerWorker); + EXPECT_EQ(processTiming.count, 1); + EXPECT_EQ(threadTiming.count, 1); + EXPECT_GT(processTiming.cpuNanos, threadTiming.cpuNanos + 10'000'000); + EXPECT_GE(processTiming.wallNanos, threadTiming.wallNanos); +} + TEST_F(CpuWallTimerTest, deltaCpuWallTimer) { CpuWallTiming timing; // Everything should be zero. diff --git a/velox/connectors/hive/CMakeLists.txt b/velox/connectors/hive/CMakeLists.txt index 871aaadee4f..22e183d6b77 100644 --- a/velox/connectors/hive/CMakeLists.txt +++ b/velox/connectors/hive/CMakeLists.txt @@ -47,6 +47,7 @@ velox_add_library( HivePartitionName.cpp HiveSplitReader.cpp PartitionIdGenerator.cpp + PartitionValue.cpp TableHandle.cpp HEADERS BufferedInputBuilder.h @@ -77,6 +78,7 @@ velox_add_library( HiveSplitReader.h IndexReader.h PartitionIdGenerator.h + PartitionValue.h TableHandle.h ) @@ -97,7 +99,7 @@ endif() velox_add_library(velox_hive_partition_function HivePartitionFunction.cpp) -velox_link_libraries(velox_hive_partition_function velox_core velox_exec) +velox_link_libraries(velox_hive_partition_function velox_core velox_exec velox_hive_hash) add_subdirectory(storage_adapters) diff --git a/velox/connectors/hive/FileConfig.cpp b/velox/connectors/hive/FileConfig.cpp index d9bff102a6d..9deda9f3532 100644 --- a/velox/connectors/hive/FileConfig.cpp +++ b/velox/connectors/hive/FileConfig.cpp @@ -45,6 +45,7 @@ const std::vector& FileConfig::registeredProperties() { VELOX_HIVE_CONFIG_REGISTER(kCacheIndexSession); VELOX_HIVE_CONFIG_REGISTER(kPinIndexSession); VELOX_HIVE_CONFIG_REGISTER(kSelectiveNimbleReaderEnabledSession); + VELOX_HIVE_CONFIG_REGISTER(kDirectBufferedInputSharedAllocationSession); VELOX_HIVE_CONFIG_REGISTER(kMaxCoalescedDistanceSession); VELOX_HIVE_CONFIG_REGISTER(kParallelUnitLoadCountSession); VELOX_HIVE_CONFIG_REGISTER(kReadTimestampUnitSession); diff --git a/velox/connectors/hive/FileConfig.h b/velox/connectors/hive/FileConfig.h index 964cd831590..b646e901a9e 100644 --- a/velox/connectors/hive/FileConfig.h +++ b/velox/connectors/hive/FileConfig.h @@ -106,6 +106,21 @@ class FileConfig { false, "Preserve dictionary encoding for Nimble string column reads.") + // TODO: Deprecate this gate and pack unconditionally once the shared + // allocation has run in production for a while. + VELOX_HIVE_CONFIG_LEGACY( + kDirectBufferedInputSharedAllocationSession, + kDirectBufferedInputSharedAllocation, + directBufferedInputSharedAllocation, + "reader.direct_buffered_input_shared_allocation", + "reader.direct-buffered-input-shared-allocation", + bool, + false, + "Pack a coalesced direct read's buffers into a single shared allocation " + "instead of one page-rounded allocation per request. Opt-in: off by " + "default, enable per cluster to roll out. Affects which bytes are " + "allocated, never which bytes are read.") + VELOX_HIVE_CONFIG_LEGACY( kNimbleLazyColumnIoSession, kNimbleLazyColumnIo, @@ -240,6 +255,7 @@ class FileConfig { "quantum per stream is always issued; subsequent quanta are loaded on " "demand. Small streams that fit in one quantum see no reduction. " "Streams coalesced with eager columns may also be loaded early.") + // --- VELOX_HIVE_CONFIG_PROPERTY properties --- VELOX_HIVE_CONFIG_PROPERTY( diff --git a/velox/connectors/hive/FileConnectorSplit.h b/velox/connectors/hive/FileConnectorSplit.h index efba64c5b8b..08544908c79 100644 --- a/velox/connectors/hive/FileConnectorSplit.h +++ b/velox/connectors/hive/FileConnectorSplit.h @@ -43,6 +43,10 @@ struct FileConnectorSplit : public ConnectorSplit { const std::unordered_map> partitionKeys; + /// Optional split-level override for table-to-file column matching. When not + /// set, the connector session property decides the mode. + const std::optional columnMappingMode; + FileConnectorSplit( const std::string& connectorId, const std::string& _filePath, @@ -53,14 +57,17 @@ struct FileConnectorSplit : public ConnectorSplit { bool cacheable = true, std::optional _properties = std::nullopt, const std::unordered_map>& - _partitionKeys = {}) + _partitionKeys = {}, + std::optional _columnMappingMode = + std::nullopt) : ConnectorSplit(connectorId, splitWeight, cacheable), filePath(_filePath), fileFormat(_fileFormat), start(_start), length(_length), properties(std::move(_properties)), - partitionKeys(_partitionKeys) {} + partitionKeys(_partitionKeys), + columnMappingMode(_columnMappingMode) {} ~FileConnectorSplit() override = default; diff --git a/velox/connectors/hive/FileConnectorUtil.cpp b/velox/connectors/hive/FileConnectorUtil.cpp index 15ea6fa886f..b0617ff53de 100644 --- a/velox/connectors/hive/FileConnectorUtil.cpp +++ b/velox/connectors/hive/FileConnectorUtil.cpp @@ -24,6 +24,7 @@ #include "velox/connectors/hive/FileConfig.h" #include "velox/connectors/hive/FileConnectorSplit.h" #include "velox/connectors/hive/FileTableHandle.h" +#include "velox/connectors/hive/PartitionValue.h" #include "velox/dwio/common/Options.h" #include "velox/dwio/common/ReaderFactory.h" #include "velox/dwio/dwrf/common/Config.h" @@ -53,6 +54,33 @@ FormatScopedConfigs makeFormatScopedConfigs( dwio::common::formatConfigPrefix(fileFormat, "_")))}; } +namespace { + +void validateColumnMappingMode( + dwio::common::ColumnMappingMode mode, + dwio::common::FileFormat fileFormat) { + // kParquetFieldId is format-specific: it matches requested columns against + // physical Parquet schema field_id metadata. Other readers don't have that + // metadata, so reject it at split setup time instead of letting a later + // reader path interpret it as a generic field-id or name/position mode. + VELOX_USER_CHECK( + mode != dwio::common::ColumnMappingMode::kParquetFieldId || + fileFormat == dwio::common::FileFormat::PARQUET, + "Column mapping mode {} is not supported for file format {}", + mode, + dwio::common::FileFormatName::toName(fileFormat)); +} + +dwio::common::ColumnMappingMode sessionColumnMappingMode( + const FileConfig& fileConfig, + const config::ConfigBase* sessionProperties) { + return fileConfig.useColumnNames(sessionProperties) + ? dwio::common::ColumnMappingMode::kName + : dwio::common::ColumnMappingMode::kPosition; +} + +} // namespace + void configureReaderOptions( const std::shared_ptr& fileConfig, const ConnectorQueryCtx* connectorQueryCtx, @@ -75,6 +103,8 @@ void configureReaderOptions( auto sessionProperties = connectorQueryCtx->sessionProperties(); VELOX_CHECK_NOT_NULL(sessionProperties, "Session properties are null"); readerOptions.setLoadQuantum(fileConfig->loadQuantum(sessionProperties)); + readerOptions.setDirectBufferedInputSharedAllocation( + fileConfig->directBufferedInputSharedAllocation(sessionProperties)); readerOptions.setMaxCoalesceBytes( fileConfig->maxCoalescedBytes(sessionProperties)); readerOptions.setMaxCoalesceDistance( @@ -94,10 +124,10 @@ void configureReaderOptions( readerOptions.setFileColumnNamesReadAsLowerCase( fileConfig->isFileColumnNamesReadAsLowerCase(sessionProperties)); readerOptions.setAllowEmptyFile(true); - readerOptions.setColumnMappingMode( - fileConfig->useColumnNames(sessionProperties) - ? dwio::common::ColumnMappingMode::kName - : dwio::common::ColumnMappingMode::kPosition); + const auto columnMappingMode = fileSplit->columnMappingMode.value_or( + sessionColumnMappingMode(*fileConfig, sessionProperties)); + validateColumnMappingMode(columnMappingMode, fileSplit->fileFormat); + readerOptions.setColumnMappingMode(columnMappingMode); readerOptions.setFileSchema(fileSchema); readerOptions.setFilePreloadThreshold(fileConfig->filePreloadThreshold()); readerOptions.setPrefetchRowGroups(fileConfig->prefetchRowGroups()); @@ -222,50 +252,14 @@ bool applyPartitionFilter( bool isPartitionDateDaysSinceEpoch, const common::Filter* filter, bool asLocalTime) { - if (type->isDate()) { - int32_t result = 0; - // days_since_epoch partition values are integers in string format. Eg. - // Iceberg partition values. - if (isPartitionDateDaysSinceEpoch) { - result = folly::to(partitionValue); - } else { - result = DATE()->toDays(partitionValue); - } - return applyFilter(*filter, result); - } - - switch (type->kind()) { - case TypeKind::BIGINT: - case TypeKind::INTEGER: - case TypeKind::SMALLINT: - case TypeKind::TINYINT: { - return applyFilter(*filter, folly::to(partitionValue)); - } - case TypeKind::REAL: - case TypeKind::DOUBLE: { - return applyFilter(*filter, folly::to(partitionValue)); - } - case TypeKind::BOOLEAN: { - return applyFilter(*filter, folly::to(partitionValue)); - } - case TypeKind::TIMESTAMP: { - VELOX_DCHECK(type->equivalent(*TIMESTAMP())); - auto result = util::fromTimestampString( - StringView(partitionValue), util::TimestampParseMode::kPrestoCast); - VELOX_CHECK(!result.hasError()); - if (asLocalTime) { - result.value().toGMT(Timestamp::defaultTimezone()); - } - return applyFilter(*filter, result.value()); - } - case TypeKind::VARCHAR: - case TypeKind::VARBINARY: { - return applyFilter(*filter, partitionValue); - } - default: - VELOX_FAIL( - "Bad type {} for partition value: {}", type->kind(), partitionValue); - } + const auto value = PartitionValue::fromString( + partitionValue, + *type, + asLocalTime ? PartitionValue::TimestampMode::kLocalTime + : PartitionValue::TimestampMode::kUtc, + isPartitionDateDaysSinceEpoch ? PartitionValue::DateMode::kDaysSinceEpoch + : PartitionValue::DateMode::kIsoString); + return applyFilter(*filter, value); } template @@ -306,22 +300,34 @@ bool testFilters( // By design, the partition key columns for Iceberg tables are included in // the data files to facilitate partition transform and partition // evolution, so we need to test both cases. - if (!rowType->containsChild(name) || iter != partitionKeys.end()) { + // + // A constant decides the column even when the file carries it: + // 'fileTypeWithId' has no subtree for it, so there are no statistics. + if (child->isConstant() || !rowType->containsChild(name) || + iter != partitionKeys.end()) { if (iter != partitionKeys.end() && iter->second.has_value()) { const auto handlesIter = partitionKeysHandle.find(name); VELOX_CHECK(handlesIter != partitionKeysHandle.end()); - // This is a non-null partition key - return applyPartitionFilter( - handlesIter->second->dataType(), - iter->second.value(), - handlesIter->second->isPartitionDateValueDaysSinceEpoch(), - child->filter(), - asLocalTime); + // A later column may still exclude the split, so keep going. + if (!applyPartitionFilter( + handlesIter->second->dataType(), + iter->second.value(), + handlesIter->second->isPartitionDateValueDaysSinceEpoch(), + child->filter(), + asLocalTime)) { + VLOG(1) << "Skipping " << filePath + << " because the partition value failed the filter for " + "column " + << name; + return false; + } + continue; } - // Column is missing from the file. If it has a constant value (e.g., - // an initial-default from schema evolution), test the filter against - // it. Otherwise treat the column as NULL. + // The value does not come from the file. Filter on the constant if + // there is one (an initial-default from schema evolution), otherwise + // treat the column as NULL. Nothing downstream re-checks it: + // testFilterOnConstant() accepts any non-null constant. bool filterMatchedConstant = false; if (child->isConstant()) { auto constantVec = child->constantValue(); diff --git a/velox/connectors/hive/FileDataSource.cpp b/velox/connectors/hive/FileDataSource.cpp index a50e606a81a..21cee2d9db0 100644 --- a/velox/connectors/hive/FileDataSource.cpp +++ b/velox/connectors/hive/FileDataSource.cpp @@ -85,6 +85,28 @@ inline void addIoLatencyMetric( } } +void addOperationStatsToRuntimeStats( + io::IoStatistics& ioStats, + std::unordered_map& res) { + for (const auto& [operation, counters] : ioStats.operationStats()) { + const auto add = [&](std::string_view counter, uint64_t value) { + if (value == 0) { + return; + } + res[fmt::format("storage.{}.{}", operation, counter)] = + RuntimeMetric(value, RuntimeCounter::Unit::kNone); + }; + add("requestCount", counters.requestCount); + add("localThrottleCount", counters.localThrottleCount); + add("globalThrottleCount", counters.globalThrottleCount); + add("resourceThrottleCount", counters.resourceThrottleCount); + add("retryCount", counters.retryCount); + // Cumulative across requests, so consumers must divide by requestCount to + // recover a per-request mean. + add("latencyInMs", counters.latencyInMs); + } +} + } // namespace void addIoStatsToRuntimeStats( @@ -629,18 +651,28 @@ void FileDataSource::addDynamicFilter( } void FileDataSource::fireScanBatchCallback(core::ScanBatchEvent event) { + // Bytes are read when the reader loads a stripe, which for small files is + // entirely inside addSplit() and for large ones is spread across next() + // calls. Reporting the delta since the previous event captures them either + // way; a window around a single next() would not. + const uint64_t totalStorageReadBytes = dataIoStats_->read().sum(); + const uint64_t storageReadBytesDelta = + totalStorageReadBytes - lastEventStorageReadBytes_; + lastEventStorageReadBytes_ = totalStorageReadBytes; if (!scanBatchCallback_) { return; } FileScanBatchEvent fileEvent; fileEvent.numRows = event.numRows; fileEvent.wallTimeMicros = event.wallTimeMicros; + fileEvent.storageReadBytes = storageReadBytesDelta; if (tableHandle_) { fileEvent.tableName = tableHandle_->name(); fileEvent.dbName = tableHandle_->dbName(); } if (split_) { fileEvent.filePath = split_->filePath; + fileEvent.fileFormat = split_->fileFormat; if (!split_->partitionKeys.empty()) { fileEvent.partitionKeys = &split_->partitionKeys; } @@ -674,6 +706,8 @@ FileDataSource::getRuntimeStats() { res.emplace(key, value); } } + + addOperationStatsToRuntimeStats(*dataIoStats_, res); return res; } diff --git a/velox/connectors/hive/FileDataSource.h b/velox/connectors/hive/FileDataSource.h index d198149817e..2e0c8c5bf65 100644 --- a/velox/connectors/hive/FileDataSource.h +++ b/velox/connectors/hive/FileDataSource.h @@ -53,6 +53,11 @@ struct FileScanBatchEvent : public core::ScanBatchEvent { /// Null when partition keys are not available. const std::unordered_map>* partitionKeys{nullptr}; + /// File format of the current split. + dwio::common::FileFormat fileFormat{dwio::common::FileFormat::UNKNOWN}; + /// Bytes fetched from storage producing this batch, including any read + /// amplification from coalescing adjacent regions. + uint64_t storageReadBytes{0}; }; class FileConfig; @@ -156,6 +161,10 @@ class FileDataSource : public DataSource { std::shared_ptr metadataIoStats_; std::shared_ptr ioStats_; + // Cumulative dataIoStats_->read().sum() as of the last scan batch event, so + // each event reports only the bytes read since the previous one. + uint64_t lastEventStorageReadBytes_{0}; + /// Column handles for the split info columns keyed on their column names. std::unordered_map infoColumns_; SpecialColumnNames specialColumns_{}; @@ -183,7 +192,7 @@ class FileDataSource : public DataSource { // post-read using the extraction chains. folly::F14FastMap extractionColumns_; - dwio::common::RuntimeStatistics runtimeStats_; + dwio::common::RuntimeStats runtimeStats_; private: // Configure extraction columns on the ScanSpec and build diff --git a/velox/connectors/hive/FileSplitReader.cpp b/velox/connectors/hive/FileSplitReader.cpp index e12ee97bc41..b179a53401a 100644 --- a/velox/connectors/hive/FileSplitReader.cpp +++ b/velox/connectors/hive/FileSplitReader.cpp @@ -21,86 +21,31 @@ #include "velox/connectors/hive/FileConfig.h" #include "velox/connectors/hive/FileConnectorSplit.h" #include "velox/connectors/hive/FileConnectorUtil.h" +#include "velox/connectors/hive/PartitionValue.h" #include "velox/dwio/common/ReaderFactory.h" -#include "velox/type/DecimalUtil.h" namespace facebook::velox::connector::hive { -namespace { -template -VectorPtr newConstantFromStringImpl( +VectorPtr newConstantFromString( const TypePtr& type, const std::optional& value, velox::memory::MemoryPool* pool, bool isLocalTimestamp, bool isDaysSinceEpoch) { - using T = typename TypeTraits::NativeType; if (!value.has_value()) { - return std::make_shared>(pool, 1, true, type, T()); - } - - if (type->isDate()) { - int32_t days = 0; - // For Iceberg, the date partition values are already in daysSinceEpoch - // form. - if (isDaysSinceEpoch) { - days = folly::to(value.value()); - } else { - days = DATE()->toDays(value.value()); - } - return std::make_shared>( - pool, 1, false, type, std::move(days)); - } - - if constexpr (std::is_same_v || std::is_same_v) { - if (type->isDecimal()) { - T decimalValue = 0; - auto [precision, scale] = getDecimalPrecisionScale(*type); - auto status = DecimalUtil::castFromString( - StringView(value.value()), precision, scale, decimalValue); - if (!status.ok()) { - VELOX_USER_FAIL(status.message()); - } - return std::make_shared>( - pool, 1, false, type, std::move(decimalValue)); - } - } - if constexpr (std::is_same_v) { - return std::make_shared>( - pool, 1, false, type, StringView(value.value())); - } else { - auto copy = velox::util::Converter::tryCast(value.value()) - .thenOrThrow(folly::identity, [&](const Status& status) { - VELOX_USER_FAIL("{}", status.message()); - }); - if constexpr (kind == TypeKind::TIMESTAMP) { - // TIMESTAMP partition value is read as local time subject to the - // 'readTimestampPartitionValueAsLocalTime' setting. TIMESTAMP_UTC - // partition value is always read as UTC. - if (type->equivalent(*TIMESTAMP()) && isLocalTimestamp) { - copy.toGMT(Timestamp::defaultTimezone()); - } - } - return std::make_shared>( - pool, 1, false, type, std::move(copy)); + return BaseVector::createNullConstant(type, 1, pool); } -} -} // namespace - -VectorPtr newConstantFromString( - const TypePtr& type, - const std::optional& value, - velox::memory::MemoryPool* pool, - bool isLocalTimestamp, - bool isDaysSinceEpoch) { - return VELOX_DYNAMIC_SCALAR_TYPE_DISPATCH_ALL( - newConstantFromStringImpl, - type->kind(), + return BaseVector::createConstant( type, - value, - pool, - isLocalTimestamp, - isDaysSinceEpoch); + PartitionValue::fromString( + value.value(), + *type, + isLocalTimestamp ? PartitionValue::TimestampMode::kLocalTime + : PartitionValue::TimestampMode::kUtc, + isDaysSinceEpoch ? PartitionValue::DateMode::kDaysSinceEpoch + : PartitionValue::DateMode::kIsoString), + 1, + pool); } std::unique_ptr FileSplitReader::create( @@ -191,7 +136,7 @@ void FileSplitReader::configureBaseReaderOptions() { void FileSplitReader::prepareSplit( std::shared_ptr metadataFilter, - dwio::common::RuntimeStatistics& runtimeStats, + dwio::common::RuntimeStats& runtimeStats, const folly::F14FastMap& fileReadOps) { createReader(fileReadOps); if (emptySplit_) { @@ -240,7 +185,7 @@ int64_t FileSplitReader::estimatedRowSize() const { } void FileSplitReader::updateRuntimeStats( - dwio::common::RuntimeStatistics& stats) const { + dwio::common::RuntimeStats& stats) const { if (baseRowReader_) { baseRowReader_->updateRuntimeStats(stats); } @@ -289,6 +234,12 @@ void FileSplitReader::createReader( if (!tableHandle_->name().empty()) { fileProperties.fileReadOps[kTableNameKey] = tableHandle_->name(); } + // Per-operation counters are attributed to the data source that opened the + // file, so they are only meaningful while file handles are not reused across + // data sources. A cached handle would both outlive these statistics and + // report one data source's reads against another's. + fileProperties.ioStatistics = + fileHandleFactory_->maxSize() == 0 ? dataIoStats_.get() : nullptr; try { fileHandleCachePtr = fileHandleFactory_->generate( @@ -340,7 +291,7 @@ RowTypePtr FileSplitReader::getAdaptedRowType() const { } bool FileSplitReader::filterOnStats( - dwio::common::RuntimeStatistics& runtimeStats) const { + dwio::common::RuntimeStats& runtimeStats) const { if (testFilters( scanSpec_.get(), baseReader_.get(), @@ -358,7 +309,7 @@ bool FileSplitReader::filterOnStats( } bool FileSplitReader::checkIfSplitIsEmpty( - dwio::common::RuntimeStatistics& runtimeStats) { + dwio::common::RuntimeStats& runtimeStats) { // emptySplit_ may already be set if the data file is not found. In this case // we don't need to test further. if (emptySplit_) { diff --git a/velox/connectors/hive/FileSplitReader.h b/velox/connectors/hive/FileSplitReader.h index e61e865c16d..71b776a91ba 100644 --- a/velox/connectors/hive/FileSplitReader.h +++ b/velox/connectors/hive/FileSplitReader.h @@ -41,7 +41,7 @@ class ConnectorQueryCtx; } // namespace facebook::velox::connector namespace facebook::velox::dwio::common { -struct RuntimeStatistics; +struct RuntimeStats; } // namespace facebook::velox::dwio::common namespace facebook::velox::memory { @@ -112,7 +112,7 @@ class FileSplitReader { /// would be called only once per incoming split virtual void prepareSplit( std::shared_ptr metadataFilter, - dwio::common::RuntimeStatistics& runtimeStats, + dwio::common::RuntimeStats& runtimeStats, const folly::F14FastMap& fileReadOps = {}); virtual uint64_t next(uint64_t size, VectorPtr& output); @@ -125,7 +125,7 @@ class FileSplitReader { int64_t estimatedRowSize() const; - void updateRuntimeStats(dwio::common::RuntimeStatistics& stats) const; + void updateRuntimeStats(dwio::common::RuntimeStats& stats) const; bool allPrefetchIssued() const; @@ -165,7 +165,7 @@ class FileSplitReader { // Check if the filters pass on the column statistics. When delta update is // present, the corresonding filter should be disabled before calling this // function. - bool filterOnStats(dwio::common::RuntimeStatistics& runtimeStats) const; + bool filterOnStats(dwio::common::RuntimeStats& runtimeStats) const; /// Check if the fileSplit_ is empty. The split is considered empty when /// 1) The data file is missing but the user chooses to ignore it @@ -173,7 +173,7 @@ class FileSplitReader { /// 3) The data in the file does not pass the filters. The test is based on /// the file metadata and partition key values /// This function needs to be called after baseReader_ is created. - bool checkIfSplitIsEmpty(dwio::common::RuntimeStatistics& runtimeStats); + bool checkIfSplitIsEmpty(dwio::common::RuntimeStats& runtimeStats); /// Create the dwio::common::RowReader object baseRowReader_, which owns the /// ColumnReaders that will be used to read the data diff --git a/velox/connectors/hive/HiveConnectorSplit.cpp b/velox/connectors/hive/HiveConnectorSplit.cpp index 125b0f1f6ac..10f3418e305 100644 --- a/velox/connectors/hive/HiveConnectorSplit.cpp +++ b/velox/connectors/hive/HiveConnectorSplit.cpp @@ -89,14 +89,7 @@ folly::dynamic HiveConnectorSplit::serialize() const { obj["infoColumns"] = infoColumnsObj; if (properties.has_value()) { - folly::dynamic propertiesObj = folly::dynamic::object; - propertiesObj["fileSize"] = properties->fileSize.has_value() - ? folly::dynamic(properties->fileSize.value()) - : nullptr; - propertiesObj["modificationTime"] = properties->modificationTime.has_value() - ? folly::dynamic(properties->modificationTime.value()) - : nullptr; - obj["properties"] = propertiesObj; + obj["properties"] = properties->serialize(); } if (rowIdProperties.has_value()) { @@ -106,6 +99,10 @@ folly::dynamic HiveConnectorSplit::serialize() const { rowIdObj["tableGuid"] = rowIdProperties->tableGuid; obj["rowIdProperties"] = rowIdObj; } + if (columnMappingMode.has_value()) { + obj["columnMappingMode"] = + dwio::common::ColumnMappingModeName::toName(*columnMappingMode); + } return obj; } @@ -173,13 +170,7 @@ std::shared_ptr HiveConnectorSplit::create( std::optional properties = std::nullopt; const auto& propertiesObj = obj.getDefault("properties", nullptr); if (propertiesObj != nullptr) { - properties = FileProperties{ - .fileSize = propertiesObj["fileSize"].isNull() - ? std::nullopt - : std::optional(propertiesObj["fileSize"].asInt()), - .modificationTime = propertiesObj["modificationTime"].isNull() - ? std::nullopt - : std::optional(propertiesObj["modificationTime"].asInt())}; + properties = FileProperties::create(propertiesObj); } std::optional rowIdProperties = std::nullopt; @@ -191,6 +182,19 @@ std::shared_ptr HiveConnectorSplit::create( .tableGuid = rowIdObj["tableGuid"].asString()}; } + std::optional columnMappingMode = + std::nullopt; + if (auto it = obj.find("columnMappingMode"); it != obj.items().end()) { + auto parsedColumnMappingMode = + dwio::common::ColumnMappingModeName::tryToColumnMappingMode( + it->second.asString()); + VELOX_USER_CHECK( + parsedColumnMappingMode.has_value(), + "Invalid HiveConnectorSplit column mapping mode: {}", + it->second.asString()); + columnMappingMode = *parsedColumnMappingMode; + } + return std::make_shared( connectorId, filePath, @@ -207,7 +211,8 @@ std::shared_ptr HiveConnectorSplit::create( infoColumns, properties, rowIdProperties, - bucketConversion); + bucketConversion, + columnMappingMode); } // static diff --git a/velox/connectors/hive/HiveConnectorSplit.h b/velox/connectors/hive/HiveConnectorSplit.h index c71ffc9b140..772dc37f271 100644 --- a/velox/connectors/hive/HiveConnectorSplit.h +++ b/velox/connectors/hive/HiveConnectorSplit.h @@ -74,6 +74,8 @@ struct HiveConnectorSplit : public FileConnectorSplit { std::optional _properties = std::nullopt, std::optional _rowIdProperties = std::nullopt, const std::optional& _bucketConversion = + std::nullopt, + std::optional _columnMappingMode = std::nullopt) : FileConnectorSplit( connectorId, @@ -84,7 +86,8 @@ struct HiveConnectorSplit : public FileConnectorSplit { splitWeight, cacheable, std::move(_properties), - _partitionKeys), + _partitionKeys, + _columnMappingMode), infoColumns(_infoColumns), serdeParameters(_serdeParameters), tableBucketNumber(_tableBucketNumber), @@ -196,6 +199,12 @@ class HiveConnectorSplitBuilder { return *this; } + HiveConnectorSplitBuilder& columnMappingMode( + dwio::common::ColumnMappingMode mode) { + columnMappingMode_ = mode; + return *this; + } + HiveConnectorSplitBuilder& batchSizeHint(int32_t hint) { batchSizeHint_ = hint; return *this; @@ -218,7 +227,8 @@ class HiveConnectorSplitBuilder { infoColumns_, fileProperties_, rowIdProperties_, - bucketConversion_); + bucketConversion_, + columnMappingMode_); split->batchSizeHint = batchSizeHint_; return split; } @@ -240,6 +250,7 @@ class HiveConnectorSplitBuilder { bool cacheable_{true}; std::optional fileProperties_; std::optional rowIdProperties_ = std::nullopt; + std::optional columnMappingMode_; int32_t batchSizeHint_{0}; }; diff --git a/velox/connectors/hive/HivePartitionFunction.cpp b/velox/connectors/hive/HivePartitionFunction.cpp index 548aff99e15..32aac03057a 100644 --- a/velox/connectors/hive/HivePartitionFunction.cpp +++ b/velox/connectors/hive/HivePartitionFunction.cpp @@ -17,128 +17,12 @@ #include +#include "velox/functions/lib/HiveHash.h" + namespace facebook::velox::connector::hive { namespace { -void mergeHash(bool mix, uint32_t oneHash, uint32_t& aggregateHash) { - aggregateHash = mix ? aggregateHash * 31 + oneHash : oneHash; -} - -int32_t hashInt64(int64_t value) { - return ((*reinterpret_cast(&value)) >> 32) ^ value; -} - -#if defined(__has_feature) -#if __has_feature(__address_sanitizer__) -__attribute__((no_sanitize("integer"))) -#endif -#endif -uint32_t hashBytes(StringView bytes, int32_t initialValue) { - uint32_t hash = initialValue; - auto* data = bytes.data(); - for (auto i = 0; i < bytes.size(); ++i) { - hash = hash * 31 + *reinterpret_cast(data + i); - } - return hash; -} - -int32_t hashTimestamp(const Timestamp& ts) { - return hashInt64((ts.getSeconds() << 30) | ts.getNanos()); -} - -template -inline uint32_t hashOne(typename TypeTraits::NativeType /* value */) { - VELOX_UNSUPPORTED( - "Hive partitioning function doesn't support {} type", - TypeTraits::name); - return 0; // Make compiler happy. -} - -template <> -inline uint32_t hashOne(bool value) { - return value ? 1 : 0; -} - -template <> -inline uint32_t hashOne(int8_t value) { - return static_cast(value); -} - -template <> -inline uint32_t hashOne(int16_t value) { - return static_cast(value); -} - -template <> -inline uint32_t hashOne(int32_t value) { - return static_cast(value); -} - -template <> -inline uint32_t hashOne(float value) { - return static_cast(*reinterpret_cast(&value)); -} - -template <> -inline uint32_t hashOne(int64_t value) { - return hashInt64(value); -} - -template <> -inline uint32_t hashOne(double value) { - return hashInt64(*reinterpret_cast(&value)); -} - -template <> -inline uint32_t hashOne(StringView value) { - return hashBytes(value, 0); -} - -template <> -inline uint32_t hashOne(StringView value) { - return hashBytes(value, 0); -} - -template <> -inline uint32_t hashOne(Timestamp value) { - return hashTimestamp(value); -} - -template <> -inline uint32_t hashOne(UnknownValue /*value*/) { - VELOX_FAIL("Unknown values cannot be non-NULL"); -} - -template -void hashPrimitive( - const DecodedVector& values, - const SelectivityVector& rows, - bool mix, - std::vector& hashes) { - if (rows.isAllSelected()) { - // The compiler seems to be a little fickle with optimizations. - // Although rows.applyToSelected should do roughly the same thing, doing - // this here along with assigning rows.size() to a variable seems to help - // the compiler to inline hashOne showing a 50% performance improvement in - // benchmarks. - vector_size_t numRows = rows.size(); - for (auto i = 0; i < numRows; ++i) { - const uint32_t hash = values.isNullAt(i) - ? 0 - : hashOne( - values.valueAt::NativeType>(i)); - mergeHash(mix, hash, hashes[i]); - } - } else { - rows.applyToSelected([&](auto row) INLINE_LAMBDA { - const uint32_t hash = values.isNullAt(row) - ? 0 - : hashOne( - values.valueAt::NativeType>(row)); - mergeHash(mix, hash, hashes[row]); - }); - } -} +using facebook::velox::functions::HiveHash; void hashPrecomputed( uint32_t precomputedHash, @@ -158,7 +42,7 @@ void HivePartitionFunction::hashTyped( bool mix, std::vector& hashes, size_t /* poolIndex */) { - hashPrimitive(values, rows, mix, hashes); + HiveHash::hashPrimitive(values, rows, mix, hashes); } template <> @@ -168,7 +52,7 @@ void HivePartitionFunction::hashTyped( bool mix, std::vector& hashes, size_t /* poolIndex */) { - hashPrimitive(values, rows, mix, hashes); + HiveHash::hashPrimitive(values, rows, mix, hashes); } template <> @@ -178,7 +62,7 @@ void HivePartitionFunction::hashTyped( bool mix, std::vector& hashes, size_t /* poolIndex */) { - hashPrimitive(values, rows, mix, hashes); + HiveHash::hashPrimitive(values, rows, mix, hashes); } template <> @@ -188,7 +72,7 @@ void HivePartitionFunction::hashTyped( bool mix, std::vector& hashes, size_t /* poolIndex */) { - hashPrimitive(values, rows, mix, hashes); + HiveHash::hashPrimitive(values, rows, mix, hashes); } template <> @@ -198,7 +82,7 @@ void HivePartitionFunction::hashTyped( bool mix, std::vector& hashes, size_t /* poolIndex */) { - hashPrimitive(values, rows, mix, hashes); + HiveHash::hashPrimitive(values, rows, mix, hashes); } template <> @@ -208,7 +92,7 @@ void HivePartitionFunction::hashTyped( bool mix, std::vector& hashes, size_t /* poolIndex */) { - hashPrimitive(values, rows, mix, hashes); + HiveHash::hashPrimitive(values, rows, mix, hashes); } template <> @@ -218,7 +102,7 @@ void HivePartitionFunction::hashTyped( bool mix, std::vector& hashes, size_t /* poolIndex */) { - hashPrimitive(values, rows, mix, hashes); + HiveHash::hashPrimitive(values, rows, mix, hashes); } template <> @@ -228,7 +112,7 @@ void HivePartitionFunction::hashTyped( bool mix, std::vector& hashes, size_t /* poolIndex */) { - hashPrimitive(values, rows, mix, hashes); + HiveHash::hashPrimitive(values, rows, mix, hashes); } template <> @@ -238,7 +122,7 @@ void HivePartitionFunction::hashTyped( bool mix, std::vector& hashes, size_t /* poolIndex */) { - hashPrimitive(values, rows, mix, hashes); + HiveHash::hashPrimitive(values, rows, mix, hashes); } template <> @@ -248,7 +132,7 @@ void HivePartitionFunction::hashTyped( bool mix, std::vector& hashes, size_t /* poolIndex */) { - hashPrimitive(values, rows, mix, hashes); + HiveHash::hashPrimitive(values, rows, mix, hashes); } template <> @@ -258,7 +142,7 @@ void HivePartitionFunction::hashTyped( bool mix, std::vector& hashes, size_t /* poolIndex */) { - hashPrimitive(values, rows, mix, hashes); + HiveHash::hashPrimitive(values, rows, mix, hashes); } template <> @@ -318,11 +202,11 @@ void HivePartitionFunction::hashTyped( const auto length = arrayVector->sizeAt(index); for (size_t i = offset; i < offset + length; ++i) { - mergeHash(true, elementsHashes[i], hash); + HiveHash::mergeHash(true, elementsHashes[i], hash); } } - mergeHash(mix, hash, hashes[row]); + HiveHash::mergeHash(mix, hash, hashes[row]); }); } @@ -389,7 +273,7 @@ void HivePartitionFunction::hashTyped( } } - mergeHash(mix, hash, hashes[row]); + HiveHash::mergeHash(mix, hash, hashes[row]); }); } @@ -429,7 +313,7 @@ void HivePartitionFunction::hashTyped( } rows.applyToSelected([&](auto row) { - mergeHash( + HiveHash::mergeHash( mix, values.isNullAt(row) ? 0 : childHashes[values.index(row)], hashes[row]); diff --git a/velox/connectors/hive/HiveSplitReader.cpp b/velox/connectors/hive/HiveSplitReader.cpp index 98a56dc1064..1744849f283 100644 --- a/velox/connectors/hive/HiveSplitReader.cpp +++ b/velox/connectors/hive/HiveSplitReader.cpp @@ -64,7 +64,7 @@ HiveSplitReader::HiveSplitReader( void HiveSplitReader::prepareSplit( std::shared_ptr metadataFilter, - dwio::common::RuntimeStatistics& runtimeStats, + dwio::common::RuntimeStats& runtimeStats, const folly::F14FastMap& fileReadOps) { validateSynthesizedColumnFilters(); FileSplitReader::prepareSplit( diff --git a/velox/connectors/hive/HiveSplitReader.h b/velox/connectors/hive/HiveSplitReader.h index e50e87ed285..f6eeaa77d91 100644 --- a/velox/connectors/hive/HiveSplitReader.h +++ b/velox/connectors/hive/HiveSplitReader.h @@ -51,7 +51,7 @@ class HiveSplitReader : public FileSplitReader { void prepareSplit( std::shared_ptr metadataFilter, - dwio::common::RuntimeStatistics& runtimeStats, + dwio::common::RuntimeStats& runtimeStats, const folly::F14FastMap& fileReadOps = {}) override; diff --git a/velox/connectors/hive/PartitionValue.cpp b/velox/connectors/hive/PartitionValue.cpp new file mode 100644 index 00000000000..f55487389ac --- /dev/null +++ b/velox/connectors/hive/PartitionValue.cpp @@ -0,0 +1,87 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/connectors/hive/PartitionValue.h" + +#include +#include + +#include + +#include "velox/common/base/Exceptions.h" +#include "velox/type/Conversions.h" +#include "velox/type/DecimalUtil.h" + +namespace facebook::velox::connector::hive { +namespace { + +template +Variant fromStringImpl( + std::string_view value, + const Type& type, + PartitionValue::TimestampMode timestampMode, + PartitionValue::DateMode dateMode) { + using NativeType = typename TypeTraits::NativeType; + + if (type.isDate()) { + const auto days = dateMode == PartitionValue::DateMode::kDaysSinceEpoch + ? folly::to(value) + : DATE()->toDays(value); + return Variant(days); + } + + if constexpr ( + std::is_same_v || + std::is_same_v) { + if (type.isDecimal()) { + NativeType decimalValue{0}; + const auto [precision, scale] = getDecimalPrecisionScale(type); + const auto status = DecimalUtil::castFromString( + StringView(value), precision, scale, decimalValue); + VELOX_USER_CHECK(status.ok(), "{}", status.message()); + return Variant::create(decimalValue); + } + } + + if constexpr (std::is_same_v) { + return Variant::create(std::string(value)); + } else { + auto converted = util::Converter::tryCast(value).thenOrThrow( + folly::identity, + [&](const Status& status) { VELOX_USER_FAIL("{}", status.message()); }); + if constexpr (kind == TypeKind::TIMESTAMP) { + if (type.equivalent(*TIMESTAMP()) && + timestampMode == PartitionValue::TimestampMode::kLocalTime) { + converted.toGMT(Timestamp::defaultTimezone()); + } + } + return Variant::create(converted); + } +} + +} // namespace + +// static +Variant PartitionValue::fromString( + std::string_view value, + const Type& type, + TimestampMode timestampMode, + DateMode dateMode) { + return VELOX_DYNAMIC_SCALAR_TYPE_DISPATCH( + fromStringImpl, type.kind(), value, type, timestampMode, dateMode); +} + +} // namespace facebook::velox::connector::hive diff --git a/velox/connectors/hive/PartitionValue.h b/velox/connectors/hive/PartitionValue.h new file mode 100644 index 00000000000..fbd71b77d61 --- /dev/null +++ b/velox/connectors/hive/PartitionValue.h @@ -0,0 +1,73 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include "velox/type/Type.h" +#include "velox/type/Variant.h" + +namespace facebook::velox::connector::hive { + +/// Converts partition key strings to typed values. TIMESTAMP and DATE have +/// more than one encoding, so the caller states which one applies. +/// +/// auto value = PartitionValue::fromString( +/// "2020-01-01 12:34:56", +/// *TIMESTAMP(), +/// PartitionValue::TimestampMode::kLocalTime, +/// PartitionValue::DateMode::kIsoString); +class PartitionValue { + public: + enum class TimestampMode { + /// Interprets the value as local time and shifts it to UTC. A + /// TIMESTAMP_UTC value is not shifted. + kLocalTime, + + /// Interprets the value as UTC. No shift. + kUtc, + }; + + enum class DateMode { + /// Parses an ISO date, for example "2020-01-02". + kIsoString, + + /// Parses an integer count of days since the epoch, for example "18263". + kDaysSinceEpoch, + }; + + /// 'value' must be non-null. Accepted input per type: + /// - BOOLEAN: t, f, 1, 0, true or false, case-insensitively. + /// - TINYINT, SMALLINT, INTEGER, BIGINT: an integer, range-checked against + /// the native type rather than widened. + /// - REAL, DOUBLE: a floating point literal. + /// - DECIMAL: a decimal literal, scaled by the type's scale. + /// - VARCHAR, VARBINARY: taken verbatim. + /// - TIMESTAMP: parsed as TimestampParseMode::kPrestoCast, then shifted per + /// 'timestampMode'. + /// - DATE: parsed per 'dateMode'. + /// + /// Fails for a non-scalar type, and for a value that does not parse as + /// 'type'. + static Variant fromString( + std::string_view value, + const Type& type, + TimestampMode timestampMode, + DateMode dateMode); +}; + +} // namespace facebook::velox::connector::hive diff --git a/velox/connectors/hive/TableHandle.cpp b/velox/connectors/hive/TableHandle.cpp index fc8b26508c6..3eff3511b9d 100644 --- a/velox/connectors/hive/TableHandle.cpp +++ b/velox/connectors/hive/TableHandle.cpp @@ -177,11 +177,10 @@ folly::dynamic HiveColumnHandle::serialize() const { return obj; } -std::string HiveColumnHandle::toString() const { +std::string HiveColumnHandle::toStringFields() const { std::ostringstream out; out << fmt::format( - "HiveColumnHandle [name: {}, columnType: {}, dataType: {},", - name_, + "columnType: {}, dataType: {},", columnTypeName(columnType_), dataType_->toString()); out << " requiredSubfields: ["; @@ -238,10 +237,14 @@ std::string HiveColumnHandle::toString() const { } out << "]"; } - out << "]"; return out.str(); } +std::string HiveColumnHandle::toString() const { + return fmt::format( + "HiveColumnHandle [name: {}, {}]", name_, toStringFields()); +} + ColumnHandlePtr HiveColumnHandle::create(const folly::dynamic& obj) { auto name = obj["hiveColumnHandleName"].asString(); auto columnType = columnTypeFromName(obj["columnType"].asString()); @@ -286,19 +289,26 @@ HiveTableHandle::HiveTableHandle( const std::unordered_map& tableParameters, std::vector filterColumnHandles, double sampleRate, - std::string dbName) + std::string dbName, + std::vector dataColumnFieldIds) : FileTableHandle(std::move(connectorId)), tableName_(tableName), subfieldFilters_(std::move(subfieldFilters)), remainingFilter_(remainingFilter), sampleRate_(sampleRate), dataColumns_(dataColumns), + dataColumnFieldIds_(std::move(dataColumnFieldIds)), indexColumns_(std::move(indexColumns)), tableParameters_(tableParameters), filterColumnHandles_(std::move(filterColumnHandles)), dbName_(std::move(dbName)) { VELOX_CHECK_GT(sampleRate_, 0.0, "Sample rate must be positive"); VELOX_CHECK_LE(sampleRate_, 1.0, "Sample rate must not exceed 1.0"); + VELOX_CHECK( + dataColumnFieldIds_.empty() || + (dataColumns_ != nullptr && + dataColumnFieldIds_.size() == dataColumns_->size()), + "Data column field IDs must be empty or aligned to data columns"); } HiveTableHandle::HiveTableHandle( @@ -380,8 +390,9 @@ std::string HiveTableHandle::toString() const { return out.str(); } -folly::dynamic HiveTableHandle::serialize() const { - folly::dynamic obj = ConnectorTableHandle::serializeBase("HiveTableHandle"); +folly::dynamic HiveTableHandle::serializeHiveFields( + const std::string& typeName) const { + folly::dynamic obj = ConnectorTableHandle::serializeBase(typeName); obj["tableName"] = tableName_; folly::dynamic subfieldFilters = folly::dynamic::array; @@ -391,8 +402,8 @@ folly::dynamic HiveTableHandle::serialize() const { pair["filter"] = filter->serialize(); subfieldFilters.push_back(pair); } - obj["subfieldFilters"] = subfieldFilters; + if (remainingFilter_) { obj["remainingFilter"] = remainingFilter_->serialize(); } @@ -404,11 +415,21 @@ folly::dynamic HiveTableHandle::serialize() const { if (dataColumns_) { obj["dataColumns"] = dataColumns_->serialize(); } + + if (!dataColumnFieldIds_.empty()) { + folly::dynamic dataColumnFieldIds = folly::dynamic::array; + for (const auto fieldId : dataColumnFieldIds_) { + dataColumnFieldIds.push_back(fieldId); + } + obj["dataColumnFieldIds"] = std::move(dataColumnFieldIds); + } + folly::dynamic tableParameters = folly::dynamic::object; for (const auto& param : tableParameters_) { tableParameters[param.first] = param.second; } obj["tableParameters"] = tableParameters; + if (!filterColumnHandles_.empty()) { folly::dynamic filterColumnHandles = folly::dynamic::array; for (const auto& handle : filterColumnHandles_) { @@ -416,6 +437,7 @@ folly::dynamic HiveTableHandle::serialize() const { } obj["filterColumnHandles"] = filterColumnHandles; } + if (!indexColumns_.empty()) { folly::dynamic indexColumns = folly::dynamic::array; for (const auto& column : indexColumns_) { @@ -431,46 +453,61 @@ folly::dynamic HiveTableHandle::serialize() const { return obj; } -ConnectorTableHandlePtr HiveTableHandle::create( +folly::dynamic HiveTableHandle::serialize() const { + return serializeHiveFields("HiveTableHandle"); +} + +// static +void HiveTableHandle::deserializeHiveFields( const folly::dynamic& obj, - void* context) { - auto connectorId = obj["connectorId"].asString(); - auto tableName = obj["tableName"].asString(); + void* context, + std::string& connectorId, + std::string& tableName, + common::SubfieldFilters& subfieldFilters, + core::TypedExprPtr& remainingFilter, + double& sampleRate, + RowTypePtr& dataColumns, + std::unordered_map& tableParameters, + std::vector& filterColumnHandles, + std::vector& indexColumns, + std::string& dbName, + std::vector& dataColumnFieldIds) { + connectorId = obj["connectorId"].asString(); + tableName = obj["tableName"].asString(); - core::TypedExprPtr remainingFilter; if (auto it = obj.find("remainingFilter"); it != obj.items().end()) { remainingFilter = ISerializable::deserialize(it->second, context); } - common::SubfieldFilters subfieldFilters; - folly::dynamic subfieldFiltersObj = obj["subfieldFilters"]; - for (const auto& subfieldFilter : subfieldFiltersObj) { - common::Subfield subfield(subfieldFilter["subfield"].asString()); - auto filter = - ISerializable::deserialize(subfieldFilter["filter"]); + for (const auto& entry : obj["subfieldFilters"]) { + common::Subfield subfield(entry["subfield"].asString()); + auto filter = ISerializable::deserialize(entry["filter"]); subfieldFilters[common::Subfield(std::move(subfield.path()))] = filter->clone(); } - double sampleRate = 1.0; + sampleRate = 1.0; if (obj.count("sampleRate")) { sampleRate = obj["sampleRate"].asDouble(); } - RowTypePtr dataColumns; if (auto it = obj.find("dataColumns"); it != obj.items().end()) { dataColumns = ISerializable::deserialize(it->second, context); } - std::unordered_map tableParameters{}; + if (auto it = obj.find("dataColumnFieldIds"); it != obj.items().end()) { + dataColumnFieldIds.reserve(it->second.size()); + for (const auto& fieldId : it->second) { + dataColumnFieldIds.push_back(fieldId.asInt()); + } + } + const auto& tableParametersObj = obj["tableParameters"]; for (const auto& key : tableParametersObj.keys()) { - const auto& value = tableParametersObj[key]; - tableParameters.emplace(key.asString(), value.asString()); + tableParameters.emplace(key.asString(), tableParametersObj[key].asString()); } - std::vector filterColumnHandles; if (auto it = obj.find("filterColumnHandles"); it != obj.items().end()) { for (const auto& handle : it->second) { filterColumnHandles.push_back( @@ -478,21 +515,50 @@ ConnectorTableHandlePtr HiveTableHandle::create( } } - std::vector indexColumns; if (auto it = obj.find("indexColumns"); it != obj.items().end()) { - for (const auto& column : it->second) { - indexColumns.push_back(column.asString()); + for (const auto& col : it->second) { + indexColumns.push_back(col.asString()); } } - std::string dbName; if (auto it = obj.find("dbName"); it != obj.items().end()) { dbName = it->second.asString(); } +} - return std::make_shared( +ConnectorTableHandlePtr HiveTableHandle::create( + const folly::dynamic& obj, + void* context) { + std::string connectorId; + std::string tableName; + std::string dbName; + common::SubfieldFilters subfieldFilters; + core::TypedExprPtr remainingFilter; + double sampleRate{1.0}; + RowTypePtr dataColumns; + std::unordered_map tableParameters; + std::vector filterColumnHandles; + std::vector indexColumns; + std::vector dataColumnFieldIds; + + deserializeHiveFields( + obj, + context, connectorId, tableName, + subfieldFilters, + remainingFilter, + sampleRate, + dataColumns, + tableParameters, + filterColumnHandles, + indexColumns, + dbName, + dataColumnFieldIds); + + return std::make_shared( + std::move(connectorId), + tableName, std::move(subfieldFilters), remainingFilter, dataColumns, @@ -500,7 +566,8 @@ ConnectorTableHandlePtr HiveTableHandle::create( tableParameters, std::move(filterColumnHandles), sampleRate, - std::move(dbName)); + std::move(dbName), + std::move(dataColumnFieldIds)); } void HiveTableHandle::registerSerDe() { diff --git a/velox/connectors/hive/TableHandle.h b/velox/connectors/hive/TableHandle.h index c0d5463c535..411510b21f8 100644 --- a/velox/connectors/hive/TableHandle.h +++ b/velox/connectors/hive/TableHandle.h @@ -159,6 +159,12 @@ class HiveColumnHandle : public FileColumnHandle { static void registerSerDe(); + protected: + // Emits the common column fields (columnType, dataType, requiredSubfields, + // extractions) without a class-name prefix. Subclasses use this to build + // their own toString() with the correct class name in the header. + std::string toStringFields() const; + private: const std::string name_; const ColumnType columnType_; @@ -188,7 +194,8 @@ class HiveTableHandle : public FileTableHandle { const std::unordered_map& tableParameters = {}, std::vector filterColumnHandles = {}, double sampleRate = 1.0, - std::string dbName = ""); + std::string dbName = "", + std::vector dataColumnFieldIds = {}); /// Legacy constructor without indexColumns parameter for backward /// compatibility. @@ -234,6 +241,12 @@ class HiveTableHandle : public FileTableHandle { return dataColumns_; } + /// Returns Iceberg field IDs aligned positionally to dataColumns(). An empty + /// vector means field IDs are unavailable. + const std::vector& dataColumnFieldIds() const { + return dataColumnFieldIds_; + } + /// Returns the names of the index columns for the table. const std::vector& indexColumns() const { return indexColumns_; @@ -269,12 +282,39 @@ class HiveTableHandle : public FileTableHandle { static void registerSerDe(); + protected: + // Serializes the Hive-common fields (tableName, subfieldFilters, + // remainingFilter, sampleRate, dataColumns, tableParameters, + // filterColumnHandles, indexColumns, dbName) into an object whose "name" + // key is set to @p typeName. Subclasses call this instead of duplicating + // the field list, then append their own keys. + folly::dynamic serializeHiveFields(const std::string& typeName) const; + + // Fills the common Hive fields from @p obj into the provided out-parameters. + // Subclasses call this from their own create() and then parse only their + // extra keys on top. + static void deserializeHiveFields( + const folly::dynamic& obj, + void* context, + std::string& connectorId, + std::string& tableName, + common::SubfieldFilters& subfieldFilters, + core::TypedExprPtr& remainingFilter, + double& sampleRate, + RowTypePtr& dataColumns, + std::unordered_map& tableParameters, + std::vector& filterColumnHandles, + std::vector& indexColumns, + std::string& dbName, + std::vector& dataColumnFieldIds); + private: const std::string tableName_; const common::SubfieldFilters subfieldFilters_; const core::TypedExprPtr remainingFilter_; const double sampleRate_; const RowTypePtr dataColumns_; + const std::vector dataColumnFieldIds_; const std::vector indexColumns_; const std::unordered_map tableParameters_; const std::vector filterColumnHandles_; diff --git a/velox/connectors/hive/iceberg/CMakeLists.txt b/velox/connectors/hive/iceberg/CMakeLists.txt index 0dbf46e6bcd..2f7581b6a8a 100644 --- a/velox/connectors/hive/iceberg/CMakeLists.txt +++ b/velox/connectors/hive/iceberg/CMakeLists.txt @@ -19,6 +19,7 @@ set( IcebergColumnHandle.cpp IcebergConfig.cpp IcebergConnector.cpp + IcebergTableHandle.cpp IcebergDataFileStatistics.cpp IcebergDataSink.cpp IcebergDataSource.cpp @@ -53,6 +54,7 @@ velox_add_library( IcebergColumnHandle.h IcebergConfig.h IcebergConnector.h + IcebergTableHandle.h IcebergDataFileStatistics.h IcebergDataSink.h IcebergDataSource.h diff --git a/velox/connectors/hive/iceberg/DeletionVectorFormat.h b/velox/connectors/hive/iceberg/DeletionVectorFormat.h index 670616e0e7b..bf0b913fd61 100644 --- a/velox/connectors/hive/iceberg/DeletionVectorFormat.h +++ b/velox/connectors/hive/iceberg/DeletionVectorFormat.h @@ -17,6 +17,7 @@ #pragma once #include +#include namespace facebook::velox::connector::hive::iceberg { @@ -35,4 +36,25 @@ inline constexpr size_t kDeletionVectorMagicSize = 4; inline constexpr size_t kDeletionVectorLengthSize = 4; inline constexpr size_t kDeletionVectorCrcSize = 4; +/// Largest Roaring64 group key (the high 32 bits of a position) the Iceberg +/// deletion-vector format can represent. +/// +/// Iceberg's RoaringPositionBitmap stores the key as a signed 32-bit int and +/// caps it one below Integer.MAX_VALUE, so a spec-compliant reader cannot load +/// a larger key back. Keys at or above 2^31 are worse still: `key << 32` +/// shifts into the sign bit of the int64 position and yields a negative row +/// ordinal. Both the writer (via DeletionVectorWriter::kMaxPosition) and the +/// reader enforce this so neither can accept a blob the other would reject. +inline constexpr uint32_t kMaxRoaring64GroupKey = 2'147'483'646; + +/// Puffin container format constants, shared by the deletion-vector writer +/// (which emits the footer) and the reader (which parses it when the manifest +/// carries no blob offset). Magic is "PFA1" and brackets both the file and the +/// footer payload. +inline constexpr char kPuffinMagic[] = {'\x50', '\x46', '\x41', '\x31'}; +inline constexpr size_t kPuffinMagicSize = 4; + +/// Blob type of an Iceberg V3 deletion vector inside a Puffin file. +inline constexpr char kDeletionVectorBlobType[] = "deletion-vector-v1"; + } // namespace facebook::velox::connector::hive::iceberg diff --git a/velox/connectors/hive/iceberg/DeletionVectorReader.cpp b/velox/connectors/hive/iceberg/DeletionVectorReader.cpp index a36c78fb0ee..4d797731569 100644 --- a/velox/connectors/hive/iceberg/DeletionVectorReader.cpp +++ b/velox/connectors/hive/iceberg/DeletionVectorReader.cpp @@ -18,6 +18,7 @@ #include "velox/connectors/hive/iceberg/DeletionVectorFormat.h" +#include #include #include @@ -41,6 +42,146 @@ uint32_t readBigEndian32(const char* data) { return folly::Endian::big(value); } +uint32_t readLittleEndian32(const char* data) { + uint32_t value; + std::memcpy(&value, data, sizeof(value)); + return folly::Endian::little(value); +} + +// Byte range of a blob inside a Puffin file. +struct BlobLocation { + uint64_t offset; + uint64_t length; +}; + +// Reads the Puffin magic that both opens the file and brackets the footer. +bool startsWithPuffinMagic(ReadFile& file) { + if (file.size() < kPuffinMagicSize) { + return false; + } + std::string magic(kPuffinMagicSize, '\0'); + file.pread(0, kPuffinMagicSize, magic.data()); + return std::string_view(magic) == + std::string_view(kPuffinMagic, kPuffinMagicSize); +} + +// Locates the deletion-vector blob by parsing the Puffin footer, for files +// whose manifest entry carried no blob offset or length. The footer is: +// Magic FooterPayload FooterPayloadSize Flags Magic +// with FooterPayloadSize and Flags each 4 bytes little-endian, read backwards +// from end of file. 'referencedDataFile' disambiguates when the file holds +// vectors for several data files; an empty value requires a single candidate. +BlobLocation locateBlobFromPuffinFooter( + ReadFile& file, + const std::string& referencedDataFile) { + // FooterPayloadSize + Flags + trailing Magic. + constexpr uint64_t kTrailerSize = 4 + 4 + kPuffinMagicSize; + // Leading Magic and the Magic that opens the footer bracket any payload. + constexpr uint64_t kMinFileSize = kTrailerSize + 2 * kPuffinMagicSize; + + const uint64_t fileSize = file.size(); + VELOX_CHECK_GE( + fileSize, + kMinFileSize, + "Puffin file is too small to contain a footer: {} bytes.", + fileSize); + + std::string trailer(kTrailerSize, '\0'); + file.pread(fileSize - kTrailerSize, kTrailerSize, trailer.data()); + const std::string_view puffinMagic(kPuffinMagic, kPuffinMagicSize); + VELOX_CHECK_EQ( + std::string_view(trailer).substr(8), + puffinMagic, + "Puffin file does not end with the expected magic."); + + const uint32_t flags = readLittleEndian32(trailer.data() + 4); + // Bit 0 of the first flag byte marks a compressed footer payload. Deletion + // vectors are always written uncompressed, so this is unreachable for files + // we produce and unsupported for files we do not. + VELOX_CHECK_EQ( + flags & 1u, 0u, "Compressed Puffin footer payloads are not supported."); + + // payloadOffset must leave room for the leading magic and the magic that + // opens the footer, so payloadSize + kMinFileSize must fit within the file. + const uint64_t payloadSize = readLittleEndian32(trailer.data()); + VELOX_CHECK_LE( + payloadSize + kMinFileSize, + fileSize, + "Puffin footer payload size {} does not fit in a {}-byte file.", + payloadSize, + fileSize); + + const uint64_t payloadOffset = fileSize - kTrailerSize - payloadSize; + std::string footerMagic(kPuffinMagicSize, '\0'); + file.pread( + payloadOffset - kPuffinMagicSize, kPuffinMagicSize, footerMagic.data()); + VELOX_CHECK_EQ( + std::string_view(footerMagic), + puffinMagic, + "Puffin footer payload is not preceded by the expected magic."); + + std::string payload(payloadSize, '\0'); + file.pread(payloadOffset, payloadSize, payload.data()); + + folly::dynamic footer; + try { + footer = folly::parseJson(payload); + } catch (const std::exception& e) { + VELOX_FAIL("Failed to parse Puffin footer payload: {}", e.what()); + } + + const auto* blobs = footer.get_ptr("blobs"); + VELOX_CHECK( + blobs != nullptr && blobs->isArray(), + "Puffin footer has no \"blobs\" array."); + + std::optional found; + for (const auto& blob : *blobs) { + const auto* type = blob.get_ptr("type"); + if (type == nullptr || !type->isString() || + type->asString() != kDeletionVectorBlobType) { + continue; + } + if (!referencedDataFile.empty()) { + const auto* properties = blob.get_ptr("properties"); + const auto* referenced = properties == nullptr + ? nullptr + : properties->get_ptr("referenced-data-file"); + if (referenced == nullptr || !referenced->isString() || + referenced->asString() != referencedDataFile) { + continue; + } + } + const auto* offset = blob.get_ptr("offset"); + const auto* length = blob.get_ptr("length"); + VELOX_CHECK( + offset != nullptr && offset->isInt() && length != nullptr && + length->isInt(), + "Puffin blob metadata is missing a numeric offset or length."); + // Both are widened to uint64_t below, where a negative value would wrap to + // a huge offset. The file-bounds check downstream would still reject it, + // but only after reporting an absurd number, so reject it here instead. + VELOX_CHECK_GE( + offset->asInt(), 0, "Puffin blob metadata has a negative offset."); + VELOX_CHECK_GE( + length->asInt(), 0, "Puffin blob metadata has a negative length."); + VELOX_CHECK( + !found.has_value(), + "Puffin file has multiple deletion vectors and no referenced data " + "file to disambiguate them."); + found = BlobLocation{ + static_cast(offset->asInt()), + static_cast(length->asInt())}; + } + + VELOX_CHECK( + found.has_value(), + "Puffin footer has no deletion-vector blob for data file: {}", + referencedDataFile.empty() ? std::string_view{"(unspecified)"} + : std::string_view{referencedDataFile}); + return found.value(); +} + // Unwraps the Iceberg deletion-vector-v1 blob frame // ([length: 4B BE][magic][bitmap][CRC-32: 4B BE]), validating the magic and // CRC-32, and returns a view of the inner roaring bitmap. If 'blob' is not @@ -120,12 +261,16 @@ void DeletionVectorReader::loadBitmap() { uint64_t blobLength = dvFile_.contentLength > 0 ? static_cast(dvFile_.contentLength) : dvFile_.fileSizeInBytes; + bool haveBlobLocation = dvFile_.contentLength > 0; if (dvFile_.contentLength == 0) { + bool haveOffset = false; + bool haveLength = false; if (auto it = dvFile_.lowerBounds.find(kDvOffsetFieldId); it != dvFile_.lowerBounds.end()) { try { blobOffset = std::stoull(it->second); + haveOffset = true; } catch (const std::exception& e) { VELOX_FAIL( "Failed to parse DV blob offset from bounds map: {}", e.what()); @@ -135,11 +280,16 @@ void DeletionVectorReader::loadBitmap() { it != dvFile_.upperBounds.end()) { try { blobLength = std::stoull(it->second); + haveLength = true; } catch (const std::exception& e) { VELOX_FAIL( "Failed to parse DV blob length from bounds map: {}", e.what()); } } + // A legacy bounds-map location is usable only when both offset and length + // are present. A partial entry falls through to Puffin-footer parsing + // rather than unframing at a wrong offset/length. + haveBlobLocation = haveOffset && haveLength; } // Pass the connector config (e.g. hive.manifold.* credentials) so @@ -156,6 +306,17 @@ void DeletionVectorReader::loadBitmap() { readFile, "Failed to open deletion vector file: {}", dvFile_.filePath); auto fileSize = readFile->size(); + + // Nothing in the manifest located the blob. A Puffin file describes its own + // blobs, so parse the footer rather than guessing. Non-Puffin inputs keep + // the legacy whole-file behaviour, which is how raw roaring blobs are read. + if (!haveBlobLocation && startsWithPuffinMagic(*readFile)) { + const auto location = + locateBlobFromPuffinFooter(*readFile, dvFile_.referencedDataFile); + blobOffset = location.offset; + blobLength = location.length; + } + VELOX_CHECK_LE( blobOffset, fileSize, @@ -232,6 +393,13 @@ void DeletionVectorReader::deserializeRoaring64Bitmap(std::string_view data) { highBits = folly::Endian::little(highBits); ptr += sizeof(uint32_t); + VELOX_CHECK_LE( + highBits, + kMaxRoaring64GroupKey, + "Roaring64Bitmap group key exceeds the maximum the Iceberg " + "deletion-vector format can represent: {}", + highBits); + int64_t highBitsOffset = static_cast(highBits) << 32; // Deserialize the 32-bit bitmap for this group. diff --git a/velox/connectors/hive/iceberg/DeletionVectorWriter.cpp b/velox/connectors/hive/iceberg/DeletionVectorWriter.cpp index 222122445ec..6e983867823 100644 --- a/velox/connectors/hive/iceberg/DeletionVectorWriter.cpp +++ b/velox/connectors/hive/iceberg/DeletionVectorWriter.cpp @@ -40,15 +40,19 @@ constexpr size_t kBitmapContainerBytes = 8'192; constexpr size_t kBitmapContainerWords = 1'024; // Puffin file format constants (per Iceberg spec). Magic is "PFA1". -constexpr char kPuffinMagic[] = {'\x50', '\x46', '\x41', '\x31'}; -constexpr size_t kPuffinMagicSize = 4; constexpr uint32_t kPuffinFooterFlags = 0; // Puffin blob metadata constants (per Iceberg V3 deletion vector spec). -constexpr char kDeletionVectorBlobType[] = "deletion-vector-v1"; -constexpr char kCompressionCodecNone[] = "none"; -// Iceberg spec: source-field-id for whole-row deletes is INT_MAX - 1. -constexpr int32_t kWholeRowDeleteFieldId = 2'147'483'646; +// The blob's "fields" list names the row-position metadata column +// (MetadataColumns.ROW_POSITION, INT_MAX - 2), matching Iceberg's own +// BaseDVFileWriter. Note this is not the positional-delete file's "pos" +// column (INT_MAX - 102), which IcebergMetadataColumns uses for a different +// purpose. +constexpr int32_t kRowPositionFieldId = 2'147'483'645; +// A deletion vector is written before its snapshot is assigned, so Iceberg +// records -1 for both. The fields are required by readers regardless. +constexpr int64_t kUnassignedSnapshotId = -1; +constexpr int64_t kUnassignedSequenceNumber = -1; void writeLittleEndian(std::string& out, uint16_t val) { val = folly::Endian::little(val); @@ -159,6 +163,12 @@ void serializeContainerData( void DeletionVectorWriter::addDeletedPosition(int64_t position) { VELOX_CHECK_GE(position, 0, "Deleted position must be non-negative."); + VELOX_CHECK_LE( + position, + kMaxPosition, + "Deleted position exceeds the maximum the Iceberg deletion-vector " + "format can represent: {}", + position); positions_.push_back(position); } @@ -257,14 +267,17 @@ std::pair writePuffinFile( uint64_t blobOffset = kPuffinMagicSize; uint64_t blobLength = framedBlob.size(); - folly::dynamic blobMeta = folly::dynamic::object( - "type", kDeletionVectorBlobType)( - "fields", - folly::dynamic::array( - folly::dynamic::object("source-field-id", kWholeRowDeleteFieldId))); + folly::dynamic blobMeta = + folly::dynamic::object("type", kDeletionVectorBlobType)( + "fields", folly::dynamic::array(kRowPositionFieldId)); + blobMeta["snapshot-id"] = kUnassignedSnapshotId; + blobMeta["sequence-number"] = kUnassignedSequenceNumber; blobMeta["offset"] = blobOffset; blobMeta["length"] = blobLength; - blobMeta["compression-codec"] = kCompressionCodecNone; + // Uncompressed blobs omit "compression-codec": Iceberg's FileMetadataParser + // writes the field only for a non-null codec, and PuffinCompressionCodec + // .forName() has no "none" entry, so emitting it literally would make a + // spec-compliant reader throw. folly::dynamic properties = folly::dynamic::object; properties["referenced-data-file"] = referencedDataFile; diff --git a/velox/connectors/hive/iceberg/DeletionVectorWriter.h b/velox/connectors/hive/iceberg/DeletionVectorWriter.h index dd133561453..a3d1f2ccae7 100644 --- a/velox/connectors/hive/iceberg/DeletionVectorWriter.h +++ b/velox/connectors/hive/iceberg/DeletionVectorWriter.h @@ -20,6 +20,8 @@ #include #include +#include "velox/connectors/hive/iceberg/DeletionVectorFormat.h" + namespace facebook::velox::memory { class MemoryPool; } // namespace facebook::velox::memory @@ -52,9 +54,25 @@ namespace facebook::velox::connector::hive::iceberg { /// std::string blob = writer.serialize(); class DeletionVectorWriter { public: + /// Largest position the Iceberg deletion-vector format can represent. + /// + /// Iceberg's RoaringPositionBitmap derives this as + /// `toPosition(MAX_KEY, Integer.MIN_VALUE)`, i.e. + /// `(key << 32) | (pos32 & 0xFFFFFFFF)` with `key = 2147483646`, and rejects + /// anything above it. The binding constraint is the Roaring64 group key: it + /// is read back as a signed 32-bit int, so a key at or above 2^31 would + /// deserialize as negative and be rejected by spec-compliant readers. + /// Matching the bound here means we never write a blob Iceberg cannot read. + static constexpr int64_t kMaxPosition = 9'223'372'030'412'324'864LL; + + static_assert( + (kMaxPosition >> 32) == kMaxRoaring64GroupKey, + "Writer position bound and reader group-key bound must agree."); + DeletionVectorWriter() = default; - /// Adds a deleted row position (0-based file row offset). + /// Adds a deleted row position (0-based file row offset). The position must + /// be in [0, kMaxPosition]. void addDeletedPosition(int64_t position); /// Adds multiple deleted row positions. diff --git a/velox/connectors/hive/iceberg/EqualityDeleteFileReader.cpp b/velox/connectors/hive/iceberg/EqualityDeleteFileReader.cpp index 6d24875f375..c37e875562f 100644 --- a/velox/connectors/hive/iceberg/EqualityDeleteFileReader.cpp +++ b/velox/connectors/hive/iceberg/EqualityDeleteFileReader.cpp @@ -146,7 +146,7 @@ EqualityDeleteFileReader::EqualityDeleteFileReader( const std::shared_ptr& fileConfig, const std::shared_ptr& ioStatistics, const std::shared_ptr& ioStats, - dwio::common::RuntimeStatistics& runtimeStats, + dwio::common::RuntimeStats& runtimeStats, const std::string& connectorId) : equalityColumnNames_(equalityColumnNames), equalityColumnTypes_(equalityColumnTypes), diff --git a/velox/connectors/hive/iceberg/EqualityDeleteFileReader.h b/velox/connectors/hive/iceberg/EqualityDeleteFileReader.h index ad3a2be359a..1efbc05d5d2 100644 --- a/velox/connectors/hive/iceberg/EqualityDeleteFileReader.h +++ b/velox/connectors/hive/iceberg/EqualityDeleteFileReader.h @@ -78,7 +78,7 @@ class EqualityDeleteFileReader { const std::shared_ptr& fileConfig, const std::shared_ptr& ioStatistics, const std::shared_ptr& ioStats, - dwio::common::RuntimeStatistics& runtimeStats, + dwio::common::RuntimeStats& runtimeStats, const std::string& connectorId); /// Applies equality deletes to the output vector by setting bits in the diff --git a/velox/connectors/hive/iceberg/IcebergColumnHandle.cpp b/velox/connectors/hive/iceberg/IcebergColumnHandle.cpp index 8a5a02f7006..819fdacb7e6 100644 --- a/velox/connectors/hive/iceberg/IcebergColumnHandle.cpp +++ b/velox/connectors/hive/iceberg/IcebergColumnHandle.cpp @@ -14,6 +14,7 @@ * limitations under the License. */ +#include #include #include #include @@ -26,6 +27,144 @@ namespace facebook::velox::connector::hive::iceberg { +namespace { + +// Produces a compact string of the set attributes in an IcebergFieldMetadata +// node. Only set fields are emitted; returns an empty string when empty(). +std::string fieldMetadataToString(const IcebergFieldMetadata& meta) { + std::ostringstream out; + bool first = true; + auto append = [&](const std::string& key, const std::string& value) { + if (!first) { + out << ", "; + } + out << key << ": " << value; + first = false; + }; + if (meta.required.has_value()) { + append("required", *meta.required ? "true" : "false"); + } + if (meta.longType.has_value()) { + append("longType", *meta.longType); + } + if (meta.timestampUnit.has_value()) { + append("timestampUnit", *meta.timestampUnit); + } + if (meta.binaryType.has_value()) { + append("binaryType", *meta.binaryType); + } + if (meta.structType.has_value()) { + append("structType", *meta.structType); + } + if (meta.length.has_value()) { + append("length", std::to_string(*meta.length)); + } + return out.str(); +} + +// Produces a compact string representation of a ParquetFieldId tree. +// Format: for leaf nodes, [, , ...] for nested nodes. +std::string fieldIdToString(const parquet::ParquetFieldId& fieldId) { + std::string result = std::to_string(fieldId.fieldId); + if (!fieldId.children.empty()) { + result += "["; + for (size_t i = 0; i < fieldId.children.size(); ++i) { + if (i > 0) { + result += ", "; + } + result += fieldIdToString(fieldId.children[i]); + } + result += "]"; + } + return result; +} + +// Serializes a ParquetFieldId tree to a folly::dynamic object. +folly::dynamic serializeFieldId(const parquet::ParquetFieldId& fieldId) { + folly::dynamic obj = folly::dynamic::object; + obj["fieldId"] = fieldId.fieldId; + folly::dynamic children = folly::dynamic::array; + for (const auto& child : fieldId.children) { + children.push_back(serializeFieldId(child)); + } + obj["children"] = children; + return obj; +} + +// Deserializes a ParquetFieldId tree from a folly::dynamic object. +parquet::ParquetFieldId deserializeFieldId(const folly::dynamic& obj) { + parquet::ParquetFieldId fieldId; + fieldId.fieldId = static_cast(obj["fieldId"].asInt()); + for (const auto& child : obj["children"]) { + fieldId.children.push_back(deserializeFieldId(child)); + } + return fieldId; +} + +// Serializes an IcebergFieldMetadata node (and its children) to a +// folly::dynamic object. Only set optional fields are emitted; an all-empty +// node is serialized as an empty object so that the children array remains +// parallel to ParquetFieldId::children. +folly::dynamic serializeFieldMetadata(const IcebergFieldMetadata& meta) { + folly::dynamic obj = folly::dynamic::object; + if (meta.required.has_value()) { + obj["required"] = *meta.required; + } + if (meta.longType.has_value()) { + obj["longType"] = *meta.longType; + } + if (meta.timestampUnit.has_value()) { + obj["timestampUnit"] = *meta.timestampUnit; + } + if (meta.binaryType.has_value()) { + obj["binaryType"] = *meta.binaryType; + } + if (meta.structType.has_value()) { + obj["structType"] = *meta.structType; + } + if (meta.length.has_value()) { + obj["length"] = *meta.length; + } + folly::dynamic children = folly::dynamic::array; + for (const auto& child : meta.children) { + children.push_back(serializeFieldMetadata(child)); + } + obj["children"] = children; + return obj; +} + +// Deserializes an IcebergFieldMetadata node (and its children) from a +// folly::dynamic object. +IcebergFieldMetadata deserializeFieldMetadata(const folly::dynamic& obj) { + IcebergFieldMetadata meta; + if (auto it = obj.find("required"); it != obj.items().end()) { + meta.required = it->second.asBool(); + } + if (auto it = obj.find("longType"); it != obj.items().end()) { + meta.longType = it->second.asString(); + } + if (auto it = obj.find("timestampUnit"); it != obj.items().end()) { + meta.timestampUnit = it->second.asString(); + } + if (auto it = obj.find("binaryType"); it != obj.items().end()) { + meta.binaryType = it->second.asString(); + } + if (auto it = obj.find("structType"); it != obj.items().end()) { + meta.structType = it->second.asString(); + } + if (auto it = obj.find("length"); it != obj.items().end()) { + meta.length = static_cast(it->second.asInt()); + } + if (auto it = obj.find("children"); it != obj.items().end()) { + for (const auto& child : it->second) { + meta.children.push_back(deserializeFieldMetadata(child)); + } + } + return meta; +} + +} // namespace + IcebergColumnHandle::IcebergColumnHandle( const std::string& name, ColumnType columnType, @@ -52,4 +191,84 @@ const parquet::ParquetFieldId& IcebergColumnHandle::field() const { return field_; } +std::string IcebergColumnHandle::toString() const { + std::string fields = HiveColumnHandle::toStringFields(); + fields += ", field: " + fieldIdToString(field_); + if (initialDefaultValue_.has_value()) { + fields += ", initialDefaultValue: " + *initialDefaultValue_; + } + if (!icebergMetadata_.empty()) { + fields += + ", icebergMetadata: {" + fieldMetadataToString(icebergMetadata_) + "}"; + } + return fmt::format("IcebergColumnHandle [name: {}, {}]", name(), fields); +} + +folly::dynamic IcebergColumnHandle::serialize() const { + folly::dynamic obj = ColumnHandle::serializeBase("IcebergColumnHandle"); + obj["hiveColumnHandleName"] = name(); + obj["columnType"] = columnTypeName(columnType()); + obj["dataType"] = dataType()->serialize(); + + folly::dynamic requiredSubfieldsArr = folly::dynamic::array; + for (const auto& subfield : requiredSubfields()) { + requiredSubfieldsArr.push_back(subfield.toString()); + } + obj["requiredSubfields"] = requiredSubfieldsArr; + + obj["field"] = serializeFieldId(field_); + + if (initialDefaultValue_.has_value()) { + obj["initialDefaultValue"] = *initialDefaultValue_; + } + + // Only serialize icebergMetadata when at least one node carries a set + // attribute, to keep the serialized form compact for callers that never + // populate V3 metadata. + if (!icebergMetadata_.empty() || !icebergMetadata_.children.empty()) { + obj["icebergMetadata"] = serializeFieldMetadata(icebergMetadata_); + } + + return obj; +} + +// static +ColumnHandlePtr IcebergColumnHandle::create(const folly::dynamic& obj) { + auto name = obj["hiveColumnHandleName"].asString(); + auto columnType = columnTypeFromName(obj["columnType"].asString()); + auto dataType = ISerializable::deserialize(obj["dataType"]); + + std::vector requiredSubfields; + for (const auto& s : obj["requiredSubfields"]) { + requiredSubfields.emplace_back(s.asString()); + } + + auto field = deserializeFieldId(obj["field"]); + + std::optional initialDefaultValue; + if (auto it = obj.find("initialDefaultValue"); it != obj.items().end()) { + initialDefaultValue = it->second.asString(); + } + + IcebergFieldMetadata icebergMetadata; + if (auto it = obj.find("icebergMetadata"); it != obj.items().end()) { + icebergMetadata = deserializeFieldMetadata(it->second); + } + + return std::make_shared( + name, + columnType, + std::move(dataType), + std::move(field), + std::move(requiredSubfields), + std::move(initialDefaultValue), + std::move(icebergMetadata)); +} + +// static +void IcebergColumnHandle::registerSerDe() { + auto& registry = DeserializationRegistryForSharedPtr(); + registry.Register("IcebergColumnHandle", IcebergColumnHandle::create); +} + } // namespace facebook::velox::connector::hive::iceberg diff --git a/velox/connectors/hive/iceberg/IcebergColumnHandle.h b/velox/connectors/hive/iceberg/IcebergColumnHandle.h index 07b69f672a6..b91957aa37d 100644 --- a/velox/connectors/hive/iceberg/IcebergColumnHandle.h +++ b/velox/connectors/hive/iceberg/IcebergColumnHandle.h @@ -52,6 +52,14 @@ class IcebergColumnHandle : public HiveColumnHandle { return initialDefaultValue_; } + std::string toString() const override; + + folly::dynamic serialize() const override; + + static ColumnHandlePtr create(const folly::dynamic& obj); + + static void registerSerDe(); + private: const parquet::ParquetFieldId field_; const std::optional initialDefaultValue_; diff --git a/velox/connectors/hive/iceberg/IcebergConnector.cpp b/velox/connectors/hive/iceberg/IcebergConnector.cpp index ab631fc6096..1b3a89a4170 100644 --- a/velox/connectors/hive/iceberg/IcebergConnector.cpp +++ b/velox/connectors/hive/iceberg/IcebergConnector.cpp @@ -26,6 +26,7 @@ #include "velox/connectors/hive/iceberg/IcebergDeletionVectorSink.h" #include "velox/connectors/hive/iceberg/IcebergMergeSink.h" #include "velox/connectors/hive/iceberg/IcebergSessionCredentials.h" +#include "velox/connectors/hive/iceberg/IcebergTableHandle.h" namespace facebook::velox::connector::hive::iceberg { @@ -164,7 +165,9 @@ std::unique_ptr IcebergConnector::createDataSink( } void IcebergConnector::registerSerDe() { + IcebergColumnHandle::registerSerDe(); IcebergFileNameGenerator::registerSerDe(); + IcebergTableHandle::registerSerDe(); } } // namespace facebook::velox::connector::hive::iceberg diff --git a/velox/connectors/hive/iceberg/IcebergSplit.cpp b/velox/connectors/hive/iceberg/IcebergSplit.cpp index 73a6a07866b..1761028d3a6 100644 --- a/velox/connectors/hive/iceberg/IcebergSplit.cpp +++ b/velox/connectors/hive/iceberg/IcebergSplit.cpp @@ -34,7 +34,10 @@ HiveIcebergSplit::HiveIcebergSplit( bool cacheable, const std::unordered_map& infoColumns, std::optional properties, - int64_t dataSequenceNumber) + int64_t dataSequenceNumber, + const std::unordered_map>& + identityPartitionKeys, + std::optional columnMappingMode) : HiveConnectorSplit( connectorId, filePath, @@ -51,8 +54,10 @@ HiveIcebergSplit::HiveIcebergSplit( infoColumns, properties, std::nullopt, - std::nullopt), - dataSequenceNumber(dataSequenceNumber) { + std::nullopt, + columnMappingMode), + dataSequenceNumber(dataSequenceNumber), + identityPartitionKeys(identityPartitionKeys) { // TODO: Deserialize _extraFileInfo to get deleteFiles; } @@ -73,6 +78,9 @@ HiveIcebergSplit::HiveIcebergSplit( const std::unordered_map& infoColumns, std::optional properties, int64_t dataSequenceNumber, + const std::unordered_map>& + identityPartitionKeys, + std::optional columnMappingMode, std::vector coalescedFiles) : HiveConnectorSplit( connectorId, @@ -90,10 +98,86 @@ HiveIcebergSplit::HiveIcebergSplit( infoColumns, properties, std::nullopt, - std::nullopt), + std::nullopt, + columnMappingMode), deleteFiles(std::move(deletes)), coalescedFiles(std::move(coalescedFiles)), - dataSequenceNumber(dataSequenceNumber) {} + dataSequenceNumber(dataSequenceNumber), + identityPartitionKeys(identityPartitionKeys) {} + +HiveIcebergSplit::HiveIcebergSplit( + const std::string& connectorId, + const std::string& filePath, + dwio::common::FileFormat fileFormat, + uint64_t start, + uint64_t length, + const std::unordered_map>& + partitionKeys, + std::optional tableBucketNumber, + const std::unordered_map& customSplitInfo, + const std::shared_ptr& extraFileInfo, + bool cacheable, + std::vector deletes, + const std::unordered_map& infoColumns, + std::optional properties, + int64_t dataSequenceNumber, + const std::unordered_map>& + identityPartitionKeys, + std::vector coalescedFiles) + : HiveIcebergSplit( + connectorId, + filePath, + fileFormat, + start, + length, + partitionKeys, + tableBucketNumber, + customSplitInfo, + extraFileInfo, + cacheable, + std::move(deletes), + infoColumns, + properties, + dataSequenceNumber, + identityPartitionKeys, + std::nullopt, + std::move(coalescedFiles)) {} + +HiveIcebergSplit::HiveIcebergSplit( + const std::string& connectorId, + const std::string& filePath, + dwio::common::FileFormat fileFormat, + uint64_t start, + uint64_t length, + const std::unordered_map>& + partitionKeys, + std::optional tableBucketNumber, + const std::unordered_map& customSplitInfo, + const std::shared_ptr& extraFileInfo, + bool cacheable, + std::vector deletes, + const std::unordered_map& infoColumns, + std::optional properties, + int64_t dataSequenceNumber, + std::vector coalescedFiles) + : HiveIcebergSplit( + connectorId, + filePath, + fileFormat, + start, + length, + partitionKeys, + tableBucketNumber, + customSplitInfo, + extraFileInfo, + cacheable, + std::move(deletes), + infoColumns, + properties, + dataSequenceNumber, + std::unordered_map>{}, + std::nullopt, + std::move(coalescedFiles)) {} std::shared_ptr IcebergSplitBuilder::build() const { return std::make_shared( @@ -110,6 +194,8 @@ std::shared_ptr IcebergSplitBuilder::build() const { deleteFiles_, infoColumns_, std::nullopt, - dataSequenceNumber_); + dataSequenceNumber_, + identityPartitionKeys_, + columnMappingMode_); } } // namespace facebook::velox::connector::hive::iceberg diff --git a/velox/connectors/hive/iceberg/IcebergSplit.h b/velox/connectors/hive/iceberg/IcebergSplit.h index 4b8ac48a991..307cfeaf523 100644 --- a/velox/connectors/hive/iceberg/IcebergSplit.h +++ b/velox/connectors/hive/iceberg/IcebergSplit.h @@ -50,6 +50,21 @@ struct HiveIcebergSplit : public connector::hive::HiveConnectorSplit { /// sequence number filtering. int64_t dataSequenceNumber{0}; + /// Partition values keyed by the Iceberg *source* column field ID, for + /// partition fields the serialized spec explicitly marks as the 'identity' + /// transform. Only an identity value equals the source column's value, so + /// only these entries may be substituted for a read of the source column. + /// + /// Transformed fields (bucket, truncate, day, void) are never present, even + /// when their derived partition-field name happens to collide with a source + /// column name. An empty map means no identity provenance was available and + /// callers must read source columns from the data file. + /// + /// Distinct from the inherited name-keyed 'partitionKeys', which carries + /// every partition field under its derived 'PartitionField.name()' and + /// therefore cannot prove a transform is identity. + std::unordered_map> identityPartitionKeys; + HiveIcebergSplit( const std::string& connectorId, const std::string& filePath, @@ -64,7 +79,11 @@ struct HiveIcebergSplit : public connector::hive::HiveConnectorSplit { bool cacheable = true, const std::unordered_map& infoColumns = {}, std::optional fileProperties = std::nullopt, - int64_t dataSequenceNumber = 0); + int64_t dataSequenceNumber = 0, + const std::unordered_map>& + identityPartitionKeys = {}, + std::optional columnMappingMode = + std::nullopt); // For tests only HiveIcebergSplit( @@ -83,7 +102,52 @@ struct HiveIcebergSplit : public connector::hive::HiveConnectorSplit { const std::unordered_map& infoColumns = {}, std::optional fileProperties = std::nullopt, int64_t dataSequenceNumber = 0, + const std::unordered_map>& + identityPartitionKeys = {}, + std::optional columnMappingMode = + std::nullopt, std::vector coalescedFiles = {}); + + // Compatibility overload for downstream callers that pass coalesced files + // immediately after identity partition keys. + HiveIcebergSplit( + const std::string& connectorId, + const std::string& filePath, + dwio::common::FileFormat fileFormat, + uint64_t start, + uint64_t length, + const std::unordered_map>& + partitionKeys, + std::optional tableBucketNumber, + const std::unordered_map& customSplitInfo, + const std::shared_ptr& extraFileInfo, + bool cacheable, + std::vector deletes, + const std::unordered_map& infoColumns, + std::optional fileProperties, + int64_t dataSequenceNumber, + const std::unordered_map>& + identityPartitionKeys, + std::vector coalescedFiles); + + // Compatibility overload for callers that predate identity partition keys. + HiveIcebergSplit( + const std::string& connectorId, + const std::string& filePath, + dwio::common::FileFormat fileFormat, + uint64_t start, + uint64_t length, + const std::unordered_map>& + partitionKeys, + std::optional tableBucketNumber, + const std::unordered_map& customSplitInfo, + const std::shared_ptr& extraFileInfo, + bool cacheable, + std::vector deletes, + const std::unordered_map& infoColumns, + std::optional fileProperties, + int64_t dataSequenceNumber, + std::vector coalescedFiles); }; /// Builds Iceberg splits with named parameters. @@ -138,6 +202,19 @@ class IcebergSplitBuilder { return *this; } + /// Sets identity-transform partition values keyed by source field ID. See + /// 'HiveIcebergSplit::identityPartitionKeys'. + IcebergSplitBuilder& identityPartitionKeys( + const std::unordered_map>& keys) { + identityPartitionKeys_ = keys; + return *this; + } + + IcebergSplitBuilder& columnMappingMode(dwio::common::ColumnMappingMode mode) { + columnMappingMode_ = mode; + return *this; + } + std::shared_ptr build() const; private: @@ -150,6 +227,9 @@ class IcebergSplitBuilder { std::unordered_map infoColumns_; std::vector deleteFiles_; int64_t dataSequenceNumber_{0}; + std::unordered_map> + identityPartitionKeys_; + std::optional columnMappingMode_; }; } // namespace facebook::velox::connector::hive::iceberg diff --git a/velox/connectors/hive/iceberg/IcebergSplitReader.cpp b/velox/connectors/hive/iceberg/IcebergSplitReader.cpp index 28b28bd2794..024ced8f8f3 100644 --- a/velox/connectors/hive/iceberg/IcebergSplitReader.cpp +++ b/velox/connectors/hive/iceberg/IcebergSplitReader.cpp @@ -81,6 +81,32 @@ void fillNullsWithInt64( namespace facebook::velox::connector::hive::iceberg { namespace { +// Indexes IcebergColumnHandles by their underlying data-column name. +// Covers output-projected handles (columnHandles) and, when tableHandle is +// non-null, filter-only handles (tableHandle->filterColumnHandles()) too. +std::unordered_map +buildIcebergHandleByName( + const ColumnHandleMap* columnHandles, + const FileTableHandle* tableHandle = nullptr) { + std::unordered_map handleByName; + const auto addHandle = [&handleByName](const auto& handle) { + if (auto* h = dynamic_cast(handle.get())) { + handleByName.emplace(h->name(), h); + } + }; + if (columnHandles) { + for (const auto& [_, handle] : *columnHandles) { + addHandle(handle); + } + } + if (tableHandle) { + for (const auto& handle : tableHandle->filterColumnHandles()) { + addHandle(handle); + } + } + return handleByName; +} + /// Returns true if a delete/update file should be skipped based on sequence /// number conflict resolution. Per the Iceberg spec (V2+): /// - Equality deletes apply when deleteSeqNum > dataSeqNum (i.e., skip when @@ -166,37 +192,63 @@ std::vector IcebergSplitReader::buildFieldIds() const { std::vector fieldIds; const auto& dataColumns = tableHandle_->dataColumns(); - if (dataColumns == nullptr || columnHandles_ == nullptr) { + if (dataColumns == nullptr) { return fieldIds; } + + const auto* hiveTableHandle = + dynamic_cast(tableHandle_.get()); + const auto* dataColumnFieldIds = hiveTableHandle != nullptr + ? &hiveTableHandle->dataColumnFieldIds() + : nullptr; + // Column handles are keyed by output alias; index them by the underlying // data-column name so we can align to dataColumns() order. - std::unordered_map handleByName; - const auto addIcebergHandle = [&handleByName](const auto& handle) { - if (auto* icebergHandle = - dynamic_cast(handle.get())) { - handleByName.emplace(icebergHandle->name(), icebergHandle); - } - }; - for (const auto& columnHandle : *columnHandles_) { - addIcebergHandle(columnHandle.second); - } - // Remaining filters can add columns to the reader output that are not - // projected by the table scan. Include those handles so filter-only columns - // are still matched by Iceberg field ID. - for (const auto& handle : tableHandle_->filterColumnHandles()) { - addIcebergHandle(handle); - } - if (handleByName.empty()) { + const auto handleByName = + buildIcebergHandleByName(columnHandles_.get(), tableHandle_.get()); + if (handleByName.empty() && + (dataColumnFieldIds == nullptr || dataColumnFieldIds->empty())) { return fieldIds; } + // Equality-delete columns absent from the user's projection have no + // IcebergColumnHandle and would otherwise get a sentinel field ID. Collect + // their real Iceberg field IDs via resolveEqualityColumns() — the single + // authoritative path for field-ID→name resolution — so the Parquet reader + // can locate those columns physically. This is purely additive: it only + // fills slots where handleByName has no entry. + std::unordered_map equalityFieldIdByName; + for (const auto& deleteFile : icebergSplit_->deleteFiles) { + if (deleteFile.content != FileContent::kEqualityDeletes || + deleteFile.recordCount == 0 || deleteFile.equalityFieldIds.empty()) { + continue; + } + if (shouldSkipBySequenceNumber( + deleteFile.dataSequenceNumber, + icebergSplit_->dataSequenceNumber, + /*isEqualityDelete=*/true)) { + continue; + } + + auto [names, types] = resolveEqualityColumns(deleteFile); + for (size_t i = 0; i < names.size(); ++i) { + equalityFieldIdByName.emplace(names[i], deleteFile.equalityFieldIds[i]); + } + } + fieldIds.reserve(dataColumns->size()); int32_t sentinelFieldId = -1; for (size_t i = 0; i < dataColumns->size(); ++i) { - auto it = handleByName.find(dataColumns->nameOf(static_cast(i))); + const auto& colName = dataColumns->nameOf(static_cast(i)); + auto it = handleByName.find(colName); if (it != handleByName.end()) { fieldIds.push_back(it->second->field()); + } else if (dataColumnFieldIds != nullptr && !dataColumnFieldIds->empty()) { + fieldIds.push_back( + dwio::common::ParquetFieldId{dataColumnFieldIds->at(i), {}}); + } else if (auto eqIt = equalityFieldIdByName.find(colName); + eqIt != equalityFieldIdByName.end()) { + fieldIds.push_back(dwio::common::ParquetFieldId{eqIt->second, {}}); } else { fieldIds.push_back(dwio::common::ParquetFieldId{sentinelFieldId--, {}}); } @@ -206,7 +258,7 @@ std::vector IcebergSplitReader::buildFieldIds() void IcebergSplitReader::prepareSplit( std::shared_ptr metadataFilter, - dwio::common::RuntimeStatistics& runtimeStats, + dwio::common::RuntimeStats& runtimeStats, const folly::F14FastMap& fileReadOps) { // Forward per-query delegated credentials into fileReadOps so delegated-auth // filesystems authorize the read as the caller rather than the service @@ -248,13 +300,7 @@ void IcebergSplitReader::prepareSplit( // IcebergColumnHandle's field-id tree (per-input-column). // Column handles are keyed by output alias; index them by physical name // to mirror buildFieldIds() and support renamed columns. - std::unordered_map handleByName; - for (const auto& [outputName, handle] : *columnHandles_) { - if (auto* icebergHandle = - dynamic_cast(handle.get())) { - handleByName.emplace(icebergHandle->name(), icebergHandle); - } - } + const auto handleByName = buildIcebergHandleByName(columnHandles_.get()); std::vector fieldIds; fieldIds.reserve(fileSchema->size()); bool allResolved = true; @@ -405,23 +451,25 @@ void IcebergSplitReader::prepareSplit( return; } - // Inject a row-number column when filters, random-skip, or positional + // Inject a row-number column when filters, random-skip, or row-skipping // deletes make the output-to-file-position mapping non-contiguous. // Check split metadata rather than positionalDeleteFileReaders_ because // the row reader must be configured before delete files are opened. Both // _row_id and $target_table_row_id need accurate file-absolute positions, - // so request injection when either is projected. - const bool hasPositionalDeletes = std::any_of( + // so request injection when either is projected. Deletion vectors skip rows + // exactly like V2 positional deletes and must be counted here too. + const bool hasRowSkippingDeletes = std::any_of( icebergSplit_->deleteFiles.begin(), icebergSplit_->deleteFiles.end(), [](const IcebergDeleteFile& deleteFile) { - return deleteFile.content == FileContent::kPositionalDeletes && + return (deleteFile.content == FileContent::kPositionalDeletes || + deleteFile.content == FileContent::kDeletionVector) && deleteFile.recordCount > 0; }); useRowNumberColumn_ = (rowIdOutputIndex_.has_value() || targetTableRowIdOutputIndex_.has_value()) && (scanSpec_->hasFilter() || baseReaderOpts_.randomSkip() != nullptr || - hasPositionalDeletes); + hasRowSkippingDeletes); if (useRowNumberColumn_) { dwio::common::RowNumberColumnInfo rowNumInfo; rowNumInfo.insertPosition = readerOutputType_->size(); @@ -560,7 +608,7 @@ void IcebergSplitReader::configureEqualityDeleteColumns() { std::vector extraNames; std::vector extraTypes; const auto& deleteFiles = icebergSplit_->deleteFiles; - const auto& splitPartitionKeys = icebergSplit_->partitionKeys; + const auto& identityPartitionKeys = icebergSplit_->identityPartitionKeys; for (const auto& deleteFile : deleteFiles) { if (deleteFile.content != FileContent::kEqualityDeletes || @@ -606,14 +654,21 @@ void IcebergSplitReader::configureEqualityDeleteColumns() { static_cast( readerOutputType_->size() + extraEqualityColumns.size())); - // For partition columns set the partition value directly as a constant - // on the scan-spec child. This is independent of whether the data file - // contains the partition column physically. With the constant set - // up-front, 'adaptColumns' does not need any special-case logic for - // augmented partition columns and the read does not depend on the - // writer's choice of including the partition column in the file. - auto partitionIt = splitPartitionKeys.find(name); - if (partitionIt != splitPartitionKeys.end()) { + // Substitute the partition value for the source column only when the + // Iceberg partition spec explicitly marks this equality field's source + // column as an identity partition field. Identity is the only transform + // whose stored partition value equals the source column value; a + // bucket, truncate, or temporal value is a transform result, and a + // 'void' value is always null while keeping the source column's name. + // Keying by the delete file's own Iceberg field ID (rather than by + // column name) also keeps this correct across column renames. + // + // Anything else -- a transformed field, a spec that could not be + // parsed, or a split with no identity metadata at all -- leaves no + // constant installed, so the column is read from the data file. + const auto identityIt = + identityPartitionKeys.find(deleteFile.equalityFieldIds[i]); + if (identityIt != identityPartitionKeys.end()) { // Iceberg encodes DATE partition values as the integer number of // days since the Unix epoch (e.g. "19345"). The standard // 'setPartitionValue' helper learns this from the planner-supplied @@ -624,7 +679,7 @@ void IcebergSplitReader::configureEqualityDeleteColumns() { const bool isDaysSinceEpoch = equalityColumnTypes[i]->isDate(); auto constant = newConstantFromString( equalityColumnTypes[i], - partitionIt->second, + identityIt->second, connectorQueryCtx_->memoryPool(), fileConfig_->readTimestampPartitionValueAsLocalTime( connectorQueryCtx_->sessionProperties()), @@ -669,20 +724,37 @@ IcebergSplitReader::resolveEqualityColumns( VELOX_CHECK( dataColumns != nullptr, "Iceberg equality delete file '{}' cannot be processed because " - "table data columns are not available in HiveTableHandle.", + "table data columns are not available in IcebergTableHandle.", deleteFile.filePath); - for (const auto& eqFieldId : deleteFile.equalityFieldIds) { - // Field IDs are 1-based sequential for non-evolved schemas. - auto colIdx = static_cast(eqFieldId - 1); + std::unordered_map columnIndexByFieldId; + if (const auto* hiveTableHandle = + dynamic_cast(tableHandle_.get())) { + const auto& dataColumnFieldIds = hiveTableHandle->dataColumnFieldIds(); + columnIndexByFieldId.reserve(dataColumnFieldIds.size()); + for (uint32_t i = 0; i < dataColumnFieldIds.size(); ++i) { + columnIndexByFieldId.emplace(dataColumnFieldIds[i], i); + } + } + + for (const auto& equalityFieldId : deleteFile.equalityFieldIds) { + VELOX_CHECK_GT( + equalityFieldId, + 0, + "Equality delete field ID must be positive: {}", + equalityFieldId); + const auto fieldIdIt = columnIndexByFieldId.find(equalityFieldId); + // Older plans and tests may not carry full-schema field IDs. Preserve the + // legacy ordinal lookup when metadata is unavailable or incomplete. + const auto columnIndex = fieldIdIt != columnIndexByFieldId.end() + ? fieldIdIt->second + : static_cast(equalityFieldId - 1); VELOX_CHECK_LT( - colIdx, + columnIndex, dataColumns->size(), - "Equality delete field ID {} out of range. This may indicate " - "schema evolution with non-sequential field IDs, which is " - "not yet supported.", - eqFieldId); - equalityColumnNames.push_back(dataColumns->nameOf(colIdx)); - equalityColumnTypes.push_back(dataColumns->childAt(colIdx)); + "Equality delete field ID cannot be resolved against table columns: {}", + equalityFieldId); + equalityColumnNames.push_back(dataColumns->nameOf(columnIndex)); + equalityColumnTypes.push_back(dataColumns->childAt(columnIndex)); } return {std::move(equalityColumnNames), std::move(equalityColumnTypes)}; } @@ -925,6 +997,12 @@ std::vector IcebergSplitReader::adaptColumns( const bool readTimestampAsLocalTime = fileConfig_->readTimestampPartitionValueAsLocalTime( connectorQueryCtx_->sessionProperties()); + + // Index all Iceberg column handles by data-column name for O(1) default-value + // lookup inside the loop. + const auto handleByName = + buildIcebergHandleByName(columnHandles_.get(), tableHandle_.get()); + // Iceberg table stores all column's data in data file. for (const auto& childSpec : childrenSpecs) { const std::string& fieldName = childSpec->fieldName(); @@ -1045,40 +1123,21 @@ std::vector IcebergSplitReader::adaptColumns( partitionIt != fileSplit_->partitionKeys.end()) { setPartitionValue(childSpec.get(), fieldName, partitionIt->second); } else { - // Check if column has an initial-default value (Iceberg V3) - bool hasDefaultValue = false; - // The columnHandles_ map is keyed by output name (which may be an - // alias). We need to find the column handle where the handle's name() - // matches fieldName. fieldName is the table column name from - // readerOutputType_. - for (const auto& [outputName, handle] : *columnHandles_) { - if (handle->name() == fieldName) { - auto icebergColumnHandle = - std::dynamic_pointer_cast(handle); - if (icebergColumnHandle && - icebergColumnHandle->initialDefaultValue().has_value()) { - // Use initial-default value for schema evolution. - auto columnType = tableSchema->findChild(fieldName); - VELOX_CHECK_NOT_NULL( - columnType, - "Column '{}' not found in table schema", - fieldName); - auto constant = newConstantFromString( - columnType, - icebergColumnHandle->initialDefaultValue().value(), - connectorQueryCtx_->memoryPool(), - readTimestampAsLocalTime, - false); - childSpec->setConstantValue(constant); - hasDefaultValue = true; - break; - } - } - } - - // Fall back to NULL if no default value - if (!hasDefaultValue) { - auto columnType = tableSchema->findChild(fieldName); + // Check if column has an initial-default value (Iceberg V3). + // Use the pre-built handleByName map that covers both output + // column handles and filter column handles. + auto it = handleByName.find(fieldName); + auto columnType = tableSchema->findChild(fieldName); + if (it != handleByName.end() && + it->second->initialDefaultValue().has_value()) { + childSpec->setConstantValue(newConstantFromString( + columnType, + it->second->initialDefaultValue().value(), + connectorQueryCtx_->memoryPool(), + /*isLocalTimestamp=*/false, + /*isDaysSinceEpoch=*/false)); + } else { + // Fall back to NULL if no default value. VELOX_CHECK_NOT_NULL( columnType, "Column '{}' not found in table schema", fieldName); childSpec->setConstantValue( diff --git a/velox/connectors/hive/iceberg/IcebergSplitReader.h b/velox/connectors/hive/iceberg/IcebergSplitReader.h index 9f9d5a3d1df..04af1ef0fa4 100644 --- a/velox/connectors/hive/iceberg/IcebergSplitReader.h +++ b/velox/connectors/hive/iceberg/IcebergSplitReader.h @@ -50,7 +50,7 @@ class IcebergSplitReader : public FileSplitReader { void prepareSplit( std::shared_ptr metadataFilter, - dwio::common::RuntimeStatistics& runtimeStats, + dwio::common::RuntimeStats& runtimeStats, const folly::F14FastMap& fileReadOps = {}) override; @@ -66,9 +66,10 @@ class IcebergSplitReader : public FileSplitReader { // Builds the requested-schema field-id trees, one per top-level column, // aligned to tableHandle_->dataColumns(). Projected and filter-only Iceberg - // handles provide field IDs. Data columns without an Iceberg handle get a + // handles provide field IDs. Data columns without an Iceberg handle fall back + // to the table handle's full-schema field IDs when available, otherwise a // negative sentinel id that matches no physical field. Returns empty when no - // Iceberg handles are available. + // Iceberg handles and no full-schema field IDs are available. std::vector buildFieldIds() const; /// Adapts the data file schema to match the table schema expected by the @@ -130,20 +131,22 @@ class IcebergSplitReader : public FileSplitReader { const RowTypePtr& fileType, const RowTypePtr& tableSchema) const override; - // Resolves the equality field IDs of an equality-delete file to the - // corresponding column names and types in the table schema. In Iceberg, - // field IDs for top-level columns are assigned sequentially starting from - // 1, matching the column order in the table schema. + // Resolves equality-delete field IDs to column names and types using the + // table handle's full-schema field IDs. Falls back to the legacy one-based + // ordinal mapping when those IDs are unavailable or incomplete. std::pair, std::vector> resolveEqualityColumns(const IcebergDeleteFile& deleteFile) const; // Discovers equality-delete columns that are not in the user's projection // and augments 'scanSpec_' and 'readerOutputType_' so they are physically - // read and made available in the output RowVector. For partition columns - // the partition value is set as a constant on the scan-spec child so the - // augmentation works regardless of whether the data file physically - // contains the partition column. Augmented columns are appended at the end - // of 'readerOutputType_' so the upstream FileDataSource's positional + // read and made available in the output RowVector. When the split proves + // the column's Iceberg field ID belongs to an identity partition field + // (see 'HiveIcebergSplit::identityPartitionKeys'), the partition value is + // set as a constant on the scan-spec child so the augmentation works + // regardless of whether the data file physically contains the column; + // every other column, including one partitioned by a transform, is read + // from the file. Augmented columns are appended at the end of + // 'readerOutputType_' so the upstream FileDataSource's positional // projection naturally drops them from the operator output. void configureEqualityDeleteColumns(); diff --git a/velox/connectors/hive/iceberg/IcebergTableHandle.cpp b/velox/connectors/hive/iceberg/IcebergTableHandle.cpp new file mode 100644 index 00000000000..5b7c70254f5 --- /dev/null +++ b/velox/connectors/hive/iceberg/IcebergTableHandle.cpp @@ -0,0 +1,193 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/connectors/hive/iceberg/IcebergTableHandle.h" + +#include +#include + +namespace facebook::velox::connector::hive::iceberg { + +IcebergTableHandle::IcebergTableHandle( + std::string connectorId, + const std::string& tableName, + common::SubfieldFilters subfieldFilters, + const core::TypedExprPtr& remainingFilter, + const RowTypePtr& dataColumns, + std::vector indexColumns, + const std::unordered_map& tableParameters, + std::vector filterColumnHandles, + double sampleRate, + std::string dbName, + std::vector dataColumnFieldIds, + bool isChangelogQuery, + std::unordered_map dataColumnHandles) + : HiveTableHandle( + std::move(connectorId), + tableName, + std::move(subfieldFilters), + remainingFilter, + dataColumns, + std::move(indexColumns), + tableParameters, + std::vector( + filterColumnHandles.begin(), + filterColumnHandles.end()), + sampleRate, + std::move(dbName), + std::move(dataColumnFieldIds)), + isChangelogQuery_(isChangelogQuery), + dataColumnHandles_(std::move(dataColumnHandles)) { + if (isChangelogQuery_) { + VELOX_USER_CHECK( + !dataColumnHandles_.empty(), + "dataColumnHandles must not be empty when isChangelogQuery is true"); + } +} + +std::string IcebergTableHandle::toString() const { + std::ostringstream out; + out << HiveTableHandle::toString(); + if (isChangelogQuery_) { + out << ", isChangelogQuery: true"; + } + if (!dataColumnHandles_.empty()) { + // Sort by name for deterministic output, mirroring HiveTableHandle's + // treatment of subfieldFilters and tableParameters. + std::map ordered; + for (const auto& [name, handle] : dataColumnHandles_) { + ordered[name] = handle.get(); + } + out << ", dataColumnHandles: ["; + bool first = true; + for (const auto& [name, handle] : ordered) { + if (!first) { + out << ", "; + } + out << name << ": " << handle->toString(); + first = false; + } + out << "]"; + } + return out.str(); +} + +folly::dynamic IcebergTableHandle::serialize() const { + // Start from the common Hive fields under the "IcebergTableHandle" type name, + // then append Iceberg-specific fields. + folly::dynamic obj = serializeHiveFields("IcebergTableHandle"); + + obj["isChangelogQuery"] = isChangelogQuery_; + + if (!dataColumnHandles_.empty()) { + folly::dynamic dataColObj = folly::dynamic::object; + for (const auto& [name, handle] : dataColumnHandles_) { + dataColObj[name] = handle->serialize(); + } + obj["dataColumnHandles"] = dataColObj; + } + + return obj; +} + +// static +ConnectorTableHandlePtr IcebergTableHandle::create( + const folly::dynamic& obj, + void* context) { + std::string connectorId; + std::string tableName; + std::string dbName; + common::SubfieldFilters subfieldFilters; + core::TypedExprPtr remainingFilter; + double sampleRate{1.0}; + RowTypePtr dataColumns; + std::unordered_map tableParameters; + // deserializeHiveFields dispatches through the SerDe registry using the + // "name" key in each handle's JSON. Because IcebergColumnHandle::serialize() + // writes "name":"IcebergColumnHandle", the registry instantiates + // IcebergColumnHandle objects, so the dynamic_cast below is always valid. + std::vector hiveFilterHandles; + std::vector indexColumns; + std::vector dataColumnFieldIds; + + deserializeHiveFields( + obj, + context, + connectorId, + tableName, + subfieldFilters, + remainingFilter, + sampleRate, + dataColumns, + tableParameters, + hiveFilterHandles, + indexColumns, + dbName, + dataColumnFieldIds); + + // Cast the already-deserialized handles to their concrete IcebergColumnHandle + // type. The dynamic_cast is guaranteed to succeed (see comment above). + std::vector filterColumnHandles; + filterColumnHandles.reserve(hiveFilterHandles.size()); + for (const auto& h : hiveFilterHandles) { + auto handle = std::dynamic_pointer_cast(h); + VELOX_CHECK_NOT_NULL( + handle, + "filterColumnHandle is not an IcebergColumnHandle during deserialization"); + filterColumnHandles.push_back(std::move(handle)); + } + + bool isChangelogQuery = false; + if (auto it = obj.find("isChangelogQuery"); it != obj.items().end()) { + isChangelogQuery = it->second.asBool(); + } + + std::unordered_map dataColumnHandles; + if (auto it = obj.find("dataColumnHandles"); it != obj.items().end()) { + for (const auto& key : it->second.keys()) { + auto name = key.asString(); + auto handle = + ISerializable::deserialize(it->second[key]); + VELOX_CHECK_NOT_NULL( + handle, + "dataColumnHandle is not an IcebergColumnHandle during deserialization"); + dataColumnHandles.emplace(std::move(name), std::move(handle)); + } + } + + return std::make_shared( + std::move(connectorId), + tableName, + std::move(subfieldFilters), + remainingFilter, + dataColumns, + std::move(indexColumns), + tableParameters, + std::move(filterColumnHandles), + sampleRate, + std::move(dbName), + std::move(dataColumnFieldIds), + isChangelogQuery, + std::move(dataColumnHandles)); +} + +// static +void IcebergTableHandle::registerSerDe() { + auto& registry = DeserializationWithContextRegistryForSharedPtr(); + registry.Register("IcebergTableHandle", IcebergTableHandle::create); +} + +} // namespace facebook::velox::connector::hive::iceberg diff --git a/velox/connectors/hive/iceberg/IcebergTableHandle.h b/velox/connectors/hive/iceberg/IcebergTableHandle.h new file mode 100644 index 00000000000..941f59707f9 --- /dev/null +++ b/velox/connectors/hive/iceberg/IcebergTableHandle.h @@ -0,0 +1,81 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +#include "velox/connectors/hive/TableHandle.h" +#include "velox/connectors/hive/iceberg/IcebergColumnHandle.h" + +namespace facebook::velox::connector::hive::iceberg { + +/// Iceberg-specific table handle that extends HiveTableHandle. +/// +/// Carries all fields inherited from HiveTableHandle (subfield filters, +/// remaining filter, data columns, table parameters, etc.) plus +/// Iceberg-specific fields: +/// - isChangelogQuery_: whether this scan is a changelog (CDC) query. +/// - dataColumnHandles_: column handles keyed by column name, used for +/// Iceberg-specific column metadata (field IDs, default values, etc.). +class IcebergTableHandle : public HiveTableHandle { + public: + IcebergTableHandle( + std::string connectorId, + const std::string& tableName, + common::SubfieldFilters subfieldFilters, + const core::TypedExprPtr& remainingFilter, + const RowTypePtr& dataColumns = nullptr, + std::vector indexColumns = {}, + const std::unordered_map& tableParameters = {}, + std::vector filterColumnHandles = {}, + double sampleRate = 1.0, + std::string dbName = "", + std::vector dataColumnFieldIds = {}, + bool isChangelogQuery = false, + std::unordered_map + dataColumnHandles = {}); + + /// Whether this scan is a changelog (CDC) query over an Iceberg table. + bool isChangelogQuery() const { + return isChangelogQuery_; + } + + /// Column handles keyed by column name, carrying Iceberg-specific metadata + /// such as field IDs and initial-default values. + const std::unordered_map& + dataColumnHandles() const { + return dataColumnHandles_; + } + + std::string toString() const override; + + folly::dynamic serialize() const override; + + static ConnectorTableHandlePtr create( + const folly::dynamic& obj, + void* context); + + static void registerSerDe(); + + private: + const bool isChangelogQuery_; + const std::unordered_map + dataColumnHandles_; +}; + +using IcebergTableHandlePtr = std::shared_ptr; + +} // namespace facebook::velox::connector::hive::iceberg diff --git a/velox/connectors/hive/iceberg/PositionalDeleteFileReader.cpp b/velox/connectors/hive/iceberg/PositionalDeleteFileReader.cpp index 4dd3100eb4a..ee7d1bd4155 100644 --- a/velox/connectors/hive/iceberg/PositionalDeleteFileReader.cpp +++ b/velox/connectors/hive/iceberg/PositionalDeleteFileReader.cpp @@ -33,7 +33,7 @@ PositionalDeleteFileReader::PositionalDeleteFileReader( const std::shared_ptr& fileConfig, const std::shared_ptr& ioStatistics, const std::shared_ptr& ioStats, - dwio::common::RuntimeStatistics& runtimeStats, + dwio::common::RuntimeStats& runtimeStats, uint64_t splitOffset, const std::string& connectorId) : deleteFile_(deleteFile), diff --git a/velox/connectors/hive/iceberg/PositionalDeleteFileReader.h b/velox/connectors/hive/iceberg/PositionalDeleteFileReader.h index 7367de7818d..6f541ecbd62 100644 --- a/velox/connectors/hive/iceberg/PositionalDeleteFileReader.h +++ b/velox/connectors/hive/iceberg/PositionalDeleteFileReader.h @@ -78,7 +78,7 @@ class PositionalDeleteFileReader { const std::shared_ptr& fileConfig, const std::shared_ptr& ioStatistics, const std::shared_ptr& ioStats, - dwio::common::RuntimeStatistics& runtimeStats, + dwio::common::RuntimeStats& runtimeStats, uint64_t splitOffset, const std::string& connectorId); diff --git a/velox/connectors/hive/iceberg/WriterOptionsAdapter.h b/velox/connectors/hive/iceberg/WriterOptionsAdapter.h index 50e31c9d713..a34d2f99ff3 100644 --- a/velox/connectors/hive/iceberg/WriterOptionsAdapter.h +++ b/velox/connectors/hive/iceberg/WriterOptionsAdapter.h @@ -58,7 +58,7 @@ class WriterOptionsAdapter { /// `icebergFieldIds` carries the per-input-column Iceberg field-id tree. /// Honored only by the NIMBLE adapter, which uses it to stamp /// `iceberg.id` (and other Iceberg V3 keys) onto each NIMBLE schema node -/// via `VeloxWriterOptions::schemaAttributes`. Pass an empty +/// via `WriterOptions::schemaAttributes`. Pass an empty /// `IcebergFieldId{}` for formats / call sites that have no field-id tree /// available (the NIMBLE adapter then produces files without /// `iceberg.id` attributes, the same wire shape as a pre-attributes diff --git a/velox/connectors/hive/iceberg/tests/CMakeLists.txt b/velox/connectors/hive/iceberg/tests/CMakeLists.txt index 3f8b853d760..a3aa3c63da3 100644 --- a/velox/connectors/hive/iceberg/tests/CMakeLists.txt +++ b/velox/connectors/hive/iceberg/tests/CMakeLists.txt @@ -44,6 +44,7 @@ if(NOT VELOX_DISABLE_GOOGLETEST) IcebergPositionalDeleteTest.cpp IcebergReadTest.cpp IcebergSplitReaderBenchmarkTest.cpp + IcebergTableHandleTest.cpp IcebergTestBase.cpp ) add_test(velox_hive_iceberg_test velox_hive_iceberg_test) @@ -172,7 +173,13 @@ if(NOT VELOX_DISABLE_GOOGLETEST) ) endif() - add_executable(velox_hive_iceberg_equality_delete_test EqualityDeleteFileReaderTest.cpp Main.cpp) + add_executable( + velox_hive_iceberg_equality_delete_test + EqualityDeleteFileReaderTest.cpp + IcebergTestBase.cpp + Main.cpp + ) + velox_add_test_headers(velox_hive_iceberg_equality_delete_test IcebergTestBase.h) add_test(velox_hive_iceberg_equality_delete_test velox_hive_iceberg_equality_delete_test) target_link_libraries( @@ -181,9 +188,19 @@ if(NOT VELOX_DISABLE_GOOGLETEST) velox_hive_iceberg_splitreader velox_exec_test_lib velox_dwio_common_test_utils + velox_vector_fuzzer GTest::gtest + GTest::gtest_main ) + if(VELOX_ENABLE_PARQUET) + target_link_libraries( + velox_hive_iceberg_equality_delete_test + velox_dwio_parquet_reader + velox_dwio_parquet_writer + ) + endif() + add_executable(velox_hive_iceberg_deletion_vector_writer_test DeletionVectorWriterTest.cpp) add_test( velox_hive_iceberg_deletion_vector_writer_test diff --git a/velox/connectors/hive/iceberg/tests/DeletionVectorReaderTest.cpp b/velox/connectors/hive/iceberg/tests/DeletionVectorReaderTest.cpp index f1dada8f455..489b3b77860 100644 --- a/velox/connectors/hive/iceberg/tests/DeletionVectorReaderTest.cpp +++ b/velox/connectors/hive/iceberg/tests/DeletionVectorReaderTest.cpp @@ -20,6 +20,8 @@ #include +#include +#include #include #include @@ -96,6 +98,201 @@ std::string serializeRoaringBitmapNoRun(const std::vector& positions) { return data; } +// Serializes a roaring bitmap in the portable no-run format, encoding any +// block whose cardinality exceeds 4096 as a 1024-word bitset container rather +// than an array container. 'serializeRoaringBitmapNoRun' above always emits +// array containers, so it cannot produce this encoding — which is the one +// Iceberg and DeletionVectorWriter use for dense blocks. +std::string serializeRoaringBitmapWithBitsetContainer( + const std::vector& positions) { + constexpr uint32_t kMaxArrayContainerCardinality = 4'096; + constexpr uint32_t kBitmapContainerBytes = 8'192; + constexpr uint32_t kCookie = 12'346; + + std::map> containers; + for (auto position : positions) { + containers[static_cast(position >> 16)].push_back( + static_cast(position & 0xFFFF)); + } + for (auto& [key, values] : containers) { + std::sort(values.begin(), values.end()); + } + + const auto numContainers = static_cast(containers.size()); + std::string data; + data.append(reinterpret_cast(&kCookie), 4); + data.append(reinterpret_cast(&numContainers), 4); + + for (const auto& [key, values] : containers) { + const auto cardMinus1 = static_cast(values.size() - 1); + data.append(reinterpret_cast(&key), 2); + data.append(reinterpret_cast(&cardMinus1), 2); + } + + uint32_t offset = 4 + 4 + numContainers * 4 + numContainers * 4; + for (const auto& [key, values] : containers) { + data.append(reinterpret_cast(&offset), 4); + offset += values.size() <= kMaxArrayContainerCardinality + ? static_cast(values.size()) * 2 + : kBitmapContainerBytes; + } + + for (const auto& [key, values] : containers) { + if (values.size() <= kMaxArrayContainerCardinality) { + for (auto value : values) { + data.append(reinterpret_cast(&value), 2); + } + } else { + std::vector words(1'024, 0); + for (auto value : values) { + words[value / 64] |= (1ULL << (value % 64)); + } + for (auto word : words) { + data.append(reinterpret_cast(&word), 8); + } + } + } + return data; +} + +// Wraps serialized 32-bit roaring bitmaps in the Roaring64 envelope: +// [numGroups: uint64] then, per group, [highBits: uint32][32-bit bitmap]. +// Reaching group N+1 requires the reader to advance exactly past group N's +// container data, so multi-group inputs exercise that skip arithmetic. +std::string wrapInRoaring64( + const std::vector>& groups) { + std::string data; + const auto numGroups = static_cast(groups.size()); + data.append(reinterpret_cast(&numGroups), 8); + for (const auto& [highBits, bitmap] : groups) { + data.append(reinterpret_cast(&highBits), 4); + data.append(bitmap); + } + return data; +} + +// Serializes a roaring bitmap whose containers may each use a different +// encoding, which none of the helpers above can express: the no-run helpers +// pick array vs bitset purely by cardinality, and the run helper makes every +// container run-encoded. +// +// Iceberg's own test suite carries an "all container types" case; this builds +// an equivalent bitmap from scratch so the reader is exercised against array, +// run, and bitset containers coexisting in one blob. +struct ContainerSpec { + enum class Encoding { kArray, kRun, kBitset }; + + uint16_t key{0}; + Encoding encoding{Encoding::kArray}; + // Values within the container's 64K block, sorted. For kRun these are still + // the expanded values; the runs are derived from them. + std::vector values; +}; + +// Collapses sorted values into (start, lengthMinus1) runs. +std::vector> toRuns( + const std::vector& values) { + std::vector> runs; + for (size_t i = 0; i < values.size();) { + size_t j = i; + while (j + 1 < values.size() && values[j + 1] == values[j] + 1) { + ++j; + } + runs.emplace_back(values[i], static_cast(j - i)); + i = j + 1; + } + return runs; +} + +std::string serializeRoaringBitmapMixed( + const std::vector& containers) { + constexpr uint32_t kSerialCookieWithRuns = 12'347; + constexpr uint32_t kRunContainersNoOffsetThreshold = 4; + constexpr size_t kBitmapContainerBytes = 8'192; + + const auto numContainers = static_cast(containers.size()); + // This helper always emits the run-bearing cookie and a run bitmap, even + // when no container is actually run-encoded, so readers always take the + // "runs present" branch. The offset section must follow the same rule or the + // two sides disagree on the header size. + const bool hasRunContainers = true; + + std::string data; + const uint32_t cookie = kSerialCookieWithRuns | ((numContainers - 1) << 16); + data.append(reinterpret_cast(&cookie), 4); + + // Run bitmap: bit i marks container i as run-encoded, LSB-first. + const uint32_t runBitmapBytes = (numContainers + 7) / 8; + std::vector runBitmap(runBitmapBytes, 0); + for (uint32_t i = 0; i < numContainers; ++i) { + if (containers[i].encoding == ContainerSpec::Encoding::kRun) { + runBitmap[i / 8] |= static_cast(1u << (i % 8)); + } + } + data.append(reinterpret_cast(runBitmap.data()), runBitmapBytes); + + for (const auto& spec : containers) { + const auto cardMinus1 = static_cast(spec.values.size() - 1); + data.append(reinterpret_cast(&spec.key), 2); + data.append(reinterpret_cast(&cardMinus1), 2); + } + + const auto containerBytes = [&](const ContainerSpec& spec) -> uint32_t { + switch (spec.encoding) { + case ContainerSpec::Encoding::kArray: + return static_cast(spec.values.size()) * 2; + case ContainerSpec::Encoding::kBitset: + return kBitmapContainerBytes; + case ContainerSpec::Encoding::kRun: + return 2 + 4 * static_cast(toRuns(spec.values).size()); + } + return 0; + }; + + // The offset section is omitted for run-bearing bitmaps below the threshold. + const bool hasOffsetSection = + !hasRunContainers || numContainers >= kRunContainersNoOffsetThreshold; + if (hasOffsetSection) { + uint32_t offset = + 4 + runBitmapBytes + 4 * numContainers + 4 * numContainers; + for (const auto& spec : containers) { + data.append(reinterpret_cast(&offset), 4); + offset += containerBytes(spec); + } + } + + for (const auto& spec : containers) { + switch (spec.encoding) { + case ContainerSpec::Encoding::kArray: + for (auto value : spec.values) { + data.append(reinterpret_cast(&value), 2); + } + break; + case ContainerSpec::Encoding::kRun: { + const auto runs = toRuns(spec.values); + const auto numRuns = static_cast(runs.size()); + data.append(reinterpret_cast(&numRuns), 2); + for (auto [start, lengthMinus1] : runs) { + data.append(reinterpret_cast(&start), 2); + data.append(reinterpret_cast(&lengthMinus1), 2); + } + break; + } + case ContainerSpec::Encoding::kBitset: { + std::vector words(1'024, 0); + for (auto value : spec.values) { + words[value / 64] |= (1ULL << (value % 64)); + } + for (auto word : words) { + data.append(reinterpret_cast(&word), 8); + } + break; + } + } + } + return data; +} + // Concatenates the deletion-vector-v1 magic bytes with a serialized roaring // bitmap. The frame's length prefix and CRC-32 both cover these bytes. std::string magicAndVectorBytes(const std::string& bitmap) { @@ -227,6 +424,87 @@ std::vector expandRuns( } // Writes binary data to a temp file and returns the path. +// Builds a Puffin file around 'blob'. Layout per the Puffin spec: +// Magic | blob | Magic | footerPayload | payloadSize(4B LE) | flags(4B LE) +// | Magic +// Built by hand rather than via writePuffinFile so the tests can produce +// footers a spec-compliant writer never would. +constexpr std::string_view kPuffinMagic{"PFA1"}; + +std::string wrapInPuffinFile( + const std::string& blob, + const folly::dynamic& footer, + uint32_t flags = 0) { + const auto appendLittleEndian32 = [](std::string& out, uint32_t value) { + const uint32_t little = folly::Endian::little(value); + out.append(reinterpret_cast(&little), sizeof(little)); + }; + + const std::string footerJson = folly::toJson(footer); + std::string file; + file.append(kPuffinMagic); + file.append(blob); + file.append(kPuffinMagic); + file.append(footerJson); + appendLittleEndian32(file, static_cast(footerJson.size())); + appendLittleEndian32(file, flags); + file.append(kPuffinMagic); + return file; +} + +// Builds a Puffin footer describing one deletion-vector-v1 blob. +folly::dynamic makeDvBlobMeta( + uint64_t blobOffset, + uint64_t blobLength, + const std::string& referencedDataFile = "") { + folly::dynamic blobMeta = + folly::dynamic::object("type", "deletion-vector-v1")( + "fields", folly::dynamic::array(2'147'483'645)); + blobMeta["snapshot-id"] = -1; + blobMeta["sequence-number"] = -1; + blobMeta["offset"] = blobOffset; + blobMeta["length"] = blobLength; + if (!referencedDataFile.empty()) { + blobMeta["properties"] = + folly::dynamic::object("referenced-data-file", referencedDataFile); + } + return blobMeta; +} + +folly::dynamic makeDvFooter( + uint64_t blobOffset, + uint64_t blobLength, + const std::string& referencedDataFile = "") { + return folly::dynamic::object( + "blobs", + folly::dynamic::array( + makeDvBlobMeta(blobOffset, blobLength, referencedDataFile))); +} + +// Creates a DV delete file that carries no blob location at all: no typed +// contentOffset/contentLength and no legacy bounds map. The reader must fall +// back to the Puffin footer. +IcebergDeleteFile makeFooterLocatedDvFile( + const std::string& filePath, + uint64_t recordCount, + uint64_t fileSize, + const std::string& referencedDataFile = "") { + IcebergDeleteFile dvFile( + FileContent::kDeletionVector, + filePath, + dwio::common::FileFormat::DWRF, + recordCount, + fileSize, + /*equalityFieldIds=*/{}, + /*lowerBounds=*/{}, + /*upperBounds=*/{}, + /*dataSequenceNumber=*/0, + /*contentOffset=*/0, + /*contentLength=*/0); + dvFile.referencedDataFile = referencedDataFile; + return dvFile; +} + std::shared_ptr writeDvFile(const std::string& bitmapData) { auto tempFile = TempFilePath::create(); // Write directly via C++ streams since TempFilePath already creates the @@ -265,6 +543,38 @@ IcebergDeleteFile makeDvDeleteFile( /*contentLength=*/contentLength); } +// Creates a DV delete file that locates its blob through the legacy +// bounds-map encoding instead of the typed 'contentOffset'/'contentLength' +// fields. Leaving 'contentLength' at 0 is what selects that fallback. +IcebergDeleteFile makeLegacyBoundsDvFile( + const std::string& filePath, + uint64_t recordCount, + uint64_t fileSize, + const std::optional& blobOffset, + const std::optional& blobLength) { + std::unordered_map lowerBounds; + std::unordered_map upperBounds; + if (blobOffset.has_value()) { + lowerBounds[DeletionVectorReader::kDvOffsetFieldId] = *blobOffset; + } + if (blobLength.has_value()) { + upperBounds[DeletionVectorReader::kDvLengthFieldId] = *blobLength; + } + + return IcebergDeleteFile( + FileContent::kDeletionVector, + filePath, + dwio::common::FileFormat::DWRF, + recordCount, + fileSize, + /*equalityFieldIds=*/{}, + lowerBounds, + upperBounds, + /*dataSequenceNumber=*/0, + /*contentOffset=*/0, + /*contentLength=*/0); +} + // Extracts which bits are set in a bitmap buffer. std::vector getSetBits(const BufferPtr& bitmap, uint64_t size) { auto* raw = bitmap->as(); @@ -521,6 +831,136 @@ TEST_F(DeletionVectorReaderTest, runContainersWithOffsetHeader) { EXPECT_TRUE(reader.noMoreData()); } +TEST_F(DeletionVectorReaderTest, bitsetContainer) { + // 5000 deletes inside a single 64K block exceeds the 4096 array-container + // threshold, so the block is stored as a 1024-word bitset. The reader picks + // the container encoding from the cardinality, so this is the only shape + // that exercises its bitset branch. + std::vector positions; + positions.reserve(5'000); + for (int64_t i = 0; i < 5'000; ++i) { + positions.push_back(i * 2); + } + + auto bitmapData = serializeRoaringBitmapWithBitsetContainer(positions); + auto tempFile = writeDvFile(bitmapData); + auto dvFile = makeDvDeleteFile( + tempFile->getPath(), + positions.size(), + static_cast(bitmapData.size())); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + + const uint64_t numRows = 10'000; + auto bitmap = allocateBitmap(numRows); + reader.readDeletePositions(0, numRows, bitmap); + + std::vector expected; + expected.reserve(positions.size()); + for (auto position : positions) { + expected.push_back(static_cast(position)); + } + EXPECT_EQ(getSetBits(bitmap, numRows), expected); + EXPECT_TRUE(reader.noMoreData()); +} + +TEST_F(DeletionVectorReaderTest, bitsetContainerMixedWithArrayContainer) { + // A dense block followed by a sparse one. The reader must advance exactly + // 8192 bytes past the bitset before reading the array container, so a wrong + // stride here corrupts the second block rather than failing outright. + std::vector positions; + positions.reserve(4'100); + for (int64_t i = 0; i < 4'100; ++i) { + positions.push_back(i); + } + positions.push_back(65'536 + 7); + positions.push_back(65'536 + 9); + + auto bitmapData = serializeRoaringBitmapWithBitsetContainer(positions); + auto tempFile = writeDvFile(bitmapData); + auto dvFile = makeDvDeleteFile( + tempFile->getPath(), + positions.size(), + static_cast(bitmapData.size())); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + + std::vector expected(positions.begin(), positions.end()); + std::sort(expected.begin(), expected.end()); + EXPECT_EQ(reader.deletedPositions(), expected); +} + +TEST_F(DeletionVectorReaderTest, runContainersAcrossRoaring64Groups) { + // Run containers wrapped in a Roaring64 envelope. Existing run-container + // coverage uses the bare 32-bit format, which returns before the 64-bit + // group loop; only a multi-group input exercises the logic that skips past + // a run container to locate the next group's header. + const std::vector< + std::pair>>> + lowGroupRuns = {{0, {{10, 4}, {100, 2}}}}; + const std::vector< + std::pair>>> + highGroupRuns = {{0, {{20, 1}}}}; + + auto bitmapData = wrapInRoaring64( + {{0, serializeRoaringBitmapWithRuns(lowGroupRuns)}, + {1, serializeRoaringBitmapWithRuns(highGroupRuns)}}); + auto tempFile = writeDvFile(bitmapData); + + std::vector expected; + for (auto position : expandRuns(lowGroupRuns)) { + expected.push_back(static_cast(position)); + } + // Group 1 positions live at highBits 1, i.e. offset 2^32. + constexpr int64_t kHighGroupBase = int64_t{1} << 32; + for (auto position : expandRuns(highGroupRuns)) { + expected.push_back(kHighGroupBase + static_cast(position)); + } + std::sort(expected.begin(), expected.end()); + + auto dvFile = makeDvDeleteFile( + tempFile->getPath(), + expected.size(), + static_cast(bitmapData.size())); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + EXPECT_EQ(reader.deletedPositions(), expected); +} + +TEST_F( + DeletionVectorReaderTest, + arrayAndBitsetContainersAcrossRoaring64Groups) { + // The multi-group skip arithmetic differs per container encoding: an array + // container advances by 2 bytes per value while a bitset always advances by + // a fixed 8192. Pair a sparse group with a dense one so both strides must be + // right for the second group's header to be found. + std::vector denseGroupPositions; + denseGroupPositions.reserve(4'200); + for (int64_t i = 0; i < 4'200; ++i) { + denseGroupPositions.push_back(i); + } + + auto bitmapData = wrapInRoaring64( + {{0, serializeRoaringBitmapWithBitsetContainer({3, 11, 65'536 + 4})}, + {1, serializeRoaringBitmapWithBitsetContainer(denseGroupPositions)}}); + auto tempFile = writeDvFile(bitmapData); + + constexpr int64_t kHighGroupBase = int64_t{1} << 32; + std::vector expected = {3, 11, 65'536 + 4}; + for (auto position : denseGroupPositions) { + expected.push_back(kHighGroupBase + position); + } + std::sort(expected.begin(), expected.end()); + + auto dvFile = makeDvDeleteFile( + tempFile->getPath(), + expected.size(), + static_cast(bitmapData.size())); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + EXPECT_EQ(reader.deletedPositions(), expected); +} + TEST_F(DeletionVectorReaderTest, largePositionsMultipleContainers) { // Positions spanning two containers: one in container 0 (key=0), one in // container 1 (key=1, i.e. pos >= 65536). @@ -568,6 +1008,168 @@ TEST_F(DeletionVectorReaderTest, blobOffset) { EXPECT_TRUE(reader.noMoreData()); } +TEST_F(DeletionVectorReaderTest, locatesBlobFromPuffinFooterWhenOffsetsAbsent) { + // With no typed contentOffset/contentLength and no legacy bounds map, the + // reader used to fall back to treating the whole file -- leading magic and + // footer included -- as the blob, which cannot parse. A Puffin file is + // self-describing, so its footer is the authoritative fallback. + const std::vector positions = {3, 7, 42, 100}; + const auto blob = serializeRoaringBitmapNoRun(positions); + const auto file = wrapInPuffinFile(blob, makeDvFooter(4, blob.size())); + auto tempFile = writeDvFile(file); + + auto dvFile = makeFooterLocatedDvFile( + tempFile->getPath(), positions.size(), file.size()); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + EXPECT_EQ(reader.deletedPositions(), positions); +} + +TEST_F(DeletionVectorReaderTest, puffinFooterSelectsBlobByReferencedDataFile) { + // One Puffin file can hold vectors for several data files. The referenced + // data file, not blob order, decides which one applies to this split. + const std::vector wantedPositions = {5, 9}; + const std::vector otherPositions = {1, 2, 3}; + const auto otherBlob = serializeRoaringBitmapNoRun(otherPositions); + const auto wantedBlob = serializeRoaringBitmapNoRun(wantedPositions); + + const uint64_t otherOffset = 4; + const uint64_t wantedOffset = otherOffset + otherBlob.size(); + folly::dynamic footer = folly::dynamic::object( + "blobs", + folly::dynamic::array( + makeDvBlobMeta(otherOffset, otherBlob.size(), "/data/other.parquet"), + makeDvBlobMeta( + wantedOffset, wantedBlob.size(), "/data/wanted.parquet"))); + + const auto file = wrapInPuffinFile(otherBlob + wantedBlob, footer); + auto tempFile = writeDvFile(file); + + auto dvFile = makeFooterLocatedDvFile( + tempFile->getPath(), + wantedPositions.size(), + file.size(), + "/data/wanted.parquet"); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + EXPECT_EQ(reader.deletedPositions(), wantedPositions); +} + +TEST_F(DeletionVectorReaderTest, rejectsAmbiguousPuffinFileWithoutReference) { + // Two candidate vectors and nothing to choose between them: picking either + // would silently apply the wrong deletes to the scan. + const auto blob = serializeRoaringBitmapNoRun({1}); + folly::dynamic footer = folly::dynamic::object( + "blobs", + folly::dynamic::array( + makeDvBlobMeta(4, blob.size()), makeDvBlobMeta(4, blob.size()))); + + const auto file = wrapInPuffinFile(blob, footer); + auto tempFile = writeDvFile(file); + auto dvFile = makeFooterLocatedDvFile(tempFile->getPath(), 1, file.size()); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + VELOX_ASSERT_THROW( + reader.deletedPositions(), + "Puffin file has multiple deletion vectors and no referenced data file"); +} + +TEST_F(DeletionVectorReaderTest, rejectsPuffinFooterWithNoMatchingBlob) { + const auto blob = serializeRoaringBitmapNoRun({1}); + const auto file = wrapInPuffinFile( + blob, makeDvFooter(4, blob.size(), "/data/other.parquet")); + auto tempFile = writeDvFile(file); + auto dvFile = makeFooterLocatedDvFile( + tempFile->getPath(), 1, file.size(), "/data/wanted.parquet"); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + VELOX_ASSERT_THROW( + reader.deletedPositions(), + "Puffin footer has no deletion-vector blob for data file"); +} + +TEST_F(DeletionVectorReaderTest, rejectsPuffinFooterWithNoBlobs) { + // A structurally valid footer that declares no blobs at all. Iceberg allows + // an empty Puffin file; for a delete file it means there is nothing to read, + // which must be an error rather than an empty (silently no-op) vector. + const auto blob = serializeRoaringBitmapNoRun({1}); + const auto file = wrapInPuffinFile( + blob, folly::dynamic::object("blobs", folly::dynamic::array())); + auto tempFile = writeDvFile(file); + auto dvFile = makeFooterLocatedDvFile(tempFile->getPath(), 1, file.size()); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + VELOX_ASSERT_THROW( + reader.deletedPositions(), + "Puffin footer has no deletion-vector blob for data file"); +} + +TEST_F(DeletionVectorReaderTest, rejectsPuffinFooterWithNegativeOffset) { + // JSON integers are signed, but the blob location is carried as uint64_t. + // A negative offset must be rejected where it is read, not allowed to wrap + // into a huge unsigned value and be reported as an absurd out-of-range read. + const auto blob = serializeRoaringBitmapNoRun({1}); + auto footer = makeDvFooter(4, blob.size()); + footer["blobs"][0]["offset"] = -8; + const auto file = wrapInPuffinFile(blob, footer); + auto tempFile = writeDvFile(file); + auto dvFile = makeFooterLocatedDvFile(tempFile->getPath(), 1, file.size()); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + VELOX_ASSERT_THROW( + reader.deletedPositions(), "Puffin blob metadata has a negative offset"); +} + +TEST_F(DeletionVectorReaderTest, rejectsPuffinFileWithoutTrailingMagic) { + const auto blob = serializeRoaringBitmapNoRun({1}); + auto file = wrapInPuffinFile(blob, makeDvFooter(4, blob.size())); + file.replace(file.size() - 4, 4, "XXXX"); + + auto tempFile = writeDvFile(file); + auto dvFile = makeFooterLocatedDvFile(tempFile->getPath(), 1, file.size()); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + VELOX_ASSERT_THROW( + reader.deletedPositions(), + "Puffin file does not end with the expected magic"); +} + +TEST_F(DeletionVectorReaderTest, rejectsPuffinFooterPayloadSizeBeyondFile) { + // A payload size larger than the file would make the payload offset + // underflow into a huge unsigned read. + const auto blob = serializeRoaringBitmapNoRun({1}); + auto file = wrapInPuffinFile(blob, makeDvFooter(4, blob.size())); + const uint32_t absurdSize = folly::Endian::little(uint32_t{1} << 30); + file.replace( + file.size() - 12, + 4, + std::string(reinterpret_cast(&absurdSize), 4)); + + auto tempFile = writeDvFile(file); + auto dvFile = makeFooterLocatedDvFile(tempFile->getPath(), 1, file.size()); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + VELOX_ASSERT_THROW( + reader.deletedPositions(), + "Puffin footer payload size 1073741824 does not fit"); +} + +TEST_F(DeletionVectorReaderTest, rejectsCompressedPuffinFooter) { + // Bit 0 of the flags marks a compressed footer payload. Deletion vectors are + // always uncompressed, so reject rather than hand zstd bytes to the parser. + const auto blob = serializeRoaringBitmapNoRun({1}); + const auto file = + wrapInPuffinFile(blob, makeDvFooter(4, blob.size()), /*flags=*/1); + + auto tempFile = writeDvFile(file); + auto dvFile = makeFooterLocatedDvFile(tempFile->getPath(), 1, file.size()); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + VELOX_ASSERT_THROW( + reader.deletedPositions(), + "Compressed Puffin footer payloads are not supported"); +} + TEST_F(DeletionVectorReaderTest, constructorRejectsWrongContentType) { auto tempFile = TempFilePath::create(); { @@ -707,3 +1309,270 @@ TEST_F(DeletionVectorReaderTest, invalidBitmapBadCookie) { reader.readDeletePositions(0, 10, bitmap), "Unknown roaring bitmap cookie"); } + +TEST_F(DeletionVectorReaderTest, legacyBoundsMapLocatesBlob) { + // Callers that predate the typed contentOffset/contentLength fields encode + // the blob's location in the delete file's bounds maps instead. Prefix the + // bitmap with padding so a wrong offset cannot accidentally still parse. + const std::vector positions = {2, 9, 70'000}; + const std::string padding(16, '\xAB'); + auto bitmapData = serializeRoaringBitmapNoRun(positions); + auto tempFile = writeDvFile(padding + bitmapData); + + auto dvFile = makeLegacyBoundsDvFile( + tempFile->getPath(), + positions.size(), + padding.size() + bitmapData.size(), + std::to_string(padding.size()), + std::to_string(bitmapData.size())); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + EXPECT_EQ(reader.deletedPositions(), positions); +} + +TEST_F(DeletionVectorReaderTest, legacyBoundsMapRejectsNonNumericOffset) { + auto bitmapData = serializeRoaringBitmapNoRun({1}); + auto tempFile = writeDvFile(bitmapData); + + auto dvFile = makeLegacyBoundsDvFile( + tempFile->getPath(), + 1, + bitmapData.size(), + "not-a-number", + std::to_string(bitmapData.size())); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + VELOX_ASSERT_THROW( + reader.deletedPositions(), "Failed to parse DV blob offset"); +} + +TEST_F(DeletionVectorReaderTest, legacyBoundsMapRejectsNonNumericLength) { + auto bitmapData = serializeRoaringBitmapNoRun({1}); + auto tempFile = writeDvFile(bitmapData); + + auto dvFile = makeLegacyBoundsDvFile( + tempFile->getPath(), 1, bitmapData.size(), "0", "not-a-number"); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + VELOX_ASSERT_THROW( + reader.deletedPositions(), "Failed to parse DV blob length"); +} + +TEST_F(DeletionVectorReaderTest, emptyBitmapYieldsNoDeletes) { + // A container-less bitmap is structurally valid but selects nothing. The + // reader must return an empty position list rather than misparse the + // header, and a subsequent read must leave the delete bitmap untouched. + auto bitmapData = serializeRoaringBitmapNoRun({}); + auto tempFile = writeDvFile(bitmapData); + + // recordCount must be positive to construct a reader at all, so this also + // covers metadata that disagrees with the blob's actual contents. + auto dvFile = makeDvDeleteFile( + tempFile->getPath(), 1, static_cast(bitmapData.size())); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + EXPECT_TRUE(reader.deletedPositions().empty()); + + auto bitmap = allocateBitmap(64); + reader.readDeletePositions(0, 64, bitmap); + EXPECT_TRUE(getSetBits(bitmap, 64).empty()); +} + +TEST_F(DeletionVectorReaderTest, emptyGroupInRoaring64IsSkipped) { + // A Roaring64 group carrying no containers contributes nothing, but the + // reader must still step over its header to reach the following group. + auto bitmapData = wrapInRoaring64( + {{0, serializeRoaringBitmapNoRun({})}, + {1, serializeRoaringBitmapNoRun({6, 12})}}); + auto tempFile = writeDvFile(bitmapData); + + constexpr int64_t kHighGroupBase = int64_t{1} << 32; + const std::vector expected = { + kHighGroupBase + 6, kHighGroupBase + 12}; + + auto dvFile = makeDvDeleteFile( + tempFile->getPath(), + expected.size(), + static_cast(bitmapData.size())); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + EXPECT_EQ(reader.deletedPositions(), expected); +} + +TEST_F(DeletionVectorReaderTest, skipsPositionsBeforeRequestedRange) { + // Reading a batch that starts past earlier deletes must advance the cursor + // over them instead of mapping them into the batch-relative bitmap. This is + // the path a split beginning at a nonzero row offset takes. + const std::vector positions = {1, 50}; + auto bitmapData = serializeRoaringBitmapNoRun(positions); + auto tempFile = writeDvFile(bitmapData); + + auto dvFile = makeDvDeleteFile( + tempFile->getPath(), + positions.size(), + static_cast(bitmapData.size())); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + + // Batch covers absolute rows [10, 60): position 1 is behind it, 50 lands at + // batch-relative index 40. + auto bitmap = allocateBitmap(50); + reader.readDeletePositions(10, 50, bitmap); + + EXPECT_EQ(getSetBits(bitmap, 50), (std::vector{40})); +} + +// The roaring serialization format is defined as little-endian. Readers that +// take a caller-configured buffer have to assert its byte order explicitly, +// because a big-endian-configured buffer would silently decode garbage. +// +// There is no equivalent footgun here — the reader takes raw bytes and +// applies an explicit folly::Endian::little conversion to every field, so the +// interpretation is fixed regardless of caller or host. What is worth pinning +// is the observable behavior that check exists to produce: bytes serialized +// big-endian must be rejected, not silently misread. +// +// They are, though via a different guard than one might expect. The +// byte-swapped cookie (0x3A300000) matches neither serial cookie, so the blob +// is treated as the 64-bit format; the byte-swapped group count is then +// absurd and trips the Roaring64 sanity limit. The assertion below pins that +// whole chain, so a future change that loosens either the cookie check or the +// group limit cannot start silently accepting byte-swapped input. +TEST_F(DeletionVectorReaderTest, rejectsBigEndianSerializedBitmap) { + auto appendBigEndian32 = [](std::string& out, uint32_t value) { + const char bytes[4] = { + static_cast((value >> 24) & 0xFF), + static_cast((value >> 16) & 0xFF), + static_cast((value >> 8) & 0xFF), + static_cast(value & 0xFF)}; + out.append(bytes, sizeof(bytes)); + }; + + // The same empty 32-bit bitmap serializeRoaringBitmapNoRun({}) produces, + // but with both header words written big-endian. + std::string bigEndianBitmap; + appendBigEndian32(bigEndianBitmap, 12'346); + appendBigEndian32(bigEndianBitmap, 0); + + auto tempFile = writeDvFile(bigEndianBitmap); + auto dvFile = makeDvDeleteFile( + tempFile->getPath(), 1, static_cast(bigEndianBitmap.size())); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + VELOX_ASSERT_THROW( + reader.deletedPositions(), + "Roaring64Bitmap group count exceeds sanity limit"); +} + +TEST_F(DeletionVectorReaderTest, enforcesRoaring64GroupKeyBound) { + // Iceberg caps the Roaring64 group key at 2147483646 because readers load it + // back as a signed 32-bit int. A key at or above 2^31 would shift into the + // sign bit of `key << 32` and surface as a negative row position. + // DeletionVectorWriter already refuses to produce such a blob; the reader + // must reject one written by anything else rather than emit garbage. + const auto readPositions = [&](uint32_t groupKey) { + auto bitmapData = + wrapInRoaring64({{groupKey, serializeRoaringBitmapNoRun({5})}}); + auto tempFile = writeDvFile(bitmapData); + auto dvFile = makeDvDeleteFile( + tempFile->getPath(), 1, static_cast(bitmapData.size())); + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + return reader.deletedPositions(); + }; + + constexpr uint32_t kMaxGroupKey = 2'147'483'646; + EXPECT_EQ( + readPositions(kMaxGroupKey), + (std::vector{(int64_t{kMaxGroupKey} << 32) + 5})); + + VELOX_ASSERT_THROW( + readPositions(kMaxGroupKey + 1), + "Roaring64Bitmap group key exceeds the maximum the Iceberg " + "deletion-vector format can represent"); +} + +// The scenarios below mirror the shapes Iceberg's own position-index tests +// cover. The bitmaps are built here rather than taken from anywhere, so they +// exercise the same encodings without depending on external fixtures. + +TEST_F(DeletionVectorReaderTest, smallAlternatingValues) { + // Sparse odd positions in a single array container — the simplest shape a + // real deletion vector takes. + const std::vector positions = {1, 3, 5, 7, 9}; + auto bitmapData = serializeRoaringBitmapNoRun(positions); + auto tempFile = writeDvFile(bitmapData); + + auto dvFile = makeDvDeleteFile( + tempFile->getPath(), + positions.size(), + static_cast(bitmapData.size())); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + EXPECT_EQ(reader.deletedPositions(), positions); +} + +TEST_F(DeletionVectorReaderTest, smallAndLargeValuesPast31Bits) { + // Two array containers far apart, the second past 2^31. The container key + // there is 32767, so a reader that sign-extends the 16-bit key or the + // resulting offset lands on a negative position. + const std::vector positions = { + 100, 101, 2'147'483'747, 2'147'483'748}; + auto bitmapData = serializeRoaringBitmapNoRun(positions); + auto tempFile = writeDvFile(bitmapData); + + auto dvFile = makeDvDeleteFile( + tempFile->getPath(), + positions.size(), + static_cast(bitmapData.size())); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + EXPECT_EQ(reader.deletedPositions(), positions); +} + +TEST_F(DeletionVectorReaderTest, allContainerEncodingsInOneBitmap) { + // Array, run, and bitset containers coexisting in a single bitmap, wrapped + // in a second Roaring64 group. Each encoding advances the read cursor by a + // different rule, so the reader has to switch strategies three times and + // still land on the next group's header. + // + // Three containers with runs present also puts this below the threshold at + // which the offset section is written, exercising the no-offset path. + std::vector denseValues; + denseValues.reserve(5'000); + for (uint16_t i = 0; i < 5'000; ++i) { + denseValues.push_back(static_cast(i * 2)); + } + + const std::vector lowGroup = { + {/*key=*/0, ContainerSpec::Encoding::kArray, {5, 7}}, + {/*key=*/1, ContainerSpec::Encoding::kRun, {1, 2, 3, 4}}, + {/*key=*/2, ContainerSpec::Encoding::kBitset, denseValues}, + }; + const std::vector highGroup = { + {/*key=*/0, ContainerSpec::Encoding::kArray, {9}}, + }; + + auto bitmapData = wrapInRoaring64( + {{0, serializeRoaringBitmapMixed(lowGroup)}, + {1, serializeRoaringBitmapMixed(highGroup)}}); + auto tempFile = writeDvFile(bitmapData); + + constexpr int64_t kHighGroupBase = int64_t{1} << 32; + std::vector expected = {5, 7}; + for (int64_t i = 1; i <= 4; ++i) { + expected.push_back(65'536 + i); + } + for (auto value : denseValues) { + expected.push_back(131'072 + value); + } + expected.push_back(kHighGroupBase + 9); + std::sort(expected.begin(), expected.end()); + + auto dvFile = makeDvDeleteFile( + tempFile->getPath(), + expected.size(), + static_cast(bitmapData.size())); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + EXPECT_EQ(reader.deletedPositions(), expected); +} diff --git a/velox/connectors/hive/iceberg/tests/DeletionVectorWriterTest.cpp b/velox/connectors/hive/iceberg/tests/DeletionVectorWriterTest.cpp index 5ddd42b9fe7..d48fdd2495a 100644 --- a/velox/connectors/hive/iceberg/tests/DeletionVectorWriterTest.cpp +++ b/velox/connectors/hive/iceberg/tests/DeletionVectorWriterTest.cpp @@ -18,10 +18,14 @@ #include +#include +#include #include #include #include +#include +#include #include "velox/common/base/BitUtil.h" #include "velox/common/base/tests/GTestUtils.h" @@ -73,6 +77,50 @@ class DeletionVectorWriterTest : public ::testing::Test { return AlignedBuffer::allocate(numBytes, pool_.get(), 0); } + /// Serializes 'positions', reads the blob back through + /// DeletionVectorReader, and returns every position the reader recovered. + /// + /// Prefer this over verifyRoundTrip for inputs spanning a wide range: + /// verifyRoundTrip walks the whole [0, maxPos] space one batch at a time, + /// which is impractical once positions reach into the 2^32 range. + std::vector roundTrip(const std::vector& positions) { + DeletionVectorWriter writer; + writer.addDeletedPositions(positions); + const auto blobData = writer.serialize(); + + auto tempFile = TempFilePath::create(); + { + std::ofstream out( + tempFile->getPath(), std::ios::binary | std::ios::trunc); + out.write(blobData.data(), static_cast(blobData.size())); + } + + IcebergDeleteFile dvFile( + FileContent::kDeletionVector, + tempFile->getPath(), + dwio::common::FileFormat::DWRF, + writer.numDistinctPositions(), + static_cast(blobData.size()), + /*equalityFieldIds=*/{}, + /*lowerBounds=*/{}, + /*upperBounds=*/{}, + /*dataSequenceNumber=*/0, + /*contentOffset=*/0, + /*contentLength=*/static_cast(blobData.size())); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + return reader.deletedPositions(); + } + + /// Returns 'positions' sorted and de-duplicated — what a round trip through + /// a roaring bitmap is expected to yield, since the bitmap is a set. + static std::vector sortedUnique(std::vector positions) { + std::sort(positions.begin(), positions.end()); + positions.erase( + std::unique(positions.begin(), positions.end()), positions.end()); + return positions; + } + /// Writes serialized bitmap to a temp file, reads it back with /// DeletionVectorReader, and verifies the positions match. void verifyRoundTrip( @@ -257,6 +305,105 @@ TEST_F(DeletionVectorWriterTest, fourOrMoreContainersWithOffsets) { verifyRoundTrip(positions, 5 * 65536 + 100); } +// Decodes the Puffin footer of a file written by writePuffinFile and returns +// the parsed footer JSON. Layout per the Puffin spec: +// Magic Blob... Footer +// Footer := Magic FooterPayload FooterPayloadSize Flags Magic +// with FooterPayloadSize and Flags each a 4-byte little-endian value. The +// trailer is located from the end of the file so nothing here depends on the +// writer's own offset arithmetic. +namespace { +folly::dynamic readPuffinFooter(const std::string& path) { + std::ifstream in(path, std::ios::binary); + const std::string bytes( + (std::istreambuf_iterator(in)), std::istreambuf_iterator()); + + constexpr size_t kMagicSize = 4; + constexpr size_t kTrailerSize = kMagicSize + 4 + 4; + EXPECT_GE(bytes.size(), kMagicSize + kTrailerSize); + EXPECT_EQ(bytes.substr(0, kMagicSize), "PFA1") << "missing leading magic"; + EXPECT_EQ(bytes.substr(bytes.size() - kMagicSize), "PFA1") + << "missing trailing magic"; + + const auto readLittleEndian32 = [&](size_t offset) { + uint32_t value; + std::memcpy(&value, bytes.data() + offset, sizeof(value)); + return folly::Endian::little(value); + }; + + const size_t flagsOffset = bytes.size() - kMagicSize - 4; + const size_t sizeOffset = flagsOffset - 4; + EXPECT_EQ(readLittleEndian32(flagsOffset), 0u) + << "flags must be 0: an uncompressed footer payload is bit 0"; + + const uint32_t payloadSize = readLittleEndian32(sizeOffset); + const size_t payloadOffset = sizeOffset - payloadSize; + EXPECT_EQ(bytes.substr(payloadOffset - kMagicSize, kMagicSize), "PFA1") + << "footer payload must be preceded by magic"; + + return folly::parseJson(bytes.substr(payloadOffset, payloadSize)); +} +} // namespace + +TEST_F(DeletionVectorWriterTest, puffinFooterIsSpecCompliant) { + // Nothing in our own read path parses the Puffin footer -- + // DeletionVectorReader locates the blob from the manifest's contentOffset -- + // so the footer is written blind. Other Iceberg engines do parse it, and + // FileMetadataParser.blobMetadataFromJson treats type, fields, snapshot-id, + // sequence-number, offset and length as required, reading fields as a list + // of integers. A footer that omits or mistypes any of them makes the whole + // deletion vector unreadable outside Velox. + DeletionVectorWriter writer; + writer.addDeletedPositions({3, 7, 42, 100}); + const auto blobData = writer.serialize(); + + auto tempDir = TempDirectoryPath::create(); + const std::string puffinPath = + std::string(tempDir->getPath()) + "/footer.puffin"; + auto sink = dwio::common::FileSink::create( + "file:" + puffinPath, {.pool = pool_.get()}); + auto [blobOffset, blobLength] = writePuffinFile( + *sink, + *pool_, + blobData, + "/data/test-data-file.parquet", + /*cardinality=*/4); + sink->close(); + + const auto footer = readPuffinFooter(puffinPath); + ASSERT_TRUE(footer.isObject()); + ASSERT_TRUE(footer["blobs"].isArray()); + ASSERT_EQ(footer["blobs"].size(), 1); + const auto& blob = footer["blobs"][0]; + + EXPECT_EQ(blob["type"].asString(), "deletion-vector-v1"); + + // Iceberg's BaseDVFileWriter sets this to the single row-position metadata + // column ID, as a list of plain integers. + constexpr int64_t kRowPositionFieldId = 2'147'483'645; + ASSERT_TRUE(blob["fields"].isArray()) << "fields must be a list"; + ASSERT_EQ(blob["fields"].size(), 1); + ASSERT_TRUE(blob["fields"][0].isInt()) + << "fields entries must be integers, not objects"; + EXPECT_EQ(blob["fields"][0].asInt(), kRowPositionFieldId); + + // Required by the parser even though a freshly written DV has no snapshot + // assigned yet; Iceberg writes -1 for both. + ASSERT_TRUE(blob.count("snapshot-id")) << "snapshot-id is required"; + ASSERT_TRUE(blob.count("sequence-number")) << "sequence-number is required"; + EXPECT_EQ(blob["snapshot-id"].asInt(), -1); + EXPECT_EQ(blob["sequence-number"].asInt(), -1); + + EXPECT_EQ(blob["offset"].asInt(), static_cast(blobOffset)); + EXPECT_EQ(blob["length"].asInt(), static_cast(blobLength)); + + const auto& properties = blob["properties"]; + EXPECT_EQ( + properties["referenced-data-file"].asString(), + "/data/test-data-file.parquet"); + EXPECT_EQ(properties["cardinality"].asString(), "4"); +} + TEST_F(DeletionVectorWriterTest, puffinFileRoundTrip) { DeletionVectorWriter writer; writer.addDeletedPositions({3, 7, 42, 100}); @@ -314,6 +461,55 @@ TEST_F(DeletionVectorWriterTest, puffinFileRoundTrip) { EXPECT_EQ(setBits, (std::vector{3, 7, 42, 100})); } +// Reads a writePuffinFile output back with NO bounds-map location, so the +// reader must parse the Puffin footer (locateBlobFromPuffinFooter) to find the +// blob. Unlike puffinFileRoundTrip -- which supplies explicit offset/length and +// therefore skips the footer parser -- this exercises the writer's footer +// layout against the reader's backwards footer parse, catching any drift in +// trailer size, magic placement, or payload-size/flags encoding. +TEST_F(DeletionVectorWriterTest, puffinFooterFallbackRoundTrip) { + DeletionVectorWriter writer; + writer.addDeletedPositions({3, 7, 42, 100}); + auto blobData = writer.serialize(); + + auto tempDir = TempDirectoryPath::create(); + const std::string puffinPath = + std::string(tempDir->getPath()) + "/test-dv-footer.puffin"; + auto sink = dwio::common::FileSink::create( + "file:" + puffinPath, {.pool = pool_.get()}); + VELOX_CHECK_NOT_NULL(sink); + writePuffinFile( + *sink, + *pool_, + blobData, + "/data/test-data-file.parquet", + /*cardinality=*/4); + sink->close(); + + std::ifstream in(puffinPath, std::ios::binary | std::ios::ate); + auto fileSize = static_cast(in.tellg()); + + // Empty bounds maps: with no offset/length the reader falls through to + // Puffin-footer parsing and selects the single deletion-vector blob. + IcebergDeleteFile dvFile( + FileContent::kDeletionVector, + puffinPath, + dwio::common::FileFormat::DWRF, + 4, + fileSize, + {}, + {}, + {}); + + DeletionVectorReader reader(dvFile, 0, pool_.get(), nullptr); + + auto bitmap = allocateBitmap(200); + reader.readDeletePositions(0, 200, bitmap); + + auto setBits = getSetBits(bitmap, 200); + EXPECT_EQ(setBits, (std::vector{3, 7, 42, 100})); +} + // Verifies the on-disk deletion-vector-v1 blob matches the Iceberg V3 spec // frame: [length: 4B BE][magic D1 D3 39 64][bitmap][CRC-32: 4B BE], where the // length and CRC-32 cover the magic + bitmap. This is what makes the DV @@ -401,3 +597,123 @@ TEST_F(DeletionVectorWriterTest, mixed32And64BitPositions) { }; verifyRoundTrip(positions, 2'048); } + +/// Verifies that duplicate positions collapse in the cardinality reported for +/// the DV blob. Seeding a writer from an existing deletion vector and then +/// adding overlapping new deletes makes 'positions_' hold duplicates, so +/// 'numPositions' overcounts and only 'numDistinctPositions' matches the +/// cardinality Iceberg expects in the blob metadata. +TEST_F(DeletionVectorWriterTest, numDistinctPositionsIgnoresDuplicates) { + DeletionVectorWriter writer; + writer.addDeletedPositions({7, 3, 7, 1, 3, 3}); + + EXPECT_EQ(writer.numPositions(), 6); + EXPECT_EQ(writer.numDistinctPositions(), 3); +} + +// Randomized round trips across sparse, dense, and mixed position sets. +// Hand-built fixtures only exercise the container shapes we thought to write +// down; random inputs cross the array/bitset thresholds and 64-bit group +// boundaries in combinations we did not enumerate. +// +// Seeds are fixed so a failure is reproducible, and reported on failure so a +// counterexample can be replayed directly. +TEST_F(DeletionVectorWriterTest, randomSparsePositionsRoundTrip) { + // Few positions spread over a wide 64-bit range: many container keys across + // several Roaring64 groups, each holding a small array container. + constexpr uint64_t kSeed = 0x1234'5678'9ABC'DEF0ULL; + constexpr int kNumPositions = 2'000; + constexpr int64_t kRange = int64_t{1} << 34; + + std::mt19937_64 rng(kSeed); + std::uniform_int_distribution positionDist(0, kRange - 1); + + std::vector positions; + positions.reserve(kNumPositions); + for (int i = 0; i < kNumPositions; ++i) { + positions.push_back(positionDist(rng)); + } + + EXPECT_EQ(roundTrip(positions), sortedUnique(positions)) << "seed=" << kSeed; +} + +TEST_F(DeletionVectorWriterTest, randomDensePositionsRoundTrip) { + // Enough positions inside one 64K block to exceed the 4096 array-container + // threshold, so the block serializes as a bitset. Deliberately includes + // duplicates: the bitmap is a set, so they must collapse. + constexpr uint64_t kSeed = 0x0FED'CBA9'8765'4321ULL; + constexpr int kNumDraws = 30'000; + constexpr int64_t kBlockBase = int64_t{7} << 16; + + std::mt19937_64 rng(kSeed); + std::uniform_int_distribution offsetDist(0, 65'535); + + std::vector positions; + positions.reserve(kNumDraws); + for (int i = 0; i < kNumDraws; ++i) { + positions.push_back(kBlockBase + offsetDist(rng)); + } + + const auto expected = sortedUnique(positions); + ASSERT_GT(expected.size(), 4'096) + << "draws should exceed the array-container threshold"; + EXPECT_EQ(roundTrip(positions), expected) << "seed=" << kSeed; +} + +TEST_F(DeletionVectorWriterTest, randomMixedPositionsRoundTrip) { + // Dense block, sparse scatter, and a contiguous run in one bitmap, split + // across two Roaring64 groups. This is the shape most likely to expose an + // offset or stride error, because the reader must step over containers of + // different encodings to find the next one. + constexpr uint64_t kSeed = 0x2468'ACE0'1357'9BDFULL; + constexpr int64_t kHighGroupBase = int64_t{1} << 32; + + std::mt19937_64 rng(kSeed); + std::vector positions; + + // Dense block in group 0. + std::uniform_int_distribution denseDist(0, 65'535); + for (int i = 0; i < 20'000; ++i) { + positions.push_back(denseDist(rng)); + } + // Sparse scatter across group 0's upper blocks. + std::uniform_int_distribution sparseDist(65'536, (int64_t{1} << 31)); + for (int i = 0; i < 500; ++i) { + positions.push_back(sparseDist(rng)); + } + // Contiguous run in group 1. + for (int64_t i = 0; i < 3'000; ++i) { + positions.push_back(kHighGroupBase + 1'000 + i); + } + // Sparse scatter in group 1. + for (int i = 0; i < 200; ++i) { + positions.push_back(kHighGroupBase + sparseDist(rng)); + } + + EXPECT_EQ(roundTrip(positions), sortedUnique(positions)) << "seed=" << kSeed; +} + +// The Roaring64 group key is read back as a signed 32-bit int, so a position +// whose high word reaches 2^31 would deserialize as a negative key and be +// rejected by spec-compliant readers. Failing at insert time keeps us from +// writing a blob Iceberg cannot read. +TEST_F(DeletionVectorWriterTest, rejectsPositionsOutsideRepresentableRange) { + DeletionVectorWriter writer; + + VELOX_ASSERT_THROW(writer.addDeletedPosition(-1), "must be non-negative"); + VELOX_ASSERT_THROW( + writer.addDeletedPosition(DeletionVectorWriter::kMaxPosition + 1), + "exceeds the maximum"); + VELOX_ASSERT_THROW( + writer.addDeletedPosition(std::numeric_limits::max()), + "exceeds the maximum"); + + // None of the rejected positions were recorded. + EXPECT_EQ(writer.numPositions(), 0); + + // The bound itself is representable and round-trips. + writer.addDeletedPosition(DeletionVectorWriter::kMaxPosition); + EXPECT_EQ( + roundTrip({DeletionVectorWriter::kMaxPosition}), + std::vector{DeletionVectorWriter::kMaxPosition}); +} diff --git a/velox/connectors/hive/iceberg/tests/EqualityDeleteFileReaderTest.cpp b/velox/connectors/hive/iceberg/tests/EqualityDeleteFileReaderTest.cpp index f8f49c07fb9..d67bf377086 100644 --- a/velox/connectors/hive/iceberg/tests/EqualityDeleteFileReaderTest.cpp +++ b/velox/connectors/hive/iceberg/tests/EqualityDeleteFileReaderTest.cpp @@ -16,133 +16,195 @@ #include -#include "velox/common/file/FileSystems.h" -#include "velox/common/testutil/TempDirectoryPath.h" -#include "velox/connectors/ConnectorRegistry.h" -#include "velox/connectors/hive/iceberg/IcebergConnector.h" -#include "velox/connectors/hive/iceberg/IcebergDeleteFile.h" -#include "velox/connectors/hive/iceberg/IcebergSplit.h" +#include "velox/connectors/hive/iceberg/tests/IcebergTestBase.h" #include "velox/exec/tests/utils/AssertQueryBuilder.h" -#include "velox/exec/tests/utils/HiveConnectorTestBase.h" -#include "velox/exec/tests/utils/PlanBuilder.h" -namespace facebook::velox::connector::hive::iceberg { - -using namespace facebook::velox::exec::test; - -namespace { - -const std::string kIcebergConnectorId = "test-iceberg-eq-delete"; - -} // namespace +namespace facebook::velox::connector::hive::iceberg::test { + +using exec::test::assertEqualResults; +using exec::test::AssertQueryBuilder; + +// --------------------------------------------------------------------------- +// FileWriteMode — parameterizes file format and whether iceberg.id field IDs +// are stamped in the written files. +// +// Valid combinations: +// {DWRF, withFieldIds=false} — plain DWRF, positional-fallback path +// {DWRF, withFieldIds=true } — DWRF with iceberg.id attrs, kFieldId path +// {PARQUET, withFieldIds=true } — Parquet with field IDs (only valid mode) +// --------------------------------------------------------------------------- +struct FileWriteMode { + dwio::common::FileFormat format; + bool withFieldIds; + + std::string toString() const { + const std::string fmt = + format == dwio::common::FileFormat::PARQUET ? "Parquet" : "Dwrf"; + return fmt + (withFieldIds ? "_FieldId" : "_Positional"); + } +}; -/// End-to-end tests for equality deletes via the IcebergSplitReader. -/// These tests write DWRF data files and delete files, then execute -/// table scans verifying that matching rows are filtered out. -class EqualityDeleteFileReaderTest : public HiveConnectorTestBase { +// --------------------------------------------------------------------------- +// EqualityDeleteFileReaderTest +// +// Non-parameterized base fixture. Partition-column and evolved-schema tests +// live here as TEST_F. The parameterized class inherits this. +// --------------------------------------------------------------------------- +class EqualityDeleteFileReaderTest : public IcebergTestBase { protected: - void SetUp() override { - HiveConnectorTestBase::SetUp(); - IcebergConnectorFactory icebergFactory; - auto icebergConnector = icebergFactory.newConnector( - kIcebergConnectorId, - std::make_shared( - std::unordered_map()), - ioExecutor_.get()); - connector::ConnectorRegistry::global().insert( - icebergConnector->connectorId(), icebergConnector); + /// Writes a file in the format described by 'mode'. When + /// Writes a file according to 'mode': + /// DWRF + withFieldIds=false — plain DWRF, no iceberg.id attributes + /// (positional fallback in DwrfReader) + /// DWRF + withFieldIds=true — DWRF with iceberg.id footer attributes + /// (kFieldId / renameByFieldId path) + /// PARQUET + withFieldIds=true — Parquet with field_id metadata stamped + /// (kParquetFieldId path) + /// PARQUET + withFieldIds=false— Parquet without field_id metadata; + /// buildFieldIds() returns empty so + /// kParquetFieldId is not activated and the + /// reader uses kPosition (ordinal binding). + std::shared_ptr writeFile( + const std::vector& data, + const std::vector& fieldIds, + const FileWriteMode& mode) { + if (mode.format == dwio::common::FileFormat::PARQUET) { +#ifdef VELOX_ENABLE_PARQUET + // Pass fieldIds only when the mode wants them stamped; empty vector means + // no field_id metadata in the Parquet schema -> kPosition mode. + return writeParquetFile( + data, mode.withFieldIds ? fieldIds : std::vector{}); +#else + VELOX_FAIL("Parquet support is not enabled"); +#endif + } + return mode.withFieldIds ? writeDwrfFileWithFieldIds(data, fieldIds) + : writeDataFile(data); } - void TearDown() override { - connector::ConnectorRegistry::global().erase(kIcebergConnectorId); - HiveConnectorTestBase::TearDown(); + /// Creates splits for the data file using the format in 'mode', attaching + /// delete files and (optionally) partition keys. + std::vector> makeSplits( + const std::string& dataFilePath, + const std::unordered_map>& + partitionKeys, + const FileWriteMode& mode, + const std::vector& deleteFiles = {}, + int64_t dataSequenceNumber = 0, + const std::unordered_map>& + identityPartitionKeys = {}) { + fileFormat_ = mode.format; // makeIcebergSplits uses fileFormat_ + return makeIcebergSplits( + dataFilePath, + deleteFiles, + partitionKeys, + /*splitCount=*/1, + /*infoColumns=*/{}, + dataSequenceNumber, + identityPartitionKeys); } - uint64_t getFileSize(const std::string& path) { - return filesystems::getFileSystem(path, nullptr) - ->openFileForRead(path) - ->size(); + /// Builds an IcebergDeleteFile descriptor. + IcebergDeleteFile makeDeleteFile( + const std::string& path, + const std::vector& equalityFieldIds, + dwio::common::FileFormat format, + int64_t recordCount = 2, + int64_t deleteSeqNum = 0) { + return IcebergDeleteFile( + FileContent::kEqualityDeletes, + path, + format, + recordCount, + getFileSize(path), + equalityFieldIds, + /*lowerBounds=*/{}, + /*upperBounds=*/{}, + deleteSeqNum); } +}; - /// Writes a DWRF data file containing the given vectors. - std::shared_ptr writeDataFile( - const std::vector& data) { - auto file = common::testutil::TempFilePath::create(); - writeToFile(file->getPath(), data); - return file; +// --------------------------------------------------------------------------- +// EqualityDeleteFileReaderTestP — parameterized over FileWriteMode. +// +// Every TEST_P runs three times: +// 1. Dwrf_Positional — DWRF, no field IDs (kPosition fallback) +// 2. Dwrf_FieldId — DWRF, iceberg.id attrs (kFieldId path) +// 3. Parquet_FieldId — Parquet field_id (kParquetFieldId path) +// --------------------------------------------------------------------------- +class EqualityDeleteFileReaderTestP + : public EqualityDeleteFileReaderTest, + public ::testing::WithParamInterface { + protected: + void SetUp() override { +#ifndef VELOX_ENABLE_PARQUET + if (GetParam().format == dwio::common::FileFormat::PARQUET) { + GTEST_SKIP() << "Parquet support not enabled"; + } +#endif + EqualityDeleteFileReaderTest::SetUp(); + fileFormat_ = GetParam().format; } - /// Writes a DWRF delete file containing the equality delete rows. - std::shared_ptr writeEqDeleteFile( - const std::vector& deleteData) { - auto file = common::testutil::TempFilePath::create(); - writeToFile(file->getPath(), deleteData); - return file; + std::shared_ptr writeDataFileP( + const std::vector& data, + const std::vector& fieldIds) { + return writeFile(data, fieldIds, GetParam()); } - /// Creates splits with equality delete files attached. - std::vector> makeSplits( + std::vector> makeSplitsP( const std::string& dataFilePath, const std::vector& deleteFiles = {}, - int64_t dataSequenceNumber = 0) { - return makeSplits( - dataFilePath, - /*partitionKeys=*/{}, - deleteFiles, - dataSequenceNumber); - } - - /// Creates splits with equality delete files and partition keys attached. - /// Use this overload to exercise the equality-delete augmentation for - /// partition columns missing from the user's projection. - std::vector> makeSplits( - const std::string& dataFilePath, + int64_t dataSequenceNumber = 0, const std::unordered_map>& - partitionKeys, - const std::vector& deleteFiles, - int64_t dataSequenceNumber = 0) { - auto fileSize = getFileSize(dataFilePath); - return {std::make_shared( - kIcebergConnectorId, + partitionKeys = {}) { + return makeSplits( dataFilePath, - dwio::common::FileFormat::DWRF, - 0, - fileSize, partitionKeys, - std::nullopt, - std::unordered_map{}, - nullptr, - /*cacheable=*/true, + GetParam(), deleteFiles, - std::unordered_map{}, - std::nullopt, - dataSequenceNumber)}; - } - - /// Builds a table scan plan node with the given schema. - core::PlanNodePtr makeTableScanPlan(const RowTypePtr& rowType) { - return makeTableScanPlan(rowType, rowType); - } - - /// Builds a table scan plan node with separate output and table column - /// schemas. Use this when the user's projection ('outputType') does not - /// contain every column referenced by an equality delete file - /// ('dataColumns' must contain the full table schema so the equality - /// column resolution can map field IDs to names). - core::PlanNodePtr makeTableScanPlan( - const RowTypePtr& outputType, - const RowTypePtr& dataColumns) { - return PlanBuilder() - .startTableScan(kIcebergConnectorId) - .outputType(outputType) - .dataColumns(dataColumns) - .endTableScan() - .planNode(); + dataSequenceNumber, + {}); } }; +// Three FileWriteMode combinations: +// +// Dwrf_Positional {DWRF, false} — no iceberg.id attributes; DwrfReader +// falls back to positional name mapping. +// +// Dwrf_FieldId {DWRF, true} — "iceberg.id" footer attributes; +// DwrfReader uses renameByFieldId. +// +// Parquet_FieldId {PARQUET, true} — Parquet field_id metadata present; +// IcebergSplitReader activates +// kParquetFieldId mapping. +// +// {PARQUET, false} does not work end-to-end: without field IDs, +// buildFieldIds() returns empty, IcebergSplitReader does not set a fileSchema +// on baseReaderOpts_, and the Parquet reader has no requested-type to match +// positions against — physical columns cannot be bound to the scan-spec names, +// yielding all-null output. All production Iceberg Parquet files carry field +// IDs; this combination is therefore not a valid use case. +INSTANTIATE_TEST_SUITE_P( + Formats, + EqualityDeleteFileReaderTestP, + ::testing::Values( + FileWriteMode{dwio::common::FileFormat::DWRF, /*withFieldIds=*/false}, + FileWriteMode{dwio::common::FileFormat::DWRF, /*withFieldIds=*/true}, + FileWriteMode{ + dwio::common::FileFormat::PARQUET, + /*withFieldIds=*/true}), + [](const ::testing::TestParamInfo& info) { + return info.param.toString(); + }); + +// =========================================================================== +// Parameterized tests +// =========================================================================== + /// Verifies that base rows matching the equality delete file are removed. -TEST_F(EqualityDeleteFileReaderTest, basicSingleColumnDelete) { +TEST_P(EqualityDeleteFileReaderTestP, basicSingleColumnDelete) { auto rowType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); auto baseData = makeRowVector( @@ -152,26 +214,16 @@ TEST_F(EqualityDeleteFileReaderTest, basicSingleColumnDelete) { makeFlatVector( {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"}), }); - auto dataFile = writeDataFile({baseData}); + auto dataFile = writeDataFileP({baseData}, {1, 2}); - // Delete rows where id == 3 or id == 7. - auto deleteData = makeRowVector( - {"id"}, - { - makeFlatVector({3, 7}), - }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto deleteData = makeRowVector({"id"}, {makeFlatVector({3, 7})}); + auto eqDeleteFile = writeDataFileP({deleteData}, {1}); - IcebergDeleteFile icebergDeleteFile( - FileContent::kEqualityDeletes, - eqDeleteFile->getPath(), - dwio::common::FileFormat::DWRF, - 2, - getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1}); // field ID 1 = column 0 = "id" + auto icebergDeleteFile = + makeDeleteFile(eqDeleteFile->getPath(), {1}, GetParam().format); - auto splits = makeSplits(dataFile->getPath(), {icebergDeleteFile}); - auto plan = makeTableScanPlan(rowType); + auto splits = makeSplitsP(dataFile->getPath(), {icebergDeleteFile}); + auto plan = makeIcebergTableScanPlan(rowType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); auto expected = makeRowVector( @@ -180,19 +232,12 @@ TEST_F(EqualityDeleteFileReaderTest, basicSingleColumnDelete) { makeFlatVector({0, 1, 2, 4, 5, 6, 8, 9}), makeFlatVector({"a", "b", "c", "e", "f", "g", "i", "j"}), }); - assertEqualResults({expected}, {result}); } -/// Regression test for the bug where IcebergSplitReader fails with -/// "Column not found in row: " when an equality-delete column is not -/// part of the user's projection. The reader must augment its scan spec to -/// physically read the equality-delete column, apply the delete, and then -/// project the column away from the output before returning to the operator. -TEST_F(EqualityDeleteFileReaderTest, equalityColumnNotInProjection) { +/// Regression test: equality-delete column absent from the user's projection. +TEST_P(EqualityDeleteFileReaderTestP, equalityColumnNotInProjection) { auto tableType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); - // The user only selects 'value'. The equality delete is on 'id', which is - // NOT in the projection — this is the case that previously failed. auto outputType = ROW({"value"}, {VARCHAR()}); auto baseData = makeRowVector( @@ -202,43 +247,26 @@ TEST_F(EqualityDeleteFileReaderTest, equalityColumnNotInProjection) { makeFlatVector( {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"}), }); - auto dataFile = writeDataFile({baseData}); + auto dataFile = writeDataFileP({baseData}, {1, 2}); - // Delete rows where id == 3 or id == 7. - auto deleteData = makeRowVector( - {"id"}, - { - makeFlatVector({3, 7}), - }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto deleteData = makeRowVector({"id"}, {makeFlatVector({3, 7})}); + auto eqDeleteFile = writeDataFileP({deleteData}, {1}); - IcebergDeleteFile icebergDeleteFile( - FileContent::kEqualityDeletes, - eqDeleteFile->getPath(), - dwio::common::FileFormat::DWRF, - 2, - getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1}); // field ID 1 = column 0 = "id" + auto icebergDeleteFile = + makeDeleteFile(eqDeleteFile->getPath(), {1}, GetParam().format); - auto splits = makeSplits(dataFile->getPath(), {icebergDeleteFile}); - auto plan = makeTableScanPlan(outputType, tableType); + auto splits = makeSplitsP(dataFile->getPath(), {icebergDeleteFile}); + auto plan = makeIcebergTableScanPlan(outputType, tableType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); - // The 'id' column must not appear in the output; only 'value' is projected. - // Rows with id=3 ("d") and id=7 ("h") are removed by the equality delete. auto expected = makeRowVector( {"value"}, - { - makeFlatVector({"a", "b", "c", "e", "f", "g", "i", "j"}), - }); - + {makeFlatVector({"a", "b", "c", "e", "f", "g", "i", "j"})}); assertEqualResults({expected}, {result}); } -/// Verifies that two equality-delete files referencing the SAME column not in -/// the user's projection only augment 'scanSpec_' once. Exercises the -/// de-duplication branch in 'IcebergSplitReader::prepareSplit'. -TEST_F(EqualityDeleteFileReaderTest, multipleDeleteFilesSameMissingColumn) { +/// Two delete files targeting the same column not in the projection. +TEST_P(EqualityDeleteFileReaderTestP, multipleDeleteFilesSameMissingColumn) { auto tableType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); auto outputType = ROW({"value"}, {VARCHAR()}); @@ -249,60 +277,31 @@ TEST_F(EqualityDeleteFileReaderTest, multipleDeleteFilesSameMissingColumn) { makeFlatVector( {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"}), }); - auto dataFile = writeDataFile({baseData}); + auto dataFile = writeDataFileP({baseData}, {1, 2}); - // Two delete files, both targeting 'id' (which is NOT in the projection). - auto deleteData1 = makeRowVector( - {"id"}, - { - makeFlatVector({2, 5}), - }); - auto eqDeleteFile1 = writeEqDeleteFile({deleteData1}); - IcebergDeleteFile icebergDeleteFile1( - FileContent::kEqualityDeletes, - eqDeleteFile1->getPath(), - dwio::common::FileFormat::DWRF, - 2, - getFileSize(eqDeleteFile1->getPath()), - /*equalityFieldIds=*/{1}); + auto deleteData1 = makeRowVector({"id"}, {makeFlatVector({2, 5})}); + auto eqDeleteFile1 = writeDataFileP({deleteData1}, {1}); + auto icebergDeleteFile1 = + makeDeleteFile(eqDeleteFile1->getPath(), {1}, GetParam().format); - auto deleteData2 = makeRowVector( - {"id"}, - { - makeFlatVector({0, 9}), - }); - auto eqDeleteFile2 = writeEqDeleteFile({deleteData2}); - IcebergDeleteFile icebergDeleteFile2( - FileContent::kEqualityDeletes, - eqDeleteFile2->getPath(), - dwio::common::FileFormat::DWRF, - 2, - getFileSize(eqDeleteFile2->getPath()), - /*equalityFieldIds=*/{1}); + auto deleteData2 = makeRowVector({"id"}, {makeFlatVector({0, 9})}); + auto eqDeleteFile2 = writeDataFileP({deleteData2}, {1}); + auto icebergDeleteFile2 = + makeDeleteFile(eqDeleteFile2->getPath(), {1}, GetParam().format); - auto splits = - makeSplits(dataFile->getPath(), {icebergDeleteFile1, icebergDeleteFile2}); - auto plan = makeTableScanPlan(outputType, tableType); + auto splits = makeSplitsP( + dataFile->getPath(), {icebergDeleteFile1, icebergDeleteFile2}); + auto plan = makeIcebergTableScanPlan(outputType, tableType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); - // Rows with id=0, 2, 5, 9 are removed (across both delete files). auto expected = makeRowVector( - {"value"}, - { - makeFlatVector({"b", "d", "e", "g", "h", "i"}), - }); - + {"value"}, {makeFlatVector({"b", "d", "e", "g", "h", "i"})}); assertEqualResults({expected}, {result}); } -/// Verifies a multi-column equality-delete file where some columns ARE in the -/// user's projection and some are NOT. Both must end up in the read output for -/// the equality probe to succeed, while only the projected columns appear in -/// the operator-visible result. -TEST_F(EqualityDeleteFileReaderTest, equalityMixedInAndOutOfProjection) { +/// Multi-column equality delete: some delete columns in projection, some not. +TEST_P(EqualityDeleteFileReaderTestP, equalityMixedInAndOutOfProjection) { auto tableType = ROW({"a", "b", "c"}, {INTEGER(), VARCHAR(), BIGINT()}); - // User selects only 'b' and 'c'. 'a' is referenced by the equality delete - // but not part of the projection. auto outputType = ROW({"b", "c"}, {VARCHAR(), BIGINT()}); auto baseData = makeRowVector( @@ -312,229 +311,276 @@ TEST_F(EqualityDeleteFileReaderTest, equalityMixedInAndOutOfProjection) { makeFlatVector({"x", "y", "z", "x", "y"}), makeFlatVector({10, 20, 30, 40, 50}), }); - auto dataFile = writeDataFile({baseData}); + auto dataFile = writeDataFileP({baseData}, {1, 2, 3}); - // Delete rows where (a=2, b="y") -- removes row 1. - // Also (a=1, b="y") -- no match (row with a=1 has b="x"). + // Delete (a=2, b="y") => row 1; (a=1, b="y") => no match. auto deleteData = makeRowVector( {"a", "b"}, { makeFlatVector({2, 1}), makeFlatVector({"y", "y"}), }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto eqDeleteFile = writeDataFileP({deleteData}, {1, 2}); - IcebergDeleteFile icebergDeleteFile( - FileContent::kEqualityDeletes, - eqDeleteFile->getPath(), - dwio::common::FileFormat::DWRF, - 3, - getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1, 2}); // field IDs 1,2 = columns "a","b" + auto icebergDeleteFile = makeDeleteFile( + eqDeleteFile->getPath(), {1, 2}, GetParam().format, /*recordCount=*/3); - auto splits = makeSplits(dataFile->getPath(), {icebergDeleteFile}); - auto plan = makeTableScanPlan(outputType, tableType); + auto splits = makeSplitsP(dataFile->getPath(), {icebergDeleteFile}); + auto plan = makeIcebergTableScanPlan(outputType, tableType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); - // Row 1 (a=2, b="y", c=20) is deleted. The remaining rows project to - // (b, c). auto expected = makeRowVector( {"b", "c"}, { makeFlatVector({"x", "z", "x", "y"}), makeFlatVector({10, 30, 40, 50}), }); + assertEqualResults({expected}, {result}); +} +/// Filter-only column upgrade: 'id' in WHERE but not SELECT, also the +/// equality-delete column. +TEST_P(EqualityDeleteFileReaderTestP, equalityFilterOnlyColumnNotInProjection) { + auto tableType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); + auto outputType = ROW({"value"}, {VARCHAR()}); + + auto baseData = makeRowVector( + {"id", "value"}, + { + makeFlatVector({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}), + makeFlatVector( + {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"}), + }); + auto dataFile = writeDataFileP({baseData}, {1, 2}); + + auto deleteData = makeRowVector({"id"}, {makeFlatVector({4, 8})}); + auto eqDeleteFile = writeDataFileP({deleteData}, {1}); + + auto icebergDeleteFile = + makeDeleteFile(eqDeleteFile->getPath(), {1}, GetParam().format); + + auto splits = makeSplitsP(dataFile->getPath(), {icebergDeleteFile}); + // WHERE id >= 3 => {3,4,5,6,7,8,9}; delete id=4,8 => {3,5,6,7,9}. + auto plan = makeIcebergTableScanPlan( + outputType, tableType, {}, /*subfieldFilters=*/{"id >= 3"}); + auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); + + auto expected = makeRowVector( + {"value"}, {makeFlatVector({"d", "f", "g", "h", "j"})}); assertEqualResults({expected}, {result}); } -/// Verifies that an equality delete on a partition column that IS in the data -/// file (Iceberg-style) but NOT in the user's projection works correctly. The -/// augmentation should set the partition value as a constant; the file-read -/// path should then leave the constant in place because the column is present -/// in 'fileType'. -TEST_F( - EqualityDeleteFileReaderTest, - equalityPartitionColumnInFileNotInProjection) { - auto tableType = ROW({"part", "value"}, {INTEGER(), VARCHAR()}); +/// Multi-column equality delete: both columns must match simultaneously. +TEST_P(EqualityDeleteFileReaderTestP, multiColumnDelete) { + auto tableType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); auto outputType = ROW({"value"}, {VARCHAR()}); - // Data file contains both 'part' and 'value', all rows in partition 2. auto baseData = makeRowVector( - {"part", "value"}, + {"id", "value"}, { - makeFlatVector({2, 2, 2, 2}), - makeFlatVector({"a", "b", "c", "d"}), + makeFlatVector({1, 2, 3, 4, 5}), + makeFlatVector({"a", "b", "c", "d", "e"}), }); - auto dataFile = writeDataFile({baseData}); + auto dataFile = writeDataFileP({baseData}, {1, 2}); auto deleteData = makeRowVector( - {"part", "value"}, + {"id", "value"}, { - makeFlatVector({2, 2}), - makeFlatVector({"b", "d"}), + makeFlatVector({3}), + makeFlatVector({"c"}), }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto eqDeleteFile = writeDataFileP({deleteData}, {1, 2}); - IcebergDeleteFile icebergDeleteFile( - FileContent::kEqualityDeletes, - eqDeleteFile->getPath(), - dwio::common::FileFormat::DWRF, - 2, - getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1, 2}); + auto icebergDeleteFile = makeDeleteFile( + eqDeleteFile->getPath(), {1, 2}, GetParam().format, /*recordCount=*/1); - auto splits = makeSplits( - dataFile->getPath(), - /*partitionKeys=*/{{"part", std::optional{"2"}}}, - {icebergDeleteFile}); - auto plan = makeTableScanPlan(outputType, tableType); + auto splits = makeSplitsP(dataFile->getPath(), {icebergDeleteFile}); + auto plan = makeIcebergTableScanPlan(outputType, tableType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); auto expected = makeRowVector( - {"value"}, - { - makeFlatVector({"a", "c"}), - }); - + {"value"}, {makeFlatVector({"a", "b", "d", "e"})}); assertEqualResults({expected}, {result}); } -/// Same as above but the partition value does NOT match the equality-delete -/// value, so no rows should be removed. -TEST_F( - EqualityDeleteFileReaderTest, - equalityPartitionColumnNonMatchingPartition) { - auto tableType = ROW({"part", "value"}, {INTEGER(), VARCHAR()}); +/// Two separate delete files both apply (multi-reader path). +TEST_P(EqualityDeleteFileReaderTestP, twoDeleteFiles) { + auto tableType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); auto outputType = ROW({"value"}, {VARCHAR()}); - // Data file holds rows in partition 2. auto baseData = makeRowVector( - {"part", "value"}, + {"id", "value"}, { - makeFlatVector({2, 2, 2}), - makeFlatVector({"a", "b", "c"}), + makeFlatVector({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}), + makeFlatVector( + {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"}), }); - auto dataFile = writeDataFile({baseData}); + auto dataFile = writeDataFileP({baseData}, {1, 2}); - // Delete (part=99, value="b"). No file row matches part=99. - auto deleteData = makeRowVector( - {"part", "value"}, + auto deleteData1 = makeRowVector({"id"}, {makeFlatVector({1, 5})}); + auto eqDeleteFile1 = writeDataFileP({deleteData1}, {1}); + auto icebergDelete1 = + makeDeleteFile(eqDeleteFile1->getPath(), {1}, GetParam().format); + + auto deleteData2 = makeRowVector({"id"}, {makeFlatVector({3, 8})}); + auto eqDeleteFile2 = writeDataFileP({deleteData2}, {1}); + auto icebergDelete2 = + makeDeleteFile(eqDeleteFile2->getPath(), {1}, GetParam().format); + + auto splits = + makeSplitsP(dataFile->getPath(), {icebergDelete1, icebergDelete2}); + auto plan = makeIcebergTableScanPlan(outputType, tableType); + auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); + + // Deleted: 1,3,5,8. Surviving: 0,2,4,6,7,9. + auto expected = makeRowVector( + {"value"}, {makeFlatVector({"a", "c", "e", "g", "h", "j"})}); + assertEqualResults({expected}, {result}); +} + +/// No rows match the delete — all rows survive. +TEST_P(EqualityDeleteFileReaderTestP, noMatchingDeletes) { + auto rowType = ROW({"id"}, {BIGINT()}); + + auto baseData = makeRowVector({"id"}, {makeFlatVector({1, 2, 3})}); + auto dataFile = writeDataFileP({baseData}, {1}); + + auto deleteData = + makeRowVector({"id"}, {makeFlatVector({100, 200})}); + auto eqDeleteFile = writeDataFileP({deleteData}, {1}); + + auto icebergDeleteFile = + makeDeleteFile(eqDeleteFile->getPath(), {1}, GetParam().format); + + auto splits = makeSplitsP(dataFile->getPath(), {icebergDeleteFile}); + auto plan = makeIcebergTableScanPlan(rowType); + auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); + + assertEqualResults( + {makeRowVector({"id"}, {makeFlatVector({1, 2, 3})})}, {result}); +} + +/// Every row is deleted. +TEST_P(EqualityDeleteFileReaderTestP, allRowsDeleted) { + auto rowType = ROW({"id"}, {BIGINT()}); + + auto baseData = makeRowVector({"id"}, {makeFlatVector({1, 2, 3})}); + auto dataFile = writeDataFileP({baseData}, {1}); + + auto deleteData = makeRowVector({"id"}, {makeFlatVector({1, 2, 3})}); + auto eqDeleteFile = writeDataFileP({deleteData}, {1}); + + auto icebergDeleteFile = makeDeleteFile( + eqDeleteFile->getPath(), {1}, GetParam().format, /*recordCount=*/3); + + auto splits = makeSplitsP(dataFile->getPath(), {icebergDeleteFile}); + auto plan = makeIcebergTableScanPlan(rowType); + auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); + + EXPECT_EQ(result->size(), 0); +} + +/// VARCHAR column equality delete. +TEST_P(EqualityDeleteFileReaderTestP, stringColumnDelete) { + auto rowType = ROW({"name", "age"}, {VARCHAR(), INTEGER()}); + + auto baseData = makeRowVector( + {"name", "age"}, { - makeFlatVector({99}), - makeFlatVector({"b"}), + makeFlatVector({"alice", "bob", "charlie", "dave"}), + makeFlatVector({25, 30, 35, 40}), }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto dataFile = writeDataFileP({baseData}, {1, 2}); - IcebergDeleteFile icebergDeleteFile( - FileContent::kEqualityDeletes, - eqDeleteFile->getPath(), - dwio::common::FileFormat::DWRF, - 1, - getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1, 2}); + auto deleteData = + makeRowVector({"name"}, {makeFlatVector({"bob", "dave"})}); + auto eqDeleteFile = writeDataFileP({deleteData}, {1}); - auto splits = makeSplits( - dataFile->getPath(), - /*partitionKeys=*/{{"part", std::optional{"2"}}}, - {icebergDeleteFile}); - auto plan = makeTableScanPlan(outputType, tableType); + auto icebergDeleteFile = + makeDeleteFile(eqDeleteFile->getPath(), {1}, GetParam().format); + + auto splits = makeSplitsP(dataFile->getPath(), {icebergDeleteFile}); + auto plan = makeIcebergTableScanPlan(rowType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); auto expected = makeRowVector( - {"value"}, + {"name", "age"}, { - makeFlatVector({"a", "b", "c"}), + makeFlatVector({"alice", "charlie"}), + makeFlatVector({25, 35}), }); - assertEqualResults({expected}, {result}); } -/// Multi-column equality delete where ONE column is a partition column not in -/// the projection and ANOTHER is a regular data column not in the projection. -/// Both must be augmented; the partition column gets a constant value, the -/// regular column is read from the file. -TEST_F( - EqualityDeleteFileReaderTest, - equalityMixedPartitionAndRegularNotInProjection) { - auto tableType = - ROW({"part", "id", "value"}, {INTEGER(), BIGINT(), VARCHAR()}); - auto outputType = ROW({"value"}, {VARCHAR()}); +/// Verifies equality deletes after a field has been dropped, leaving sparse +/// top-level field IDs. +TEST_F(EqualityDeleteFileReaderTest, nonSequentialEqualityFieldId) { + auto tableType = ROW({"id", "category"}, {BIGINT(), VARCHAR()}); + const std::vector fieldIds{1, 3}; - // Data file contains all three columns, all rows in partition 7. auto baseData = makeRowVector( - {"part", "id", "value"}, + {"id", "dropped", "category"}, { - makeFlatVector({7, 7, 7, 7}), - makeFlatVector({10, 20, 30, 40}), - makeFlatVector({"a", "b", "c", "d"}), + makeFlatVector({1, 2, 3, 4, 5}), + makeFlatVector({10, 20, 30, 40, 50}), + makeFlatVector({"A", "B", "A", "C", "B"}), }); - auto dataFile = writeDataFile({baseData}); + auto dataFile = writeDwrfFileWithFieldIds({baseData}, {1, 2, 3}); - // Delete (part=7, id=20) and (part=7, id=40). Should remove "b" and "d". auto deleteData = makeRowVector( - {"part", "id"}, + {"category"}, { - makeFlatVector({7, 7}), - makeFlatVector({20, 40}), + makeFlatVector({"B"}), }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto eqDeleteFile = writeDataFile({deleteData}); IcebergDeleteFile icebergDeleteFile( FileContent::kEqualityDeletes, eqDeleteFile->getPath(), dwio::common::FileFormat::DWRF, - 2, + 1, getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1, 2}); // field IDs 1,2 = part, id + /*equalityFieldIds=*/{3}); auto splits = makeSplits( dataFile->getPath(), - /*partitionKeys=*/{{"part", std::optional{"7"}}}, + {}, + {dwio::common::FileFormat::DWRF, false}, {icebergDeleteFile}); - auto plan = makeTableScanPlan(outputType, tableType); + auto plan = makeIcebergTableScanPlan(tableType, tableType, fieldIds); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); auto expected = makeRowVector( - {"value"}, + {"id", "category"}, { - makeFlatVector({"a", "c"}), + makeFlatVector({1, 3, 4}), + makeFlatVector({"A", "A", "C"}), }); - assertEqualResults({expected}, {result}); } -/// Verifies equality delete on a DATE partition column not in projection. -/// Iceberg encodes DATE partition values as days-since-epoch (e.g. "19345"). -/// This exercises the type-derived 'isDaysSinceEpoch' flag in -/// 'configureEqualityDeleteColumns' — for DATE columns the partition string -/// must be parsed as an integer day count, NOT as an ISO-8601 date string. TEST_F( EqualityDeleteFileReaderTest, - equalityDatePartitionColumnNotInProjection) { - auto tableType = ROW({"part_date", "value"}, {DATE(), VARCHAR()}); - auto outputType = ROW({"value"}, {VARCHAR()}); + nonSequentialEqualityFieldIdNotInProjection) { + auto tableType = ROW({"id", "category"}, {BIGINT(), VARCHAR()}); + auto outputType = ROW({"id"}, {BIGINT()}); + const std::vector fieldIds{1, 3}; - // 19345 days since 1970-01-01 == 2022-12-22. All file rows belong to that - // partition. - constexpr int32_t kPartitionDays = 19345; auto baseData = makeRowVector( - {"part_date", "value"}, + {"id", "dropped", "category"}, { - makeFlatVector( - {kPartitionDays, kPartitionDays, kPartitionDays}, DATE()), - makeFlatVector({"a", "b", "c"}), + makeFlatVector({1, 2, 3, 4, 5}), + makeFlatVector({10, 20, 30, 40, 50}), + makeFlatVector({"A", "B", "A", "C", "B"}), }); - auto dataFile = writeDataFile({baseData}); + auto dataFile = writeDwrfFileWithFieldIds({baseData}, {1, 2, 3}); - // Delete (part_date=19345, value="b"). auto deleteData = makeRowVector( - {"part_date", "value"}, + {"category"}, { - makeFlatVector({kPartitionDays}, DATE()), - makeFlatVector({"b"}), + makeFlatVector({"B"}), }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto eqDeleteFile = writeDataFile({deleteData}); IcebergDeleteFile icebergDeleteFile( FileContent::kEqualityDeletes, @@ -542,306 +588,246 @@ TEST_F( dwio::common::FileFormat::DWRF, 1, getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1, 2}); + /*equalityFieldIds=*/{3}); auto splits = makeSplits( dataFile->getPath(), - /*partitionKeys=*/ - {{"part_date", - std::optional{std::to_string(kPartitionDays)}}}, + /*partitionKeys=*/{}, + {dwio::common::FileFormat::DWRF, false}, {icebergDeleteFile}); - auto plan = makeTableScanPlan(outputType, tableType); + auto plan = makeIcebergTableScanPlan(outputType, tableType, fieldIds); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); auto expected = makeRowVector( - {"value"}, + {"id"}, { - makeFlatVector({"a", "c"}), + makeFlatVector({1, 3, 4}), }); - assertEqualResults({expected}, {result}); } -/// Exercises the filter-only column upgrade path in -/// 'configureEqualityDeleteColumns'. The equality-delete column 'id' is -/// referenced by a WHERE predicate (so the planner installs a scan-spec -/// child with 'projectOut=false') but is NOT in the user's SELECT -/// projection. The augmentation must upgrade the existing scan-spec child -/// from filter-only to 'projectOut=true' and assign a non-conflicting -/// channel so the equality-delete reader can probe by name. -TEST_F(EqualityDeleteFileReaderTest, equalityFilterOnlyColumnNotInProjection) { - auto tableType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); - auto outputType = ROW({"value"}, {VARCHAR()}); +/// Verifies the ordinal fallback on a non-first column when full-schema field +/// IDs are unavailable. +TEST_P(EqualityDeleteFileReaderTestP, deleteOnSecondColumn) { + auto rowType = ROW({"id", "category"}, {BIGINT(), VARCHAR()}); auto baseData = makeRowVector( - {"id", "value"}, + {"id", "category"}, { - makeFlatVector({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}), - makeFlatVector( - {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"}), + makeFlatVector({1, 2, 3, 4, 5}), + makeFlatVector({"A", "B", "A", "C", "B"}), }); - auto dataFile = writeDataFile({baseData}); + auto dataFile = writeDataFileP({baseData}, {1, 2}); - // Equality delete removes id == 4 and id == 8. - auto deleteData = makeRowVector( - {"id"}, - { - makeFlatVector({4, 8}), - }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto deleteData = + makeRowVector({"category"}, {makeFlatVector({"B"})}); + auto eqDeleteFile = writeDataFileP({deleteData}, {2}); - IcebergDeleteFile icebergDeleteFile( - FileContent::kEqualityDeletes, - eqDeleteFile->getPath(), - dwio::common::FileFormat::DWRF, - 2, - getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1}); // field ID 1 = "id" - - auto splits = makeSplits(dataFile->getPath(), {icebergDeleteFile}); - // WHERE id >= 3 keeps rows {3,4,5,6,7,8,9} from the file; the equality - // delete then removes id=4 and id=8, leaving values {d->skipped} no, we - // expect surviving values for ids {3,5,6,7,9}, projected as 'value' only. - auto plan = PlanBuilder() - .startTableScan(kIcebergConnectorId) - .outputType(outputType) - .dataColumns(tableType) - .subfieldFilter("id >= 3") - .endTableScan() - .planNode(); + auto icebergDeleteFile = makeDeleteFile( + eqDeleteFile->getPath(), {2}, GetParam().format, /*recordCount=*/1); + + auto splits = makeSplitsP(dataFile->getPath(), {icebergDeleteFile}); + auto plan = makeIcebergTableScanPlan(rowType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); auto expected = makeRowVector( - {"value"}, + {"id", "category"}, { - makeFlatVector({"d", "f", "g", "h", "j"}), + makeFlatVector({1, 3, 4}), + makeFlatVector({"A", "A", "C"}), }); - assertEqualResults({expected}, {result}); } -/// Regression test for the schema-evolution + -/// partition-column-not-in-projection scenario surfaced by the Presto Iceberg -/// integration test 'testEqualityDeleteWithPartitionColumnMissingInSelect'. -/// -/// Setup (mirrors the Presto test for the older data file): -/// - Full table schema: (a, b, c, d) with a, c, d as partition columns. -/// - The data file under test was written BEFORE 'd' was added, so it -/// physically contains only (a, b, c). -/// - User projection: (a, b, d). 'd' must be NULL-filled (not in file); -/// 'c' is NOT in the projection but IS referenced by the equality -/// delete and IS the file's identity-partition column. -/// -/// The equality delete (a=6, c=2, b=1006) targets the file row -/// (6, '1006', 2). The augmentation must: -/// 1. Add 'c' to 'scanSpec_' / 'readerOutputType_' so the eq-delete -/// probe can find it by name. -/// 2. Leave 'd' alone — 'd' is in the user projection and gets the -/// standard schema-evolution NULL-fill from 'adaptColumns'. -/// 3. Honour the existing partition-key constant on 'c' regardless of -/// whether the file physically contains 'c'. -TEST_F( - EqualityDeleteFileReaderTest, - equalityPartitionColumnNotInProjectionWithEvolvedSchema) { - // Full evolved table schema (after 'ALTER TABLE ADD COLUMN d'). - auto tableType = - ROW({"a", "b", "c", "d"}, {INTEGER(), VARCHAR(), INTEGER(), VARCHAR()}); - // User selects 'a', 'b', 'd'. Note 'c' is NOT projected. - auto outputType = ROW({"a", "b", "d"}, {INTEGER(), VARCHAR(), VARCHAR()}); +/// Delete applies when deleteSeqNum > dataSeqNum. +TEST_P(EqualityDeleteFileReaderTestP, sequenceNumberDeleteApplies) { + auto rowType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); - // Data file contains only (a, b, c) — written before 'd' was added. - // Both rows are in the (a=6, c=2) partition. auto baseData = makeRowVector( - {"a", "b", "c"}, + {"id", "value"}, { - makeFlatVector({6, 6}), - makeFlatVector({"1006", "1009"}), - makeFlatVector({2, 2}), + makeFlatVector({1, 2, 3, 4, 5}), + makeFlatVector({"a", "b", "c", "d", "e"}), }); - auto dataFile = writeDataFile({baseData}); + auto dataFile = writeDataFileP({baseData}, {1, 2}); - // Equality delete on (a, b, c) with values (6, '1006', 2). Field IDs - // are in field-id order = [1, 2, 3]. - auto deleteData = makeRowVector( - {"a", "b", "c"}, - { - makeFlatVector({6}), - makeFlatVector({"1006"}), - makeFlatVector({2}), - }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto deleteData = makeRowVector({"id"}, {makeFlatVector({2, 4})}); + auto eqDeleteFile = writeDataFileP({deleteData}, {1}); - IcebergDeleteFile icebergDeleteFile( - FileContent::kEqualityDeletes, - eqDeleteFile->getPath(), - dwio::common::FileFormat::DWRF, - 1, - getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1, 2, 3}); + // deleteSeqNum=5 > dataSeqNum=3 => applies. + auto icebergDeleteFile = makeDeleteFile( + eqDeleteFile->getPath(), {1}, GetParam().format, 2, /*deleteSeqNum=*/5); - auto splits = makeSplits( - dataFile->getPath(), - /*partitionKeys=*/ - {{"a", std::optional{"6"}}, - {"c", std::optional{"2"}}}, - {icebergDeleteFile}); - auto plan = makeTableScanPlan(outputType, tableType); + auto splits = makeSplitsP( + dataFile->getPath(), {icebergDeleteFile}, /*dataSequenceNumber=*/3); + auto plan = makeIcebergTableScanPlan(rowType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); - // Row (6, '1006', 2) is deleted by the equality delete; (6, '1009', 2) - // survives. 'd' is NULL because the data file was written before 'd' - // was added. auto expected = makeRowVector( - {"a", "b", "d"}, + {"id", "value"}, { - makeFlatVector({6}), - makeFlatVector({"1009"}), - makeNullableFlatVector({std::nullopt}), + makeFlatVector({1, 3, 5}), + makeFlatVector({"a", "c", "e"}), }); - assertEqualResults({expected}, {result}); } -/// Verifies multi-column equality deletes (both columns must match). -TEST_F(EqualityDeleteFileReaderTest, multiColumnDelete) { - auto rowType = ROW({"a", "b", "c"}, {INTEGER(), VARCHAR(), BIGINT()}); +/// Delete skipped when deleteSeqNum <= dataSeqNum. +TEST_P(EqualityDeleteFileReaderTestP, sequenceNumberDeleteSkipped) { + auto rowType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); auto baseData = makeRowVector( - {"a", "b", "c"}, - { - makeFlatVector({1, 2, 3, 4, 5}), - makeFlatVector({"x", "y", "z", "x", "y"}), - makeFlatVector({10, 20, 30, 40, 50}), - }); - auto dataFile = writeDataFile({baseData}); - - // Delete rows where (a=2, b="y") — matches row index 1. - // Also (a=5, b="y") — matches row index 4. - // But (a=1, b="y") — no match (a=1 has b="x"). - auto deleteData = makeRowVector( - {"a", "b"}, + {"id", "value"}, { - makeFlatVector({2, 5, 1}), - makeFlatVector({"y", "y", "y"}), + makeFlatVector({1, 2, 3}), + makeFlatVector({"a", "b", "c"}), }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto dataFile = writeDataFileP({baseData}, {1, 2}); - IcebergDeleteFile icebergDeleteFile( - FileContent::kEqualityDeletes, - eqDeleteFile->getPath(), - dwio::common::FileFormat::DWRF, - 3, - getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1, 2}); // field IDs 1,2 = columns "a","b" + auto deleteData = makeRowVector({"id"}, {makeFlatVector({1, 2, 3})}); + auto eqDeleteFile = writeDataFileP({deleteData}, {1}); + + // deleteSeqNum=2 <= dataSeqNum=5 => skipped. + auto icebergDeleteFile = makeDeleteFile( + eqDeleteFile->getPath(), {1}, GetParam().format, 3, /*deleteSeqNum=*/2); - auto splits = makeSplits(dataFile->getPath(), {icebergDeleteFile}); - auto plan = makeTableScanPlan(rowType); + auto splits = makeSplitsP( + dataFile->getPath(), {icebergDeleteFile}, /*dataSequenceNumber=*/5); + auto plan = makeIcebergTableScanPlan(rowType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); - // Rows 0, 2, 3 survive (rows 1 and 4 deleted). auto expected = makeRowVector( - {"a", "b", "c"}, + {"id", "value"}, { - makeFlatVector({1, 3, 4}), - makeFlatVector({"x", "z", "x"}), - makeFlatVector({10, 30, 40}), + makeFlatVector({1, 2, 3}), + makeFlatVector({"a", "b", "c"}), }); - assertEqualResults({expected}, {result}); } -/// Verifies that when no rows match, all rows survive. -TEST_F(EqualityDeleteFileReaderTest, noMatchingDeletes) { - auto rowType = ROW({"id"}, {BIGINT()}); +/// Delete skipped when deleteSeqNum == dataSeqNum (edge case of <=). +TEST_P(EqualityDeleteFileReaderTestP, sequenceNumberEqualSkipped) { + auto rowType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); auto baseData = makeRowVector( - {"id"}, + {"id", "value"}, { makeFlatVector({1, 2, 3}), + makeFlatVector({"a", "b", "c"}), }); - auto dataFile = writeDataFile({baseData}); + auto dataFile = writeDataFileP({baseData}, {1, 2}); - // Delete file has values not present in base data. - auto deleteData = makeRowVector( - {"id"}, - { - makeFlatVector({100, 200}), - }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto deleteData = makeRowVector({"id"}, {makeFlatVector({1, 2, 3})}); + auto eqDeleteFile = writeDataFileP({deleteData}, {1}); - IcebergDeleteFile icebergDeleteFile( - FileContent::kEqualityDeletes, - eqDeleteFile->getPath(), - dwio::common::FileFormat::DWRF, - 2, - getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1}); + // deleteSeqNum=5 == dataSeqNum=5 => skipped. + auto icebergDeleteFile = makeDeleteFile( + eqDeleteFile->getPath(), {1}, GetParam().format, 3, /*deleteSeqNum=*/5); - auto splits = makeSplits(dataFile->getPath(), {icebergDeleteFile}); - auto plan = makeTableScanPlan(rowType); + auto splits = makeSplitsP( + dataFile->getPath(), {icebergDeleteFile}, /*dataSequenceNumber=*/5); + auto plan = makeIcebergTableScanPlan(rowType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); auto expected = makeRowVector( - {"id"}, + {"id", "value"}, { makeFlatVector({1, 2, 3}), + makeFlatVector({"a", "b", "c"}), }); - assertEqualResults({expected}, {result}); } -/// Verifies that all rows are deleted when every base row matches. -TEST_F(EqualityDeleteFileReaderTest, allRowsDeleted) { +/// deleteSeqNum=0 disables filtering — delete always applies. +TEST_P(EqualityDeleteFileReaderTestP, sequenceNumberZeroAlwaysApplies) { auto rowType = ROW({"id"}, {BIGINT()}); + auto baseData = makeRowVector({"id"}, {makeFlatVector({1, 2, 3})}); + auto dataFile = writeDataFileP({baseData}, {1}); + + auto deleteData = makeRowVector({"id"}, {makeFlatVector({2})}); + auto eqDeleteFile = writeDataFileP({deleteData}, {1}); + + // deleteSeqNum=0 => filtering disabled, applies despite dataSeqNum=10. + auto icebergDeleteFile = makeDeleteFile( + eqDeleteFile->getPath(), {1}, GetParam().format, 1, /*deleteSeqNum=*/0); + + auto splits = makeSplitsP( + dataFile->getPath(), {icebergDeleteFile}, /*dataSequenceNumber=*/10); + auto plan = makeIcebergTableScanPlan(rowType); + auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); + + assertEqualResults( + {makeRowVector({"id"}, {makeFlatVector({1, 3})})}, {result}); +} + +/// Only delete files with higher sequence numbers than the data file apply. +TEST_P(EqualityDeleteFileReaderTestP, mixedSequenceNumbers) { + auto rowType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); + auto baseData = makeRowVector( - {"id"}, + {"id", "value"}, { - makeFlatVector({1, 2, 3}), + makeFlatVector({1, 2, 3, 4, 5}), + makeFlatVector({"a", "b", "c", "d", "e"}), }); - auto dataFile = writeDataFile({baseData}); + auto dataFile = writeDataFileP({baseData}, {1, 2}); - auto deleteData = makeRowVector( - {"id"}, - { - makeFlatVector({1, 2, 3}), - }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + // seqNum=10 > dataSeqNum=5 => applied. + auto deleteData1 = makeRowVector({"id"}, {makeFlatVector({2})}); + auto eqDeleteFile1 = writeDataFileP({deleteData1}, {1}); + auto icebergDeleteFile1 = makeDeleteFile( + eqDeleteFile1->getPath(), {1}, GetParam().format, 1, /*deleteSeqNum=*/10); - IcebergDeleteFile icebergDeleteFile( - FileContent::kEqualityDeletes, - eqDeleteFile->getPath(), - dwio::common::FileFormat::DWRF, - 3, - getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1}); + // seqNum=3 <= dataSeqNum=5 => skipped. + auto deleteData2 = makeRowVector({"id"}, {makeFlatVector({4})}); + auto eqDeleteFile2 = writeDataFileP({deleteData2}, {1}); + auto icebergDeleteFile2 = makeDeleteFile( + eqDeleteFile2->getPath(), {1}, GetParam().format, 1, /*deleteSeqNum=*/3); - auto splits = makeSplits(dataFile->getPath(), {icebergDeleteFile}); - auto plan = makeTableScanPlan(rowType); + auto splits = makeSplitsP( + dataFile->getPath(), + {icebergDeleteFile1, icebergDeleteFile2}, + /*dataSequenceNumber=*/5); + auto plan = makeIcebergTableScanPlan(rowType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); - EXPECT_EQ(result->size(), 0); + // id=2 deleted; id=4 survives. + auto expected = makeRowVector( + {"id", "value"}, + { + makeFlatVector({1, 3, 4, 5}), + makeFlatVector({"a", "c", "d", "e"}), + }); + assertEqualResults({expected}, {result}); } -/// Verifies equality deletes with VARCHAR columns. -TEST_F(EqualityDeleteFileReaderTest, stringColumnDelete) { - auto rowType = ROW({"name", "age"}, {VARCHAR(), INTEGER()}); +// =========================================================================== +// Non-parameterized tests -- partition columns and evolved schema (DWRF only). +// =========================================================================== + +/// Equality delete on a partition column in the data file but not projected. +TEST_F( + EqualityDeleteFileReaderTest, + equalityPartitionColumnInFileNotInProjection) { + auto tableType = ROW({"part", "value"}, {INTEGER(), VARCHAR()}); + auto outputType = ROW({"value"}, {VARCHAR()}); auto baseData = makeRowVector( - {"name", "age"}, + {"part", "value"}, { - makeFlatVector({"alice", "bob", "charlie", "dave"}), - makeFlatVector({25, 30, 35, 40}), + makeFlatVector({2, 2, 2, 2}), + makeFlatVector({"a", "b", "c", "d"}), }); auto dataFile = writeDataFile({baseData}); - // Delete rows where name is "bob" or "dave". auto deleteData = makeRowVector( - {"name"}, + {"part", "value"}, { - makeFlatVector({"bob", "dave"}), + makeFlatVector({2, 2}), + makeFlatVector({"b", "d"}), }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto eqDeleteFile = writeDataFile({deleteData}); IcebergDeleteFile icebergDeleteFile( FileContent::kEqualityDeletes, @@ -849,41 +835,46 @@ TEST_F(EqualityDeleteFileReaderTest, stringColumnDelete) { dwio::common::FileFormat::DWRF, 2, getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1}); // field ID 1 = "name" + /*equalityFieldIds=*/{1, 2}); - auto splits = makeSplits(dataFile->getPath(), {icebergDeleteFile}); - auto plan = makeTableScanPlan(rowType); + auto splits = makeSplits( + dataFile->getPath(), + /*partitionKeys=*/{}, + {dwio::common::FileFormat::DWRF, false}, + {icebergDeleteFile}, + /*dataSequenceNumber=*/0, + // Source field ID 1 ('part') is an explicit identity partition field. + /*identityPartitionKeys=*/{{1, std::optional{"2"}}}); + auto plan = makeIcebergTableScanPlan(outputType, tableType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); - auto expected = makeRowVector( - {"name", "age"}, - { - makeFlatVector({"alice", "charlie"}), - makeFlatVector({25, 35}), - }); - - assertEqualResults({expected}, {result}); + assertEqualResults( + {makeRowVector({"value"}, {makeFlatVector({"a", "c"})})}, + {result}); } -/// Verifies equality deletes on a non-first column (field ID 2). -TEST_F(EqualityDeleteFileReaderTest, deleteOnSecondColumn) { - auto rowType = ROW({"id", "category"}, {BIGINT(), VARCHAR()}); +/// Same as above but partition value does not match -- no rows deleted. +TEST_F( + EqualityDeleteFileReaderTest, + equalityPartitionColumnNonMatchingPartition) { + auto tableType = ROW({"part", "value"}, {INTEGER(), VARCHAR()}); + auto outputType = ROW({"value"}, {VARCHAR()}); auto baseData = makeRowVector( - {"id", "category"}, + {"part", "value"}, { - makeFlatVector({1, 2, 3, 4, 5}), - makeFlatVector({"A", "B", "A", "C", "B"}), + makeFlatVector({2, 2, 2}), + makeFlatVector({"a", "b", "c"}), }); auto dataFile = writeDataFile({baseData}); - // Delete rows where category == "B". auto deleteData = makeRowVector( - {"category"}, + {"part", "value"}, { - makeFlatVector({"B"}), + makeFlatVector({99}), + makeFlatVector({"b"}), }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto eqDeleteFile = writeDataFile({deleteData}); IcebergDeleteFile icebergDeleteFile( FileContent::kEqualityDeletes, @@ -891,296 +882,293 @@ TEST_F(EqualityDeleteFileReaderTest, deleteOnSecondColumn) { dwio::common::FileFormat::DWRF, 1, getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{2}); // field ID 2 = column 1 = "category" + /*equalityFieldIds=*/{1, 2}); - auto splits = makeSplits(dataFile->getPath(), {icebergDeleteFile}); - auto plan = makeTableScanPlan(rowType); + auto splits = makeSplits( + dataFile->getPath(), + /*partitionKeys=*/{}, + {dwio::common::FileFormat::DWRF, false}, + {icebergDeleteFile}, + /*dataSequenceNumber=*/0, + /*identityPartitionKeys=*/{{1, std::optional{"2"}}}); + auto plan = makeIcebergTableScanPlan(outputType, tableType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); - // Rows with category="B" (indices 1,4) deleted. - auto expected = makeRowVector( - {"id", "category"}, - { - makeFlatVector({1, 3, 4}), - makeFlatVector({"A", "A", "C"}), - }); - - assertEqualResults({expected}, {result}); + assertEqualResults( + {makeRowVector( + {"value"}, {makeFlatVector({"a", "b", "c"})})}, + {result}); } -/// Verifies that equality deletes apply when the delete file has a higher -/// sequence number than the data file (per the Iceberg V2+ spec). -TEST_F(EqualityDeleteFileReaderTest, sequenceNumberDeleteApplies) { - auto rowType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); +/// One partition column + one regular column, both absent from the projection. +TEST_F( + EqualityDeleteFileReaderTest, + equalityMixedPartitionAndRegularNotInProjection) { + auto tableType = + ROW({"part", "id", "value"}, {INTEGER(), BIGINT(), VARCHAR()}); + auto outputType = ROW({"value"}, {VARCHAR()}); auto baseData = makeRowVector( - {"id", "value"}, + {"part", "id", "value"}, { - makeFlatVector({1, 2, 3, 4, 5}), - makeFlatVector({"a", "b", "c", "d", "e"}), + makeFlatVector({7, 7, 7, 7}), + makeFlatVector({10, 20, 30, 40}), + makeFlatVector({"a", "b", "c", "d"}), }); auto dataFile = writeDataFile({baseData}); auto deleteData = makeRowVector( - {"id"}, + {"part", "id"}, { - makeFlatVector({2, 4}), + makeFlatVector({7, 7}), + makeFlatVector({20, 40}), }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto eqDeleteFile = writeDataFile({deleteData}); - // Delete file has sequence number 5, data file has sequence number 3. - // Since deleteSeq (5) > dataSeq (3), the delete should apply. IcebergDeleteFile icebergDeleteFile( FileContent::kEqualityDeletes, eqDeleteFile->getPath(), dwio::common::FileFormat::DWRF, 2, getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1}, - /*lowerBounds=*/{}, - /*upperBounds=*/{}, - /*dataSequenceNumber=*/5); + /*equalityFieldIds=*/{1, 2}); auto splits = makeSplits( dataFile->getPath(), + {}, + {dwio::common::FileFormat::DWRF, false}, {icebergDeleteFile}, - /*dataSequenceNumber=*/3); - auto plan = makeTableScanPlan(rowType); + /*dataSequenceNumber=*/0, + /*identityPartitionKeys=*/{{1, std::optional{"7"}}}); + auto plan = makeIcebergTableScanPlan(outputType, tableType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); - // Rows with id=2 and id=4 are deleted. - auto expected = makeRowVector( - {"id", "value"}, - { - makeFlatVector({1, 3, 5}), - makeFlatVector({"a", "c", "e"}), - }); - - assertEqualResults({expected}, {result}); + assertEqualResults( + {makeRowVector({"value"}, {makeFlatVector({"a", "c"})})}, + {result}); } -/// Verifies that equality deletes are skipped when the delete file has a -/// lower or equal sequence number compared to the data file. -TEST_F(EqualityDeleteFileReaderTest, sequenceNumberDeleteSkipped) { - auto rowType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); +/// DATE partition column not in projection (days-since-epoch encoding). +TEST_F( + EqualityDeleteFileReaderTest, + equalityDatePartitionColumnNotInProjection) { + auto tableType = ROW({"part_date", "value"}, {DATE(), VARCHAR()}); + auto outputType = ROW({"value"}, {VARCHAR()}); + constexpr int32_t kPartitionDays = 19345; // 2022-12-22 auto baseData = makeRowVector( - {"id", "value"}, + {"part_date", "value"}, { - makeFlatVector({1, 2, 3}), + makeFlatVector( + {kPartitionDays, kPartitionDays, kPartitionDays}, DATE()), makeFlatVector({"a", "b", "c"}), }); auto dataFile = writeDataFile({baseData}); auto deleteData = makeRowVector( - {"id"}, + {"part_date", "value"}, { - makeFlatVector({1, 2, 3}), + makeFlatVector({kPartitionDays}, DATE()), + makeFlatVector({"b"}), }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto eqDeleteFile = writeDataFile({deleteData}); - // Delete file has sequence number 2, data file has sequence number 5. - // Since deleteSeq (2) <= dataSeq (5), the delete should be skipped. IcebergDeleteFile icebergDeleteFile( FileContent::kEqualityDeletes, eqDeleteFile->getPath(), dwio::common::FileFormat::DWRF, - 3, + 1, getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1}, - /*lowerBounds=*/{}, - /*upperBounds=*/{}, - /*dataSequenceNumber=*/2); + /*equalityFieldIds=*/{1, 2}); auto splits = makeSplits( dataFile->getPath(), + /*partitionKeys=*/{}, + {dwio::common::FileFormat::DWRF, false}, {icebergDeleteFile}, - /*dataSequenceNumber=*/5); - auto plan = makeTableScanPlan(rowType); + /*dataSequenceNumber=*/0, + /*identityPartitionKeys=*/ + {{1, std::optional{std::to_string(kPartitionDays)}}}); + auto plan = makeIcebergTableScanPlan(outputType, tableType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); - // All rows survive because the delete file is skipped. - auto expected = makeRowVector( - {"id", "value"}, - { - makeFlatVector({1, 2, 3}), - makeFlatVector({"a", "b", "c"}), - }); - - assertEqualResults({expected}, {result}); + assertEqualResults( + {makeRowVector({"value"}, {makeFlatVector({"a", "c"})})}, + {result}); } -/// Verifies that equality deletes are skipped when the delete file has the -/// same sequence number as the data file (edge case of the <= check). -TEST_F(EqualityDeleteFileReaderTest, sequenceNumberEqualSkipped) { - auto rowType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); +/// Regression for transformed partition fields whose derived partition-field +/// name collides with a real source column name. +/// +/// A partition field may be named anything, so a 'bucket[4]' field on source +/// column 'id' can itself be named "id". The split's name-keyed +/// 'partitionKeys' then maps "id" to the *bucket ordinal*, not to any row's +/// 'id'. Substituting that value for the source column would corrupt the +/// equality-delete probe. Only a field the spec marks 'identity' may be +/// substituted, and no identity metadata is supplied here, so the reader must +/// read the physical 'id' column from the data file. +TEST_F( + EqualityDeleteFileReaderTest, + equalityBucketPartitionNameCollisionNotInProjection) { + auto tableType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); + auto outputType = ROW({"value"}, {VARCHAR()}); auto baseData = makeRowVector( {"id", "value"}, { - makeFlatVector({1, 2, 3}), - makeFlatVector({"a", "b", "c"}), + makeFlatVector({10, 20, 30, 40}), + makeFlatVector({"a", "b", "c", "d"}), }); auto dataFile = writeDataFile({baseData}); + // Equality delete removes the row whose physical id is 20. auto deleteData = makeRowVector( {"id"}, { - makeFlatVector({1, 2, 3}), + makeFlatVector({20}), }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto eqDeleteFile = writeDataFile({deleteData}); - // Delete file and data file have the same sequence number (5). - // Since deleteSeq (5) <= dataSeq (5), the delete should be skipped. IcebergDeleteFile icebergDeleteFile( FileContent::kEqualityDeletes, eqDeleteFile->getPath(), dwio::common::FileFormat::DWRF, - 3, + 1, getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1}, - /*lowerBounds=*/{}, - /*upperBounds=*/{}, - /*dataSequenceNumber=*/5); + /*equalityFieldIds=*/{1}); auto splits = makeSplits( dataFile->getPath(), + // The bucket ordinal, stored under a name that collides with the + // source column. + /*partitionKeys=*/{{"id", std::optional{"2"}}}, + {dwio::common::FileFormat::DWRF, false}, {icebergDeleteFile}, - /*dataSequenceNumber=*/5); - auto plan = makeTableScanPlan(rowType); + /*dataSequenceNumber=*/0, + // 'bucket[4]' is not identity, so nothing is substitutable. + /*identityPartitionKeys=*/{}); + auto plan = makeIcebergTableScanPlan(outputType, tableType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); - // All rows survive because the delete file is skipped (equal seq#). + // Only the physically matching row is deleted. Substituting the bucket + // ordinal would make every row's id 2 and delete nothing. auto expected = makeRowVector( - {"id", "value"}, + {"value"}, { - makeFlatVector({1, 2, 3}), - makeFlatVector({"a", "b", "c"}), + makeFlatVector({"a", "c", "d"}), }); assertEqualResults({expected}, {result}); } -/// Verifies that when either sequence number is 0 (unassigned/legacy V1), -/// the delete file is always applied (filtering is disabled). -TEST_F(EqualityDeleteFileReaderTest, sequenceNumberZeroAlwaysApplies) { - auto rowType = ROW({"id"}, {BIGINT()}); +/// Regression for the 'void' transform, which always stores a null partition +/// value and — unlike bucket/truncate/temporal — keeps the source column's +/// name by default. A name-keyed lookup therefore finds an entry for "id" +/// and would install a constant null over every row, making the null +/// equality-delete key match everything. With identity metadata absent, the +/// physical non-null 'id' values must survive instead. +TEST_F(EqualityDeleteFileReaderTest, equalityVoidPartitionNotInProjection) { + auto tableType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); + auto outputType = ROW({"value"}, {VARCHAR()}); auto baseData = makeRowVector( - {"id"}, + {"id", "value"}, { - makeFlatVector({1, 2, 3}), + makeFlatVector({10, 20, 30}), + makeFlatVector({"a", "b", "c"}), }); auto dataFile = writeDataFile({baseData}); + // The delete key is null, matching the null the void transform stores. auto deleteData = makeRowVector( {"id"}, { - makeFlatVector({2}), + makeNullableFlatVector({std::nullopt}), }); - auto eqDeleteFile = writeEqDeleteFile({deleteData}); + auto eqDeleteFile = writeDataFile({deleteData}); - // Delete file has sequence number 0 (legacy), data file has sequence 10. - // Since deleteSeq is 0, filtering is disabled and the delete applies. IcebergDeleteFile icebergDeleteFile( FileContent::kEqualityDeletes, eqDeleteFile->getPath(), dwio::common::FileFormat::DWRF, 1, getFileSize(eqDeleteFile->getPath()), - /*equalityFieldIds=*/{1}, - /*lowerBounds=*/{}, - /*upperBounds=*/{}, - /*dataSequenceNumber=*/0); + /*equalityFieldIds=*/{1}); auto splits = makeSplits( dataFile->getPath(), + /*partitionKeys=*/{{"id", std::nullopt}}, + {dwio::common::FileFormat::DWRF, false}, {icebergDeleteFile}, - /*dataSequenceNumber=*/10); - auto plan = makeTableScanPlan(rowType); + /*dataSequenceNumber=*/0, + // 'void' is not identity, so its null is not substitutable. + /*identityPartitionKeys=*/{}); + auto plan = makeIcebergTableScanPlan(outputType, tableType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); - // Row id=2 is deleted because sequence number filtering is disabled. auto expected = makeRowVector( - {"id"}, + {"value"}, { - makeFlatVector({1, 3}), + makeFlatVector({"a", "b", "c"}), }); assertEqualResults({expected}, {result}); } -/// Verifies that when multiple delete files have different sequence numbers, -/// only those with higher sequence numbers than the data file are applied. -TEST_F(EqualityDeleteFileReaderTest, mixedSequenceNumbers) { - auto rowType = ROW({"id", "value"}, {BIGINT(), VARCHAR()}); +/// Schema evolution + partition column not in projection (Presto regression). +TEST_F( + EqualityDeleteFileReaderTest, + equalityPartitionColumnNotInProjectionWithEvolvedSchema) { + auto tableType = + ROW({"a", "b", "c", "d"}, {INTEGER(), VARCHAR(), INTEGER(), VARCHAR()}); + auto outputType = ROW({"a", "b", "d"}, {INTEGER(), VARCHAR(), VARCHAR()}); + // Data file was written before 'd' was added -- contains only (a, b, c). auto baseData = makeRowVector( - {"id", "value"}, + {"a", "b", "c"}, { - makeFlatVector({1, 2, 3, 4, 5}), - makeFlatVector({"a", "b", "c", "d", "e"}), + makeFlatVector({6, 6}), + makeFlatVector({"1006", "1009"}), + makeFlatVector({2, 2}), }); auto dataFile = writeDataFile({baseData}); - // First delete file: seqNum=10 (higher than data seqNum=5) → applied. - auto deleteData1 = makeRowVector( - {"id"}, - { - makeFlatVector({2}), - }); - auto eqDeleteFile1 = writeEqDeleteFile({deleteData1}); - IcebergDeleteFile icebergDeleteFile1( - FileContent::kEqualityDeletes, - eqDeleteFile1->getPath(), - dwio::common::FileFormat::DWRF, - 1, - getFileSize(eqDeleteFile1->getPath()), - /*equalityFieldIds=*/{1}, - /*lowerBounds=*/{}, - /*upperBounds=*/{}, - /*dataSequenceNumber=*/10); - - // Second delete file: seqNum=3 (lower than data seqNum=5) → skipped. - auto deleteData2 = makeRowVector( - {"id"}, + auto deleteData = makeRowVector( + {"a", "b", "c"}, { - makeFlatVector({4}), + makeFlatVector({6}), + makeFlatVector({"1006"}), + makeFlatVector({2}), }); - auto eqDeleteFile2 = writeEqDeleteFile({deleteData2}); - IcebergDeleteFile icebergDeleteFile2( + auto eqDeleteFile = writeDataFile({deleteData}); + + IcebergDeleteFile icebergDeleteFile( FileContent::kEqualityDeletes, - eqDeleteFile2->getPath(), + eqDeleteFile->getPath(), dwio::common::FileFormat::DWRF, 1, - getFileSize(eqDeleteFile2->getPath()), - /*equalityFieldIds=*/{1}, - /*lowerBounds=*/{}, - /*upperBounds=*/{}, - /*dataSequenceNumber=*/3); + getFileSize(eqDeleteFile->getPath()), + /*equalityFieldIds=*/{1, 2, 3}); auto splits = makeSplits( dataFile->getPath(), - {icebergDeleteFile1, icebergDeleteFile2}, - /*dataSequenceNumber=*/5); - auto plan = makeTableScanPlan(rowType); + {{"a", std::optional{"6"}}, + {"c", std::optional{"2"}}}, + {dwio::common::FileFormat::DWRF, false}, + {icebergDeleteFile}, + 0); + auto plan = makeIcebergTableScanPlan(outputType, tableType); auto result = AssertQueryBuilder(plan).splits(splits).copyResults(pool()); - // Only id=2 is deleted (from delete file 1 with seqNum=10). - // id=4 survives because delete file 2 (seqNum=3) is skipped. + // (6,'1006',2) deleted; (6,'1009',2) survives. 'd' is NULL-filled. auto expected = makeRowVector( - {"id", "value"}, + {"a", "b", "d"}, { - makeFlatVector({1, 3, 4, 5}), - makeFlatVector({"a", "c", "d", "e"}), + makeFlatVector({6}), + makeFlatVector({"1009"}), + makeNullableFlatVector({std::nullopt}), }); - assertEqualResults({expected}, {result}); } -// TODO: Add a Parquet-format equality delete test. Currently all equality -// delete tests use DWRF because writeToFile() (from HiveConnectorTestBase) -// only supports DWRF. Adding a Parquet test requires adding Parquet writer -// dependencies to this test target's BUCK file and a Parquet write helper. - -} // namespace facebook::velox::connector::hive::iceberg +} // namespace facebook::velox::connector::hive::iceberg::test diff --git a/velox/connectors/hive/iceberg/tests/IcebergConnectorTest.cpp b/velox/connectors/hive/iceberg/tests/IcebergConnectorTest.cpp index 018b555d449..f34f6fb81a2 100644 --- a/velox/connectors/hive/iceberg/tests/IcebergConnectorTest.cpp +++ b/velox/connectors/hive/iceberg/tests/IcebergConnectorTest.cpp @@ -16,9 +16,13 @@ #include "velox/connectors/hive/iceberg/IcebergConnector.h" #include +#include "velox/common/io/IoStatistics.h" #include "velox/connectors/ConnectorRegistry.h" +#include "velox/connectors/hive/FileConfig.h" +#include "velox/connectors/hive/FileConnectorUtil.h" #include "velox/connectors/hive/HiveConfig.h" #include "velox/connectors/hive/iceberg/IcebergColumnHandle.h" +#include "velox/connectors/hive/iceberg/IcebergSplit.h" #include "velox/connectors/hive/iceberg/tests/IcebergTestBase.h" #include "velox/type/Type.h" @@ -69,6 +73,46 @@ TEST_F(IcebergConnectorTest, connectorProperties) { ASSERT_NE(icebergConnector->ioExecutor(), nullptr); } +TEST_F(IcebergConnectorTest, splitColumnMappingMode) { + auto split = IcebergSplitBuilder("/tmp/testfile") + .connectorId(test::kIcebergConnectorId) + .fileFormat(dwio::common::FileFormat::PARQUET) + .columnMappingMode(dwio::common::ColumnMappingMode::kName) + .build(); + + ASSERT_TRUE(split->columnMappingMode.has_value()); + EXPECT_EQ( + split->columnMappingMode.value(), dwio::common::ColumnMappingMode::kName); +} + +TEST_F(IcebergConnectorTest, splitMappingOverridesSession) { + auto fileConfig = std::make_shared( + std::make_shared( + std::unordered_map{}), + "hive."); + setConnectorSessionProperty(FileConfig::kUseColumnNamesSession, "false"); + auto split = IcebergSplitBuilder("/tmp/testfile") + .connectorId(test::kIcebergConnectorId) + .fileFormat(dwio::common::FileFormat::PARQUET) + .columnMappingMode(dwio::common::ColumnMappingMode::kName) + .build(); + + dwio::common::ReaderOptions readerOptions(pool_.get()); + readerOptions.setDataIoStats(std::make_shared()); + readerOptions.setMetadataIoStats(std::make_shared()); + configureReaderOptions( + fileConfig, + connectorQueryCtx_.get(), + /*fileSchema=*/nullptr, + split, + /*tableParameters=*/{}, + readerOptions); + + EXPECT_EQ( + readerOptions.columnMappingMode(), + dwio::common::ColumnMappingMode::kName); +} + TEST_F(IcebergConnectorTest, columnHandleForwardsPostProcessor) { auto called = std::make_shared(false); std::function postProcessor = [called](VectorPtr&) { diff --git a/velox/connectors/hive/iceberg/tests/IcebergDeletionVectorReadTest.cpp b/velox/connectors/hive/iceberg/tests/IcebergDeletionVectorReadTest.cpp index 779ea0f59ce..26a4c6ba531 100644 --- a/velox/connectors/hive/iceberg/tests/IcebergDeletionVectorReadTest.cpp +++ b/velox/connectors/hive/iceberg/tests/IcebergDeletionVectorReadTest.cpp @@ -21,7 +21,12 @@ #include +#include "velox/common/file/FileSystems.h" +#include "velox/connectors/hive/iceberg/IcebergMetadataColumns.h" +#include "velox/connectors/hive/iceberg/IcebergSplit.h" #include "velox/dwio/common/FileSink.h" +#include "velox/dwio/dwrf/reader/ReaderBase.h" +#include "velox/dwio/dwrf/writer/FlushPolicy.h" #include "velox/exec/tests/utils/AssertQueryBuilder.h" #include "velox/exec/tests/utils/PlanBuilder.h" @@ -116,5 +121,123 @@ TEST_F(IcebergDeletionVectorReadTest, deletionVectorFiltersDeletedRowsInScan) { exec::test::AssertQueryBuilder(plan).splits(splits).assertResults(expected); } +// A deletion vector removes rows from the scan output, so an output row's +// index no longer matches its position in the data file. _row_id must still +// report file-absolute positions. +// +// The row-number column that carries those positions used to be injected only +// for V2 positional delete files, so a split whose only delete file was a V3 +// deletion vector fell back to deriving positions from the output index. A +// DELETE with no WHERE clause hit this: with no filter to force the row-number +// column on, the rewritten deletion vector recorded 0..N-1 instead of the true +// positions and left the trailing rows of the file undeleted. +TEST_F(IcebergDeletionVectorReadTest, rowIdIsFileAbsoluteWithDeletionVector) { + auto dataFile = TempFilePath::create(); + writeToFile( + dataFile->getPath(), + {makeRowVector({makeFlatVector({10, 20, 30, 40, 50})})}); + + auto [puffinDir, dvFile] = writeDeletionVector(dataFile->getPath(), {1, 3}); + + const std::unordered_map infoColumns{ + {IcebergMetadataColumn::kFirstRowIdInfoColumn, "200"}, + {IcebergMetadataColumn::kDataSequenceNumberInfoColumn, "42"}}; + + const std::vector outputNames{ + "c0", "_row_id", "_last_updated_sequence_number"}; + auto plan = exec::test::PlanBuilder() + .startTableScan(test::kIcebergConnectorId) + .outputType( + ROW({outputNames[0], outputNames[1], outputNames[2]}, + {BIGINT(), BIGINT(), BIGINT()})) + .dataColumns(ROW({"c0"}, {BIGINT()})) + .endTableScan() + .planNode(); + + // Positions 1 and 3 are deleted, so the surviving rows keep file positions + // 0, 2 and 4 and _row_id is firstRowId plus those positions. + auto expected = makeRowVector( + outputNames, + { + makeFlatVector({10, 30, 50}), + makeFlatVector({200, 202, 204}), + makeFlatVector({42, 42, 42}), + }); + exec::test::AssertQueryBuilder(plan) + .splits({makeIcebergSplitWithInfoColumns( + dataFile->getPath(), infoColumns, {dvFile})}) + .assertResults(expected); +} + +// Deletion-vector positions are absolute row ordinals in the base data file, +// while a split's 'start'/'length' are a byte range over that file. This +// pins down the conversion between those two coordinate systems for a split +// that begins at a nonzero byte offset. +// +// IcebergSplitReader resolves the two by asking the format reader, not by +// arithmetic on the byte offset: after creating the row reader it sets +// 'splitOffset_ = baseRowReader_->nextRowNumber()', which for DWRF/ORC is the +// cumulative row count of the preceding stripes (row groups for Parquet), and +// each batch then recovers its absolute range as +// 'splitOffset_ + baseReadOffset_'. A DV position preceding the split must be +// ignored rather than shifted onto a row this split actually reads. +TEST_F(IcebergDeletionVectorReadTest, deletionVectorAppliesToNonZeroByteSplit) { + // Two DWRF stripes of five rows each: values 0..4 then 5..9, where each + // value equals its absolute row position. + auto dataFile = TempFilePath::create(); + writeToFile( + dataFile->getPath(), + { + makeRowVector({makeFlatVector({0, 1, 2, 3, 4})}), + makeRowVector({makeFlatVector({5, 6, 7, 8, 9})}), + }, + std::make_shared(), + []() { + return std::make_unique([]() { return true; }); + }); + + auto readFile = filesystems::getFileSystem(dataFile->getPath(), nullptr) + ->openFileForRead(dataFile->getPath()); + const uint64_t fileSize = readFile->size(); + dwio::common::ReaderOptions readerOptions{pool()}; + auto reader = std::make_unique( + readerOptions, + std::make_unique( + std::shared_ptr(std::move(readFile)), *pool())); + reader->loadCache(); + + // The premise of the test: the second stripe really does start at a nonzero + // byte offset and at absolute row 5, so the byte offset and the row offset + // are different numbers and cannot be confused for one another. + ASSERT_EQ(reader->footer().stripesSize(), 2); + const uint64_t secondStripeByteOffset = reader->footer().stripes(1).offset(); + ASSERT_GT(secondStripeByteOffset, 0); + ASSERT_EQ(reader->footer().stripes(0).numberOfRows(), 5); + + // Absolute data-file positions: 1 lives in the first stripe and is outside + // this split, 7 is the third row of the second stripe. + auto [puffinDir, dvFile] = writeDeletionVector(dataFile->getPath(), {1, 7}); + + auto split = IcebergSplitBuilder(dataFile->getPath()) + .connectorId(test::kIcebergConnectorId) + .fileFormat(fileFormat_) + .start(secondStripeByteOffset) + .length(fileSize - secondStripeByteOffset) + .deleteFiles({dvFile}) + .build(); + + auto plan = exec::test::PlanBuilder() + .startTableScan(test::kIcebergConnectorId) + .outputType(ROW({"c0"}, {BIGINT()})) + .endTableScan() + .planNode(); + + // Position 7 is removed from the second stripe. Position 1 precedes the + // split and must not delete anything; in particular it must not be + // misread as the second row of this split (value 6). + auto expected = makeRowVector({makeFlatVector({5, 6, 8, 9})}); + exec::test::AssertQueryBuilder(plan).splits({split}).assertResults(expected); +} + } // namespace } // namespace facebook::velox::connector::hive::iceberg diff --git a/velox/connectors/hive/iceberg/tests/IcebergDeletionVectorSinkTest.cpp b/velox/connectors/hive/iceberg/tests/IcebergDeletionVectorSinkTest.cpp index a284abe7fda..86ca83dc6f3 100644 --- a/velox/connectors/hive/iceberg/tests/IcebergDeletionVectorSinkTest.cpp +++ b/velox/connectors/hive/iceberg/tests/IcebergDeletionVectorSinkTest.cpp @@ -684,3 +684,125 @@ TEST_F(IcebergDeletionVectorSinkTest, dictionaryWrappedRowIdStructFlatPath) { pool_.get()); EXPECT_EQ(deleted, (std::vector{1, 3})); } + +TEST_F(IcebergDeletionVectorSinkTest, rejectsUnsupportedInputType) { + auto tempDir = TempDirectoryPath::create(); + auto handle = makeDeletionVectorHandle(tempDir->getPath()); + + // Neither the flat (file_path, pos) shape nor a ROW whose first two fields + // are (VARCHAR, BIGINT), so the sink cannot locate the row-id. + VELOX_ASSERT_USER_THROW( + IcebergDeletionVectorSink( + ROW({"a", "b"}, {BIGINT(), BIGINT()}), + handle, + connectorQueryCtx_.get(), + connector::CommitStrategy::kNoCommit, + hiveConfig_), + "IcebergDeletionVectorSink expects a two-column (file_path, pos) input"); +} + +TEST_F(IcebergDeletionVectorSinkTest, appendDataIgnoresNullAndEmptyPages) { + auto tempDir = TempDirectoryPath::create(); + auto handle = makeDeletionVectorHandle(tempDir->getPath()); + + IcebergDeletionVectorSink sink( + ROW({"file_path", "pos"}, {VARCHAR(), BIGINT()}), + handle, + connectorQueryCtx_.get(), + connector::CommitStrategy::kNoCommit, + hiveConfig_); + + sink.appendData(nullptr); + sink.appendData(makePositionDeleteRows({}, {})); + + // Neither page created per-file state, so nothing is written. + EXPECT_TRUE(sink.finish()); + EXPECT_TRUE(sink.close().empty()); + EXPECT_EQ(sink.stats().numWrittenFiles, 0u); +} + +TEST_F(IcebergDeletionVectorSinkTest, rowIdStructSkipsNullRowId) { + auto tempDir = TempDirectoryPath::create(); + auto handle = makeDeletionVectorHandle(tempDir->getPath()); + const std::string dataFile = tempDir->getPath() + "/A.parquet"; + + IcebergDeletionVectorSink sink( + ROW({"$row_id"}, + {ROW( + {"_file", "_pos", "_spec_id", "partition_data"}, + {VARCHAR(), BIGINT(), INTEGER(), VARCHAR()})}), + handle, + connectorQueryCtx_.get(), + connector::CommitStrategy::kNoCommit, + hiveConfig_); + + // A null row-id carries no (file_path, pos) to delete and must be skipped + // rather than dereferenced. + auto page = makeRowIdStructDeletePage( + dataFile, {10, 20, 30}, {0, 1, 2}, /*filePathConstant=*/true); + auto rowIdWithNull = page->childAt(0); + rowIdWithNull->setNull(1, true); + + sink.appendData(page); + EXPECT_TRUE(sink.finish()); + + auto messages = sink.close(); + ASSERT_EQ(messages.size(), 1); + const auto parsed = folly::parseJson(messages[0]); + // Positions 10 and 30 survive; the null row at index 1 contributes nothing. + EXPECT_EQ(parsed["metrics"]["recordCount"].asInt(), 2); +} + +TEST_F(IcebergDeletionVectorSinkTest, finishAndAbortAreIdempotent) { + auto tempDir = TempDirectoryPath::create(); + auto handle = makeDeletionVectorHandle(tempDir->getPath()); + + IcebergDeletionVectorSink sink( + ROW({"file_path", "pos"}, {VARCHAR(), BIGINT()}), + handle, + connectorQueryCtx_.get(), + connector::CommitStrategy::kNoCommit, + hiveConfig_); + + sink.appendData( + makePositionDeleteRows({tempDir->getPath() + "/A.parquet"}, {5})); + + EXPECT_TRUE(sink.finish()); + const auto statsAfterFirstFinish = sink.stats(); + + // A second finish() must not re-write the Puffin file or duplicate the + // commit message. + EXPECT_TRUE(sink.finish()); + EXPECT_EQ( + sink.stats().numWrittenFiles, statsAfterFirstFinish.numWrittenFiles); + EXPECT_EQ(sink.close().size(), 1); + + // abort() after finish() is a no-op rather than clearing committed state. + sink.abort(); + sink.abort(); + EXPECT_EQ(sink.close().size(), 1); +} + +TEST_F(IcebergDeletionVectorSinkTest, puffinPathFallsBackToTargetPath) { + auto tempDir = TempDirectoryPath::create(); + auto handle = makeDeletionVectorHandle(tempDir->getPath()); + + IcebergDeletionVectorSink sink( + ROW({"file_path", "pos"}, {VARCHAR(), BIGINT()}), + handle, + connectorQueryCtx_.get(), + connector::CommitStrategy::kNoCommit, + hiveConfig_); + + // A data-file name with no '/' has no parent directory to co-locate with, + // so the Puffin lands in the location handle's target path instead. + sink.appendData(makePositionDeleteRows({"bare-name.parquet"}, {1})); + EXPECT_TRUE(sink.finish()); + + auto messages = sink.close(); + ASSERT_EQ(messages.size(), 1); + const auto parsed = folly::parseJson(messages[0]); + const auto puffinPath = parsed["path"].asString(); + EXPECT_EQ(puffinPath.rfind(tempDir->getPath() + "/dv-", 0), 0) + << "puffin path should sit under the target path: " << puffinPath; +} diff --git a/velox/connectors/hive/iceberg/tests/IcebergDwrfInsertTest.cpp b/velox/connectors/hive/iceberg/tests/IcebergDwrfInsertTest.cpp index 7677da9db90..d4492a6d1f9 100644 --- a/velox/connectors/hive/iceberg/tests/IcebergDwrfInsertTest.cpp +++ b/velox/connectors/hive/iceberg/tests/IcebergDwrfInsertTest.cpp @@ -240,6 +240,128 @@ TEST_F(IcebergDwrfInsertTest, partitioned) { exec::test::AssertQueryBuilder(plan).splits(splits).assertResults(vectors); } +/// IcebergFileNameGenerator participates in plan serialization +/// (IcebergConnector::registerSerDe registers it), so its serialize / +/// deserialize / toString methods are reachable from any serialized plan even +/// though the write tests never round-trip a plan. +TEST_F(IcebergDwrfInsertTest, fileNameGeneratorSerDeRoundTrip) { + IcebergFileNameGenerator::registerSerDe(); + + const IcebergFileNameGenerator generator; + EXPECT_EQ(generator.toString(), "IcebergFileNameGenerator"); + + const auto serialized = generator.serialize(); + EXPECT_EQ(serialized["name"].asString(), "IcebergFileNameGenerator"); + + const auto deserialized = + IcebergFileNameGenerator::deserialize(serialized, /*context=*/nullptr); + ASSERT_NE(deserialized, nullptr); + EXPECT_EQ(deserialized->toString(), generator.toString()); +} + +/// Exercises writer file rotation. FileDataSink rotates once a writer's +/// current file reaches 'maxTargetFileBytes', which HiveDataSink derives from +/// the format-specific max-target-file-size setting (0 = unlimited, the +/// default, so no other test reaches this path). IcebergDataSink overrides +/// rotateWriter() to capture per-file statistics before the writer is reset; +/// without rotation coverage that override and the post-rotation null-writer +/// handling in closeInternal() never run. +TEST_F(IcebergDwrfInsertTest, writerRotationProducesMultipleFiles) { + // 1 byte forces a rotation after essentially every batch. + setConnectorSessionProperty( + connector::hive::HiveConfig::kOrcMaxTargetFileSizeSession, "1B"); + + auto rowType = ROW({"c1", "c2"}, {BIGINT(), VARCHAR()}); + const auto outputDirectory = TempDirectoryPath::create(); + const auto dataPath = outputDirectory->getPath(); + const auto vectors = createTestData(rowType, 3, 50); + + const auto dataSink = createDataSinkAndAppendData(vectors, dataPath); + const auto commitTasks = dataSink->close(); + + // Rotation means one unpartitioned writer emits several files, each with its + // own commit task and its own record count. + ASSERT_GT(commitTasks.size(), 1) + << "expected rotation to split the write across multiple files"; + EXPECT_EQ(listFiles(dataPath).size(), commitTasks.size()); + + int64_t totalRecords = 0; + for (const auto& task : commitTasks) { + const auto taskJson = folly::parseJson(task); + const auto records = taskJson["metrics"]["recordCount"].asInt(); + // Per-file counts are deltas, not the running total, so none may be zero + // and they must sum to the rows written. + EXPECT_GT(records, 0); + totalRecords += records; + } + EXPECT_EQ(totalRecords, 3 * 50); + + // The rotated files still read back as the original data. + auto splits = createSplitsForDirectory(dataPath); + ASSERT_EQ(splits.size(), commitTasks.size()); + auto plan = exec::test::PlanBuilder() + .startTableScan(test::kIcebergConnectorId) + .outputType(rowType) + .endTableScan() + .planNode(); + exec::test::AssertQueryBuilder(plan).splits(splits).assertResults(vectors); +} + +/// Partition values are serialized into the commit message by a per-type +/// dispatch. VARBINARY and TIMESTAMP take dedicated specializations — +/// base64-encoding and micros-since-epoch respectively — that the BIGINT and +/// VARCHAR partition tests never reach. +TEST_F(IcebergDwrfInsertTest, varbinaryPartitionValue) { + auto rowType = ROW({"c1", "c2"}, {VARBINARY(), BIGINT()}); + const auto outputDirectory = TempDirectoryPath::create(); + const auto dataPath = outputDirectory->getPath(); + const auto vectors = createTestData(rowType, 1, 20); + + std::vector partitionTransforms = { + {0, TransformType::kIdentity, std::nullopt}}; + const auto dataSink = + createDataSinkAndAppendData(vectors, dataPath, partitionTransforms); + const auto commitTasks = dataSink->close(); + ASSERT_GT(commitTasks.size(), 0); + + for (const auto& task : commitTasks) { + const auto taskJson = folly::parseJson(task); + ASSERT_GT(taskJson.count("partitionDataJson"), 0); + const auto partitionData = + folly::parseJson(taskJson["partitionDataJson"].asString()); + ASSERT_EQ(partitionData["partitionValues"].size(), 1); + // Binary partition values are base64 strings, never raw bytes. + const auto& value = partitionData["partitionValues"][0]; + EXPECT_TRUE(value.isString() || value.isNull()); + } +} + +TEST_F(IcebergDwrfInsertTest, timestampPartitionValue) { + auto rowType = ROW({"c1", "c2"}, {TIMESTAMP(), BIGINT()}); + const auto outputDirectory = TempDirectoryPath::create(); + const auto dataPath = outputDirectory->getPath(); + const auto vectors = createTestData(rowType, 1, 20); + + std::vector partitionTransforms = { + {0, TransformType::kIdentity, std::nullopt}}; + const auto dataSink = + createDataSinkAndAppendData(vectors, dataPath, partitionTransforms); + const auto commitTasks = dataSink->close(); + ASSERT_GT(commitTasks.size(), 0); + + for (const auto& task : commitTasks) { + const auto taskJson = folly::parseJson(task); + ASSERT_GT(taskJson.count("partitionDataJson"), 0); + const auto partitionData = + folly::parseJson(taskJson["partitionDataJson"].asString()); + ASSERT_EQ(partitionData["partitionValues"].size(), 1); + // Iceberg stores timestamps as micros since epoch, so the serialized + // partition value must be an integer rather than a formatted string. + const auto& value = partitionData["partitionValues"][0]; + EXPECT_TRUE(value.isInt() || value.isNull()); + } +} + /// Regression test for the isPartitioned() guard added to ensureWriter(). /// Without the guard, calling ensureWriter() on a non-partitioned table /// invoked makeCommitPartitionValue(), which dereferences diff --git a/velox/connectors/hive/iceberg/tests/IcebergMergeSinkTest.cpp b/velox/connectors/hive/iceberg/tests/IcebergMergeSinkTest.cpp index c68a8f54f92..d3d245bf5c1 100644 --- a/velox/connectors/hive/iceberg/tests/IcebergMergeSinkTest.cpp +++ b/velox/connectors/hive/iceberg/tests/IcebergMergeSinkTest.cpp @@ -468,6 +468,42 @@ TEST_F(IcebergMergeSinkTest, dataInputTypeNamesComeFromHandleNotSource) { // VELOX_ENABLE_PARQUET guard removed: rely on the iceberg_connector // target's unconditional Parquet writer dependency. +TEST_F(IcebergMergeSinkTest, appendDataIgnoresNullAndEmptyPages) { + auto tempDir = TempDirectoryPath::create(); + auto sink = makeSink(tempDir->getPath()); + + // Neither page carries rows, so neither sub-sink is driven and no operation + // byte is validated. + sink->appendData(nullptr); + sink->appendData(makeInput( + /*ids=*/std::vector>{}, + /*names=*/std::vector>{}, + /*operations=*/std::vector{}, + /*filePaths=*/std::vector>{}, + /*positions=*/std::vector>{}, + /*insertFromUpdate=*/std::vector{})); + + EXPECT_TRUE(sink->finish()); + EXPECT_TRUE(sink->close().empty()); +} + +TEST_F(IcebergMergeSinkTest, abortIsIdempotent) { + auto tempDir = TempDirectoryPath::create(); + auto sink = makeSink(tempDir->getPath()); + + sink->appendData(makeInput( + /*ids=*/{1}, + /*names=*/{{std::string("a")}}, + /*operations=*/{IMS::kInsertOperationNumber}, + /*filePaths=*/{std::nullopt}, + /*positions=*/{std::nullopt}, + /*insertFromUpdate=*/{0})); + + // The second abort must not re-enter the sub-sinks, which are not safe to + // abort twice. + sink->abort(); + sink->abort(); +} } // namespace } // namespace facebook::velox::connector::hive::iceberg::test diff --git a/velox/connectors/hive/iceberg/tests/IcebergNimbleInsertTest.cpp b/velox/connectors/hive/iceberg/tests/IcebergNimbleInsertTest.cpp index 92e8ad84672..e3a33479693 100644 --- a/velox/connectors/hive/iceberg/tests/IcebergNimbleInsertTest.cpp +++ b/velox/connectors/hive/iceberg/tests/IcebergNimbleInsertTest.cpp @@ -15,16 +15,17 @@ */ // This end-to-end test exercises the batch NimbleReader/NimbleWriter factories, -// which live under dwio/nimble/.../fb/ and are internal-only (not shipped to -// OSS). Guard the whole test so the OSS build (VELOX_ENABLE_NIMBLE off) does -// not try to include the fb/ headers. Mirrors WriterOptionsAdapterTest.cpp. +// both now under velox/dwio/nimble/ and neither behind an fb/ segment: both are +// OSS-exportable. Guard the whole test so the OSS build (VELOX_ENABLE_NIMBLE +// off) does not try to build the batch reader, which stays internal-only. +// Mirrors WriterOptionsAdapterTest.cpp. #ifdef VELOX_ENABLE_NIMBLE -#include "dwio/nimble/velox/reader/fb/NimbleReader.h" -#include "dwio/nimble/writer/fb/NimbleWriter.h" #include "velox/connectors/hive/HiveConfig.h" #include "velox/connectors/hive/iceberg/IcebergColumnHandle.h" #include "velox/connectors/hive/iceberg/tests/IcebergTestBase.h" +#include "velox/dwio/nimble/velox/reader/NimbleReaderFactory.h" +#include "velox/dwio/nimble/writer/WriterFactory.h" #include "velox/exec/tests/utils/AssertQueryBuilder.h" #include "velox/exec/tests/utils/PlanBuilder.h" diff --git a/velox/connectors/hive/iceberg/tests/IcebergParquetStatsTest.cpp b/velox/connectors/hive/iceberg/tests/IcebergParquetStatsTest.cpp index cf336ffad77..0a02c688568 100644 --- a/velox/connectors/hive/iceberg/tests/IcebergParquetStatsTest.cpp +++ b/velox/connectors/hive/iceberg/tests/IcebergParquetStatsTest.cpp @@ -592,7 +592,7 @@ TEST_F(IcebergParquetStatsTest, mixedDoubleFloat) { EXPECT_FLOAT_EQ(maxFloatVal, -std::numeric_limits::infinity()); } -TEST_F(IcebergParquetStatsTest, NaN) { +TEST_F(IcebergParquetStatsTest, nan) { constexpr vector_size_t size = 1'000; constexpr int32_t expectedNulls = 500; constexpr int32_t doubleColId = 1; diff --git a/velox/connectors/hive/iceberg/tests/IcebergReadTest.cpp b/velox/connectors/hive/iceberg/tests/IcebergReadTest.cpp index d412f545bf3..22dad1d4b72 100644 --- a/velox/connectors/hive/iceberg/tests/IcebergReadTest.cpp +++ b/velox/connectors/hive/iceberg/tests/IcebergReadTest.cpp @@ -391,9 +391,9 @@ TEST_F(IcebergReadTest, readParquetFlatSchemaEvolutionByFieldId) { ROW({"enabled", "id"}, {BOOLEAN(), BIGINT()}), ROW({"id", "flag", "status"}, {BIGINT(), BOOLEAN(), VARCHAR()}), { - {"id", "id", BIGINT(), makeFieldId(1)}, - {"enabled", "flag", BOOLEAN(), makeFieldId(2)}, - {"status", "status", VARCHAR(), makeFieldId(3)}, + {"id", "id", BIGINT(), makeFieldId(1), {}}, + {"enabled", "flag", BOOLEAN(), makeFieldId(2), {}}, + {"status", "status", VARCHAR(), makeFieldId(3), {}}, }, makeRowVector( {"enabled", "id"}, @@ -407,10 +407,10 @@ TEST_F(IcebergReadTest, readParquetFlatSchemaEvolutionByFieldId) { ROW({"id", "flag", "status", "score"}, {BIGINT(), BOOLEAN(), VARCHAR(), INTEGER()}), { - {"id", "id", BIGINT(), makeFieldId(1)}, - {"flag", "flag", BOOLEAN(), makeFieldId(2)}, - {"status", "status", VARCHAR(), makeFieldId(3)}, - {"score", "score", INTEGER(), makeFieldId(4)}, + {"id", "id", BIGINT(), makeFieldId(1), {}}, + {"flag", "flag", BOOLEAN(), makeFieldId(2), {}}, + {"status", "status", VARCHAR(), makeFieldId(3), {}}, + {"score", "score", INTEGER(), makeFieldId(4), {}}, }, makeRowVector( {"id", "flag", "status", "score"}, @@ -425,8 +425,8 @@ TEST_F(IcebergReadTest, readParquetFlatSchemaEvolutionByFieldId) { ROW({"id", "status"}, {BIGINT(), VARCHAR()}), ROW({"id", "status"}, {BIGINT(), VARCHAR()}), { - {"id", "id", BIGINT(), makeFieldId(1)}, - {"status", "status", VARCHAR(), makeFieldId(3)}, + {"id", "id", BIGINT(), makeFieldId(1), {}}, + {"status", "status", VARCHAR(), makeFieldId(3), {}}, }, makeRowVector( {"id", "status"}, @@ -438,7 +438,7 @@ TEST_F(IcebergReadTest, readParquetFlatSchemaEvolutionByFieldId) { ROW({"score"}, {INTEGER()}), ROW({"score"}, {INTEGER()}), { - {"score", "score", INTEGER(), makeFieldId(4)}, + {"score", "score", INTEGER(), makeFieldId(4), {}}, }, makeRowVector( {"score"}, @@ -450,8 +450,8 @@ TEST_F(IcebergReadTest, readParquetFlatSchemaEvolutionByFieldId) { ROW({"id", "status"}, {BIGINT(), BOOLEAN()}), ROW({"id", "status"}, {BIGINT(), BOOLEAN()}), { - {"id", "id", BIGINT(), makeFieldId(1)}, - {"status", "status", BOOLEAN(), makeFieldId(4)}, + {"id", "id", BIGINT(), makeFieldId(1), {}}, + {"status", "status", BOOLEAN(), makeFieldId(4), {}}, }, makeRowVector( {"id", "status"}, @@ -481,7 +481,7 @@ TEST_F(IcebergReadTest, readParquetFilterOnlyColumnByFieldId) { .outputType(ROW({"id"}, {BIGINT()})) .dataColumns(testData.writeType) .assignments(makeFieldIdAssignments({ - {"id", "id", BIGINT(), makeFieldId(1)}, + {"id", "id", BIGINT(), makeFieldId(1), {}}, })) .filterColumnHandles({ makeIcebergHandle("status", VARCHAR(), makeFieldId(3)), @@ -573,8 +573,8 @@ TEST_F(IcebergReadTest, readParquetArrayByFieldId) { {"items", "items", nestedReadType->childAt(0), - makeFieldId( - 1, {makeFieldId(2, {makeFieldId(4), makeFieldId(3)})})}, + makeFieldId(1, {makeFieldId(2, {makeFieldId(4), makeFieldId(3)})}), + {}}, }, {nestedExpected}); } @@ -596,7 +596,8 @@ TEST_F(IcebergReadTest, readParquetArrayByFieldId) { {"items", "items", nestedProjectedReadType->childAt(0), - makeFieldId(1, {makeFieldId(2, {makeFieldId(4)})})}, + makeFieldId(1, {makeFieldId(2, {makeFieldId(4)})}), + {}}, }, {nestedProjectedExpected}); } @@ -645,7 +646,8 @@ TEST_F(IcebergReadTest, readParquetMapByFieldId) { makeFieldId( 1, {makeFieldId(2), - makeFieldId(3, {makeFieldId(5), makeFieldId(4)})})}, + makeFieldId(3, {makeFieldId(5), makeFieldId(4)})}), + {}}, }, {mapExpected}); } @@ -671,8 +673,8 @@ TEST_F(IcebergReadTest, readParquetMapByFieldId) { {"attributes", "attributes", mapProjectedReadType->childAt(0), - makeFieldId( - 1, {makeFieldId(2), makeFieldId(3, {makeFieldId(5)})})}, + makeFieldId(1, {makeFieldId(2), makeFieldId(3, {makeFieldId(5)})}), + {}}, }, {mapProjectedExpected}); } @@ -809,12 +811,7 @@ TEST_F(IcebergReadTest, addColumnWithDefaultAllTypes) { Timestamp(1705314600, 0)})})}; assertDefaultValues( - newRowType, - newRowType, - assignments, - dataVectors, - expectedVectors, - {{HiveConfig::kReadTimestampPartitionValueAsLocalTimeSession, "false"}}); + newRowType, newRowType, assignments, dataVectors, expectedVectors); } TEST_F(IcebergReadTest, addColumnWithInvalidDefault) { @@ -1389,6 +1386,94 @@ TEST_F(IcebergReadTest, targetTableRowIdSynthesis) { .assertResults({expected}); } +// Info columns arrive as strings on the split and are parsed at read time. +// A value the coordinator could not have produced means the split metadata is +// corrupt, so the reader must fail loudly rather than silently substituting a +// default: $first_row_id and $data_sequence_number both feed V3 row lineage, +// and a wrong value there mislabels every row in the file. +class IcebergInfoColumnValidationTest : public IcebergReadTest { + protected: + // Runs a scan of a single BIGINT column with 'infoColumns' attached to the + // split, projecting 'outputType' (which decides whether the row-lineage or + // MERGE row-id parsing paths run at all). + void assertScanFails( + const std::unordered_map& infoColumns, + const RowTypePtr& outputType, + const std::string& expectedMessage) { + std::vector inputVectors = { + makeRowVector({"c0"}, {makeFlatVector({10, 20, 30})})}; + auto dataFilePath = TempFilePath::create(); + writeToFile(dataFilePath->getPath(), inputVectors); + + auto plan = exec::test::PlanBuilder() + .startTableScan(test::kIcebergConnectorId) + .outputType(outputType) + .dataColumns(ROW({"c0"}, {BIGINT()})) + .endTableScan() + .planNode(); + + VELOX_ASSERT_THROW( + exec::test::AssertQueryBuilder(plan) + .splits({makeIcebergSplitWithInfoColumns( + dataFilePath->getPath(), infoColumns)}) + .copyResults(pool()), + expectedMessage); + } + + // Projecting _row_id is what makes the reader parse $first_row_id. + RowTypePtr rowLineageOutputType() const { + return ROW( + {"c0", IcebergMetadataColumn::kRowIdColumnName}, {BIGINT(), BIGINT()}); + } + + // Projecting $target_table_row_id is what makes the reader parse $spec_id. + RowTypePtr targetRowIdOutputType() const { + return ROW( + {"c0", IcebergMetadataColumn::kTargetTableRowIdColumnName}, + {BIGINT(), + ROW({"file_path", "row_position", "spec_id", "partition_data"}, + {VARCHAR(), BIGINT(), INTEGER(), VARCHAR()})}); + } +}; + +TEST_F(IcebergInfoColumnValidationTest, rejectsNonNumericFirstRowId) { + assertScanFails( + {{IcebergMetadataColumn::kFirstRowIdInfoColumn, "not-a-number"}}, + rowLineageOutputType(), + "Invalid $first_row_id value in split info columns"); +} + +TEST_F(IcebergInfoColumnValidationTest, rejectsNegativeFirstRowId) { + // Parses cleanly but is out of range: row ids are file-absolute offsets. + assertScanFails( + {{IcebergMetadataColumn::kFirstRowIdInfoColumn, "-1"}}, + rowLineageOutputType(), + "First row ID must be non-negative"); +} + +TEST_F(IcebergInfoColumnValidationTest, rejectsNonNumericDataSequenceNumber) { + assertScanFails( + {{IcebergMetadataColumn::kFirstRowIdInfoColumn, "0"}, + {IcebergMetadataColumn::kDataSequenceNumberInfoColumn, "abc"}}, + rowLineageOutputType(), + "Invalid $data_sequence_number value in split info columns"); +} + +TEST_F(IcebergInfoColumnValidationTest, rejectsNegativeDataSequenceNumber) { + assertScanFails( + {{IcebergMetadataColumn::kFirstRowIdInfoColumn, "0"}, + {IcebergMetadataColumn::kDataSequenceNumberInfoColumn, "-5"}}, + rowLineageOutputType(), + "Data sequence number must be non-negative"); +} + +TEST_F(IcebergInfoColumnValidationTest, rejectsNonNumericSpecId) { + assertScanFails( + {{IcebergMetadataColumn::kSpecIdInfoColumn, "spec-seven"}}, + targetRowIdOutputType(), + "Invalid $spec_id value in split info columns"); +} + TEST_F(IcebergReadTest, flatMapAsStruct) { // Write a DWRF file with a MAP column. auto mapType = MAP(BIGINT(), DOUBLE()); @@ -1443,5 +1528,114 @@ TEST_F(IcebergReadTest, flatMapAsStruct) { .assertResults({expected}); } +TEST_F(IcebergReadTest, filterPushdownWithInitialDefaultInFilterColumnHandles) { + // Test's a scenario where filter on default value column is used in the query + // TABLE = [id int , country varchar(defaultValue='IN')] + // QUERY = SELECT id FROM table WHERE country = 'IN' + // When filter pushdown is enabled, column handle for 'country' is present + // only in filterColumnHandles_ of HiveTableHandle. adaptColumns was searching + // only in columnHandles_ for the default value, missing it and creating null + // vector. + + // Old data file: only the 'id' column present. + std::vector dataVectors = { + makeRowVector({makeFlatVector({1, 2, 3})})}; + auto dataFilePath = TempFilePath::create(); + writeToFile(dataFilePath->getPath(), dataVectors); + + auto outputType = ROW({"id"}, {BIGINT()}); + + ColumnHandleMap assignments; + assignments["id"] = makeIcebergHandle("id", BIGINT(), 1); + + // filterColumnHandles carries country WITH initialDefaultValue="IN". + std::vector filterHandles = { + makeIcebergHandle("country", VARCHAR(), 2, "IN")}; + + // Expected: all 3 rows (country filter passes via default constant). + std::vector allRowsIN = { + makeRowVector(outputType->names(), {makeFlatVector({1, 2, 3})})}; + + // Full schema used for subfieldFilter expression parsing (country must be + // reachable even though it is not in outputType). + auto fullSchema = ROW({"id", "country"}, {BIGINT(), VARCHAR()}); + + auto assertFilter = + [&](const std::string& subfieldFilter, + const std::vector& expected, + const std::vector>& splits, + int32_t numSplitsSkipped = 0) { + auto plan = exec::test::PlanBuilder() + .startTableScan() + .connectorId(test::kIcebergConnectorId) + .outputType(outputType) + .dataColumns(fullSchema) + .assignments(assignments) + .filterColumnHandles(filterHandles) + .subfieldFilter(subfieldFilter) + .endTableScan() + .planNode(); + auto task = + exec::test::AssertQueryBuilder(plan).splits(splits).assertResults( + expected); + ASSERT_EQ( + task->taskStats() + .pipelineStats[0] + .operatorStats[0] + .runtimeStats["skippedSplits"] + .sum, + numSplitsSkipped) + << "Unexpected skipped splits for filter: " << subfieldFilter; + }; + + // Bug scenario: country absent from file, assignments handle has no default. + // Without fix: adaptColumns finds no default → sets NULL → testFilters skips + // file (NULL != 'IN') → 0 rows, numSplitsSkipped=1. WRONG. + // With fix: adaptColumns finds default 'IN' in filterColumnHandles → + // constant 'IN' → testFilters passes → 3 rows. CORRECT. + // Note: splits must be recreated for each assertFilter call — ConnectorSplit + // objects have their dataSource set during execution and cannot be reused. + assertFilter( + "country = 'IN'", + allRowsIN, + makeIcebergSplits(dataFilePath->getPath()), + /*numSplitsSkipped=*/0); + + // Non-matching default: constant 'IN' != 'US' → file skipped regardless. + assertFilter( + "country = 'US'", + {}, + makeIcebergSplits(dataFilePath->getPath()), + /*numSplitsSkipped=*/1); + + // New file written AFTER ALTER TABLE: country physically present = 'US'. + // Output still only has {id} — country is filter-only. + std::vector newData = {makeRowVector( + {"id", "country"}, + {makeFlatVector({4, 5}), + makeFlatVector({"US", "US"})})}; + auto newFilePath = TempFilePath::create(); + writeToFile(newFilePath->getPath(), newData); + + auto makeTwoSplits = [&]() { + auto s1 = makeIcebergSplits(dataFilePath->getPath()); + auto s2 = makeIcebergSplits(newFilePath->getPath()); + s1.insert(s1.end(), s2.begin(), s2.end()); + return s1; + }; + + // country='IN': old file passes (constant 'IN'), new file skipped ('US'). + // 1 split skipped, rows {1,2,3} from old file. + assertFilter( + "country = 'IN'", allRowsIN, makeTwoSplits(), /*numSplitsSkipped=*/1); + + // country='US': old file skipped (constant 'IN'!='US'), new file passes. + // 1 split skipped, rows {4,5} from new file. + std::vector newRowsUS = { + makeRowVector(outputType->names(), {makeFlatVector({4, 5})})}; + assertFilter( + "country = 'US'", newRowsUS, makeTwoSplits(), /*numSplitsSkipped=*/1); +} + } // namespace } // namespace facebook::velox::connector::hive::iceberg diff --git a/velox/connectors/hive/iceberg/tests/IcebergSplitReaderBenchmark.cpp b/velox/connectors/hive/iceberg/tests/IcebergSplitReaderBenchmark.cpp index 7c514460a37..33d18933989 100644 --- a/velox/connectors/hive/iceberg/tests/IcebergSplitReaderBenchmark.cpp +++ b/velox/connectors/hive/iceberg/tests/IcebergSplitReaderBenchmark.cpp @@ -18,6 +18,7 @@ #include #include "velox/connectors/hive/HiveConfig.h" +#include "velox/connectors/hive/iceberg/IcebergTableHandle.h" using namespace facebook::velox; using namespace facebook::velox::dwio; @@ -235,7 +236,7 @@ int IcebergSplitReaderBenchmark::read( const RowTypePtr& rowType, uint32_t nextSize, std::unique_ptr icebergSplitReader) { - runtimeStats_ = RuntimeStatistics(); + runtimeStats_ = RuntimeStats(); icebergSplitReader->resetFilterCaches(); int resultSize = 0; auto result = BaseVector::create(rowType, 0, leafPool_.get()); @@ -282,8 +283,8 @@ void IcebergSplitReaderBenchmark::readSingleColumn( core::TypedExprPtr remainingFilterExpr; - std::shared_ptr hiveTableHandle = - std::make_shared( + const std::shared_ptr icebergTableHandle = + std::make_shared( "kHiveConnectorId", "tableName", std::move(filters), @@ -339,7 +340,7 @@ void IcebergSplitReaderBenchmark::readSingleColumn( std::unique_ptr icebergSplitReader = std::make_unique( icebergSplit, - hiveTableHandle, + icebergTableHandle, nullptr, connectorQueryCtx_.get(), hiveConfig, diff --git a/velox/connectors/hive/iceberg/tests/IcebergSplitReaderBenchmark.h b/velox/connectors/hive/iceberg/tests/IcebergSplitReaderBenchmark.h index 5d96cc21005..f808ffed990 100644 --- a/velox/connectors/hive/iceberg/tests/IcebergSplitReaderBenchmark.h +++ b/velox/connectors/hive/iceberg/tests/IcebergSplitReaderBenchmark.h @@ -119,7 +119,7 @@ class IcebergSplitReaderBenchmark { std::shared_ptr rootPool_; std::shared_ptr leafPool_; std::unique_ptr writer_; - dwio::common::RuntimeStatistics runtimeStats_; + dwio::common::RuntimeStats runtimeStats_; dwio::common::FileFormat fileFormat_{dwio::common::FileFormat::DWRF}; const std::string kHiveConnectorId = "hive-iceberg"; diff --git a/velox/connectors/hive/iceberg/tests/IcebergTableHandleTest.cpp b/velox/connectors/hive/iceberg/tests/IcebergTableHandleTest.cpp new file mode 100644 index 00000000000..f6e97dcad07 --- /dev/null +++ b/velox/connectors/hive/iceberg/tests/IcebergTableHandleTest.cpp @@ -0,0 +1,472 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/connectors/hive/iceberg/IcebergTableHandle.h" + +#include +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/connectors/hive/TableHandle.h" +#include "velox/connectors/hive/iceberg/IcebergColumnHandle.h" +#include "velox/dwio/common/ParquetFieldId.h" +#include "velox/type/Type.h" + +// facebook::velox +using facebook::velox::BIGINT; +using facebook::velox::ISerializable; +using facebook::velox::ROW; +using facebook::velox::Type; +using facebook::velox::TypePtr; +using facebook::velox::VARCHAR; +using facebook::velox::dwio::common::ParquetFieldId; + +// facebook::velox::common +using facebook::velox::common::Subfield; +using facebook::velox::common::SubfieldFilters; + +// facebook::velox::connector::hive +using facebook::velox::connector::hive::FileColumnHandle; +using facebook::velox::connector::hive::HiveColumnHandle; +using facebook::velox::connector::hive::HiveTableHandle; + +// facebook::velox::connector::hive::iceberg +using facebook::velox::connector::hive::iceberg::IcebergColumnHandle; +using facebook::velox::connector::hive::iceberg::IcebergColumnHandlePtr; +using facebook::velox::connector::hive::iceberg::IcebergFieldMetadata; +using facebook::velox::connector::hive::iceberg::IcebergTableHandle; + +namespace { + +// Registers all SerDe entries needed to round-trip IcebergTableHandle and +// IcebergColumnHandle. +void registerAll() { + Type::registerSerDe(); + HiveColumnHandle::registerSerDe(); + IcebergColumnHandle::registerSerDe(); + HiveTableHandle::registerSerDe(); + IcebergTableHandle::registerSerDe(); +} + +// Builds a minimal IcebergColumnHandle for use in table handle tests. +IcebergColumnHandlePtr makeIcebergCol( + const std::string& name, + const TypePtr& type, + int32_t fieldId = 1) { + return std::make_shared( + name, + FileColumnHandle::ColumnType::kRegular, + type, + ParquetFieldId{fieldId, {}}); +} + +// Builds a minimal IcebergTableHandle with default Iceberg fields. +std::shared_ptr makeMinimal( + const std::string& connectorId = "test-iceberg", + const std::string& tableName = "test_table") { + return std::make_shared( + connectorId, + tableName, + /*subfieldFilters=*/SubfieldFilters{}, + /*remainingFilter=*/nullptr); +} + +} // namespace + +// --------------------------------------------------------------------------- +// Field accessors +// --------------------------------------------------------------------------- + +TEST(IcebergTableHandleTest, defaultFields) { + registerAll(); + + auto handle = makeMinimal(); + + ASSERT_EQ(handle->tableName(), "test_table"); + ASSERT_EQ(handle->name(), "test_table"); + ASSERT_FALSE(handle->isChangelogQuery()); + ASSERT_TRUE(handle->dataColumnHandles().empty()); + ASSERT_TRUE(handle->subfieldFilters().empty()); + ASSERT_EQ(handle->remainingFilter(), nullptr); + ASSERT_EQ(handle->sampleRate(), 1.0); + ASSERT_EQ(handle->dataColumns(), nullptr); + ASSERT_TRUE(handle->dbName().empty()); +} + +// isChangelogQuery=true with empty dataColumnHandles must throw. +TEST(IcebergTableHandleTest, changelogQueryRequiresDataColumnHandles) { + registerAll(); + + VELOX_ASSERT_THROW( + std::make_shared( + "test-iceberg", + "cdc_table", + SubfieldFilters{}, + /*remainingFilter=*/nullptr, + /*dataColumns=*/nullptr, + /*indexColumns=*/std::vector{}, + /*tableParameters=*/std::unordered_map{}, + /*filterColumnHandles=*/std::vector{}, + /*sampleRate=*/1.0, + /*dbName=*/"", + /*dataColumnFieldIds=*/std::vector{}, + /*isChangelogQuery=*/true, + /*dataColumnHandles=*/ + std::unordered_map{}), + "dataColumnHandles must not be empty when isChangelogQuery is true"); +} + +// --------------------------------------------------------------------------- +// toString — fully-populated IcebergColumnHandle inside dataColumnHandles. +// Covers: nested fieldId children, initialDefaultValue, icebergMetadata +// (empty, partial, fully-populated), sorted map order. +// --------------------------------------------------------------------------- + +TEST(IcebergTableHandleTest, toString) { + registerAll(); + + // Three data columns to exercise sorted output ("id" < "payload" < "score") + // and all three icebergMetadata states: partial, empty, fully-populated. + + // "id": partial icebergMetadata (required + longType only). + IcebergFieldMetadata icebergFieldMetadata; + icebergFieldMetadata.required = false; + icebergFieldMetadata.longType = "LONG"; + auto idCol = std::make_shared( + "id", + FileColumnHandle::ColumnType::kRegular, + BIGINT(), + ParquetFieldId{3, {}}, + /*requiredSubfields=*/std::vector{}, + /*initialDefaultValue=*/std::nullopt, + icebergFieldMetadata); + + // "payload": nested fieldId=10[11, 12], initialDefaultValue, no metadata. + ParquetFieldId nestedField{ + 10, {ParquetFieldId{11, {}}, ParquetFieldId{12, {}}}}; + auto payloadCol = std::make_shared( + "payload", + FileColumnHandle::ColumnType::kRegular, + ROW({{"key", BIGINT()}, {"value", VARCHAR()}}), + nestedField, + /*requiredSubfields=*/std::vector{}, + /*initialDefaultValue=*/std::optional{"{}"}); // empty struct + + // "score": fully-populated icebergMetadata, all six attributes set. + IcebergFieldMetadata fullMeta; + fullMeta.required = true; + fullMeta.longType = "LONG"; + fullMeta.timestampUnit = "MICROS"; + fullMeta.binaryType = "UUID"; + fullMeta.structType = "VariantStruct"; + fullMeta.length = 16; + auto scoreCol = std::make_shared( + "score", + FileColumnHandle::ColumnType::kRegular, + BIGINT(), + ParquetFieldId{7, {}}, + /*requiredSubfields=*/std::vector{}, + /*initialDefaultValue=*/std::nullopt, + fullMeta); + + std::unordered_map dataColumnHandles = { + {"id", idCol}, + {"payload", payloadCol}, + {"score", scoreCol}, + }; + + auto dataColumns = + ROW({{"c0", BIGINT()}, {"c1", VARCHAR()}, {"c2", BIGINT()}}); + auto handle = std::make_shared( + "test-iceberg", + "cdc_table", + SubfieldFilters{}, + /*remainingFilter=*/nullptr, + dataColumns, + /*indexColumns=*/std::vector{}, + /*tableParameters=*/ + std::unordered_map{{"format", "parquet"}}, + /*filterColumnHandles=*/std::vector{}, + /*sampleRate=*/0.5, + /*dbName=*/"analytics", + /*dataColumnFieldIds=*/std::vector{}, + /*isChangelogQuery=*/true, + dataColumnHandles); + + // dataColumnHandles_ is sorted by name: "id" < "payload" < "score". + ASSERT_EQ( + handle->toString(), + "table: cdc_table" + ", sample rate: 0.5" + ", data columns: ROW" + ", table parameters: [format:parquet]" + ", isChangelogQuery: true" + ", dataColumnHandles: [" + "id: IcebergColumnHandle [name: id, columnType: Regular," + " dataType: BIGINT, requiredSubfields: [ ], field: 3," + " icebergMetadata: {required: false, longType: LONG}]" + ", payload: IcebergColumnHandle [name: payload, columnType: Regular," + " dataType: ROW, requiredSubfields: [ ]," + " field: 10[11, 12], initialDefaultValue: {}]" + ", score: IcebergColumnHandle [name: score, columnType: Regular," + " dataType: BIGINT, requiredSubfields: [ ], field: 7," + " icebergMetadata: {required: true, longType: LONG," + " timestampUnit: MICROS, binaryType: UUID," + " structType: VariantStruct, length: 16}]" + "]"); +} + +// --------------------------------------------------------------------------- +// SerDe round-trips +// --------------------------------------------------------------------------- + +// Minimal handle: only connectorId, tableName, subfieldFilters, +// remainingFilter. All optional fields must deserialize to their defaults. +TEST(IcebergTableHandleTest, serdeMinimal) { + registerAll(); + + auto handle = makeMinimal(); + auto clone = ISerializable::deserialize( + handle->serialize(), /*context=*/nullptr); + + ASSERT_EQ(clone->connectorId(), handle->connectorId()); + ASSERT_EQ(clone->tableName(), handle->tableName()); + ASSERT_TRUE(clone->subfieldFilters().empty()); + ASSERT_EQ(clone->remainingFilter(), nullptr); + ASSERT_EQ(clone->dataColumns(), nullptr); + ASSERT_TRUE(clone->dataColumnFieldIds().empty()); + ASSERT_TRUE(clone->indexColumns().empty()); + ASSERT_TRUE(clone->tableParameters().empty()); + ASSERT_DOUBLE_EQ(clone->sampleRate(), 1.0); + ASSERT_TRUE(clone->dbName().empty()); + ASSERT_FALSE(clone->isChangelogQuery()); + ASSERT_TRUE(clone->dataColumnHandles().empty()); + ASSERT_TRUE(clone->hiveFilterColumnHandles().empty()); +} + +// Fully-populated round-trip: every serialized field is set to a non-default +// value and verified to survive the serialize/deserialize cycle. +TEST(IcebergTableHandleTest, serdeFullyPopulated) { + registerAll(); + + // dataColumnHandles: leaf + nested field, one with initialDefaultValue. + const ParquetFieldId nestedField{ + 10, {ParquetFieldId{11, {}}, ParquetFieldId{12, {}}}}; + auto payloadCol = std::make_shared( + "payload", + FileColumnHandle::ColumnType::kRegular, + ROW({{"key", BIGINT()}, {"value", VARCHAR()}}), + nestedField, + /*requiredSubfields=*/std::vector{}, + /*initialDefaultValue=*/std::optional{"active"}); + + const std::unordered_map + dataColumnHandles = { + {"id", makeIcebergCol("id", BIGINT(), /*fieldId=*/7)}, + {"payload", payloadCol}, + }; + + // filterColumnHandles: two Iceberg-typed handles. + std::vector filterHandles = { + makeIcebergCol("partition_date", VARCHAR(), /*fieldId=*/5), + makeIcebergCol("region_id", BIGINT(), /*fieldId=*/6), + }; + + auto dataColumns = ROW({{"c0", BIGINT()}, {"c1", VARCHAR()}}); + auto handle = std::make_shared( + "test-iceberg", + "cdc_table", + SubfieldFilters{}, + /*remainingFilter=*/nullptr, + dataColumns, + /*indexColumns=*/std::vector{"id", "event_ts"}, + /*tableParameters=*/ + std::unordered_map{ + {"format", "parquet"}, + {"write.target-file-size-bytes", "134217728"}, + }, + /*filterColumnHandles=*/filterHandles, + /*sampleRate=*/0.1, + /*dbName=*/"warehouse", + /*dataColumnFieldIds=*/std::vector{10, 20}, + /*isChangelogQuery=*/true, + dataColumnHandles); + + auto clone = ISerializable::deserialize( + handle->serialize(), /*context=*/nullptr); + + // HiveTableHandle fields. + ASSERT_EQ(clone->connectorId(), handle->connectorId()); + ASSERT_EQ(clone->tableName(), handle->tableName()); + ASSERT_EQ(clone->dataColumns()->toString(), dataColumns->toString()); + ASSERT_EQ(clone->dataColumnFieldIds(), handle->dataColumnFieldIds()); + ASSERT_EQ(clone->indexColumns(), handle->indexColumns()); + ASSERT_EQ(clone->tableParameters(), handle->tableParameters()); + ASSERT_DOUBLE_EQ(clone->sampleRate(), handle->sampleRate()); + ASSERT_EQ(clone->dbName(), handle->dbName()); + + // filterColumnHandles: concrete type survives deserialization. + const auto& restoredFilters = clone->hiveFilterColumnHandles(); + ASSERT_EQ(restoredFilters.size(), 2); + const auto* fh0 = + dynamic_cast(restoredFilters[0].get()); + const auto* fh1 = + dynamic_cast(restoredFilters[1].get()); + ASSERT_NE(fh0, nullptr); + ASSERT_NE(fh1, nullptr); + ASSERT_EQ(fh0->name(), "partition_date"); + ASSERT_EQ(fh0->field().fieldId, 5); + ASSERT_EQ(fh1->name(), "region_id"); + ASSERT_EQ(fh1->field().fieldId, 6); + + // Iceberg-specific fields. + ASSERT_TRUE(clone->isChangelogQuery()); + ASSERT_EQ( + clone->dataColumnHandles().size(), handle->dataColumnHandles().size()); + + // "id": leaf field. + const auto& cloneId = clone->dataColumnHandles().at("id"); + ASSERT_EQ(cloneId->field().fieldId, 7); + ASSERT_EQ(*cloneId->dataType(), *BIGINT()); + + // "payload": nested field + initialDefaultValue. + const auto& clonePayload = clone->dataColumnHandles().at("payload"); + ASSERT_EQ(clonePayload->field().fieldId, nestedField.fieldId); + ASSERT_EQ(clonePayload->field().children.size(), 2); + ASSERT_EQ(clonePayload->field().children[0].fieldId, 11); + ASSERT_EQ(clonePayload->field().children[1].fieldId, 12); + ASSERT_EQ( + clonePayload->initialDefaultValue(), + std::optional{"active"}); +} + +// --------------------------------------------------------------------------- +// IcebergColumnHandle icebergMetadata_ SerDe +// --------------------------------------------------------------------------- + +// All V3 attributes and recursive children survive serialization. +TEST(IcebergTableHandleTest, serdeIcebergMetadata) { + registerAll(); + + IcebergFieldMetadata childMeta; + childMeta.required = false; + childMeta.longType = "LONG"; + + IcebergFieldMetadata meta; + meta.required = true; + meta.longType = "LONG"; + meta.timestampUnit = "MICROS"; + meta.binaryType = "UUID"; + meta.structType = "VariantStruct"; + meta.length = 16; + meta.children = {childMeta}; + + auto col = std::make_shared( + "ts", + FileColumnHandle::ColumnType::kRegular, + BIGINT(), + ParquetFieldId{99, {ParquetFieldId{100, {}}}}, + /*requiredSubfields=*/std::vector{}, + /*initialDefaultValue=*/std::nullopt, + meta); + + auto obj = col->serialize(); + ASSERT_TRUE(obj.count("icebergMetadata")); + + auto restored = ISerializable::deserialize(obj); + const auto& rm = restored->icebergMetadata(); + + ASSERT_EQ(rm.required, meta.required); + ASSERT_EQ(rm.longType, meta.longType); + ASSERT_EQ(rm.timestampUnit, meta.timestampUnit); + ASSERT_EQ(rm.binaryType, meta.binaryType); + ASSERT_EQ(rm.structType, meta.structType); + ASSERT_EQ(rm.length, meta.length); + ASSERT_EQ(rm.children.size(), meta.children.size()); + ASSERT_EQ(rm.children[0].required, meta.children[0].required); + ASSERT_EQ(rm.children[0].longType, meta.children[0].longType); +} + +// Empty icebergMetadata is not written to serialized form. +TEST(IcebergTableHandleTest, serdeIcebergMetadataOmittedWhenEmpty) { + registerAll(); + + auto col = makeIcebergCol("id", BIGINT(), /*fieldId=*/1); + auto obj = col->serialize(); + // No V3 metadata was set — the key must be absent. + ASSERT_EQ(obj.count("icebergMetadata"), 0); + + auto restored = ISerializable::deserialize(obj); + ASSERT_TRUE(restored->icebergMetadata().empty()); + ASSERT_TRUE(restored->icebergMetadata().children.empty()); +} + +// --------------------------------------------------------------------------- +// Negative tests — malformed / missing Iceberg metadata +// --------------------------------------------------------------------------- + +// Deserializing an IcebergTableHandle JSON that is missing the required +// "tableName" key must throw with a clear diagnostic. +TEST(IcebergTableHandleTest, deserializeMissingTableName) { + registerAll(); + + // Build a valid handle, serialize it, then surgically remove "tableName". + auto obj = makeMinimal()->serialize(); + obj.erase("tableName"); + + try { + ISerializable::deserialize(obj, /*context=*/nullptr); + FAIL() << "Expected std::out_of_range for missing 'tableName'"; + } catch (const std::out_of_range& e) { + EXPECT_NE(std::string(e.what()).find("tableName"), std::string::npos) + << "Exception message was: " << e.what(); + } +} + +// Deserializing an IcebergColumnHandle JSON that is missing the required +// "field" key (the ParquetFieldId) must throw with a clear diagnostic. +TEST(IcebergTableHandleTest, deserializeMissingFieldKey) { + registerAll(); + + auto col = makeIcebergCol("id", BIGINT(), /*fieldId=*/1); + auto obj = col->serialize(); + obj.erase("field"); + + try { + ISerializable::deserialize(obj); + FAIL() << "Expected std::out_of_range for missing 'field'"; + } catch (const std::out_of_range& e) { + EXPECT_NE(std::string(e.what()).find("field"), std::string::npos) + << "Exception message was: " << e.what(); + } +} + +// Deserializing a ParquetFieldId JSON that is missing the required +// "fieldId" key inside the "field" object must throw with a clear diagnostic. +TEST(IcebergTableHandleTest, deserializeMissingFieldId) { + registerAll(); + + auto col = makeIcebergCol("id", BIGINT(), /*fieldId=*/42); + auto obj = col->serialize(); + // Remove "fieldId" from the nested "field" object. + obj["field"].erase("fieldId"); + + try { + ISerializable::deserialize(obj); + FAIL() << "Expected std::out_of_range for missing 'fieldId'"; + } catch (const std::out_of_range& e) { + EXPECT_NE(std::string(e.what()).find("fieldId"), std::string::npos) + << "Exception message was: " << e.what(); + } +} diff --git a/velox/connectors/hive/iceberg/tests/IcebergTestBase.cpp b/velox/connectors/hive/iceberg/tests/IcebergTestBase.cpp index 897c47f4d39..407205d4564 100644 --- a/velox/connectors/hive/iceberg/tests/IcebergTestBase.cpp +++ b/velox/connectors/hive/iceberg/tests/IcebergTestBase.cpp @@ -333,7 +333,9 @@ std::vector> IcebergTestBase::makeIcebergSplits( partitionKeys, uint32_t splitCount, const std::unordered_map& infoColumns, - int64_t dataSequenceNumber) { + int64_t dataSequenceNumber, + const std::unordered_map>& + identityPartitionKeys) { VELOX_CHECK_GT(splitCount, 0); std::vector> splits; const auto fileSize = getFileSize(dataFilePath); @@ -350,6 +352,7 @@ std::vector> IcebergTestBase::makeIcebergSplits( .deleteFiles(deleteFiles) .infoColumns(infoColumns) .dataSequenceNumber(dataSequenceNumber) + .identityPartitionKeys(identityPartitionKeys) .build()); } @@ -368,6 +371,177 @@ IcebergTestBase::makeIcebergSplitWithInfoColumns( return splits.front(); } +std::shared_ptr IcebergTestBase::writeDataFile( + const std::vector& data) { + auto file = common::testutil::TempFilePath::create(); + writeToFile(file->getPath(), data); + return file; +} + +std::shared_ptr +IcebergTestBase::writeDwrfFileWithFieldIds( + const std::vector& data, + const std::vector& icebergFieldIds) { + VELOX_CHECK(!data.empty()); + const uint32_t numCols = data[0]->type()->size(); + VELOX_CHECK_EQ(icebergFieldIds.size(), numCols); + + // Build schemaAttributes: DWRF pre-order node 0 is the root struct (no + // iceberg.id); nodes 1..numCols are the top-level columns. + std::unordered_map>> + attrs; + for (uint32_t i = 0; i < numCols; ++i) { + attrs[i + 1] = {{"iceberg.id", std::to_string(icebergFieldIds[i])}}; + } + + auto file = common::testutil::TempFilePath::create(); + auto fs = filesystems::getFileSystem(file->getPath(), {}); + auto writeFile = fs->openFileForWrite( + file->getPath(), + {.shouldCreateParentDirectories = true, + .shouldThrowOnFileAlreadyExists = false}); + auto sink = std::make_unique( + std::move(writeFile), file->getPath()); + dwio::common::WriterOptions writerOptions; + auto dwrfOptions = std::make_shared(); + dwrfOptions->schemaAttributes = std::move(attrs); + writerOptions.formatSpecificOptions = dwrfOptions; + writerOptions.schema = data[0]->type(); + auto childPool = + rootPool_->addAggregateChild("writeDwrfFileWithFieldIds.writer"); + writerOptions.memoryPool = childPool.get(); + dwrf::Writer writer{std::move(sink), writerOptions}; + for (const auto& batch : data) { + writer.write(batch); + } + writer.close(); + return file; +} + +#ifdef VELOX_ENABLE_PARQUET +std::shared_ptr +IcebergTestBase::writeParquetFile( + const std::vector& data, + const std::vector& icebergFieldIds) { + VELOX_CHECK(!data.empty()); + auto file = common::testutil::TempFilePath::create(); + auto writeFile = + std::make_unique(file->getPath(), true, false); + auto sink = std::make_unique( + std::move(writeFile), file->getPath()); + dwio::common::WriterOptions writerOptions; + writerOptions.memoryPool = rootPool_.get(); + parquet::ParquetWriterOptions parquetOptions; + if (!icebergFieldIds.empty()) { + VELOX_CHECK_EQ(icebergFieldIds.size(), data[0]->type()->size()); + parquetOptions.parquetFieldIds.reserve(icebergFieldIds.size()); + for (int32_t id : icebergFieldIds) { + parquetOptions.parquetFieldIds.push_back(parquet::ParquetFieldId{id, {}}); + } + } + writerOptions.formatSpecificOptions = + std::make_shared( + std::move(parquetOptions)); + auto writer = std::make_unique( + std::move(sink), writerOptions, asRowType(data[0]->type())); + for (const auto& batch : data) { + writer->write(batch); + } + writer->close(); + return file; +} +#endif // VELOX_ENABLE_PARQUET + +core::PlanNodePtr IcebergTestBase::makeIcebergTableScanPlan( + const RowTypePtr& outputType, + const RowTypePtr& dataColumns, + const std::vector& dataColumnFieldIds, + const std::vector& subfieldFilters, + const std::string& remainingFilter) { + VELOX_CHECK_NOT_NULL(dataColumns); + + // Build IcebergColumnHandle assignments for each output-projected column. + // The Iceberg field ID is taken from dataColumnFieldIds when available, + // otherwise it defaults to the 1-based ordinal position in dataColumns. + connector::ColumnHandleMap assignments; + assignments.reserve(outputType->size()); + for (uint32_t i = 0; i < outputType->size(); ++i) { + const auto& name = outputType->nameOf(i); + const auto& type = outputType->childAt(i); + auto tableIdx = dataColumns->getChildIdxIfExists(name); + VELOX_CHECK( + tableIdx.has_value(), + "Output column '{}' not found in dataColumns.", + name); + const int32_t fieldId = !dataColumnFieldIds.empty() + ? dataColumnFieldIds[*tableIdx] + : static_cast(*tableIdx + 1); + assignments.emplace( + name, + std::make_shared( + name, + FileColumnHandle::ColumnType::kRegular, + type, + parquet::ParquetFieldId{fieldId, {}})); + } + + // Build filter-only IcebergColumnHandles for columns referenced by pushed- + // down filters but absent from the output projection. These are needed so + // buildIcebergHandleByName() can resolve their Iceberg field IDs and + // configureEqualityDeleteColumns() can promote them to projected columns + // when they also serve as equality-delete keys. + std::vector filterHandles; + if (!subfieldFilters.empty() || !remainingFilter.empty()) { + for (uint32_t i = 0; i < dataColumns->size(); ++i) { + const auto& name = dataColumns->nameOf(i); + if (assignments.count(name)) { + continue; // Already in the output projection. + } + // Include this column as a filter handle if any subfield filter names it. + bool usedInFilter = std::any_of( + subfieldFilters.begin(), + subfieldFilters.end(), + [&name](const std::string& f) { + return f.find(name) != std::string::npos; + }); + // Also include it if it appears in the remainingFilter expression. + if (!usedInFilter && !remainingFilter.empty()) { + usedInFilter = remainingFilter.find(name) != std::string::npos; + } + if (!usedInFilter) { + continue; + } + const auto& type = dataColumns->childAt(i); + const int32_t fieldId = !dataColumnFieldIds.empty() + ? dataColumnFieldIds[i] + : static_cast(i + 1); + filterHandles.push_back( + std::make_shared( + name, + FileColumnHandle::ColumnType::kRegular, + type, + parquet::ParquetFieldId{fieldId, {}})); + } + } + + return exec::test::PlanBuilder() + .startTableScan(kIcebergConnectorId) + .outputType(outputType) + .dataColumns(dataColumns) + .subfieldFilters(subfieldFilters) + .remainingFilter(remainingFilter) + .dataColumnFieldIds(dataColumnFieldIds) + .filterColumnHandles(std::move(filterHandles)) + .assignments(assignments) + .endTableScan() + .planNode(); +} + +core::PlanNodePtr IcebergTestBase::makeIcebergTableScanPlan( + const RowTypePtr& rowType) { + return makeIcebergTableScanPlan(rowType, rowType); +} + ColumnHandleMap IcebergTestBase::makeColumnHandles( const RowTypePtr& rowType, const std::unordered_set& partitionIndices) { diff --git a/velox/connectors/hive/iceberg/tests/IcebergTestBase.h b/velox/connectors/hive/iceberg/tests/IcebergTestBase.h index f4ff9a8ac9a..cfe047ee326 100644 --- a/velox/connectors/hive/iceberg/tests/IcebergTestBase.h +++ b/velox/connectors/hive/iceberg/tests/IcebergTestBase.h @@ -25,14 +25,21 @@ #include #include "velox/common/testutil/TempDirectoryPath.h" +#include "velox/connectors/hive/iceberg/IcebergColumnHandle.h" #include "velox/connectors/hive/iceberg/IcebergConfig.h" #include "velox/connectors/hive/iceberg/IcebergDataSink.h" #include "velox/connectors/hive/iceberg/IcebergDeleteFile.h" +#include "velox/connectors/hive/iceberg/IcebergSplit.h" +#include "velox/dwio/common/FileSink.h" +#include "velox/dwio/dwrf/writer/Writer.h" #include "velox/exec/tests/utils/HiveConnectorTestBase.h" +#include "velox/exec/tests/utils/PlanBuilder.h" #include "velox/vector/fuzzer/VectorFuzzer.h" #ifdef VELOX_ENABLE_PARQUET +#include "velox/common/file/LocalFile.h" #include "velox/dwio/parquet/RegisterParquetWriter.h" #include "velox/dwio/parquet/reader/ParquetReader.h" +#include "velox/dwio/parquet/writer/Writer.h" #endif namespace facebook::velox::connector::hive::iceberg::test { @@ -86,7 +93,9 @@ class IcebergTestBase : public exec::test::HiveConnectorTestBase { partitionKeys = {}, uint32_t splitCount = 1, const std::unordered_map& infoColumns = {}, - int64_t dataSequenceNumber = 0); + int64_t dataSequenceNumber = 0, + const std::unordered_map>& + identityPartitionKeys = {}); /// Creates one Iceberg connector split for a full data file with info /// columns. @@ -96,6 +105,41 @@ class IcebergTestBase : public exec::test::HiveConnectorTestBase { const std::vector& deleteFiles = {}, int64_t dataSequenceNumber = 0); + /// Writes a DWRF data file with no iceberg.id footer attributes. + /// The DWRF reader falls back to positional name mapping for these files. + std::shared_ptr writeDataFile( + const std::vector& data); + + /// Writes a DWRF file stamping "iceberg.id" footer attributes on each + /// top-level column. 'icebergFieldIds[i]' is the Iceberg field ID for the + /// i-th column; DWRF pre-order node IDs: 0=root, 1=first column, etc. + std::shared_ptr writeDwrfFileWithFieldIds( + const std::vector& data, + const std::vector& icebergFieldIds); + +#ifdef VELOX_ENABLE_PARQUET + /// Writes a Parquet file. 'icebergFieldIds[i]' is stamped as the Parquet + /// field ID for column i so the reader resolves columns by field ID under + /// kParquetFieldId mode. Pass an empty vector to omit field IDs. + std::shared_ptr writeParquetFile( + const std::vector& data, + const std::vector& icebergFieldIds = {}); +#endif + + /// Builds an Iceberg table scan plan. + /// Field IDs are derived from each output column's 1-based position in + /// 'dataColumns' (the full table schema), which is the authoritative source + /// for Iceberg field IDs regardless of file format or projection. + core::PlanNodePtr makeIcebergTableScanPlan( + const RowTypePtr& outputType, + const RowTypePtr& dataColumns, + const std::vector& dataColumnFieldIds = {}, + const std::vector& subfieldFilters = {}, + const std::string& remainingFilter = ""); + + /// Convenience overload: outputType == dataColumns (full-projection scan). + core::PlanNodePtr makeIcebergTableScanPlan(const RowTypePtr& rowType); + /// Creates Hive column handles for all columns in 'rowType', marking /// specified columns as partition keys. ColumnHandleMap makeColumnHandles( diff --git a/velox/connectors/hive/iceberg/tests/WriterOptionsAdapterTest.cpp b/velox/connectors/hive/iceberg/tests/WriterOptionsAdapterTest.cpp index feffabda89b..07c042cc42c 100644 --- a/velox/connectors/hive/iceberg/tests/WriterOptionsAdapterTest.cpp +++ b/velox/connectors/hive/iceberg/tests/WriterOptionsAdapterTest.cpp @@ -18,7 +18,7 @@ #include #ifdef VELOX_ENABLE_NIMBLE -#include "dwio/nimble/writer/fb/NimbleWriter.h" +#include "velox/dwio/nimble/writer/WriterFactory.h" #endif #include "velox/common/base/tests/GTestUtils.h" #include "velox/connectors/hive/iceberg/IcebergColumnHandle.h" diff --git a/velox/connectors/hive/paimon/PaimonSplitReader.cpp b/velox/connectors/hive/paimon/PaimonSplitReader.cpp index 446e47897dc..00ca6693ab2 100644 --- a/velox/connectors/hive/paimon/PaimonSplitReader.cpp +++ b/velox/connectors/hive/paimon/PaimonSplitReader.cpp @@ -76,7 +76,7 @@ PaimonSplitReader::PaimonSplitReader( void PaimonSplitReader::prepareSplit( std::shared_ptr metadataFilter, - dwio::common::RuntimeStatistics& runtimeStats, + dwio::common::RuntimeStats& runtimeStats, const folly::F14FastMap& /*fileReadOps*/) { // Save for re-use when advancing to subsequent files. metadataFilter_ = std::move(metadataFilter); diff --git a/velox/connectors/hive/paimon/PaimonSplitReader.h b/velox/connectors/hive/paimon/PaimonSplitReader.h index c21e660ee2b..966f78f437e 100644 --- a/velox/connectors/hive/paimon/PaimonSplitReader.h +++ b/velox/connectors/hive/paimon/PaimonSplitReader.h @@ -65,7 +65,7 @@ class PaimonSplitReader : public FileSplitReader { /// ensureFileSplitReader(). File validation is done in the constructor. void prepareSplit( std::shared_ptr metadataFilter, - dwio::common::RuntimeStatistics& runtimeStats, + dwio::common::RuntimeStats& runtimeStats, const folly::F14FastMap& fileReadOps = {}) override; @@ -99,7 +99,7 @@ class PaimonSplitReader : public FileSplitReader { // Saved from prepareSplit() for re-use when advancing to subsequent files. std::shared_ptr metadataFilter_; - dwio::common::RuntimeStatistics* runtimeStats_{nullptr}; + dwio::common::RuntimeStats* runtimeStats_{nullptr}; }; } // namespace facebook::velox::connector::hive::paimon diff --git a/velox/connectors/hive/storage_adapters/abfs/tests/AbfsFileSystemTest.cpp b/velox/connectors/hive/storage_adapters/abfs/tests/AbfsFileSystemTest.cpp index 35b1d9fc780..985da9b70e8 100644 --- a/velox/connectors/hive/storage_adapters/abfs/tests/AbfsFileSystemTest.cpp +++ b/velox/connectors/hive/storage_adapters/abfs/tests/AbfsFileSystemTest.cpp @@ -167,7 +167,13 @@ class AbfsFileSystemTest : public testing::Test { } void TearDown() override { - azuriteServer_->stop(); + // azuriteServer_ is left null if SetUp() threw before it could be + // constructed (e.g. the azurite-blob executable wasn't found). + // TearDown() runs unconditionally after SetUp(), even on failure, so + // it must not assume construction succeeded. + if (azuriteServer_ != nullptr) { + azuriteServer_->stop(); + } } static std::string generateRandomData(int size) { diff --git a/velox/connectors/hive/storage_adapters/hdfs/HdfsFileSystem.cpp b/velox/connectors/hive/storage_adapters/hdfs/HdfsFileSystem.cpp index c8b408f1a8f..52269b85ba7 100644 --- a/velox/connectors/hive/storage_adapters/hdfs/HdfsFileSystem.cpp +++ b/velox/connectors/hive/storage_adapters/hdfs/HdfsFileSystem.cpp @@ -26,10 +26,20 @@ std::string_view HdfsFileSystem::kViewfsScheme("viewfs://"); class HdfsFileSystem::Impl { public: - // Keep config here for possible use in the future. explicit Impl( const config::ConfigBase* config, const HdfsServiceEndpoint& endpoint) { + // Read-retry policy. Defaults preserve the original fail-fast behavior: + // maxReadAttempts == 1 means no retries. Retries are opt-in because the + // JNI-backed libhdfs.so already retries and fails over internally, whereas + // libhdfs3 does not. config may be null (getFileSystem can be called with a + // null config), in which case the member defaults are kept. + if (config != nullptr) { + maxReadAttempts_ = config->get("hive.hdfs.read-max-attempts", 1); + retryBaseDelayMs_ = + config->get("hive.hdfs.read-retry-delay-ms", 100); + } + auto status = filesystems::arrow::io::internal::ConnectLibHdfs(&driver_); VELOX_CHECK( status.ok(), "Failed to connect to libhdfs: {}", status.ToString()); @@ -89,9 +99,19 @@ class HdfsFileSystem::Impl { return driver_; } + int maxReadAttempts() const { + return maxReadAttempts_; + } + + int retryBaseDelayMs() const { + return retryBaseDelayMs_; + } + private: hdfsFS hdfsClient_{nullptr}; filesystems::arrow::io::internal::LibHdfsShim* driver_{nullptr}; + int maxReadAttempts_{1}; + int retryBaseDelayMs_{100}; bool closed_{false}; }; @@ -117,7 +137,11 @@ std::unique_ptr HdfsFileSystem::openFileForRead( } } return std::make_unique( - impl_->hdfsShim(), impl_->hdfsClient(), path); + impl_->hdfsShim(), + impl_->hdfsClient(), + path, + impl_->maxReadAttempts(), + impl_->retryBaseDelayMs()); } std::unique_ptr HdfsFileSystem::openFileForWrite( diff --git a/velox/connectors/hive/storage_adapters/hdfs/HdfsReadFile.cpp b/velox/connectors/hive/storage_adapters/hdfs/HdfsReadFile.cpp index 94adb31a742..e532841d42a 100644 --- a/velox/connectors/hive/storage_adapters/hdfs/HdfsReadFile.cpp +++ b/velox/connectors/hive/storage_adapters/hdfs/HdfsReadFile.cpp @@ -15,9 +15,23 @@ */ #include "HdfsReadFile.h" + +#include +#include +#include +#include + #include "velox/external/hdfs/ArrowHdfsInternal.h" namespace facebook::velox { +namespace { +// Upper bound on the exponential backoff between read retries. Without a cap +// the delay doubles unbounded as the attempt count grows (and the shift would +// eventually overflow), so a large maxReadAttempts_ could stall a read for +// hours. 30s is long enough to ride out a transient DataNode blip while keeping +// the worst-case wait bounded. +constexpr int64_t kMaxRetryDelayMs = 30'000; +} // namespace struct HdfsFile { filesystems::arrow::io::internal::LibHdfsShim* driver_; @@ -31,6 +45,12 @@ struct HdfsFile { } } + // Owns a raw libhdfs handle, so it is not copyable or movable. + HdfsFile(const HdfsFile&) = delete; + HdfsFile& operator=(const HdfsFile&) = delete; + HdfsFile(HdfsFile&&) = delete; + HdfsFile& operator=(HdfsFile&&) = delete; + void open( filesystems::arrow::io::internal::LibHdfsShim* driver, hdfsFS client, @@ -53,10 +73,35 @@ struct HdfsFile { driver_->GetLastExceptionRootCause()); } + // Close the current handle (if any) and reopen the file. Used by + // preadInternal to recover a thread-local handle whose stream went bad after + // a transient read failure. file_ is a folly::ThreadLocal in the owning Impl, + // so each thread reopens its own handle; the shared HDFS client is untouched. + void reopen(const std::string& path) { + if (handle_) { + // Ignore the close result: the stream is already in a bad state and is + // being discarded regardless. + driver_->CloseFile(client_, handle_); + handle_ = nullptr; + } + handle_ = driver_->OpenFile(client_, path.data(), O_RDONLY, 0, 0, 0); + VELOX_CHECK_NOT_NULL( + handle_, + "Unable to reopen file {}. got error: {}", + path, + driver_->GetLastExceptionRootCause()); + } + + // Returns the raw libhdfs3 result, including non-positive values on failure. + // The caller (preadInternal) decides whether a non-positive result is + // retriable. int32_t read(char* pos, uint64_t length) const { - auto bytesRead = driver_->Read(client_, handle_, pos, length); - VELOX_CHECK(bytesRead >= 0, "Read failure in HDFSReadFile::preadInternal."); - return bytesRead; + // hdfsRead takes a signed tSize; cap the request so the unsigned length + // never narrows into a negative value. preadInternal loops on a short read, + // so servicing a large request in tSize-sized chunks is fine. + const auto chunk = static_cast(std::min( + length, static_cast(std::numeric_limits::max()))); + return driver_->Read(client_, handle_, pos, chunk); } }; @@ -65,8 +110,22 @@ class HdfsReadFile::Impl { Impl( filesystems::arrow::io::internal::LibHdfsShim* driver, hdfsFS hdfs, - const std::string_view path) - : driver_(driver), hdfsClient_(hdfs), filePath_(path) { + const std::string_view path, + int maxReadAttempts, + int retryBaseDelayMs) + : driver_(driver), + hdfsClient_(hdfs), + filePath_(path), + maxReadAttempts_(maxReadAttempts), + retryBaseDelayMs_(retryBaseDelayMs) { + // maxReadAttempts_ counts the initial read plus any retries, so it must be + // at least 1. Reject non-positive values up front; otherwise the retry + // budget check below would produce a confusing "after 0 attempts" message. + VELOX_USER_CHECK_GE( + maxReadAttempts_, + 1, + "hive.hdfs.read-max-attempts must be at least 1, got {}", + maxReadAttempts_); fileInfo_ = driver_->GetPathInfo(hdfsClient_, filePath_.data()); if (fileInfo_ == nullptr) { auto error = fmt::format( @@ -96,8 +155,46 @@ class HdfsReadFile::Impl { } file_->seek(offset); uint64_t totalBytesRead = 0; + // attempt counts the read attempts for this pread. maxReadAttempts_ == 1 + // means fail-fast with no retries; the budget spans the whole pread. + int attempt = 1; while (totalBytesRead < length) { auto bytesRead = file_->read(pos, length - totalBytesRead); + // checkFileReadParameters guarantees offset + length stays within the + // file, so we never legitimately hit EOF here. A non-positive result is + // therefore always a failure to make progress: a negative value is an + // explicit libhdfs3 error, and a zero means the read stalled without + // advancing. Both are treated as transient and retried; leaving the zero + // case out would spin this loop forever. + if (bytesRead <= 0) { + VELOX_CHECK_LT( + attempt, + maxReadAttempts_, + "Read failure in HDFSReadFile::preadInternal after {} attempts, " + "file: {}, offset: {}, length: {}, root cause: {}", + maxReadAttempts_, + filePath_, + offset, + length, + driver_->GetLastExceptionRootCause()); + LOG(WARNING) << "Transient HDFS read failure on " << filePath_ + << " (offset=" << offset + totalBytesRead << ", attempt " + << attempt << "/" << maxReadAttempts_ << "), root cause: " + << driver_->GetLastExceptionRootCause(); + // Exponential backoff, capped at kMaxRetryDelayMs. The shift is clamped + // (and done in 64-bit) so a large attempt count can neither overflow + // nor produce an absurdly long sleep. + const int shift = std::min(attempt - 1, 20); + const int64_t delayMs = std::min( + int64_t{retryBaseDelayMs_} << shift, kMaxRetryDelayMs); + std::this_thread::sleep_for(std::chrono::milliseconds(delayMs)); + ++attempt; + // Rebuild the handle and reposition to the first unread byte; already + // read bytes are kept. + file_->reopen(filePath_); + file_->seek(offset + totalBytesRead); + continue; + } totalBytesRead += bytesRead; pos += bytesRead; } @@ -153,14 +250,24 @@ class HdfsReadFile::Impl { hdfsFS hdfsClient_; std::string filePath_; hdfsFileInfo* fileInfo_; + const int maxReadAttempts_; + const int retryBaseDelayMs_; folly::ThreadLocal file_; }; HdfsReadFile::HdfsReadFile( filesystems::arrow::io::internal::LibHdfsShim* driver, hdfsFS hdfs, - const std::string_view path) - : pImpl(std::make_unique(driver, hdfs, path)) {} + const std::string_view path, + int maxReadAttempts, + int retryBaseDelayMs) + : pImpl( + std::make_unique( + driver, + hdfs, + path, + maxReadAttempts, + retryBaseDelayMs)) {} HdfsReadFile::~HdfsReadFile() = default; diff --git a/velox/connectors/hive/storage_adapters/hdfs/HdfsReadFile.h b/velox/connectors/hive/storage_adapters/hdfs/HdfsReadFile.h index 6208702eab5..0c049c3d92b 100644 --- a/velox/connectors/hive/storage_adapters/hdfs/HdfsReadFile.h +++ b/velox/connectors/hive/storage_adapters/hdfs/HdfsReadFile.h @@ -20,7 +20,7 @@ namespace facebook::velox { namespace filesystems::arrow::io::internal { -class LibHdfsShim; +struct LibHdfsShim; } /** @@ -28,10 +28,22 @@ class LibHdfsShim; */ class HdfsReadFile final : public ReadFile { public: + /// @param maxReadAttempts Total number of read attempts for a transient + /// failure, including the initial one. 1 (the default) means fail-fast with + /// no retries, preserving the original behavior; set >1 to enable retries. + /// This is opt-in because backends differ: the JNI-backed libhdfs.so already + /// retries and fails over internally via DFSInputStream, whereas libhdfs3 + /// does not. + /// @param retryBaseDelayMs Base delay for the exponential backoff between + /// retries: the Nth retry waits retryBaseDelayMs * 2^(N-1) milliseconds, + /// capped at an internal maximum so a large maxReadAttempts cannot stall a + /// read for an unbounded amount of time. explicit HdfsReadFile( filesystems::arrow::io::internal::LibHdfsShim* driver, hdfsFS hdfs, - std::string_view path); + std::string_view path, + int maxReadAttempts = 1, + int retryBaseDelayMs = 100); ~HdfsReadFile() override; std::string_view pread( diff --git a/velox/connectors/hive/storage_adapters/hdfs/HdfsWriteFile.h b/velox/connectors/hive/storage_adapters/hdfs/HdfsWriteFile.h index fb311b1a6c3..de055fef4d5 100644 --- a/velox/connectors/hive/storage_adapters/hdfs/HdfsWriteFile.h +++ b/velox/connectors/hive/storage_adapters/hdfs/HdfsWriteFile.h @@ -21,7 +21,7 @@ namespace facebook::velox { namespace filesystems::arrow::io::internal { -class LibHdfsShim; +struct LibHdfsShim; } /// Implementation of hdfs write file. Nothing written to the file should be diff --git a/velox/connectors/hive/storage_adapters/hdfs/tests/CMakeLists.txt b/velox/connectors/hive/storage_adapters/hdfs/tests/CMakeLists.txt index 7a6b0832765..c7a554f2e2f 100644 --- a/velox/connectors/hive/storage_adapters/hdfs/tests/CMakeLists.txt +++ b/velox/connectors/hive/storage_adapters/hdfs/tests/CMakeLists.txt @@ -32,6 +32,25 @@ target_link_libraries( target_compile_options(velox_hdfs_file_test PRIVATE -Wno-deprecated-declarations) +# Unit test for HdfsReadFile's retry path. Drives a stub LibHdfsShim, so it needs +# no Hadoop MiniCluster and is safe to run in parallel with the cluster-based +# tests above. +add_executable(velox_hdfs_read_file_test HdfsReadFileTest.cpp) + +add_test(velox_hdfs_read_file_test velox_hdfs_read_file_test) +target_link_libraries( + velox_hdfs_read_file_test + velox_file + velox_hdfs + velox_core + velox_exec + GTest::gtest + GTest::gtest_main + GTest::gmock +) + +target_compile_options(velox_hdfs_read_file_test PRIVATE -Wno-deprecated-declarations) + add_executable(velox_hdfs_insert_test HdfsInsertTest.cpp HdfsMiniCluster.cpp HdfsUtilTest.cpp) velox_add_test_headers(velox_hdfs_insert_test HdfsMiniCluster.h) diff --git a/velox/connectors/hive/storage_adapters/hdfs/tests/HdfsReadFileTest.cpp b/velox/connectors/hive/storage_adapters/hdfs/tests/HdfsReadFileTest.cpp new file mode 100644 index 00000000000..2f5a5a3aac0 --- /dev/null +++ b/velox/connectors/hive/storage_adapters/hdfs/tests/HdfsReadFileTest.cpp @@ -0,0 +1,201 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/connectors/hive/storage_adapters/hdfs/HdfsReadFile.h" + +#include +#include + +#include "gtest/gtest.h" +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/external/hdfs/ArrowHdfsInternal.h" + +// Deterministic coverage for the transient-read-failure retry path in +// HdfsReadFile::Impl::preadInternal, without a live HDFS cluster. +// +// LibHdfsShim dispatches every libhdfs3 call through a plain function pointer +// (e.g. this->hdfsRead(...)), and HdfsReadFile takes the shim plus an opaque +// hdfsFS by pointer. So we can hand it a shim whose pointers target the stubs +// below and drive Read() to fail a fixed number of times before succeeding -- +// something a real MiniCluster (which can only be up or down) cannot do +// reproducibly. +// +// Retries are opt-in via the maxReadAttempts constructor argument. The tests +// that exercise the retry loop pass a small non-default value; a fast base +// delay (kTestRetryDelayMs) keeps the backoff sleeps negligible. + +namespace facebook::velox { +namespace { + +using filesystems::arrow::io::internal::LibHdfsShim; + +// hdfsFS / hdfsFile are opaque pointers; the stubs never dereference them, so +// any non-null value works. +hdfsFS kFakeFs = reinterpret_cast(0x1); +hdfsFile kFakeHandle = reinterpret_cast(0x2); + +// Small backoff base so the retry tests don't actually sleep for long. +constexpr int kTestRetryDelayMs = 1; + +// libhdfs3 uses C function pointers, which cannot capture state, so the stub +// behaviour is driven through translation-unit-local variables. The fixture +// resets all of them before every test. +tSize gFileSize; +int gReadCalls; // number of times stubRead was entered +int gFailCount; // first gFailCount reads fail (return gFailReturn) +tSize gFailReturn; // -1 (libhdfs3 error) or 0 (no progress) +tSize gMaxChunk; // >0 caps the bytes returned by a successful read (short read) +hdfsFileInfo gFileInfo; + +hdfsFileInfo* stubGetPathInfo(hdfsFS, const char*) { + gFileInfo = {}; + gFileInfo.mSize = gFileSize; + gFileInfo.mBlockSize = gFileSize; + return &gFileInfo; +} + +// gFileInfo is a static, so there is nothing to free. +void stubFreeFileInfo(hdfsFileInfo*, int) {} + +hdfsFile stubOpenFile(hdfsFS, const char*, int, int, short, tSize) { // NOLINT + return kFakeHandle; +} + +int stubCloseFile(hdfsFS, hdfsFile) { + return 0; +} + +int stubSeek(hdfsFS, hdfsFile, tOffset) { + return 0; +} + +tSize stubRead(hdfsFS, hdfsFile, void* buffer, tSize length) { + if (gReadCalls++ < gFailCount) { + return gFailReturn; + } + const tSize n = (gMaxChunk > 0 && gMaxChunk < length) ? gMaxChunk : length; + std::memset(buffer, 'x', n); + return n; +} + +char* stubGetLastExceptionRootCause() { + return strdup("mock transient failure"); +} + +class HdfsReadFileRetryTest : public testing::Test { + protected: + void SetUp() override { + gFileSize = 1024; + gReadCalls = 0; + gFailCount = 0; + gFailReturn = -1; + gMaxChunk = 0; + + shim_.Initialize(); + shim_.hdfsGetPathInfo = stubGetPathInfo; + shim_.hdfsFreeFileInfo = stubFreeFileInfo; + shim_.hdfsOpenFile = stubOpenFile; + shim_.hdfsCloseFile = stubCloseFile; + shim_.hdfsSeek = stubSeek; + shim_.hdfsRead = stubRead; + shim_.hdfsGetLastExceptionRootCause = stubGetLastExceptionRootCause; + } + + LibHdfsShim shim_; +}; + +// With retries disabled (maxReadAttempts == 1, the default), the very first +// transient failure throws immediately and Read() is attempted exactly once. +// This pins "default == original fail-fast behavior" as a regression guard: the +// default behavior is a no-op unless a caller explicitly opts in. +TEST_F(HdfsReadFileRetryTest, failFastByDefault) { + gFailCount = 1; + gFailReturn = -1; + + HdfsReadFile readFile(&shim_, kFakeFs, "/mock"); + VELOX_ASSERT_THROW(readFile.pread(0, gFileSize), "after 1 attempts"); + EXPECT_EQ(gReadCalls, 1); +} + +// A non-positive maxReadAttempts (e.g. a misconfigured +// hive.hdfs.read-max-attempts) is rejected at construction with a user error, +// rather than silently producing a confusing "after 0 attempts" message on the +// first read failure. +TEST_F(HdfsReadFileRetryTest, rejectsNonPositiveMaxAttempts) { + VELOX_ASSERT_THROW( + HdfsReadFile(&shim_, kFakeFs, "/mock", /*maxReadAttempts=*/0), + "hive.hdfs.read-max-attempts must be at least 1"); +} + +// Two transient errors, then success: the read recovers and returns the full +// payload. Exactly three Read() calls (2 failed + 1 success) confirms the retry +// loop re-issued the read rather than giving up or double-counting. +TEST_F(HdfsReadFileRetryTest, recoversAfterTransientNegativeFailures) { + gFailCount = 2; + gFailReturn = -1; + + HdfsReadFile readFile( + &shim_, kFakeFs, "/mock", /*maxReadAttempts=*/4, kTestRetryDelayMs); + const auto data = readFile.pread(0, gFileSize); + + EXPECT_EQ(data, std::string(gFileSize, 'x')); + EXPECT_EQ(gReadCalls, 3); +} + +// A zero return means the read made no progress. It must be retried just like a +// negative return; treating it as success would spin the while loop forever. +// This pins the `bytesRead <= 0` predicate (not `< 0`). +TEST_F(HdfsReadFileRetryTest, retriesOnZeroReturn) { + gFailCount = 2; + gFailReturn = 0; + + HdfsReadFile readFile( + &shim_, kFakeFs, "/mock", /*maxReadAttempts=*/4, kTestRetryDelayMs); + const auto data = readFile.pread(0, gFileSize); + + EXPECT_EQ(data, std::string(gFileSize, 'x')); + EXPECT_EQ(gReadCalls, 3); +} + +// A persistent failure throws after the retry budget is exhausted. With +// maxReadAttempts == 4, Read() is attempted 4 times in total (1 initial + 3 +// retries) before the final throw, and the message reports it. +TEST_F(HdfsReadFileRetryTest, throwsAfterExhaustion) { + gFailCount = 1000; // always fails + gFailReturn = -1; + + HdfsReadFile readFile( + &shim_, kFakeFs, "/mock", /*maxReadAttempts=*/4, kTestRetryDelayMs); + VELOX_ASSERT_THROW(readFile.pread(0, gFileSize), "after 4 attempts"); + EXPECT_EQ(gReadCalls, 4); +} + +// A successful but short read is not a failure: preadInternal must keep the +// bytes and loop until the request is filled, without consuming any retries. +// maxReadAttempts == 1 proves short reads never touch the retry budget. +TEST_F(HdfsReadFileRetryTest, shortReadsAccumulateWithoutRetry) { + gFailCount = 0; + gMaxChunk = 256; // 1024 / 256 = 4 short reads + + HdfsReadFile readFile(&shim_, kFakeFs, "/mock"); + const auto data = readFile.pread(0, gFileSize); + + EXPECT_EQ(data, std::string(gFileSize, 'x')); + EXPECT_EQ(gReadCalls, 4); +} + +} // namespace +} // namespace facebook::velox diff --git a/velox/connectors/hive/storage_adapters/s3fs/RegisterS3FileSystem.cpp b/velox/connectors/hive/storage_adapters/s3fs/RegisterS3FileSystem.cpp index d04569dbb4c..1654e12267b 100644 --- a/velox/connectors/hive/storage_adapters/s3fs/RegisterS3FileSystem.cpp +++ b/velox/connectors/hive/storage_adapters/s3fs/RegisterS3FileSystem.cpp @@ -32,14 +32,14 @@ using FileSystemMap = folly::Synchronized< std::unordered_map>>; /// Multiple S3 filesystems are supported. -/// Key is the endpoint value specified in the config using hive.s3.endpoint. +/// Key is the endpoint value specified in the config using s3.endpoint. /// If the endpoint is empty, it will default to AWS S3 Library. /// Different S3 buckets can be accessed with different client configurations. /// This allows for different endpoints, data read and write strategies. -/// The bucket specific option is set by replacing the hive.s3. prefix on an -/// option with hive.s3.bucket.BUCKETNAME., where BUCKETNAME is the name of the +/// The bucket specific option is set by replacing the s3. prefix on an +/// option with s3.bucket.BUCKETNAME., where BUCKETNAME is the name of the /// bucket. When connecting to a bucket, all options explicitly set will -/// override the base hive.s3. values. +/// override the base s3. values. FileSystemMap& fileSystems() { static FileSystemMap instances; @@ -82,10 +82,10 @@ std::shared_ptr fileSystemGenerator( } auto logLevel = - properties->get(S3Config::kS3LogLevel, std::string("FATAL")); - std::optional logLocation = - static_cast>( - properties->get(S3Config::kS3LogLocation)); + S3Config::configValue(*properties, S3Config::kS3LogLevel) + .value_or("FATAL"); + auto logLocation = + S3Config::configValue(*properties, S3Config::kS3LogLocation); initializeS3(logLevel, logLocation); std::shared_ptr fs; if (fileSystemFactory) { diff --git a/velox/connectors/hive/storage_adapters/s3fs/S3Config.cpp b/velox/connectors/hive/storage_adapters/s3fs/S3Config.cpp index 906bf0f9b17..3e0c1bb72af 100644 --- a/velox/connectors/hive/storage_adapters/s3fs/S3Config.cpp +++ b/velox/connectors/hive/storage_adapters/s3fs/S3Config.cpp @@ -24,18 +24,33 @@ namespace facebook::velox::filesystems { static constexpr size_t kMinimumMultipartMinPartSize = 5U << 20; // 5MB static constexpr size_t kMaximumMultipartMinPartSize = 5U << 30; // 5GB +std::optional S3Config::configValue( + const config::ConfigBase& config, + std::string_view configKey) { + if (auto value = config.get(std::string(configKey))) { + return value; + } + // Fall back to the deprecated "hive.s3." prefix. + VELOX_CHECK( + configKey.substr(0, std::string_view(kS3Prefix).size()) == kS3Prefix, + "S3 config key must be prefixed with '{}': {}", + kS3Prefix, + configKey); + const auto suffix = configKey.substr(std::string_view(kS3Prefix).size()); + return config.get( + fmt::format("{}{}", kS3DeprecatedPrefix, suffix)); +} + std::string S3Config::cacheKey( std::string_view bucket, std::shared_ptr config) { - auto bucketEndpoint = bucketConfigKey(Keys::kEndpoint, bucket); - if (config->valueExists(bucketEndpoint)) { - return fmt::format( - "{}-{}", config->get(bucketEndpoint).value(), bucket); + if (auto bucketEndpoint = + configValue(*config, bucketConfigKey(Keys::kEndpoint, bucket))) { + return fmt::format("{}-{}", bucketEndpoint.value(), bucket); } - auto baseEndpoint = baseConfigKey(Keys::kEndpoint); - if (config->valueExists(baseEndpoint)) { - return fmt::format( - "{}-{}", config->get(baseEndpoint).value(), bucket); + if (auto baseEndpoint = + configValue(*config, baseConfigKey(Keys::kEndpoint))) { + return fmt::format("{}-{}", baseEndpoint.value(), bucket); } return std::string(bucket); } @@ -49,32 +64,22 @@ S3Config::S3Config( key++) { auto s3Key = static_cast(key); auto value = S3Config::configTraits().find(s3Key)->second; - auto configSuffix = value.first; auto configDefault = value.second; - // Set bucket S3 config "hive.s3.bucket.*" if present. - std::stringstream bucketConfig; - bucketConfig << kS3BucketPrefix << bucket << "." << configSuffix; - auto configVal = static_cast>( - properties->get(bucketConfig.str())); - if (configVal.has_value()) { + // Prefer the bucket-specific "s3.bucket.*" config, then the base "s3.*" + // config, then the default. Each lookup falls back to the deprecated + // "hive.s3." prefix when the canonical key is absent. + if (auto configVal = + configValue(*properties, bucketConfigKey(s3Key, bucket))) { config_[s3Key] = configVal.value(); + } else if (auto baseVal = configValue(*properties, baseConfigKey(s3Key))) { + config_[s3Key] = baseVal.value(); } else { - // Set base config "hive.s3.*" if present. - std::stringstream baseConfig; - baseConfig << kS3Prefix << configSuffix; - configVal = static_cast>( - properties->get(baseConfig.str())); - if (configVal.has_value()) { - config_[s3Key] = configVal.value(); - } else { - // Set the default value. - config_[s3Key] = configDefault; - } + config_[s3Key] = configDefault; } } payloadSigningPolicy_ = - properties->get(kS3PayloadSigningPolicy, "Never"); + configValue(*properties, kS3PayloadSigningPolicy).value_or("Never"); VELOX_CHECK_GE( minPartSize(), diff --git a/velox/connectors/hive/storage_adapters/s3fs/S3Config.h b/velox/connectors/hive/storage_adapters/s3fs/S3Config.h index ab1252f242f..d26091db8d5 100644 --- a/velox/connectors/hive/storage_adapters/s3fs/S3Config.h +++ b/velox/connectors/hive/storage_adapters/s3fs/S3Config.h @@ -15,8 +15,12 @@ */ #pragma once +#include #include +#include #include +#include +#include #include "velox/common/base/Exceptions.h" namespace facebook::velox::config { @@ -26,36 +30,43 @@ class ConfigBase; namespace facebook::velox::filesystems { /// Build config required to initialize an S3FileSystem instance. -/// All hive.s3 options can be set on a per-bucket basis. -/// The bucket-specific option is set by replacing the hive.s3. prefix on an -/// option with hive.s3.bucket.BUCKETNAME., where BUCKETNAME is the name of the -/// bucket. +/// All s3 options can be set on a per-bucket basis. +/// The bucket-specific option is set by replacing the s3. prefix on an option +/// with s3.bucket.BUCKETNAME., where BUCKETNAME is the name of the bucket. /// When connecting to a bucket, all options explicitly set will override the -/// base hive.s3. values. +/// base s3. values. /// These semantics are similar to the Apache Hadoop-Aws module. /// https://hadoop.apache.org/docs/current/hadoop-aws/tools/hadoop-aws/index.html +/// +/// The canonical config prefix is "s3." and is usable from any connector. The +/// legacy "hive.s3." prefix is accepted as a deprecated fallback: a canonical +/// key takes precedence over its "hive.s3." counterpart when both are set. class S3Config { public: S3Config() = delete; /// S3 config prefix. - static constexpr const char* kS3Prefix = "hive.s3."; + static constexpr const char* kS3Prefix = "s3."; + + /// Deprecated S3 config prefix, accepted as a fallback for the canonical + /// kS3Prefix. Prefer kS3Prefix in new configurations. + static constexpr const char* kS3DeprecatedPrefix = "hive.s3."; /// S3 bucket config prefix - static constexpr const char* kS3BucketPrefix = "hive.s3.bucket."; + static constexpr const char* kS3BucketPrefix = "s3.bucket."; /// Log granularity of AWS C++ SDK. - static constexpr const char* kS3LogLevel = "hive.s3.log-level"; + static constexpr const char* kS3LogLevel = "s3.log-level"; /// Payload signing policy. static constexpr const char* kS3PayloadSigningPolicy = - "hive.s3.payload-signing-policy"; + "s3.payload-signing-policy"; /// S3FileSystem default identity. static constexpr const char* kDefaultS3Identity = "s3-default-identity"; /// Log location of AWS C++ SDK. - static constexpr const char* kS3LogLocation = "hive.s3.log-location"; + static constexpr const char* kS3LogLocation = "s3.log-location"; /// Keys to identify the config. enum class Keys { @@ -155,6 +166,14 @@ class S3Config { return buffer.str(); } + /// Returns the value for the canonical 'configKey', which must be prefixed + /// with kS3Prefix. When the canonical key is absent, falls back to the + /// deprecated kS3DeprecatedPrefix form. Returns std::nullopt when neither is + /// set. + static std::optional configValue( + const config::ConfigBase& config, + std::string_view configKey); + /// The S3 storage endpoint server. This can be used to connect to an /// S3-compatible storage system instead of AWS. std::optional endpoint() const { diff --git a/velox/connectors/hive/storage_adapters/s3fs/S3FileSystem.cpp b/velox/connectors/hive/storage_adapters/s3fs/S3FileSystem.cpp index a39f68cc1a6..7c386a3bac4 100644 --- a/velox/connectors/hive/storage_adapters/s3fs/S3FileSystem.cpp +++ b/velox/connectors/hive/storage_adapters/s3fs/S3FileSystem.cpp @@ -416,7 +416,7 @@ class S3FileSystem::Impl { VELOX_USER_CHECK_GE( maxAttempts.value(), 0, - "Invalid configuration: specified 'hive.s3.max-attempts' value {} is < 0.", + "Invalid configuration: specified S3 'max-attempts' value {} is < 0.", maxAttempts.value()); return std::make_shared( maxAttempts.value()); @@ -429,7 +429,7 @@ class S3FileSystem::Impl { VELOX_USER_CHECK_GE( maxAttempts.value(), 0, - "Invalid configuration: specified 'hive.s3.max-attempts' value {} is < 0.", + "Invalid configuration: specified S3 'max-attempts' value {} is < 0.", maxAttempts.value()); return std::make_shared( maxAttempts.value()); @@ -442,7 +442,7 @@ class S3FileSystem::Impl { VELOX_USER_CHECK_GE( maxAttempts.value(), 0, - "Invalid configuration: specified 'hive.s3.max-attempts' value {} is < 0.", + "Invalid configuration: specified S3 'max-attempts' value {} is < 0.", maxAttempts.value()); return std::make_shared( maxAttempts.value()); @@ -512,7 +512,7 @@ S3FileSystem::S3FileSystem( std::string_view bucketName, const std::shared_ptr config) : FileSystem(config) { - auto s3Config = std::make_shared(bucketName, config); + auto s3Config = std::make_shared(bucketName, config_); impl_ = std::make_shared(std::move(s3Config)); } diff --git a/velox/connectors/hive/storage_adapters/s3fs/S3Util.h b/velox/connectors/hive/storage_adapters/s3fs/S3Util.h index 9265d24387e..7aa277976ec 100644 --- a/velox/connectors/hive/storage_adapters/s3fs/S3Util.h +++ b/velox/connectors/hive/storage_adapters/s3fs/S3Util.h @@ -188,7 +188,7 @@ inline std::string getRequestID( error.ShouldRetry()) { \ errMsg.append( \ fmt::format( \ - " Request failed after retrying {} times. Try increasing the value of 'hive.s3.max-attempts'.", \ + " Request failed after retrying {} times. Try increasing the value of the S3 'max-attempts' configuration.", \ outcome.GetRetryCount())); \ } \ if (error.GetResponseCode() == Aws::Http::HttpResponseCode::NOT_FOUND) { \ diff --git a/velox/connectors/hive/storage_adapters/s3fs/tests/MinioServer.h b/velox/connectors/hive/storage_adapters/s3fs/tests/MinioServer.h index 164bb590518..27bc2c7bf68 100644 --- a/velox/connectors/hive/storage_adapters/s3fs/tests/MinioServer.h +++ b/velox/connectors/hive/storage_adapters/s3fs/tests/MinioServer.h @@ -56,15 +56,15 @@ class MinioServer { return tempPath_->getPath(); } - std::shared_ptr hiveConfig( + std::shared_ptr s3Config( const std::unordered_map configOverride = {}) const { std::unordered_map config({ - {"hive.s3.aws-access-key", accessKey_}, - {"hive.s3.aws-secret-key", secretKey_}, - {"hive.s3.endpoint", connectionString_}, - {"hive.s3.ssl.enabled", "false"}, - {"hive.s3.path-style-access", "true"}, + {"s3.aws-access-key", accessKey_}, + {"s3.aws-secret-key", secretKey_}, + {"s3.endpoint", connectionString_}, + {"s3.ssl.enabled", "false"}, + {"s3.path-style-access", "true"}, }); // Update the default config map with the supplied configOverride map diff --git a/velox/connectors/hive/storage_adapters/s3fs/tests/S3ConfigTest.cpp b/velox/connectors/hive/storage_adapters/s3fs/tests/S3ConfigTest.cpp index f0a1e6e377c..f03e881315a 100644 --- a/velox/connectors/hive/storage_adapters/s3fs/tests/S3ConfigTest.cpp +++ b/velox/connectors/hive/storage_adapters/s3fs/tests/S3ConfigTest.cpp @@ -22,6 +22,7 @@ namespace facebook::velox::filesystems { namespace { + TEST(S3ConfigTest, defaultConfig) { auto config = std::make_shared( std::unordered_map()); @@ -50,7 +51,7 @@ TEST(S3ConfigTest, overrideConfig) { {S3Config::baseConfigKey(S3Config::Keys::kSSLEnabled), "false"}, {S3Config::baseConfigKey(S3Config::Keys::kUseInstanceCredentials), "true"}, - {"hive.s3.payload-signing-policy", "RequestDependent"}, + {S3Config::kS3PayloadSigningPolicy, "RequestDependent"}, {S3Config::baseConfigKey(S3Config::Keys::kEndpoint), "endpoint"}, {S3Config::baseConfigKey(S3Config::Keys::kEndpointRegion), "region"}, {S3Config::baseConfigKey(S3Config::Keys::kAccessKey), "access"}, @@ -99,7 +100,7 @@ TEST(S3ConfigTest, overrideBucketConfig) { {S3Config::baseConfigKey(S3Config::Keys::kAccessKey), "access"}, {S3Config::bucketConfigKey(S3Config::Keys::kAccessKey, bucket), "bucket-access"}, - {"hive.s3.payload-signing-policy", "Always"}, + {S3Config::kS3PayloadSigningPolicy, "Always"}, {S3Config::baseConfigKey(S3Config::Keys::kSecretKey), "secret"}, {S3Config::bucketConfigKey(S3Config::Keys::kSecretKey, bucket), "bucket-secret"}, @@ -138,6 +139,54 @@ TEST(S3ConfigTest, overrideBucketConfig) { ASSERT_EQ(s3Config.minPartSize(), 20971520); } +TEST(S3ConfigTest, deprecatedPrefixFallback) { + std::string_view bucket = "bucket"; + // Configure entirely through the deprecated "hive.s3." prefix. + std::unordered_map configFromFile = { + {"hive.s3.endpoint", "endpoint"}, + {"hive.s3.aws-access-key", "access"}, + {"hive.s3.aws-secret-key", "secret"}, + {"hive.s3.bucket.bucket.aws-access-key", "bucket-access"}, + {"hive.s3.payload-signing-policy", "Always"}, + {"hive.s3.log-level", "Info"}, + {"hive.s3.log-location", "/tmp/logs"}, + }; + auto configBase = + std::make_shared(std::move(configFromFile)); + auto s3Config = S3Config(bucket, configBase); + ASSERT_EQ(s3Config.endpoint(), std::optional("endpoint")); + ASSERT_EQ(s3Config.accessKey(), std::optional("bucket-access")); + ASSERT_EQ(s3Config.secretKey(), std::optional("secret")); + ASSERT_EQ(s3Config.payloadSigningPolicy(), "Always"); + // cacheKey also honors the deprecated endpoint key. + ASSERT_EQ(s3Config.cacheKey(bucket, configBase), "endpoint-bucket"); + // The log settings that RegisterS3FileSystem reads honor the fallback too. + ASSERT_EQ( + S3Config::configValue(*configBase, S3Config::kS3LogLevel), + std::optional("Info")); + ASSERT_EQ( + S3Config::configValue(*configBase, S3Config::kS3LogLocation), + std::optional("/tmp/logs")); +} + +TEST(S3ConfigTest, canonicalPrefixWins) { + // When both prefixes are set, the canonical "s3." value takes precedence. + std::unordered_map configFromFile = { + {S3Config::baseConfigKey(S3Config::Keys::kEndpoint), "canonical"}, + {"hive.s3.endpoint", "deprecated"}, + }; + auto configBase = + std::make_shared(std::move(configFromFile)); + auto s3Config = S3Config("bucket", configBase); + ASSERT_EQ(s3Config.endpoint(), std::optional("canonical")); + ASSERT_EQ(s3Config.cacheKey("bucket", configBase), "canonical-bucket"); + + ASSERT_EQ( + S3Config::configValue( + *configBase, S3Config::baseConfigKey(S3Config::Keys::kEndpoint)), + std::optional("canonical")); +} + TEST(S3ConfigTest, minPartSizeValidation) { // Test that setting min-part-size below 5MB throws an error. std::unordered_map configFromFile = { diff --git a/velox/connectors/hive/storage_adapters/s3fs/tests/S3FileSystemMetricsTest.cpp b/velox/connectors/hive/storage_adapters/s3fs/tests/S3FileSystemMetricsTest.cpp index 907acf92783..651abc1c9b5 100644 --- a/velox/connectors/hive/storage_adapters/s3fs/tests/S3FileSystemMetricsTest.cpp +++ b/velox/connectors/hive/storage_adapters/s3fs/tests/S3FileSystemMetricsTest.cpp @@ -147,8 +147,8 @@ TEST_F(S3FileSystemMetricsTest, metrics) { const auto file = "test.txt"; const auto filename = localPath(bucketName) + "/" + file; const auto s3File = s3URI(bucketName, file); - auto hiveConfig = minioServer_->hiveConfig(); - S3FileSystem s3fs(bucketName, hiveConfig); + auto s3Config = minioServer_->s3Config(); + S3FileSystem s3fs(bucketName, s3Config); auto pool = memory::memoryManager()->addLeafPool("S3FileSystemMetricsTest"); auto writeFile = diff --git a/velox/connectors/hive/storage_adapters/s3fs/tests/S3FileSystemRegistrationTest.cpp b/velox/connectors/hive/storage_adapters/s3fs/tests/S3FileSystemRegistrationTest.cpp index 6b98a1d290a..7558bdd7a27 100644 --- a/velox/connectors/hive/storage_adapters/s3fs/tests/S3FileSystemRegistrationTest.cpp +++ b/velox/connectors/hive/storage_adapters/s3fs/tests/S3FileSystemRegistrationTest.cpp @@ -23,7 +23,7 @@ namespace { std::string cacheKeyFunc( std::shared_ptr config, std::string_view path) { - return config->get("hive.s3.endpoint").value(); + return config->get("s3.endpoint").value(); } class CustomS3FileSystem : public S3FileSystem { @@ -63,9 +63,9 @@ TEST_F(S3FileSystemRegistrationTest, readViaRegistry) { LocalWriteFile writeFile(filename); writeData(&writeFile); } - auto hiveConfig = minioServer_->hiveConfig(); + auto s3Config = minioServer_->s3Config(); { - auto s3fs = filesystems::getFileSystem(s3File, hiveConfig); + auto s3fs = filesystems::getFileSystem(s3File, s3Config); auto readFile = s3fs->openFileForRead(s3File); readData(readFile.get()); } @@ -81,34 +81,34 @@ TEST_F(S3FileSystemRegistrationTest, fileHandle) { LocalWriteFile writeFile(filename); writeData(&writeFile); } - auto hiveConfig = minioServer_->hiveConfig(); + auto s3Config = minioServer_->s3Config(); FileHandleFactory factory( std::make_unique>(1000), - std::make_unique(hiveConfig)); + std::make_unique(s3Config)); FileHandleKey key{s3File}; auto fileHandleCachePtr = factory.generate(key); readData(fileHandleCachePtr->file.get()); } TEST_F(S3FileSystemRegistrationTest, cacheKey) { - auto hiveConfig = minioServer_->hiveConfig(); - auto s3fs = filesystems::getFileSystem(kDummyPath, hiveConfig); + auto s3Config = minioServer_->s3Config(); + auto s3fs = filesystems::getFileSystem(kDummyPath, s3Config); std::string_view kDummyPath2 = "s3://dummy2/foo.txt"; - auto s3fs_new = filesystems::getFileSystem(kDummyPath2, hiveConfig); + auto s3fs_new = filesystems::getFileSystem(kDummyPath2, s3Config); // The cacheKeyFunc function allows fs caching based on the endpoint value. ASSERT_EQ(s3fs, s3fs_new); } TEST_F(S3FileSystemRegistrationTest, customFileSystemFactory) { - auto hiveConfig = minioServer_->hiveConfig(); - auto s3fs = filesystems::getFileSystem(kDummyPath, hiveConfig); + auto s3Config = minioServer_->s3Config(); + auto s3fs = filesystems::getFileSystem(kDummyPath, s3Config); auto customS3fs = std::dynamic_pointer_cast(s3fs); VELOX_CHECK_NOT_NULL(customS3fs); } TEST_F(S3FileSystemRegistrationTest, finalize) { - auto hiveConfig = minioServer_->hiveConfig(); - auto s3fs = filesystems::getFileSystem(kDummyPath, hiveConfig); + auto s3Config = minioServer_->s3Config(); + auto s3fs = filesystems::getFileSystem(kDummyPath, s3Config); VELOX_ASSERT_THROW( filesystems::finalizeS3FileSystem(), "Cannot finalize S3FileSystem while in use"); diff --git a/velox/connectors/hive/storage_adapters/s3fs/tests/S3FileSystemTest.cpp b/velox/connectors/hive/storage_adapters/s3fs/tests/S3FileSystemTest.cpp index 3c71b02c6aa..54dde14466e 100644 --- a/velox/connectors/hive/storage_adapters/s3fs/tests/S3FileSystemTest.cpp +++ b/velox/connectors/hive/storage_adapters/s3fs/tests/S3FileSystemTest.cpp @@ -35,7 +35,6 @@ class S3FileSystemTest : public S3Test { void SetUp() override { S3Test::SetUp(); - auto hiveConfig = minioServer_->hiveConfig({}); filesystems::initializeS3("Info", kLogLocation_); } @@ -70,8 +69,8 @@ TEST_F(S3FileSystemTest, writeAndRead) { LocalWriteFile writeFile(filename); writeData(&writeFile); } - auto hiveConfig = minioServer_->hiveConfig(); - filesystems::S3FileSystem s3fs(bucketName, hiveConfig); + auto s3Config = minioServer_->s3Config(); + filesystems::S3FileSystem s3fs(bucketName, s3Config); auto readFile = s3fs.openFileForRead(s3File); readData(readFile.get()); } @@ -79,68 +78,68 @@ TEST_F(S3FileSystemTest, writeAndRead) { TEST_F(S3FileSystemTest, invalidCredentialsConfig) { { std::unordered_map config( - {{"hive.s3.aws-session-token", "dummy-token"}}); - auto hiveConfig = + {{"s3.aws-session-token", "dummy-token"}}); + auto s3Config = std::make_shared(std::move(config)); VELOX_ASSERT_THROW( - filesystems::S3FileSystem("", hiveConfig), + filesystems::S3FileSystem("", s3Config), "Invalid configuration: session token requires both access key and secret key"); } { std::unordered_map config( - {{"hive.s3.use-instance-credentials", "true"}, - {"hive.s3.iam-role", "dummy-iam-role"}}); - auto hiveConfig = + {{"s3.use-instance-credentials", "true"}, + {"s3.iam-role", "dummy-iam-role"}}); + auto s3Config = std::make_shared(std::move(config)); // Both instance credentials and iam-role cannot be specified VELOX_ASSERT_THROW( - filesystems::S3FileSystem("", hiveConfig), + filesystems::S3FileSystem("", s3Config), "Invalid configuration: specify only one among 'access/secret keys', 'use instance credentials', 'IAM role'"); } { std::unordered_map config( - {{"hive.s3.aws-secret-key", "dummy-key"}, - {"hive.s3.aws-access-key", "dummy-key"}, - {"hive.s3.iam-role", "dummy-iam-role"}}); - auto hiveConfig = + {{"s3.aws-secret-key", "dummy-key"}, + {"s3.aws-access-key", "dummy-key"}, + {"s3.iam-role", "dummy-iam-role"}}); + auto s3Config = std::make_shared(std::move(config)); // Both access/secret keys and iam-role cannot be specified VELOX_ASSERT_THROW( - filesystems::S3FileSystem("", hiveConfig), + filesystems::S3FileSystem("", s3Config), "Invalid configuration: specify only one among 'access/secret keys', 'use instance credentials', 'IAM role'"); } { std::unordered_map config( - {{"hive.s3.aws-secret-key", "dummy"}, - {"hive.s3.aws-access-key", "dummy"}, - {"hive.s3.use-instance-credentials", "true"}}); - auto hiveConfig = + {{"s3.aws-secret-key", "dummy"}, + {"s3.aws-access-key", "dummy"}, + {"s3.use-instance-credentials", "true"}}); + auto s3Config = std::make_shared(std::move(config)); // Both access/secret keys and instance credentials cannot be specified VELOX_ASSERT_THROW( - filesystems::S3FileSystem("", hiveConfig), + filesystems::S3FileSystem("", s3Config), "Invalid configuration: specify only one among 'access/secret keys', 'use instance credentials', 'IAM role'"); } { std::unordered_map config( - {{"hive.s3.aws-secret-key", "dummy"}}); - auto hiveConfig = + {{"s3.aws-secret-key", "dummy"}}); + auto s3Config = std::make_shared(std::move(config)); // Both access key and secret key must be specified VELOX_ASSERT_THROW( - filesystems::S3FileSystem("", hiveConfig), + filesystems::S3FileSystem("", s3Config), "Invalid configuration: both access key and secret key must be specified"); } } TEST_F(S3FileSystemTest, temporaryCredentials) { - auto hiveConfig = std::make_shared( + auto s3Config = std::make_shared( std::unordered_map{ - {"hive.s3.aws-access-key", "access"}, - {"hive.s3.aws-secret-key", "secret"}, - {"hive.s3.aws-session-token", "token"}}); - filesystems::S3FileSystem s3fs("", hiveConfig); + {"s3.aws-access-key", "access"}, + {"s3.aws-secret-key", "secret"}, + {"s3.aws-session-token", "token"}}); + filesystems::S3FileSystem s3fs("", s3Config); const auto credentials = s3fs.getCredentialSnapshot(); EXPECT_EQ(credentials.accessKeyId, "access"); @@ -153,8 +152,8 @@ TEST_F(S3FileSystemTest, missingFile) { const char* file = "i-do-not-exist.txt"; const std::string s3File = s3URI(bucketName, file); addBucket(bucketName); - auto hiveConfig = minioServer_->hiveConfig(); - filesystems::S3FileSystem s3fs(bucketName, hiveConfig); + auto s3Config = minioServer_->s3Config(); + filesystems::S3FileSystem s3fs(bucketName, s3Config); VELOX_ASSERT_RUNTIME_THROW_CODE( s3fs.openFileForRead(s3File), error_code::kFileNotFound, @@ -162,8 +161,8 @@ TEST_F(S3FileSystemTest, missingFile) { } TEST_F(S3FileSystemTest, missingBucket) { - auto hiveConfig = minioServer_->hiveConfig(); - filesystems::S3FileSystem s3fs("", hiveConfig); + auto s3Config = minioServer_->s3Config(); + filesystems::S3FileSystem s3fs("", s3Config); VELOX_ASSERT_RUNTIME_THROW_CODE( s3fs.openFileForRead(kDummyPath), error_code::kFileNotFound, @@ -171,9 +170,8 @@ TEST_F(S3FileSystemTest, missingBucket) { } TEST_F(S3FileSystemTest, invalidAccessKey) { - auto hiveConfig = - minioServer_->hiveConfig({{"hive.s3.aws-access-key", "dummy-key"}}); - filesystems::S3FileSystem s3fs("", hiveConfig); + auto s3Config = minioServer_->s3Config({{"s3.aws-access-key", "dummy-key"}}); + filesystems::S3FileSystem s3fs("", s3Config); // Minio credentials are wrong and this should throw VELOX_ASSERT_THROW( s3fs.openFileForRead(kDummyPath), @@ -181,9 +179,8 @@ TEST_F(S3FileSystemTest, invalidAccessKey) { } TEST_F(S3FileSystemTest, invalidSecretKey) { - auto hiveConfig = - minioServer_->hiveConfig({{"hive.s3.aws-secret-key", "dummy-key"}}); - filesystems::S3FileSystem s3fs("", hiveConfig); + auto s3Config = minioServer_->s3Config({{"s3.aws-secret-key", "dummy-key"}}); + filesystems::S3FileSystem s3fs("", s3Config); // Minio credentials are wrong and this should throw. VELOX_ASSERT_THROW( s3fs.openFileForRead("s3://dummy/foo.txt"), @@ -191,9 +188,8 @@ TEST_F(S3FileSystemTest, invalidSecretKey) { } TEST_F(S3FileSystemTest, noBackendServer) { - auto hiveConfig = - minioServer_->hiveConfig({{"hive.s3.aws-secret-key", "dummy-key"}}); - filesystems::S3FileSystem s3fs("", hiveConfig); + auto s3Config = minioServer_->s3Config({{"s3.aws-secret-key", "dummy-key"}}); + filesystems::S3FileSystem s3fs("", s3Config); // Stop Minio and check error. minioServer_->stop(); VELOX_ASSERT_THROW( @@ -217,7 +213,7 @@ TEST_F(S3FileSystemTest, logLevel) { // S3 log level is set once during initialization. // It does not change with a new config. - config["hive.s3.log-level"] = "Trace"; + config["s3.log-level"] = "Trace"; checkLogLevelName("INFO"); } @@ -238,7 +234,7 @@ TEST_F(S3FileSystemTest, logLocation) { // S3 log location is set once during initialization. // It does not change with a new config. - config["hive.s3.log-location"] = "/home/foobar"; + config["s3.log-location"] = "/home/foobar"; checkLogPrefix(expected); } @@ -248,8 +244,8 @@ TEST_F(S3FileSystemTest, mkdirAndRename) { const auto s3File = s3URI(bucketName, file); addBucket(bucketName); - auto hiveConfig = minioServer_->hiveConfig(); - filesystems::S3FileSystem s3fs(bucketName, hiveConfig); + auto s3Config = minioServer_->s3Config(); + filesystems::S3FileSystem s3fs(bucketName, s3Config); ASSERT_FALSE(s3fs.exists(s3File)); s3fs.mkdir(s3File); @@ -269,9 +265,9 @@ TEST_F(S3FileSystemTest, writeFileAndRead) { const auto filename = localPath(bucketName) + "/" + file; const auto s3File = s3URI(bucketName, file); - auto hiveConfig = - minioServer_->hiveConfig({{"hive.s3.multipart-upload-threads", "4"}}); - filesystems::S3FileSystem s3fs(bucketName, hiveConfig); + auto s3Config = + minioServer_->s3Config({{"s3.multipart-upload-threads", "4"}}); + filesystems::S3FileSystem s3fs(bucketName, s3Config); auto pool = memory::memoryManager()->addLeafPool("S3FileSystemTest"); auto writeFile = s3fs.openFileForWrite(s3File, {{}, pool.get(), std::nullopt}); @@ -348,9 +344,9 @@ TEST_F(S3FileSystemTest, writeFileAndRead) { TEST_F(S3FileSystemTest, abortWrite) { const auto bucketName = "abortwrite"; addBucket(bucketName); - auto hiveConfig = - minioServer_->hiveConfig({{"hive.s3.multipart-upload-threads", "4"}}); - filesystems::S3FileSystem s3fs(bucketName, hiveConfig); + auto s3Config = + minioServer_->s3Config({{"s3.multipart-upload-threads", "4"}}); + filesystems::S3FileSystem s3fs(bucketName, s3Config); auto pool = memory::memoryManager()->addLeafPool("S3AbortWriteTest"); for (const auto& [file, size] : std::vector>{ @@ -371,14 +367,13 @@ TEST_F(S3FileSystemTest, abortWrite) { } TEST_F(S3FileSystemTest, invalidConnectionSettings) { - auto hiveConfig = - minioServer_->hiveConfig({{"hive.s3.connect-timeout", "400"}}); + auto s3Config = minioServer_->s3Config({{"s3.connect-timeout", "400"}}); VELOX_ASSERT_THROW( - filesystems::S3FileSystem("", hiveConfig), "Invalid duration"); + filesystems::S3FileSystem("", s3Config), "Invalid duration"); - hiveConfig = minioServer_->hiveConfig({{"hive.s3.socket-timeout", "abc"}}); + s3Config = minioServer_->s3Config({{"s3.socket-timeout", "abc"}}); VELOX_ASSERT_THROW( - filesystems::S3FileSystem("", hiveConfig), "Invalid duration"); + filesystems::S3FileSystem("", s3Config), "Invalid duration"); } TEST_F(S3FileSystemTest, registerCredentialProviderFactories) { @@ -389,15 +384,15 @@ TEST_F(S3FileSystemTest, registerCredentialProviderFactories) { return std::make_shared(); }); - auto hiveConfig = minioServer_->hiveConfig( - {{"hive.s3.aws-credentials-provider", credentialsProvider}}); - ASSERT_NO_THROW(filesystems::S3FileSystem("", hiveConfig)); + auto s3Config = minioServer_->s3Config( + {{"s3.aws-credentials-provider", credentialsProvider}}); + ASSERT_NO_THROW(filesystems::S3FileSystem("", s3Config)); // Configure with unregistered credential provider. - hiveConfig = minioServer_->hiveConfig( - {{"hive.s3.aws-credentials-provider", invalidCredentialsProvider}}); + s3Config = minioServer_->s3Config( + {{"s3.aws-credentials-provider", invalidCredentialsProvider}}); VELOX_ASSERT_THROW( - filesystems::S3FileSystem({"", hiveConfig}), + filesystems::S3FileSystem({"", s3Config}), "CredentialsProviderFactory for 'invalid-credentials-provider' not registered"); // Register invalid credentials provider name. diff --git a/velox/connectors/hive/storage_adapters/s3fs/tests/S3InsertTest.cpp b/velox/connectors/hive/storage_adapters/s3fs/tests/S3InsertTest.cpp index 0861df83d48..eb51de73bff 100644 --- a/velox/connectors/hive/storage_adapters/s3fs/tests/S3InsertTest.cpp +++ b/velox/connectors/hive/storage_adapters/s3fs/tests/S3InsertTest.cpp @@ -37,7 +37,7 @@ class S3InsertTest : public S3Test, public test::InsertTest { void SetUp() override { S3Test::SetUp(); - InsertTest::SetUp(minioServer_->hiveConfig(), ioExecutor_.get()); + InsertTest::SetUp(minioServer_->s3Config(), ioExecutor_.get()); } void TearDown() override { diff --git a/velox/connectors/hive/storage_adapters/s3fs/tests/S3MultipleEndpointsTest.cpp b/velox/connectors/hive/storage_adapters/s3fs/tests/S3MultipleEndpointsTest.cpp index 3c90521bb56..6d483616771 100644 --- a/velox/connectors/hive/storage_adapters/s3fs/tests/S3MultipleEndpointsTest.cpp +++ b/velox/connectors/hive/storage_adapters/s3fs/tests/S3MultipleEndpointsTest.cpp @@ -66,11 +66,11 @@ class S3MultipleEndpoints : public S3Test, public ::test::VectorTestBase { connector::hive::HiveConnectorFactory factory; auto hiveConnector1 = factory.newConnector( std::string(connectorId1), - minioServer_->hiveConfig(config1Override), + minioServer_->s3Config(config1Override), ioExecutor_.get()); auto hiveConnector2 = factory.newConnector( std::string(connectorId2), - minioSecondServer_->hiveConfig(config2Override), + minioSecondServer_->s3Config(config2Override), ioExecutor_.get()); connector::ConnectorRegistry::global().insert( hiveConnector1->connectorId(), hiveConnector1); @@ -204,19 +204,19 @@ TEST_F(S3MultipleEndpoints, bucketEndpoints) { auto configOverride = [](std::shared_ptr config) { return std::unordered_map{ - {"hive.s3.bucket.writedata.endpoint", - config->get("hive.s3.endpoint").value()}, - {"hive.s3.bucket.writedata.aws-access-key", - config->get("hive.s3.aws-access-key").value()}, - {"hive.s3.bucket.writedata.aws-secret-key", - config->get("hive.s3.aws-secret-key").value()}, - {"hive.s3.endpoint", "fail"}, - {"hive.s3.aws-access-key", "fail"}, - {"hive.s3.aws-secret-key", "fail"}, + {"s3.bucket.writedata.endpoint", + config->get("s3.endpoint").value()}, + {"s3.bucket.writedata.aws-access-key", + config->get("s3.aws-access-key").value()}, + {"s3.bucket.writedata.aws-secret-key", + config->get("s3.aws-secret-key").value()}, + {"s3.endpoint", "fail"}, + {"s3.aws-access-key", "fail"}, + {"s3.aws-secret-key", "fail"}, }; }; - auto config1 = configOverride(minioServer_->hiveConfig()); - auto config2 = configOverride(minioSecondServer_->hiveConfig()); + auto config1 = configOverride(minioServer_->s3Config()); + auto config2 = configOverride(minioSecondServer_->s3Config()); registerConnectors(kConnectorId1, kConnectorId2, config1, config2); testJoin(kExpectedRows, outputDirectory, kConnectorId1, kConnectorId2); diff --git a/velox/connectors/hive/storage_adapters/s3fs/tests/S3ReadTest.cpp b/velox/connectors/hive/storage_adapters/s3fs/tests/S3ReadTest.cpp index 9062eed288a..8c9fc1e46ce 100644 --- a/velox/connectors/hive/storage_adapters/s3fs/tests/S3ReadTest.cpp +++ b/velox/connectors/hive/storage_adapters/s3fs/tests/S3ReadTest.cpp @@ -43,7 +43,7 @@ class S3ReadTest : public S3Test, public ::test::VectorTestBase { filesystems::registerS3FileSystem(); connector::hive::HiveConnectorFactory factory; auto hiveConnector = - factory.newConnector(kHiveConnectorId, minioServer_->hiveConfig()); + factory.newConnector(kHiveConnectorId, minioServer_->s3Config()); connector::ConnectorRegistry::global().insert( hiveConnector->connectorId(), hiveConnector); parquet::registerParquetReaderFactory(); diff --git a/velox/connectors/hive/tests/CMakeLists.txt b/velox/connectors/hive/tests/CMakeLists.txt index f698f19945c..c72fe355704 100644 --- a/velox/connectors/hive/tests/CMakeLists.txt +++ b/velox/connectors/hive/tests/CMakeLists.txt @@ -28,6 +28,7 @@ add_executable( HivePartitionNameTest.cpp HiveSplitTest.cpp PartitionIdGeneratorTest.cpp + PartitionValueTest.cpp TableHandleTest.cpp ) add_test(velox_hive_connector_test velox_hive_connector_test) diff --git a/velox/connectors/hive/tests/FileConfigTest.cpp b/velox/connectors/hive/tests/FileConfigTest.cpp index 18c714d242a..8b1d6ae5506 100644 --- a/velox/connectors/hive/tests/FileConfigTest.cpp +++ b/velox/connectors/hive/tests/FileConfigTest.cpp @@ -56,6 +56,7 @@ TEST(FileConfigTest, defaultConfig) { EXPECT_FALSE(config.nimbleStringDecoderZeroCopy(emptySession.get())); EXPECT_FALSE(config.nimblePreserveDictionaryEncoding(emptySession.get())); EXPECT_FALSE(config.nimbleLazyColumnIo(emptySession.get())); + EXPECT_FALSE(config.directBufferedInputSharedAllocation(emptySession.get())); } TEST(FileConfigTest, overrideConfig) { @@ -79,6 +80,7 @@ TEST(FileConfigTest, overrideConfig) { {FileConfig::kNimbleStringDecoderZeroCopy, "true"}, {FileConfig::kNimblePreserveDictionaryEncoding, "true"}, {FileConfig::kNimbleLazyColumnIo, "true"}, + {FileConfig::kDirectBufferedInputSharedAllocation, "true"}, }; FileConfig config( std::make_shared(std::move(configFromFile)), "hive."); @@ -105,6 +107,10 @@ TEST(FileConfigTest, overrideConfig) { EXPECT_TRUE(config.nimbleStringDecoderZeroCopy(emptySession.get())); EXPECT_TRUE(config.nimblePreserveDictionaryEncoding(emptySession.get())); EXPECT_TRUE(config.nimbleLazyColumnIo(emptySession.get())); + // The catalog key is the only way this gate can be enabled in production: the + // session key contains a dot, so the Presto CLI cannot parse it, and there is + // no HiveSessionProperties.java entry for it. + EXPECT_TRUE(config.directBufferedInputSharedAllocation(emptySession.get())); } TEST(FileConfigTest, connectorScopedReaderOptions) { @@ -144,6 +150,7 @@ TEST(FileConfigTest, overrideSession) { {FileConfig::kNimbleStringDecoderZeroCopySession, "true"}, {FileConfig::kNimblePreserveDictionaryEncodingSession, "true"}, {FileConfig::kNimbleLazyColumnIoSession, "true"}, + {FileConfig::kDirectBufferedInputSharedAllocationSession, "true"}, }; const auto session = std::make_unique(std::move(sessionOverride)); @@ -165,6 +172,7 @@ TEST(FileConfigTest, overrideSession) { EXPECT_TRUE(config.nimbleStringDecoderZeroCopy(session.get())); EXPECT_TRUE(config.nimblePreserveDictionaryEncoding(session.get())); EXPECT_TRUE(config.nimbleLazyColumnIo(session.get())); + EXPECT_TRUE(config.directBufferedInputSharedAllocation(session.get())); } TEST(FileConfigTest, nullConfig) { diff --git a/velox/connectors/hive/tests/FileConnectorUtilTest.cpp b/velox/connectors/hive/tests/FileConnectorUtilTest.cpp index 2ebb6100e40..95b6bfcc07a 100644 --- a/velox/connectors/hive/tests/FileConnectorUtilTest.cpp +++ b/velox/connectors/hive/tests/FileConnectorUtilTest.cpp @@ -77,7 +77,9 @@ class FileConnectorUtilTest : public exec::test::HiveConnectorTestBase { std::shared_ptr makeSplit( dwio::common::FileFormat format = dwio::common::FileFormat::DWRF, const std::string& path = "/tmp/testfile", - bool cacheable = true) { + bool cacheable = true, + std::optional columnMappingMode = + std::nullopt) { return std::make_shared( "testConnectorId", path, @@ -85,7 +87,10 @@ class FileConnectorUtilTest : public exec::test::HiveConnectorTestBase { /*_start=*/0, /*_length=*/std::numeric_limits::max(), /*splitWeight=*/0, - cacheable); + cacheable, + std::nullopt, + std::unordered_map>{}, + columnMappingMode); } std::string writeDataFile(const RowVectorPtr& data) { @@ -164,6 +169,53 @@ TEST_F(FileConnectorUtilTest, configureReaderOptions) { EXPECT_EQ(readerOptions.footerSpeculativeIoSize(), 128UL << 10); } + // Split-level column mapping mode overrides the shared session property. + { + auto holder = makeConnectorQueryCtx( + {{hive::FileConfig::kUseColumnNamesSession, "true"}}); + auto split = makeSplit( + dwio::common::FileFormat::ORC, + "/tmp/testfile", + true, + dwio::common::ColumnMappingMode::kPosition); + dwio::common::ReaderOptions readerOptions(pool_.get()); + readerOptions.setDataIoStats(dataIoStats_); + readerOptions.setMetadataIoStats(metadataIoStats_); + hive::configureReaderOptions( + fileConfig, + holder.ctx.get(), + /*fileSchema=*/nullptr, + split, + /*tableParameters=*/{}, + readerOptions); + + EXPECT_EQ( + readerOptions.columnMappingMode(), + dwio::common::ColumnMappingMode::kPosition); + } + + // Parquet field-id matching is only valid for Parquet files. + { + auto holder = makeConnectorQueryCtx(); + auto split = makeSplit( + dwio::common::FileFormat::ORC, + "/tmp/testfile", + true, + dwio::common::ColumnMappingMode::kParquetFieldId); + dwio::common::ReaderOptions readerOptions(pool_.get()); + readerOptions.setDataIoStats(dataIoStats_); + readerOptions.setMetadataIoStats(metadataIoStats_); + VELOX_ASSERT_THROW( + hive::configureReaderOptions( + fileConfig, + holder.ctx.get(), + /*fileSchema=*/nullptr, + split, + /*tableParameters=*/{}, + readerOptions), + "not supported for file format orc"); + } + // Test format mismatch throws. { auto holder = makeConnectorQueryCtx(); @@ -516,6 +568,102 @@ TEST_F(FileConnectorUtilTest, testFiltersPartitionKeyFails) { /*asLocalTime=*/false)); } +// A constant decides the filter even when the data file carries the column. +// The reader built from the scan spec has no 'typeWithId()' subtree for it, so +// reading its statistics dereferences null. +TEST_F(FileConnectorUtilTest, testFiltersConstantOverridingFileColumn) { + auto batch = + makeRowVector({"c0"}, {makeFlatVector(100, folly::identity)}); + auto filePath = writeDataFile(batch); + + // Tests 'filter' on a scan that returns 'constant' for 'c0'. + auto keepsSplit = [&](const VectorPtr& constant, + std::shared_ptr filter) { + auto scanSpec = std::make_shared(""); + auto* child = scanSpec->addField("c0", 0); + child->setConstantValue(constant); + child->setFilter(std::move(filter)); + + dwio::common::ReaderOptions readerOpts(pool_.get()); + readerOpts.setFileFormat(dwio::common::FileFormat::DWRF); + readerOpts.setScanSpec(scanSpec); + auto reader = dwrf::DwrfReader::create( + std::make_unique( + std::make_shared(filePath), readerOpts.memoryPool()), + readerOpts); + + return hive::testFilters( + scanSpec.get(), + reader.get(), + filePath, + /*partitionKeys=*/{}, + /*partitionKeysHandle=*/{}, + /*asLocalTime=*/false); + }; + + // The file holds 0 through 99: statistics would keep the split for both. + auto constant = makeConstant(1'000, 1); + EXPECT_FALSE(keepsSplit( + constant, std::make_shared(0, 99, false))); + EXPECT_TRUE(keepsSplit( + constant, std::make_shared(1'000, 1'000, false))); + + // A null constant decides from testNull() alone. The file holds no null, so + // statistics would answer the other way for both. + auto nullConstant = BaseVector::createNullConstant(BIGINT(), 1, pool_.get()); + EXPECT_FALSE(keepsSplit(nullConstant, std::make_shared())); + EXPECT_TRUE(keepsSplit(nullConstant, std::make_shared())); +} + +// Every filtered partition key has to be tested, not just the first. Nothing +// downstream re-checks them: testFilterOnConstant() accepts any non-null +// constant. +TEST_F(FileConnectorUtilTest, testFiltersSecondPartitionKeyFails) { + auto batch = + makeRowVector({"c0"}, {makeFlatVector(100, folly::identity)}); + auto filePath = writeDataFile(batch); + auto reader = makeReader(filePath); + + auto scanSpec = std::make_shared(""); + scanSpec->addField("c0", 0); + // 'ds' comes first and passes, so stopping at the first key misses 'hour'. + scanSpec->addField("ds", 1)->setFilter( + std::make_unique( + std::vector{"2024-01-01"}, false)); + scanSpec->addField("hour", 2)->setFilter( + std::make_unique(0, 11, false)); + + const std::unordered_map> + partitionKeys = { + {"ds", "2024-01-01"}, + {"hour", "23"}, + }; + const std::unordered_map + partitionKeysHandle = { + {"ds", + std::make_shared( + "ds", + hive::HiveColumnHandle::ColumnType::kPartitionKey, + VARCHAR(), + VARCHAR())}, + {"hour", + std::make_shared( + "hour", + hive::HiveColumnHandle::ColumnType::kPartitionKey, + BIGINT(), + BIGINT())}, + }; + + EXPECT_FALSE( + hive::testFilters( + scanSpec.get(), + reader.get(), + filePath, + partitionKeys, + partitionKeysHandle, + /*asLocalTime=*/false)); +} + TEST_F(FileConnectorUtilTest, testFiltersNullPartitionKeyRejectsNotNull) { auto rowType = ROW({"c0"}, {BIGINT()}); auto batch = diff --git a/velox/connectors/hive/tests/HiveConnectorSerDeTest.cpp b/velox/connectors/hive/tests/HiveConnectorSerDeTest.cpp index 9bec9c0c5a7..d5f37bc332c 100644 --- a/velox/connectors/hive/tests/HiveConnectorSerDeTest.cpp +++ b/velox/connectors/hive/tests/HiveConnectorSerDeTest.cpp @@ -105,9 +105,21 @@ class HiveConnectorSerDeTest : public exec::test::HiveConnectorTestBase { ASSERT_EQ( split.properties->modificationTime, clone->properties->modificationTime); + ASSERT_EQ( + split.properties->readRangeHint, clone->properties->readRangeHint); + if (split.properties->extraFileInfo != nullptr) { + ASSERT_NE(clone->properties->extraFileInfo, nullptr); + ASSERT_EQ( + *split.properties->extraFileInfo, + *clone->properties->extraFileInfo); + } else { + ASSERT_EQ(clone->properties->extraFileInfo, nullptr); + } + ASSERT_EQ(split.properties->fileReadOps, clone->properties->fileReadOps); } else { ASSERT_FALSE(clone->properties.has_value()); } + ASSERT_EQ(split.columnMappingMode, clone->columnMappingMode); } }; @@ -270,7 +282,11 @@ TEST_F(HiveConnectorSerDeTest, hiveConnectorSplit) { const std::unordered_map infoColumns{ {"c0", "0"}, {"c1", "1"}}; FileProperties fileProperties{ - .fileSize = 2048, .modificationTime = std::nullopt}; + .fileSize = 2048, + .modificationTime = std::nullopt, + .readRangeHint = std::nullopt, + .extraFileInfo = nullptr, + .fileReadOps = {}}; const auto properties = std::optional(fileProperties); RowIdProperties rowIdProperties{ .metadataVersion = 2, .partitionId = 3, .tableGuid = "test"}; @@ -289,7 +305,9 @@ TEST_F(HiveConnectorSerDeTest, hiveConnectorSplit) { cacheable, infoColumns, properties, - rowIdProperties); + rowIdProperties, + std::nullopt, + dwio::common::ColumnMappingMode::kName); ASSERT_EQ(split1.cacheable, cacheable); testSerde(split1); @@ -319,5 +337,36 @@ TEST_F(HiveConnectorSerDeTest, hiveConnectorSplit) { testSerde(split3); } +TEST_F(HiveConnectorSerDeTest, hiveConnectorSplitFileProperties) { + std::string descriptor; + for (int i = 0; i < 256; ++i) { + descriptor.push_back(static_cast(i)); + } + + auto split = HiveConnectorSplit( + "testSerde", "/testSerde/p", dwio::common::FileFormat::DWRF); + split.properties = FileProperties{ + .fileSize = 2048, + .modificationTime = 1024, + .readRangeHint = 4096, + .extraFileInfo = std::make_shared(descriptor), + .fileReadOps = {{"o0", "0"}, {"o1", "1"}}}; + testSerde(split); + + // A producer that predates these keys omits them rather than writing nulls. + auto obj = split.serialize(); + obj["properties"].erase("readRangeHint"); + obj["properties"].erase("extraFileInfo"); + obj["properties"].erase("fileReadOps"); + + const auto clone = ISerializable::deserialize(obj); + ASSERT_TRUE(clone->properties.has_value()); + EXPECT_EQ(clone->properties->fileSize, 2048); + EXPECT_EQ(clone->properties->modificationTime, 1024); + EXPECT_FALSE(clone->properties->readRangeHint.has_value()); + EXPECT_EQ(clone->properties->extraFileInfo, nullptr); + EXPECT_TRUE(clone->properties->fileReadOps.empty()); +} + } // namespace } // namespace facebook::velox::connector::hive::test diff --git a/velox/connectors/hive/tests/HiveConnectorTest.cpp b/velox/connectors/hive/tests/HiveConnectorTest.cpp index 180e347b4b6..daf1efd09f5 100644 --- a/velox/connectors/hive/tests/HiveConnectorTest.cpp +++ b/velox/connectors/hive/tests/HiveConnectorTest.cpp @@ -631,10 +631,9 @@ TEST_F(HiveConnectorTest, extractFiltersFromRemainingFilter) { extractFiltersFromRemainingFilter(expr, &evaluator, filters, sampleRate); ASSERT_EQ(sampleRate, 1); ASSERT_GT(filters.count(Subfield("c2")), 0); - // Change these once HUGEINT filter merge is fixed. - ASSERT_TRUE(remaining); - ASSERT_EQ( - remaining->toString(), "not(lt(ROW[\"c2\"],cast(0 as DECIMAL(20, 0))))"); + auto expectedFilter = exec::betweenHugeint(0, 1); + ASSERT_TRUE(expectedFilter->testingEquals(*filters.at(Subfield("c2")))); + ASSERT_FALSE(remaining); // parseExpr gives AND/OR with 2 arguments. We need to construct the node // manually to have more than 2. diff --git a/velox/connectors/hive/tests/HiveDataSinkTest.cpp b/velox/connectors/hive/tests/HiveDataSinkTest.cpp index 583ea12787a..f38fc9f4ee7 100644 --- a/velox/connectors/hive/tests/HiveDataSinkTest.cpp +++ b/velox/connectors/hive/tests/HiveDataSinkTest.cpp @@ -2145,10 +2145,16 @@ TEST_F(HiveDataSinkTest, sessionParquetConfigsMergeIntoProvidedFormatOptions) { dwio::common::FileFormat::PARQUET, parquet::ParquetConfig::kWriterBatchSizeSession), "97"); + connectorSessionProperties_->set( + dwio::common::formatSessionProperty( + dwio::common::FileFormat::PARQUET, + parquet::ParquetConfig::kWriterRowGroupSizeSession), + "2MB"); auto writerOptions = std::make_shared(); auto parquetOptions = std::make_shared(); parquetOptions->batchSize = 11; + parquetOptions->rowGroupSizeBytes = 1 << 20; parquetOptions->bufferGrowRatio = 1.7; writerOptions->formatSpecificOptions = parquetOptions; @@ -2164,6 +2170,7 @@ TEST_F(HiveDataSinkTest, sessionParquetConfigsMergeIntoProvidedFormatOptions) { dataSink->appendData(createVectors(10, 1).front()); EXPECT_EQ(parquetOptions->batchSize, 97); + EXPECT_EQ(parquetOptions->rowGroupSizeBytes, 2 << 20); EXPECT_EQ(parquetOptions->bufferGrowRatio, 1.7); } #endif diff --git a/velox/connectors/hive/tests/PartitionValueTest.cpp b/velox/connectors/hive/tests/PartitionValueTest.cpp new file mode 100644 index 00000000000..954f44b743b --- /dev/null +++ b/velox/connectors/hive/tests/PartitionValueTest.cpp @@ -0,0 +1,172 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/connectors/hive/PartitionValue.h" + +#include + +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/type/Filter.h" +#include "velox/type/TimestampConversion.h" + +namespace facebook::velox::connector::hive { +namespace { + +using TimestampMode = PartitionValue::TimestampMode; +using DateMode = PartitionValue::DateMode; + +Variant toVariant( + std::string_view value, + const TypePtr& type, + TimestampMode timestampMode = TimestampMode::kUtc, + DateMode dateMode = DateMode::kIsoString) { + return PartitionValue::fromString(value, *type, timestampMode, dateMode); +} + +Timestamp parseTimestamp(std::string_view value) { + return util::fromTimestampString( + value.data(), value.size(), util::TimestampParseMode::kPrestoCast) + .value(); +} + +TEST(PartitionValueTest, scalarTypes) { + EXPECT_EQ(toVariant("true", BOOLEAN()), Variant(true)); + EXPECT_EQ(toVariant("-1", TINYINT()), Variant(int8_t{-1})); + EXPECT_EQ(toVariant("-2", SMALLINT()), Variant(int16_t{-2})); + EXPECT_EQ(toVariant("-3", INTEGER()), Variant(int32_t{-3})); + EXPECT_EQ(toVariant("-4", BIGINT()), Variant(int64_t{-4})); + EXPECT_EQ(toVariant("1.25", REAL()), Variant(1.25f)); + EXPECT_EQ(toVariant("2.5", DOUBLE()), Variant(2.5)); + EXPECT_EQ(toVariant("hello", VARCHAR()), Variant(std::string("hello"))); + EXPECT_EQ(toVariant("binary", VARBINARY()), Variant::binary("binary")); +} + +TEST(PartitionValueTest, booleanAcceptedStrings) { + EXPECT_EQ(toVariant("t", BOOLEAN()), Variant(true)); + EXPECT_EQ(toVariant("0", BOOLEAN()), Variant(false)); + EXPECT_EQ(toVariant("FALSE", BOOLEAN()), Variant(false)); + + VELOX_ASSERT_USER_THROW( + toVariant("yes", BOOLEAN()), "Cannot cast yes to BOOLEAN"); + VELOX_ASSERT_USER_THROW( + toVariant("off", BOOLEAN()), "Cannot cast off to BOOLEAN"); +} + +// Each integer type is range-checked against its own native type. +TEST(PartitionValueTest, outOfRangeNarrowInteger) { + VELOX_ASSERT_USER_THROW( + toVariant("99999999999", INTEGER()), "Overflow during conversion"); + VELOX_ASSERT_USER_THROW( + toVariant("40000", SMALLINT()), "Overflow during conversion"); + VELOX_ASSERT_USER_THROW( + toVariant("300", TINYINT()), "Overflow during conversion"); +} + +TEST(PartitionValueTest, decimalTypes) { + EXPECT_EQ( + toVariant("12.34", DECIMAL(10, 2)).value(), 1'234); + EXPECT_EQ( + toVariant("12345678901234567890.12", DECIMAL(25, 2)) + .value(), + HugeInt::parse("1234567890123456789012")); +} + +TEST(PartitionValueTest, dateEncodings) { + EXPECT_EQ(toVariant("2020-01-02", DATE()).value(), 18'263); + EXPECT_EQ( + toVariant("18263", DATE(), TimestampMode::kUtc, DateMode::kDaysSinceEpoch) + .value(), + 18'263); +} + +TEST(PartitionValueTest, timestampModes) { + const auto unshifted = parseTimestamp("2020-01-01 12:34:56"); + auto shifted = unshifted; + shifted.toGMT(Timestamp::defaultTimezone()); + + EXPECT_EQ( + toVariant("2020-01-01 12:34:56", TIMESTAMP(), TimestampMode::kLocalTime) + .value(), + shifted); + EXPECT_EQ( + toVariant("2020-01-01 12:34:56", TIMESTAMP(), TimestampMode::kUtc) + .value(), + unshifted); + EXPECT_EQ( + toVariant( + "2020-01-01 12:34:56", TIMESTAMP_UTC(), TimestampMode::kLocalTime) + .value(), + unshifted); +} + +TEST(PartitionValueTest, filterOnConvertedValue) { + const common::BigintRange bigintRange(10, 20, /*nullAllowed=*/false); + EXPECT_TRUE(applyFilter(bigintRange, toVariant("15", BIGINT()))); + EXPECT_FALSE(applyFilter(bigintRange, toVariant("25", BIGINT()))); + + // A REAL value is tested as a float, not as a double. + const common::FloatRange floatRange( + 1.0f, + /*lowerUnbounded=*/false, + /*lowerExclusive=*/false, + 2.0f, + /*upperUnbounded=*/false, + /*upperExclusive=*/false, + /*nullAllowed=*/false); + EXPECT_TRUE(applyFilter(floatRange, toVariant("1.25", REAL()))); + EXPECT_FALSE(applyFilter(floatRange, toVariant("2.5", REAL()))); + + const common::BigintRange dateRange(18'263, 18'263, /*nullAllowed=*/false); + EXPECT_TRUE(applyFilter(dateRange, toVariant("2020-01-02", DATE()))); + + const common::BoolValue boolFilter(true, /*nullAllowed=*/false); + EXPECT_TRUE(applyFilter(boolFilter, toVariant("true", BOOLEAN()))); + EXPECT_FALSE(applyFilter(boolFilter, toVariant("false", BOOLEAN()))); +} + +TEST(PartitionValueTest, filterOnDecimal) { + const common::BigintRange shortRange(1'234, 1'234, /*nullAllowed=*/false); + EXPECT_TRUE(applyFilter(shortRange, toVariant("12.34", DECIMAL(10, 2)))); + EXPECT_FALSE(applyFilter(shortRange, toVariant("12.35", DECIMAL(10, 2)))); + + const auto longValue = HugeInt::parse("1234567890123456789012"); + const common::HugeintRange longRange( + longValue, longValue, /*nullAllowed=*/false); + EXPECT_TRUE(applyFilter( + longRange, toVariant("12345678901234567890.12", DECIMAL(25, 2)))); +} + +TEST(PartitionValueTest, filterUsesTimestampMode) { + const auto unshifted = parseTimestamp("2020-01-01 12:34:56"); + auto shifted = unshifted; + shifted.toGMT(Timestamp::defaultTimezone()); + + const common::TimestampRange shiftedFilter( + shifted, shifted, /*nullAllowed=*/false); + EXPECT_TRUE(applyFilter( + shiftedFilter, + toVariant( + "2020-01-01 12:34:56", TIMESTAMP(), TimestampMode::kLocalTime))); + + const common::TimestampRange unshiftedFilter( + unshifted, unshifted, /*nullAllowed=*/false); + EXPECT_TRUE(applyFilter( + unshiftedFilter, + toVariant("2020-01-01 12:34:56", TIMESTAMP(), TimestampMode::kUtc))); +} + +} // namespace +} // namespace facebook::velox::connector::hive diff --git a/velox/connectors/hive/tests/TableHandleTest.cpp b/velox/connectors/hive/tests/TableHandleTest.cpp index c342add7f4e..c2e8693f3ba 100644 --- a/velox/connectors/hive/tests/TableHandleTest.cpp +++ b/velox/connectors/hive/tests/TableHandleTest.cpp @@ -101,6 +101,40 @@ TEST(TableHandleTest, hiveTableHandleDbName) { ASSERT_TRUE(cloneNoDb->dbName().empty()); } +TEST(TableHandleTest, hiveTableHandleDataColumnFieldIds) { + Type::registerSerDe(); + connector::hive::HiveTableHandle::registerSerDe(); + + const auto dataColumns = ROW({"a", "c"}, {BIGINT(), VARCHAR()}); + const std::vector fieldIds{1, 3}; + auto handle = std::make_shared( + "test-connector", + "test_table", + common::SubfieldFilters{}, + /*remainingFilter=*/nullptr, + dataColumns, + /*indexColumns=*/std::vector{}, + /*tableParameters=*/std::unordered_map{}, + /*filterColumnHandles=*/ + std::vector{}, + /*sampleRate=*/1.0, + /*dbName=*/"", + fieldIds); + + EXPECT_EQ(handle->dataColumnFieldIds(), fieldIds); + + auto serialized = handle->serialize(); + auto clone = ISerializable::deserialize( + serialized, /*context=*/nullptr); + EXPECT_EQ(clone->dataColumnFieldIds(), fieldIds); + + serialized.erase("dataColumnFieldIds"); + auto legacyClone = + ISerializable::deserialize( + serialized, /*context=*/nullptr); + EXPECT_TRUE(legacyClone->dataColumnFieldIds().empty()); +} + TEST(TableHandleTest, hiveTableHandleIndexSupport) { // Test HiveTableHandle without index columns. auto tableHandleWithoutIndex = diff --git a/velox/core/PlanNode.cpp b/velox/core/PlanNode.cpp index 18945fd4773..475f64ac730 100644 --- a/velox/core/PlanNode.cpp +++ b/velox/core/PlanNode.cpp @@ -408,6 +408,20 @@ std::vector deserializeStrings(const folly::dynamic& array) { return ISerializable::deserialize>(array); } +std::vector> deserializeOptionalStrings( + const folly::dynamic& array) { + std::vector> names; + names.reserve(array.size()); + for (const auto& name : array) { + if (name.isNull()) { + names.emplace_back(std::nullopt); + } else { + names.emplace_back(name.asString()); + } + } + return names; +} + RowTypePtr deserializeRowType(const folly::dynamic& obj) { return ISerializable::deserialize(obj); } @@ -1369,7 +1383,7 @@ UnnestNode::UnnestNode( const PlanNodeId& id, std::vector replicateVariables, std::vector unnestVariables, - std::vector unnestNames, + std::vector> unnestNames, std::optional ordinalityName, std::optional markerName, const PlanNodePtr& source) @@ -1387,7 +1401,7 @@ UnnestNode::UnnestNode( const PlanNodeId& id, std::vector replicateVariables, std::vector unnestVariables, - std::vector unnestNames, + std::vector> unnestNames, std::optional ordinalityName, std::optional markerName, std::optional splitOutput, @@ -1428,16 +1442,25 @@ UnnestNode::UnnestNode( int unnestIndex = 0; for (const auto& variable : unnestVariables_) { if (variable->type()->isArray()) { - names.emplace_back(unnestNames_[unnestIndex++]); - types.emplace_back(variable->type()->asArray().elementType()); + if (unnestNames_[unnestIndex].has_value()) { + names.emplace_back(unnestNames_[unnestIndex].value()); + types.emplace_back(variable->type()->asArray().elementType()); + } + ++unnestIndex; } else if (variable->type()->isMap()) { const auto& mapType = variable->type()->asMap(); - names.emplace_back(unnestNames_[unnestIndex++]); - types.emplace_back(mapType.keyType()); + if (unnestNames_[unnestIndex].has_value()) { + names.emplace_back(unnestNames_[unnestIndex].value()); + types.emplace_back(mapType.keyType()); + } + ++unnestIndex; - names.emplace_back(unnestNames_[unnestIndex++]); - types.emplace_back(mapType.valueType()); + if (unnestNames_[unnestIndex].has_value()) { + names.emplace_back(unnestNames_[unnestIndex].value()); + types.emplace_back(mapType.valueType()); + } + ++unnestIndex; } else { VELOX_FAIL( "Unexpected type of unnest variable. Expected ARRAY or MAP, but got {}.", @@ -1466,7 +1489,13 @@ folly::dynamic UnnestNode::serialize() const { auto obj = PlanNode::serialize(); obj["replicateVariables"] = ISerializable::serialize(replicateVariables_); obj["unnestVariables"] = ISerializable::serialize(unnestVariables_); - obj["unnestNames"] = ISerializable::serialize(unnestNames_); + folly::dynamic unnestNames = folly::dynamic::array; + for (const auto& name : unnestNames_) { + unnestNames.push_back( + name.has_value() ? folly::dynamic(name.value()) + : folly::dynamic(nullptr)); + } + obj["unnestNames"] = std::move(unnestNames); if (ordinalityName_.has_value()) { obj["ordinalityName"] = ordinalityName_.value(); @@ -1492,7 +1521,7 @@ PlanNodePtr UnnestNode::create(const folly::dynamic& obj, void* context) { auto replicateVariables = deserializeFields(obj["replicateVariables"], context); auto unnestVariables = deserializeFields(obj["unnestVariables"], context); - auto unnestNames = deserializeStrings(obj["unnestNames"]); + auto unnestNames = deserializeOptionalStrings(obj["unnestNames"]); std::optional ordinalityName = std::nullopt; if (obj.count("ordinalityName")) { ordinalityName = obj["ordinalityName"].asString(); @@ -2177,7 +2206,9 @@ bool NestedLoopJoinNode::isSupported(JoinType joinType) { case JoinType::kLeft: case JoinType::kRight: case JoinType::kFull: + case JoinType::kLeftSemiFilter: case JoinType::kLeftSemiProject: + case JoinType::kAnti: return true; default: @@ -3067,6 +3098,18 @@ void validateGroupingKeys( } } // namespace +InsertTableHandle::InsertTableHandle( + const std::string& connectorId, + const connector::ConnectorInsertTableHandlePtr& connectorInsertTableHandle, + folly::F14FastSet notNullColumns) + : connectorId_(connectorId), + connectorInsertTableHandle_(connectorInsertTableHandle), + notNullColumns_(std::move(notNullColumns)) { + for (const auto& name : notNullColumns_) { + VELOX_USER_CHECK(!name.empty(), "NOT NULL column name must not be empty"); + } +} + TableWriteNode::TableWriteNode( const PlanNodeId& id, const RowTypePtr& columns, @@ -3095,6 +3138,17 @@ TableWriteNode::TableWriteNode( "Column not found in TableWrite input: {}", column); } + const auto& notNullColumns = insertTableHandle_->notNullColumns(); + if (!notNullColumns.empty()) { + const folly::F14FastSet columnNameSet( + columnNames_.begin(), columnNames_.end()); + for (const auto& name : notNullColumns) { + VELOX_USER_CHECK( + columnNameSet.contains(name), + "NOT NULL column is not in the table schema: {}", + name); + } + } if (columnStatsSpec_.has_value()) { VELOX_USER_CHECK( columnStatsSpec_->aggregationStep == AggregationNode::Step::kSingle || @@ -3144,8 +3198,14 @@ void addStatsSpecDetails( } // namespace void TableWriteNode::addDetails(std::stringstream& stream) const { - stream << insertTableHandle_->connectorId() << ", " - << folly::join(", ", columnNames_); + stream << insertTableHandle_->connectorId(); + const auto& notNullColumns = insertTableHandle_->notNullColumns(); + for (const auto& columnName : columnNames_) { + stream << ", " << columnName; + if (notNullColumns.contains(columnName)) { + stream << " not null"; + } + } if (columnStatsSpec_.has_value()) { stream << ", "; addStatsSpecDetails(stream, columnStatsSpec_); @@ -3225,6 +3285,14 @@ folly::dynamic TableWriteNode::serialize() const { obj["outputType"] = outputType_->serialize(); obj["commitStrategy"] = std::string(connector::CommitStrategyName::toName(commitStrategy_)); + const auto& notNullColumns = insertTableHandle_->notNullColumns(); + if (!notNullColumns.empty()) { + // Sorted to keep the serialized form stable across runs. + std::vector sortedNotNullColumns( + notNullColumns.begin(), notNullColumns.end()); + std::sort(sortedNotNullColumns.begin(), sortedNotNullColumns.end()); + obj["notNullColumns"] = ISerializable::serialize(sortedNotNullColumns); + } return obj; } @@ -3252,13 +3320,19 @@ PlanNodePtr TableWriteNode::create(const folly::dynamic& obj, void* context) { if (obj.count("columnStatsSpec") != 0) { columnStatsSpec = ColumnStatsSpec::create(obj["columnStatsSpec"], context); } + folly::F14FastSet notNullColumns; + if (obj.count("notNullColumns") != 0) { + const auto names = ISerializable::deserialize>( + obj["notNullColumns"]); + notNullColumns.insert(names.begin(), names.end()); + } return std::make_shared( id, columns, columnNames, std::move(columnStatsSpec), std::make_shared( - connectorId, connectorInsertTableHandle), + connectorId, connectorInsertTableHandle, std::move(notNullColumns)), hasPartitioningScheme, outputType, commitStrategy, @@ -3440,6 +3514,7 @@ PartitionedOutputNode::PartitionedOutputNode( RowTypePtr outputType, std::string serdeKind, std::string transportKind, + std::string transportOptions, PlanNodePtr source) : PlanNode(id), kind_(kind), @@ -3450,6 +3525,7 @@ PartitionedOutputNode::PartitionedOutputNode( partitionFunctionSpec_(std::move(partitionFunctionSpec)), serdeKind_(std::move(serdeKind)), transportKind_(std::move(transportKind)), + transportOptions_(std::move(transportOptions)), outputType_(std::move(outputType)) { VELOX_USER_CHECK_GT(numPartitions_, 0); if (numPartitions_ == 1) { @@ -3603,6 +3679,7 @@ folly::dynamic PartitionedOutputNode::serialize() const { obj["partitionFunctionSpec"] = partitionFunctionSpec_->serialize(); obj["serdeKind"] = serdeKind_; obj["transportKind"] = transportKind_; + obj["transportOptions"] = transportOptions_; obj["outputType"] = outputType_->serialize(); return obj; } @@ -3629,6 +3706,7 @@ PlanNodePtr PartitionedOutputNode::create( obj["serdeKind"].asString(), obj.getDefault("transportKind", std::string{TransportKind::kInMemory}) .asString(), + obj.getDefault("transportOptions", "").asString(), deserializeSingleSource(obj, context)); } diff --git a/velox/core/PlanNode.h b/velox/core/PlanNode.h index b5ebb046c88..06cd8a54fdb 100644 --- a/velox/core/PlanNode.h +++ b/velox/core/PlanNode.h @@ -16,6 +16,7 @@ #pragma once #include +#include #include @@ -45,6 +46,9 @@ struct TransportKind { /// delivered is decided by the layer above -- read locally, fetched over /// HTTP, or written to a shuffle service. static constexpr std::string_view kInMemory{"in-memory"}; + /// Materialized output buffering backed by an application-provided durable + /// exchange implementation. + static constexpr std::string_view kMaterialized{"materialized"}; /// UCX-based RDMA exchange for high-bandwidth GPU transfers between workers. static constexpr std::string_view kUcx{"UCX"}; /// Deprecated source-compat alias for kInMemory; prefer kInMemory. @@ -54,12 +58,22 @@ struct TransportKind { /// Generic representation of InsertTable struct InsertTableHandle { public: + /// @param notNullColumns Throws a user error if any name is empty. + InsertTableHandle( + const std::string& connectorId, + const connector::ConnectorInsertTableHandlePtr& + connectorInsertTableHandle, + folly::F14FastSet notNullColumns); + +#ifdef VELOX_ENABLE_BACKWARD_COMPATIBILITY + /// Legacy constructor. Prefer the overload above, which takes the NOT NULL + /// columns. Removed once all callers have migrated. InsertTableHandle( const std::string& connectorId, const connector::ConnectorInsertTableHandlePtr& connectorInsertTableHandle) - : connectorId_(connectorId), - connectorInsertTableHandle_(connectorInsertTableHandle) {} + : InsertTableHandle(connectorId, connectorInsertTableHandle, {}) {} +#endif // VELOX_ENABLE_BACKWARD_COMPATIBILITY const std::string& connectorId() const { return connectorId_; @@ -70,12 +84,19 @@ struct InsertTableHandle { return connectorInsertTableHandle_; } + /// Target columns that must not contain nulls. Empty if unconstrained. + const folly::F14FastSet& notNullColumns() const { + return notNullColumns_; + } + private: // Connector ID const std::string connectorId_; // Write request to a DataSink of that connector type const connector::ConnectorInsertTableHandlePtr connectorInsertTableHandle_; + + const folly::F14FastSet notNullColumns_; }; class SortOrder { @@ -1556,7 +1577,8 @@ class TableWriteNode : public PlanNode { /// - grouping keys must be a subset of 'columns' (partition columns). /// - grouping keys must not contain duplicates. /// @param insertTableHandle Connector-specific handle identifying the - /// target table and write operation. + /// target table and write operation. Its notNullColumns() must be a subset + /// of 'columnNames'. /// @param hasPartitioningScheme Whether a partitioning scheme is configured /// for shuffles. Controls which query config determines the number of /// writer operator instances: 'task_partitioned_writer_count' if true, @@ -2726,8 +2748,33 @@ class PartitionedOutputNode : public PlanNode { RowTypePtr outputType, std::string serdeKind, std::string transportKind, + std::string transportOptions, PlanNodePtr source); + PartitionedOutputNode( + const PlanNodeId& id, + Kind kind, + const std::vector& keys, + int numPartitions, + bool replicateNullsAndAny, + PartitionFunctionSpecPtr partitionFunctionSpec, + RowTypePtr outputType, + std::string serdeKind, + std::string transportKind, + PlanNodePtr source) + : PartitionedOutputNode( + id, + kind, + keys, + numPartitions, + replicateNullsAndAny, + std::move(partitionFunctionSpec), + std::move(outputType), + std::move(serdeKind), + std::move(transportKind), + {}, + std::move(source)) {} + // Backward-compatible ctor without an explicit transport; defaults to the // in-memory transport. Prefer the ctor above. PartitionedOutputNode( @@ -2831,6 +2878,7 @@ class PartitionedOutputNode : public PlanNode { outputType_ = other.outputType(); serdeKind_ = other.serdeKind(); transportKind_ = other.transportKind(); + transportOptions_ = other.transportOptions(); VELOX_CHECK_EQ(other.sources().size(), 1); source_ = other.sources()[0]; } @@ -2880,6 +2928,11 @@ class PartitionedOutputNode : public PlanNode { return *this; } + Builder& transportOptions(std::string transportOptions) { + transportOptions_ = std::move(transportOptions); + return *this; + } + Builder& source(PlanNodePtr source) { source_ = std::move(source); return *this; @@ -2921,6 +2974,7 @@ class PartitionedOutputNode : public PlanNode { outputType_.value(), serdeKind_.value(), transportKind_.value(), + transportOptions_.value_or(std::string{}), source_.value()); } @@ -2934,6 +2988,7 @@ class PartitionedOutputNode : public PlanNode { std::optional outputType_; std::optional serdeKind_; std::optional transportKind_; + std::optional transportOptions_; std::optional source_; }; @@ -2985,6 +3040,11 @@ class PartitionedOutputNode : public PlanNode { return transportKind_; } + /// Opaque configuration interpreted by the selected output transport. + const std::string& transportOptions() const { + return transportOptions_; + } + /// Returns true if an arbitrary row and all rows with null keys must be /// replicated to all destinations. This is used to ensure correct results /// for anti-join which requires all nodes to know whether combined build @@ -3021,6 +3081,7 @@ class PartitionedOutputNode : public PlanNode { const PartitionFunctionSpecPtr partitionFunctionSpec_; const std::string serdeKind_; const std::string transportKind_; + const std::string transportOptions_; const RowTypePtr outputType_; }; @@ -4822,7 +4883,9 @@ class UnnestNode : public PlanNode { /// or MAP. /// @param unnestNames Names to use for unnested outputs: one name for each /// array (element); two names for each map (key and value). The output - /// names must appear in the same order as unnestVariables. + /// names must appear in the same order as unnestVariables. A std::nullopt + /// entry prunes the corresponding output column (not emitted, not + /// materialized). /// @param ordinalityName Optional name for the ordinality columns. If not /// present, ordinality column is not produced. /// @param markerName Optional name for column which indicates whether an @@ -4839,7 +4902,7 @@ class UnnestNode : public PlanNode { const PlanNodeId& id, std::vector replicateVariables, std::vector unnestVariables, - std::vector unnestNames, + std::vector> unnestNames, std::optional ordinalityName, std::optional markerName, const PlanNodePtr& source); @@ -4848,12 +4911,35 @@ class UnnestNode : public PlanNode { const PlanNodeId& id, std::vector replicateVariables, std::vector unnestVariables, - std::vector unnestNames, + std::vector> unnestNames, std::optional ordinalityName, std::optional markerName, std::optional splitOutput, const PlanNodePtr& source); +#ifdef VELOX_ENABLE_BACKWARD_COMPATIBILITY + /// Deprecated. Use the std::vector> overload. + UnnestNode( + const PlanNodeId& id, + std::vector replicateVariables, + std::vector unnestVariables, + std::vector unnestNames, + std::optional ordinalityName, + std::optional markerName, + const PlanNodePtr& source) + : UnnestNode( + id, + std::move(replicateVariables), + std::move(unnestVariables), + std::vector>( + unnestNames.begin(), + unnestNames.end()), + std::move(ordinalityName), + std::move(markerName), + std::nullopt, + source) {} +#endif + class Builder { public: Builder() = default; @@ -4887,7 +4973,7 @@ class UnnestNode : public PlanNode { return *this; } - Builder& unnestNames(std::vector unnestNames) { + Builder& unnestNames(std::vector> unnestNames) { unnestNames_ = std::move(unnestNames); return *this; } @@ -4939,7 +5025,7 @@ class UnnestNode : public PlanNode { std::optional id_; std::optional> replicateVariables_; std::optional> unnestVariables_; - std::optional> unnestNames_; + std::optional>> unnestNames_; std::optional ordinalityName_; std::optional markerName_; std::optional source_; @@ -4972,7 +5058,7 @@ class UnnestNode : public PlanNode { return unnestVariables_; } - const std::vector& unnestNames() const { + const std::vector>& unnestNames() const { return unnestNames_; } @@ -5009,7 +5095,7 @@ class UnnestNode : public PlanNode { const std::vector replicateVariables_; const std::vector unnestVariables_; - const std::vector unnestNames_; + const std::vector> unnestNames_; const std::optional ordinalityName_; const std::optional markerName_; const std::optional splitOutput_; @@ -5952,7 +6038,8 @@ using MarkSortedNodePtr = std::shared_ptr; /// Optimized version of a WindowNode for a single row_number, rank or /// dense_rank function with a limit over sorted partitions. The output of this /// node contains all input columns followed by an optional -/// 'rowNumberColumnName' BIGINT column. +/// 'rowNumberColumnName' BIGINT column, with rows within each partition emitted +/// in ascending order of sorting keys (matching WindowNode). /// TODO: This node will be renamed to TopNRank or TopNRowNode once all the /// support for handling rank and dense_rank is committed to Velox. class TopNRowNumberNode : public PlanNode { @@ -6484,45 +6571,6 @@ class RPCNode : public PlanNode { rpc::RPCStreamingMode streamingMode = rpc::RPCStreamingMode::kPerRow, int32_t dispatchBatchSize = 0); -#ifdef VELOX_ENABLE_BACKWARD_COMPATIBILITY - /// Legacy constructor. Prefer the CallTypedExpr constructor above. - /// - /// Accepts the flattened call fields and builds the CallTypedExpr internally: - /// each argument becomes a FieldAccessTypedExpr referencing - /// argumentColumns[i] (type argumentTypes[i]), except positions with a - /// non-null constantInputs[i], which become a ConstantTypedExpr wrapping that - /// constant vector. Defined inline (header-only) so it compiles in the - /// read-only-synced Prestissimo build, whose Buck targets define - /// VELOX_ENABLE_BACKWARD_COMPATIBILITY; velox and open-source builds never - /// do. Removed in the CONTRACT step once all callers use the CallTypedExpr - /// constructor. - RPCNode( - const PlanNodeId& id, - PlanNodePtr source, - std::string functionName, - TypePtr functionResultType, - std::string outputColumn, - RowTypePtr outputType, - std::vector argumentColumns, - std::vector argumentTypes, - std::vector constantInputs, - rpc::RPCStreamingMode streamingMode = rpc::RPCStreamingMode::kPerRow, - int32_t dispatchBatchSize = 0) - : RPCNode( - id, - std::move(source), - rpcCallFromLegacyFields( - std::move(functionName), - std::move(functionResultType), - argumentColumns, - argumentTypes, - constantInputs), - std::move(outputColumn), - std::move(outputType), - streamingMode, - dispatchBatchSize) {} -#endif // VELOX_ENABLE_BACKWARD_COMPATIBILITY - const PlanNodePtr& source() const { return sources_[0]; } @@ -6540,48 +6588,6 @@ class RPCNode : public PlanNode { return call_->type(); } -#ifdef VELOX_ENABLE_BACKWARD_COMPATIBILITY - /// Legacy accessors over the folded call, each derived from call()->inputs(). - /// Retained so pre-migration callers (e.g. the presto-cpp conversion test) - /// keep compiling; removed in the CONTRACT step once every caller uses - /// call()->inputs() directly. argumentColumns() yields the FieldAccess name - /// for column arguments and an empty string for constants; constantInputs() - /// yields the constant vector for constant arguments and nullptr for columns. - /// Defined inline (header-only) for the read-only-synced Prestissimo build. - std::vector argumentColumns() const { - std::vector columns; - columns.reserve(call_->inputs().size()); - for (const auto& input : call_->inputs()) { - if (auto* field = input->asUnchecked()) { - columns.push_back(field->name()); - } else { - columns.emplace_back(); - } - } - return columns; - } - std::vector argumentTypes() const { - std::vector types; - types.reserve(call_->inputs().size()); - for (const auto& input : call_->inputs()) { - types.push_back(input->type()); - } - return types; - } - std::vector constantInputs() const { - std::vector constants; - constants.reserve(call_->inputs().size()); - for (const auto& input : call_->inputs()) { - if (auto* constant = input->asUnchecked()) { - constants.push_back(constant->valueVector()); - } else { - constants.push_back(nullptr); - } - } - return constants; - } -#endif // VELOX_ENABLE_BACKWARD_COMPATIBILITY - const std::string& outputColumn() const { return outputColumn_; } @@ -6611,45 +6617,6 @@ class RPCNode : public PlanNode { static PlanNodePtr create(const folly::dynamic& obj, void* context); private: -#ifdef VELOX_ENABLE_BACKWARD_COMPATIBILITY - // Builds the RPC CallTypedExpr from the legacy flattened call fields, for the - // legacy constructor above. Each argument becomes a FieldAccessTypedExpr on - // argumentColumns[i], except positions with a non-null constantInputs[i], - // which become a ConstantTypedExpr. Header-only for the read-only-synced - // Prestissimo build; removed in the CONTRACT step. - static core::CallTypedExprPtr rpcCallFromLegacyFields( - std::string functionName, - TypePtr functionResultType, - const std::vector& argumentColumns, - const std::vector& argumentTypes, - const std::vector& constantInputs) { - VELOX_CHECK_EQ( - argumentColumns.size(), - argumentTypes.size(), - "RPCNode argumentColumns and argumentTypes must have the same size"); - VELOX_CHECK_EQ( - argumentColumns.size(), - constantInputs.size(), - "RPCNode argumentColumns and constantInputs must have the same size"); - std::vector callInputs; - callInputs.reserve(argumentColumns.size()); - for (size_t i = 0; i < argumentColumns.size(); ++i) { - if (constantInputs[i] != nullptr) { - callInputs.push_back( - std::make_shared(constantInputs[i])); - } else { - callInputs.push_back( - std::make_shared( - argumentTypes[i], argumentColumns[i])); - } - } - return std::make_shared( - std::move(functionResultType), - std::move(callInputs), - std::move(functionName)); - } -#endif // VELOX_ENABLE_BACKWARD_COMPATIBILITY - void addDetails(std::stringstream& stream) const override; std::vector sources_; diff --git a/velox/core/QueryConfig.cpp b/velox/core/QueryConfig.cpp index 80a5b5c934b..a522a1e7db5 100644 --- a/velox/core/QueryConfig.cpp +++ b/velox/core/QueryConfig.cpp @@ -155,6 +155,8 @@ const std::vector& QueryConfig::registeredProperties() { VELOX_REGISTER_QUERY_CONFIG(kHashProbeDynamicFilterPushdownEnabled); VELOX_REGISTER_QUERY_CONFIG(kHashProbeStringDynamicFilterPushdownEnabled); VELOX_REGISTER_QUERY_CONFIG(kHashProbeBloomFilterPushdownMaxSize); + VELOX_REGISTER_QUERY_CONFIG(kBypassHashProbeBloomFilterMinRows); + VELOX_REGISTER_QUERY_CONFIG(kBypassHashProbeBloomFilterMinPct); VELOX_REGISTER_QUERY_CONFIG(kMinTableRowsForParallelJoinBuild); // Debug and validation. diff --git a/velox/core/QueryConfig.h b/velox/core/QueryConfig.h index 92d7f19f1c8..dd8bc3969e2 100644 --- a/velox/core/QueryConfig.h +++ b/velox/core/QueryConfig.h @@ -989,6 +989,30 @@ class QueryConfig { 0, "Maximum byte size of Bloom filter from hash probe. 0 disables.") + /// Number of probe rows used to decide whether to bypass the build-side + /// Bloom filter for left joins and non-null-aware left semi-project and left + /// anti joins. 0 disables local Bloom filter probing. + VELOX_QUERY_CONFIG( + kBypassHashProbeBloomFilterMinRows, + bypassHashProbeBloomFilterMinRows, + "bypass_hash_probe_bloom_filter_min_rows", + int32_t, + 0, + "Number of probe rows used to decide whether to bypass the build-side " + "Bloom filter for left joins and non-null-aware left semi-project and " + "left anti joins. 0 disables local Bloom filter probing.") + + /// Bypass the build-side Bloom filter if its acceptance percentage meets + /// or exceeds this value. 0 bypasses the Bloom filter without sampling. + VELOX_QUERY_CONFIG( + kBypassHashProbeBloomFilterMinPct, + bypassHashProbeBloomFilterMinPct, + "bypass_hash_probe_bloom_filter_min_pct", + int32_t, + 85, + "Bypass the build-side Bloom filter if its acceptance percentage meets " + "or exceeds this value. 0 bypasses the Bloom filter without sampling.") + /// The minimum number of table rows that can trigger the parallel hash join /// table build. VELOX_QUERY_CONFIG( @@ -1547,46 +1571,46 @@ class QueryConfig { "admission-controlled dispatch this ceiling now actually bounds in-flight " "rows, so it must be sized for the backend's healthy concurrency.") - /// Enables the adaptive per-tier RPC rate limiter (RPCRateLimiter). + /// Enables AIMD adaptation of each backend's rate-limit capacity. VELOX_QUERY_CONFIG( kRpcRateLimiterAdaptiveEnabled, rpcRateLimiterAdaptiveEnabled, "rpc.ratelimiter.adaptive_enabled", bool, true, - "When true (default), the process-global per-tier RPC rate limiter adapts " - "its max-pending cap via AIMD driven by the backend overload signal " + "When true (default), each backend's rate limiter adapts its capacity " + "via AIMD driven by the backend overload signal " "(rate-limit/timeout): multiplicative-decrease on an overload-classified " "drain, additive-increase on a clean drain. On by default because it is " "the protective behavior for shared, rate-limited inference backends; set " "false to keep a static cap. Unlike the per-driver congestion window, this " "coordinates all drivers on the worker and reacts to the rate-limit signal " - "directly, not to RTT.") + "directly, not to RTT. The first query to reach a backend fixes its policy for the life of the worker process; later queries contribute their outcomes to the adaptation but cannot change the setting. ") - /// Floor for the adaptive per-tier RPC rate limiter's max-pending cap. + /// Floor the adaptive rate-limit capacity may shrink to. VELOX_QUERY_CONFIG( kRpcRateLimiterMinLimit, rpcRateLimiterMinLimit, "rpc.ratelimiter.min_limit", int64_t, 50, - "Floor the adaptive RPC rate limiter's per-tier max-pending cap may " + "Floor that a backend's adaptive rate-limit capacity may " "shrink to under sustained overload. Default 50 (a floor of 1 can stall a " "workload under sustained throttling). Only used when " - "rpc.ratelimiter.adaptive_enabled is true.") + "rpc.ratelimiter.adaptive_enabled is true. The first query to reach a backend fixes its policy for the life of the worker process; later queries contribute their outcomes to the adaptation but cannot change the setting. ") - /// Multiplicative-decrease factor for the adaptive RPC rate limiter. + /// Multiplicative-decrease factor for the adaptive rate-limit capacity. VELOX_QUERY_CONFIG( kRpcRateLimiterDecreaseFactor, rpcRateLimiterDecreaseFactor, "rpc.ratelimiter.decrease_factor", double, 0.5, - "Factor applied to the adaptive RPC rate limiter's per-tier max-pending " - "cap on each overload-classified drain. Default 0.5 (halve). Clamped to " - "(0, 1). Only used when rpc.ratelimiter.adaptive_enabled is true.") + "Factor applied to a backend's adaptive rate-limit capacity " + "on each overload-classified drain. Default 0.5 (halve). Clamped to " + "(0, 1). Only used when rpc.ratelimiter.adaptive_enabled is true. The first query to reach a backend fixes its policy for the life of the worker process; later queries contribute their outcomes to the adaptation but cannot change the setting. ") - /// Ceiling for the per-tier RPC rate-limiter max-pending cap. + /// Ceiling for a backend's rate-limit capacity. VELOX_QUERY_CONFIG( kRpcRateLimiterMaxLimit, rpcRateLimiterMaxLimit, @@ -1594,11 +1618,16 @@ class QueryConfig { int64_t, 200, "Ceiling (and, with adaptive enabled, the starting value) for the " - "process-global per-tier RPC rate-limiter max-pending cap. Default 200 " + "per-backend rate-limit capacity, shared across drivers. Default 200 " "(validated for LLM-inference backends); 0 falls back to the built-in 20. " - "With admission-controlled dispatch this cap actually bounds process-wide " - "in-flight rows per tier; the adaptive limiter shrinks from here toward " - "rpc.ratelimiter.min_limit under overload.") + "With admission-controlled dispatch this cap bounds in-flight work " + "against that backend across every driver on the worker; the adaptive " + "limiter shrinks from here toward rpc.ratelimiter.min_limit under " + "overload. Any positive value here overrides a ceiling the function " + "asked for through its own options; set 0 to defer to that. The first " + "query to reach a backend fixes its policy for the life of the worker " + "process; later queries contribute their outcomes to the adaptation but " + "cannot change the setting.") // --- Hand-written accessors for properties that need custom logic --- diff --git a/velox/core/tests/PlanNodeBuilderTest.cpp b/velox/core/tests/PlanNodeBuilderTest.cpp index 17d2c9ca1e6..fddc420a7b9 100644 --- a/velox/core/tests/PlanNodeBuilderTest.cpp +++ b/velox/core/tests/PlanNodeBuilderTest.cpp @@ -370,8 +370,10 @@ TEST_F(PlanNodeBuilderTest, tableWriteNode) { std::vector{"sum(c0)"}); const auto outputType = TableWriteTraits::outputType(statsSpec); - const auto insertTableHandle = - std::make_shared("connector_id", nullptr); + const auto insertTableHandle = std::make_shared( + "connector_id", + nullptr, + /*notNullColumns=*/folly::F14FastSet{}); const auto verify = [&](const std::shared_ptr& node) { EXPECT_EQ(node->id(), id); @@ -402,6 +404,24 @@ TEST_F(PlanNodeBuilderTest, tableWriteNode) { verify(node2); } +TEST_F(PlanNodeBuilderTest, tableWriteNodeNotNullColumnOutsideSchema) { + const auto insertTableHandle = std::make_shared( + "connector_id", nullptr, folly::F14FastSet{"c1"}); + + VELOX_ASSERT_USER_THROW( + TableWriteNode::Builder() + .id("test_id") + .columns(ROW({"c0"}, {INTEGER()})) + .columnNames({"c0"}) + .insertTableHandle(insertTableHandle) + .hasPartitioningScheme(false) + .outputType(TableWriteTraits::outputType(std::nullopt)) + .commitStrategy(connector::CommitStrategy::kNoCommit) + .source(source_) + .build(), + "NOT NULL column is not in the table schema: c1"); +} + TEST_F(PlanNodeBuilderTest, tableWriteMergeNode) { const PlanNodeId id = "table_write_merge_node_id"; @@ -614,6 +634,7 @@ TEST_F(PlanNodeBuilderTest, partitionedOutputNode) { std::make_shared(); const RowTypePtr outputType = ROW({"c0"}, {BIGINT()}); const auto serdeKind = "Presto"; + const std::string transportOptions = R"({"exchangeId":"test"})"; const auto verify = [&](const std::shared_ptr& node) { @@ -624,6 +645,7 @@ TEST_F(PlanNodeBuilderTest, partitionedOutputNode) { EXPECT_EQ(node->isReplicateNullsAndAny(), replicateNullsAndAny); EXPECT_EQ(node->outputType(), outputType); EXPECT_EQ(node->serdeKind(), serdeKind); + EXPECT_EQ(node->transportOptions(), transportOptions); EXPECT_EQ(node->partitionFunctionSpecPtr(), partitionFunctionSpec); EXPECT_EQ(node->sources(), std::vector{source_}); }; @@ -638,12 +660,15 @@ TEST_F(PlanNodeBuilderTest, partitionedOutputNode) { .outputType(outputType) .serdeKind(serdeKind) .transportKind(std::string{TransportKind::kInMemory}) + .transportOptions(transportOptions) .source(source_) .build(); verify(node); const auto node2 = PartitionedOutputNode::Builder(*node).build(); verify(node2); + + EXPECT_EQ(node->serialize()["transportOptions"], transportOptions); } TEST_F(PlanNodeBuilderTest, hashJoinNode) { @@ -1005,7 +1030,7 @@ TEST_F(PlanNodeBuilderTest, unnestNode) { std::make_shared(BIGINT(), "a")}; std::vector unnestVariables{ std::make_shared(ARRAY(BIGINT()), "b")}; - std::vector unnestNames{"b"}; + std::vector> unnestNames{"b"}; std::optional ordinalityName = std::make_optional("ord"); std::optional splitOutput = false; @@ -1028,7 +1053,7 @@ TEST_F(PlanNodeBuilderTest, unnestNode) { expectedNames.push_back(variable->name()); } for (const auto& name : unnestNames) { - expectedNames.push_back(name); + expectedNames.push_back(name.value()); } if (ordinalityName.has_value()) { expectedNames.push_back(ordinalityName.value()); diff --git a/velox/core/tests/StringTest.cpp b/velox/core/tests/StringTest.cpp index 118ecbbc457..c1551f63ee5 100644 --- a/velox/core/tests/StringTest.cpp +++ b/velox/core/tests/StringTest.cpp @@ -27,7 +27,7 @@ StringWriter createStringWriter(const std::string& str) { } } // namespace -TEST(String, StringWriter) { +TEST(String, stringWriter) { // Default constructor & empty string copy. { const std::string emptyString; diff --git a/velox/docs/configs.rst b/velox/docs/configs.rst index 2e076cbd307..c097d7ada48 100644 --- a/velox/docs/configs.rst +++ b/velox/docs/configs.rst @@ -173,6 +173,18 @@ Generic Configuration probe. When set to 0, no Bloom filter will be generated. To achieve optimal performance, this should not be too larger than the CPU cache size on the host. + * - bypass_hash_probe_bloom_filter_min_rows + - integer + - 0 + - The number of probe rows used to decide whether to bypass the build-side + Bloom filter for left joins and non-null-aware left semi-project and left + anti joins. When set to 0, local Bloom filter probing is disabled. + * - bypass_hash_probe_bloom_filter_min_pct + - integer + - 85 + - Bypass the build-side Bloom filter if its acceptance percentage meets + or exceeds this value. When set to 0, the Bloom filter is bypassed + without sampling. * - debug.validate_output_from_operators - bool - false @@ -947,9 +959,12 @@ Common Options * - ``use-column-names`` - bool - false - - Map table fields to file fields using names instead of indices for all - file formats. The connector property is scoped by connector ID, for - example ``hive.use-column-names`` or ``iceberg.use-column-names``. + - Map table fields to file fields using names instead of indices. This is + a connector/session-level default column matching policy and applies + uniformly to all file formats read by the connector, for example ORC, + DWRF, and Parquet. Split-specific column mapping mode, when present, + takes precedence over this setting. The connector property is scoped by + connector ID, for example ``hive.use-column-names`` or ``iceberg.use-column-names``. Session: ``use_column_names``. ORC Options (prefix ``hive.orc.``) @@ -1058,7 +1073,7 @@ Parquet Options (prefix ``hive.parquet.``) silently over-consuming. The reservation shrinks as row groups are skipped by filterRowGroups and is released in full when the reader is destroyed. When tracking engages, the estimate is also surfaced - per scan via the runtime stat ``parquetFooterEstimatedBytes`` so + per scan via the runtime stat ``parquet.footerEstimatedBytes`` so operators can compare it against actual pool usage. Session: ``parquet_footer_memory_tracking_threshold``. * - ``writer.max-target-file-size`` - capacity @@ -1070,12 +1085,10 @@ Parquet Options (prefix ``hive.parquet.``) Session: ``parquet_writer_max_target_file_size``. - Row-group sizing is independent of this setting and is not user-configurable: a row group is - flushed at a 128MB byte target or 1,048,576 rows, whichever comes first. The byte target is - soft - a row group may slightly exceed it, since the writer flushes only after buffered bytes - reach the target; the row count is a hard cap. When ``writer.max-target-file-size`` is set, - the writer may flush the current row group early so the accumulated file size is visible and - rotation can occur. + Row-group sizing is independent of this setting. The soft byte target is configured by + ``writer.row-group-size``, while the row count is capped at 1,048,576. When + ``writer.max-target-file-size`` is set, the writer may flush the current row group early so + the accumulated file size is visible and rotation can occur. * - ``writer.enable-dictionary`` - bool - true @@ -1105,6 +1118,12 @@ Parquet Options (prefix ``hive.parquet.``) - integer - 1024 - Batch size used when writing into Parquet through Arrow bridge. Session: ``hive.parquet.writer.batch_size``. + * - ``writer.row-group-size`` + - string + - 128MB + - Soft target for the serialized row group size. The estimate includes compressed bytes for + encoded pages and estimated bytes for data that has not yet been serialized into pages, so a + row group may slightly exceed the target. Session: ``parquet_writer_row_group_size``. * - ``writer.created-by`` - string - parquet-cpp-velox version 0.0.0 @@ -1115,6 +1134,13 @@ Parquet Options (prefix ``hive.parquet.``) - Whether to store DECIMAL values using integer physical types (INT32/INT64) when precision allows. When false, all DECIMAL values are stored as FIXED_LEN_BYTE_ARRAY regardless of precision. Session: ``hive.parquet.writer.enable_store_decimal_as_integer``. + * - ``writer.enable-page-index`` + - bool + - false + - Whether to write the Parquet page index (column index and offset index) when writing into + Parquet through the Arrow bridge. When enabled, per-page statistics are stored in the page + index instead of the data page headers, letting readers skip pages that cannot match a filter. + Session: ``hive.parquet.writer.enable_page_index``. Nimble Options (prefix ``hive.nimble.``) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -1162,6 +1188,12 @@ Nimble Options (prefix ``hive.nimble.``) ``Amazon S3 Configuration`` ^^^^^^^^^^^^^^^^^^^^^^^^^^^ +The S3 configuration keys below use the ``s3.`` prefix and are shared by all +connectors that read from or write to S3, such as Hive, Iceberg, and Paimon. +The legacy ``hive.s3.`` prefix is still accepted as a deprecated fallback; a +key set with the ``s3.`` prefix takes precedence over the same key set with the +``hive.s3.`` prefix. New configurations should use the ``s3.`` prefix. + .. list-table:: :widths: 30 10 10 70 :header-rows: 1 @@ -1170,98 +1202,98 @@ Nimble Options (prefix ``hive.nimble.``) - Type - Default Value - Description - * - hive.s3.use-instance-credentials + * - ``s3.use-instance-credentials`` - bool - false - Use the EC2 metadata service to retrieve API credentials. This works with IAM roles in EC2. - * - hive.s3.aws-access-key + * - ``s3.aws-access-key`` - string - - Default AWS access key to use. - * - hive.s3.aws-secret-key + * - ``s3.aws-secret-key`` - string - - Default AWS secret key to use. - * - hive.s3.endpoint + * - ``s3.endpoint`` - string - - The S3 storage endpoint server. This can be used to connect to an S3-compatible storage system instead of AWS. - * - hive.s3.endpoint.region + * - ``s3.endpoint.region`` - string - us-east-1 - The S3 storage endpoint server region. Default is set by the AWS SDK. If not configured, region will be attempted - to be parsed from the hive.s3.endpoint value. - * - hive.s3.path-style-access + to be parsed from the ``s3.endpoint`` value. + * - ``s3.path-style-access`` - bool - false - Use path-style access for all requests to the S3-compatible storage. This is for S3-compatible storage that doesn't support virtual-hosted-style access. - * - hive.s3.ssl.enabled + * - ``s3.ssl.enabled`` - bool - true - Use HTTPS to communicate with the S3 API. - * - hive.s3.log-level + * - ``s3.log-level`` - string - FATAL - **Allowed values:** "OFF", "FATAL", "ERROR", "WARN", "INFO", "DEBUG", "TRACE". Granularity of logging generated by the AWS C++ SDK library. - * - hive.s3.log-location + * - ``s3.log-location`` - string - "" - Specifies the path where the log files are created. Generated log files start with "aws_sdk\_" and use the default AWS S3 logger. Example: setting "/tmp" results in files "/tmp/aws_sdk_*". - * - hive.s3.payload-signing-policy + * - ``s3.payload-signing-policy`` - string - Never - **Allowed values:** "Always", "RequestDependent", "Never". When set to "Always", the payload checksum is included in the signature calculation. When set to "RequestDependent", the payload checksum is included based on the value returned by "AmazonWebServiceRequest::SignBody()". - * - hive.s3.iam-role + * - ``s3.iam-role`` - string - - IAM role to assume. - * - hive.s3.iam-role-session-name + * - ``s3.iam-role-session-name`` - string - velox-session - Session name associated with the IAM role. - * - hive.s3.use-proxy-from-env + * - ``s3.use-proxy-from-env`` - bool - false - Utilize the configuration of the environment variables http_proxy, https_proxy, and no_proxy for use with the S3 API. - * - hive.s3.connect-timeout + * - ``s3.connect-timeout`` - string - - Socket connect timeout. - * - hive.s3.socket-timeout + * - ``s3.socket-timeout`` - string - - Socket read timeout. - * - hive.s3.max-connections + * - ``s3.max-connections`` - integer - - Maximum concurrent TCP connections for a single http client. - * - hive.s3.max-attempts + * - ``s3.max-attempts`` - integer - - Maximum attempts for connections to a single http client, work together with retry-mode. By default, it's 3 for standard/adaptive mode and 10 for legacy mode. - * - hive.s3.retry-mode + * - ``s3.retry-mode`` - string - - **Allowed values:** "standard", "adaptive", "legacy". By default it's empty, S3 client will be created with RetryStrategy. Legacy mode only enables throttled retry for transient errors. Standard mode is built on top of legacy mode and has throttled retry enabled for throttling errors apart from transient errors. Adaptive retry mode dynamically limits the rate of AWS requests to maximize success rate. - * - hive.s3.aws-credentials-provider + * - ``s3.aws-credentials-provider`` - string - - A custom credential provider, if specified, will be used to create the client in favor of other authentication mechanisms. The provider must be registered using "registerAWSCredentialsProvider" before it can be used. - * - hive.s3.aws-imds-enabled + * - ``s3.aws-imds-enabled`` - bool - true - AWS Instance Metadata Service (IMDS) is an AWS EC2 instance component used by applications to securely access metadata. We must disable it on other instances to avoid high first-time read latency from S3 compatible object storages. - * - hive.s3.min-part-size + * - ``s3.min-part-size`` - string - 10MB - Minimum multi-part upload part size. The smallest allowed value is 5MB. The largest allowed value is 5GB. @@ -1272,10 +1304,14 @@ Nimble Options (prefix ``hive.nimble.``) Bucket Level Configuration """""""""""""""""""""""""" -All "hive.s3.*" config (except "hive.s3.log-level") can be set on a per-bucket basis. The bucket-specific option is set by -replacing the "hive.s3." prefix on a config with "hive.s3.bucket.BUCKETNAME.", where BUCKETNAME is the name of the -bucket. e.g. the endpoint for a bucket named "velox" can be specified by the config "hive.s3.bucket.velox.endpoint". -When connecting to a bucket, all options explicitly set will override the base "hive.s3." values. +All ``s3.*`` config except ``s3.log-level`` can be set on a per-bucket basis. +The bucket-specific option is set by replacing the ``s3.`` prefix on a config +with ``s3.bucket.BUCKETNAME.``, where BUCKETNAME is the name of the bucket. For +example, the endpoint for a bucket named "velox" can be specified by the config +``s3.bucket.velox.endpoint``. When connecting to a bucket, all options +explicitly set will override the base ``s3.`` values. The deprecated +``hive.s3.`` prefix follows the same rules (for example +``hive.s3.bucket.velox.endpoint``). These semantics are similar to the `Apache Hadoop-Aws module `_. ``Google Cloud Storage Configuration`` @@ -1440,6 +1476,13 @@ Spark-specific Configuration - bool - true - If true, Spark ``collect_list`` aggregate function ignores nulls in the input. + * - spark.decimal_to_float_high_precision_cast_enabled + - bool + - false + - If true, casts from ``DECIMAL`` to ``REAL``/``DOUBLE`` use a high-precision conversion + (via an intermediate string) for values that cannot be represented exactly by floating + point arithmetic, aligning the result with Spark. Disabled by default due to the + significant performance regression; users sensitive to precision loss can enable it. Tracing -------- diff --git a/velox/docs/develop/aggregations.rst b/velox/docs/develop/aggregations.rst index 46c04818914..59747bc6947 100644 --- a/velox/docs/develop/aggregations.rst +++ b/velox/docs/develop/aggregations.rst @@ -367,3 +367,7 @@ Many aggregate functions implement toIntermediate() fast path. Some examples inc Runtime statistic `abandonedPartialAggregationRows` counts rows that bypassed partial aggregation after it was abandoned. A value greater than 0 indicates that partial aggregation was abandoned. + +Runtime statistic `toIntermediateFastPathCalls` counts calls to aggregate +functions that use the `toIntermediate()` fast path after partial aggregation +is abandoned. diff --git a/velox/docs/develop/connectors.rst b/velox/docs/develop/connectors.rst index b4c4275cdb8..5a618b8e6f5 100644 --- a/velox/docs/develop/connectors.rst +++ b/velox/docs/develop/connectors.rst @@ -107,9 +107,8 @@ By default, the C++ AWS S3 client does not honor the configuration of the environment variables http_proxy, https_proxy, and no_proxy. The Java AWS S3 client supports this. The environment variables can be specified as lower case, upper case or both. -In order to enable the use of a proxy the hive connector configuration variable -`hive.s3.use-proxy-from-env` must be set to `true`. By default, the value -is `false`. +In order to enable the use of a proxy, set the S3 configuration variable +``s3.use-proxy-from-env`` to ``true``. By default, the value is ``false``. This is the behavior when the proxy settings are enabled: diff --git a/velox/docs/develop/operators.rst b/velox/docs/develop/operators.rst index 76a971357a4..2e7ae6a3cbe 100644 --- a/velox/docs/develop/operators.rst +++ b/velox/docs/develop/operators.rst @@ -678,7 +678,7 @@ output rows with empty unnest values are not produced. * - unnestVariables - Input columns of type array or map to expand. * - unnestNames - - Names to use for expanded columns. One name per array column. Two names per map column. + - Names to use for expanded columns. One name per array column. Two names per map column. A name may be absent (null) to prune the corresponding expanded column, which is then neither emitted nor materialized. * - ordinalityName - Optional name for the ordinality column. * - emptyUnnestValueName @@ -709,7 +709,7 @@ the written file paths on storage and the collected column stats. * - aggregationNode - Optional Aggregation plan node used to collect column stats for the data written to storage. * - insertTableHandle - - Connector-specific description of the destination table. + - Connector-specific description of the destination table. Its notNullColumns is a subset of columnNames; writing a null into one of them fails the query. * - outputType - A list of output columns containing the metadata of the data written storage. diff --git a/velox/docs/functions/presto/binary.rst b/velox/docs/functions/presto/binary.rst index 1ce77a306f9..4b882105b56 100644 --- a/velox/docs/functions/presto/binary.rst +++ b/velox/docs/functions/presto/binary.rst @@ -176,3 +176,13 @@ Binary Functions :noindex: Computes the xxhash64 hash of ``binary`` with ``bigint`` seed. + +.. function:: xxhash128(binary) -> varbinary + + Computes the XXH3 128-bit hash of ``binary``, returned as the 16-byte + big-endian canonical representation. + +.. function:: xxhash128(binary, bigint) -> varbinary + :noindex: + + Computes the XXH3 128-bit hash of ``binary`` with ``bigint`` seed. diff --git a/velox/docs/functions/spark/conversion.rst b/velox/docs/functions/spark/conversion.rst index 02960fb77f9..469c2f9536c 100644 --- a/velox/docs/functions/spark/conversion.rst +++ b/velox/docs/functions/spark/conversion.rst @@ -53,8 +53,11 @@ Integral types include bigint, integer, smallint, and tinyint. From integral types ^^^^^^^^^^^^^^^^^^^ +*(ANSI compliant)* + Casting one integral type to another is allowed. When the input value exceeds the range of result type, -a value of the result type is created forcedly with the input value. +a value of the result type is created forcedly with the input value when ANSI mode is disabled; +throws an error when ANSI mode is enabled. Valid examples: @@ -62,14 +65,17 @@ Valid examples: SELECT cast(1234567 as bigint); -- 1234567 SELECT cast(12 as tinyint); -- 12 - SELECT cast(1234 as tinyint); -- -46 - SELECT cast(1234567 as smallint); -- -10617 + SELECT cast(1234 as tinyint); -- -46 (ANSI OFF) / ERROR (ANSI ON) + SELECT cast(1234567 as smallint); -- -10617 (ANSI OFF) / ERROR (ANSI ON) From floating-point types ^^^^^^^^^^^^^^^^^^^^^^^^^ +*(ANSI compliant)* + Casting from floating-point input to an integral type truncates the input value. -It is allowed when the truncated result exceeds the range of result type. +It is allowed when the truncated result exceeds the range of result type when ANSI mode +is disabled; throws an error when ANSI mode is enabled. Valid examples @@ -79,12 +85,12 @@ Valid examples SELECT cast(12345.67 as bigint); -- 12345 SELECT cast(127.1 as tinyint); -- 127 SELECT cast(127.8 as tinyint); -- 127 - SELECT cast(1234567.89 as smallint); -- -10617 - SELECT cast(cast('inf' as double) as bigint); -- 9223372036854775807 - SELECT cast(cast('nan' as double) as integer); -- 0 - SELECT cast(cast('nan' as double) as smallint); -- 0 - SELECT cast(cast('nan' as double) as tinyint); -- 0 - SELECT cast(cast('nan' as double) as bigint); -- 0 + SELECT cast(1234567.89 as smallint); -- -10617 (ANSI OFF) / ERROR (ANSI ON) + SELECT cast(cast('inf' as double) as bigint); -- 9223372036854775807 (ANSI OFF) / ERROR (ANSI ON) + SELECT cast(cast('nan' as double) as integer); -- 0 (ANSI OFF) / ERROR (ANSI ON) + SELECT cast(cast('nan' as double) as smallint); -- 0 (ANSI OFF) / ERROR (ANSI ON) + SELECT cast(cast('nan' as double) as tinyint); -- 0 (ANSI OFF) / ERROR (ANSI ON) + SELECT cast(cast('nan' as double) as bigint); -- 0 (ANSI OFF) / ERROR (ANSI ON) From strings ^^^^^^^^^^^^ @@ -175,6 +181,43 @@ Invalid examples SELECT cast(cast('2025-02-25 08:00:26.88' as timestamp) as smallint); -- NULL (ANSI OFF) / ERROR (ANSI ON) SELECT cast(cast('2025-02-25 08:00:26.88' as timestamp) as tinyint); -- NULL (ANSI OFF) / ERROR (ANSI ON) +Cast to Floating-Point Types +---------------------------- + +From strings +^^^^^^^^^^^^ + +*(ANSI compliant)* + +Casting a string to ``REAL`` or ``DOUBLE`` accepts decimal and scientific +notation, an optional leading sign, and the case-insensitive special literals +``nan``, ``inf``, ``infinity`` (optionally signed). Values that overflow the +target type produce ``Infinity`` rather than an error, matching Spark. + +Casting from invalid strings returns NULL when ANSI mode is disabled; throws an +error otherwise. + +Valid examples + +:: + + SELECT cast('1.5' as double); -- 1.5 + SELECT cast('-3.14E-2' as real); -- -0.0314 + SELECT cast('1.5e10' as double); -- 1.5E10 + SELECT cast('nan' as double); -- NaN (case insensitive) + SELECT cast('-Infinity' as double); -- -Infinity (case insensitive) + SELECT cast('1e39' as real); -- Infinity (overflow) + SELECT cast('1e309' as double); -- Infinity (overflow) + +Invalid examples + +:: + + SELECT cast('abc' as double); -- NULL (ANSI OFF) / ERROR (ANSI ON) + SELECT cast('1.2a' as double); -- NULL (ANSI OFF) / ERROR (ANSI ON) + SELECT cast('1.2.3' as real); -- NULL (ANSI OFF) / ERROR (ANSI ON) + SELECT cast('' as double); -- NULL (ANSI OFF) / ERROR (ANSI ON) + Cast to Boolean --------------- @@ -331,6 +374,75 @@ Invalid examples SELECT cast('2012/10/23' as date); -- NULL // Invalid argument SELECT cast('2012.10.23' as date); -- NULL // Invalid argument +From TIMESTAMP_UTC +^^^^^^^^^^^^^^^^^^ + +*(ANSI compliant)* + +Casting a timestamp_utc to date extracts the date from the stored timestamp +fields without applying the session timezone. + +``cast`` throws when the value is too far from the epoch to fit in a date +(regardless of ANSI mode); ``try_cast`` returns NULL instead. + +Valid examples + +:: + + SELECT cast(TIMESTAMP_NTZ '2020-01-01 15:30:00' as date); -- 2020-01-01 + SELECT cast(TIMESTAMP_NTZ '2020-01-01 00:00:00' as date); -- 2020-01-01 + +Under session timezone ``America/Los_Angeles`` (UTC-8): :: + + SELECT cast(TIMESTAMP_NTZ '2020-01-01 00:00:00' as date); -- 2020-01-01 + +Cast to Time +------------ + +.. note:: + The TIME type was introduced in Apache Spark 4.1.0. + +From strings +^^^^^^^^^^^^ + +*(ANSI compliant)* + +Supported format is ``H:m[:s[.SSSSSS]]`` where: + + * ``H`` is hour (0-23) + * ``m`` is minute (0-59) + * ``s`` is an optional second (0-59); omitted seconds default to zero + * ``SSSSSS`` is optional fractional seconds (0-999999, up to microseconds) + +All leading and trailing UTF8 white-spaces are trimmed before casting. +Velox represents Spark ``TIME`` using ``TIME MICRO UTC``, whose values are +stored as microseconds since midnight (0 to 86,399,999,999). + +**ANSI mode behavior:** + + * **ANSI ON**: Invalid time strings throw an error. + * **ANSI OFF**: Invalid time strings return NULL. + +Valid examples + +:: + + SELECT cast('00:00:00' as time); -- 0 (midnight) + SELECT cast('12:30' as time); -- 45000000000 (seconds default to zero) + SELECT cast('12:30:45' as time); -- 45045000000 (12:30:45 in microseconds) + SELECT cast('23:59:59' as time); -- 86399000000 + SELECT cast('12:03:17.123' as time); -- 43397123000 (with milliseconds) + SELECT cast('12:03:17.123456' as time); -- 43397123456 (with microseconds) + SELECT cast(' 12:30:45 ' as time); -- 45045000000 (whitespace trimmed) + +Invalid examples + +:: + + SELECT cast('24:00:00' as time); -- NULL / throws error (hour out of range) + SELECT cast('12:60:00' as time); -- NULL / throws error (minute out of range) + SELECT cast('12:30:60' as time); -- NULL / throws error (second out of range) + Cast to Decimal --------------- @@ -395,6 +507,34 @@ Invalid examples SELECT cast(cast(100 as integer) as decimal(17, 16)); -- NULL (ANSI OFF) / ERROR (ANSI ON) // Value too large SELECT cast(cast(-100 as bigint) as decimal(17, 16)); -- NULL (ANSI OFF) / ERROR (ANSI ON) // Value too large +From decimal types +^^^^^^^^^^^^^^^^^^ + +*(ANSI compliant)* + +Casting a decimal value to a decimal of a different precision and scale is +allowed. + +When ANSI mode is enabled, casting a value that overflows the target precision +and scale throws an error. Otherwise, such casts return NULL. + +Valid examples + +:: + + SELECT cast(cast(-0.03 as decimal(2, 2)) as decimal(4, 4)); -- -0.0300 + SELECT cast(cast(1.05 as decimal(20, 2)) as decimal(10, 5)); -- 1.05000 + SELECT cast(cast(55.00 as decimal(6, 2)) as decimal(20, 10)); -- 55.0000000000 + SELECT cast(cast(1.2345 as decimal(6, 4)) as decimal(20, 1)); -- 1.2 + +Invalid examples + +:: + + SELECT cast(cast(-1000.000 as decimal(20, 3)) as decimal(6, 4)); -- NULL (ANSI OFF) / ERROR (ANSI ON) // Value too large + SELECT cast(cast(99999999999999999999999999999999999999 as decimal(38, 0)) as decimal(38, 1)); -- NULL (ANSI OFF) / ERROR (ANSI ON) // Value too large + SELECT cast(cast(-99999999999999999999999999999999999999 as decimal(38, 0)) as decimal(38, 1)); -- NULL (ANSI OFF) / ERROR (ANSI ON) // Value too large + From floating-point types ^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -433,6 +573,36 @@ Invalid examples SELECT cast(cast('inf' as double) as decimal(38, 2)); -- NULL (ANSI OFF) / ERROR (ANSI ON) // Value is not finite SELECT cast(cast('nan' as double) as decimal(38, 2)); -- NULL (ANSI OFF) / ERROR (ANSI ON) // Value is not finite +From boolean +^^^^^^^^^^^^ + +*(ANSI compliant)* + +Casting a boolean to a decimal of given precision and scale is allowed. +``true`` becomes 1 and ``false`` becomes 0. + +When ANSI mode is enabled, casting a value that overflows the target precision +and scale throws an error. Otherwise, such casts return NULL. Only ``true`` can +overflow, and only when the target has no integer digits (precision equals +scale), since 1 cannot be represented there. ``false`` becomes 0, which fits any +precision and scale. + +Valid examples + +:: + + SELECT cast(true as decimal(6, 2)); -- 1.00 + SELECT cast(false as decimal(6, 2)); -- 0.00 + SELECT cast(true as decimal(20, 10)); -- 1.0000000000 + SELECT cast(false as decimal(1, 1)); -- 0.0 + +Invalid examples + +:: + + SELECT cast(true as decimal(1, 1)); -- NULL (ANSI OFF) / ERROR (ANSI ON) // Value too large + SELECT cast(true as decimal(38, 38)); -- NULL (ANSI OFF) / ERROR (ANSI ON) // Value too large + Cast to Varbinary ----------------- @@ -474,8 +644,24 @@ Valid example From floating-point types ^^^^^^^^^^^^^^^^^^^^^^^^^ +*(ANSI compliant)* + Casting from floating-point input to timestamp type is allowed. -The input value is treated as the number of seconds since the epoch (1970-01-01 00:00:00 UTC) and converted to microseconds by truncating the fractional part. +The input value is treated as the number of seconds since the epoch +(``1970-01-01 00:00:00 UTC``) and converted to microseconds by truncating +the fractional part. + +When ANSI mode is disabled: + +* overflow is allowed and the result is saturated to the minimum or maximum + representable timestamp +* ``NaN`` and ``Infinity`` return NULL + +When ANSI mode is enabled: + +* overflow throws an error +* malformed floating-point values such as ``NaN`` and ``Infinity`` throw an + error Valid examples @@ -484,11 +670,21 @@ Valid examples SELECT cast(0.0 as timestamp); -- 1970-01-01 00:00:00 SELECT cast(1727181032.0 as timestamp); -- 2024-09-24 12:30:32 SELECT cast(-1727181032.0 as timestamp); -- 1915-04-09 11:29:28 - SELECT cast(cast(9223372036855.999 as double) as timestamp); -- 294247-01-10 04:00:54.775807 - SELECT cast(cast(-9223372036856.999 as double) as timestamp); -- -290308-12-21 19:59:05.224192 - SELECT cast(cast(1.79769e+308 as double) as timestamp); -- 294247-01-10 04:00:54.775807 - SELECT cast(cast('inf' as double) as timestamp); -- NULL - SELECT cast(cast('nan' as double) as timestamp); -- NULL + +Overflow examples + +:: + + SELECT cast(cast(9223372036855.999 as double) as timestamp); -- 294247-01-10 04:00:54.775807 (ANSI OFF) / ERROR (ANSI ON) + SELECT cast(cast(-9223372036856.999 as double) as timestamp); -- -290308-12-21 19:59:05.224192 (ANSI OFF) / ERROR (ANSI ON) + SELECT cast(cast(1.79769e+308 as double) as timestamp); -- 294247-01-10 04:00:54.775807 (ANSI OFF) / ERROR (ANSI ON) + +Malformed examples + +:: + + SELECT cast(cast('inf' as double) as timestamp); -- NULL (ANSI OFF) / ERROR (ANSI ON) + SELECT cast(cast('nan' as double) as timestamp); -- NULL (ANSI OFF) / ERROR (ANSI ON) From strings ^^^^^^^^^^^^ @@ -497,7 +693,10 @@ From strings Casting from strings to timestamp uses Spark-compatible timestamp parsing. The parser accepts date-only values, both ``' '`` and ``'T'`` as date-time -separators, fractional seconds, and leading or trailing spaces. +separators, fractional seconds, and leading or trailing spaces. Both ``' '`` +and ``'T'`` date-time separators must be followed immediately by a digit. +Outer whitespace is trimmed before the parser runs, so a bare trailing +separator such as ``"2015-03-18 "`` is handled as a date-only value. Casting from invalid strings returns NULL when ANSI mode is disabled and throws an error when ANSI mode is enabled. @@ -518,6 +717,9 @@ Invalid examples SELECT cast('INVALID' as timestamp); -- NULL (ANSI OFF) / ERROR (ANSI ON) SELECT cast('2012-Oct-01' as timestamp); -- NULL (ANSI OFF) / ERROR (ANSI ON) + SELECT cast('2015-03-18T' as timestamp); -- NULL (ANSI OFF) / ERROR (ANSI ON) + SELECT cast('2015-03-18T 12:00:00' as timestamp); -- NULL (ANSI OFF) / ERROR (ANSI ON) + SELECT cast('2015-03-18 Z' as timestamp); -- NULL (ANSI OFF) / ERROR (ANSI ON) From boolean ^^^^^^^^^^^^ @@ -588,3 +790,26 @@ Valid examples SELECT cast('2015-03-18 12:03:17.123' as timestamp_ntz); -- 2015-03-18 12:03:17.123 SELECT cast('1970-01-01 00:00:00-08:00' as timestamp_ntz); -- 1970-01-01 00:00:00 SELECT cast('2015-03-18T12:03:17Z' as timestamp_ntz); -- 2015-03-18 12:03:17 + +From DATE +^^^^^^^^^ + +*(ANSI compliant)* + +Casting a date to timestamp_utc returns midnight of the given date, not +subject to the session timezone. + +``cast`` throws when the date is too far from the epoch to fit in a +timestamp_utc (regardless of ANSI mode); ``try_cast`` returns NULL instead. + +Valid examples + +Under session timezone UTC: :: + + SELECT cast(DATE '2020-01-01' as timestamp_ntz); -- 2020-01-01 00:00:00 + SELECT cast(DATE '1970-01-01' as timestamp_ntz); -- 1970-01-01 00:00:00 + +Under session timezone ``America/Los_Angeles`` (UTC-8): :: + + SELECT cast(DATE '2020-01-01' as timestamp_ntz); -- 2020-01-01 00:00:00 + SELECT cast(DATE '1970-01-01' as timestamp_ntz); -- 1970-01-01 00:00:00 diff --git a/velox/docs/functions/spark/coverage.rst b/velox/docs/functions/spark/coverage.rst index f8fe4e00fd2..3158458b7fe 100644 --- a/velox/docs/functions/spark/coverage.rst +++ b/velox/docs/functions/spark/coverage.rst @@ -16,6 +16,7 @@ Here is a list of all scalar, aggregate, and window functions from Spark, with f table.coverage tr:nth-child(2) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(2) td:nth-child(9) {background-color: #6BA81E;} table.coverage tr:nth-child(3) td:nth-child(1) {background-color: #6BA81E;} + table.coverage tr:nth-child(3) td:nth-child(7) {background-color: #6BA81E;} table.coverage tr:nth-child(4) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(5) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(5) td:nth-child(2) {background-color: #6BA81E;} @@ -45,8 +46,10 @@ Here is a list of all scalar, aggregate, and window functions from Spark, with f table.coverage tr:nth-child(13) td:nth-child(7) {background-color: #6BA81E;} table.coverage tr:nth-child(14) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(14) td:nth-child(2) {background-color: #6BA81E;} + table.coverage tr:nth-child(14) td:nth-child(5) {background-color: #6BA81E;} table.coverage tr:nth-child(15) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(15) td:nth-child(2) {background-color: #6BA81E;} + table.coverage tr:nth-child(15) td:nth-child(5) {background-color: #6BA81E;} table.coverage tr:nth-child(16) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(16) td:nth-child(2) {background-color: #6BA81E;} table.coverage tr:nth-child(17) td:nth-child(1) {background-color: #6BA81E;} @@ -60,8 +63,10 @@ Here is a list of all scalar, aggregate, and window functions from Spark, with f table.coverage tr:nth-child(19) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(19) td:nth-child(2) {background-color: #6BA81E;} table.coverage tr:nth-child(19) td:nth-child(4) {background-color: #6BA81E;} + table.coverage tr:nth-child(19) td:nth-child(5) {background-color: #6BA81E;} table.coverage tr:nth-child(20) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(20) td:nth-child(2) {background-color: #6BA81E;} + table.coverage tr:nth-child(20) td:nth-child(4) {background-color: #6BA81E;} table.coverage tr:nth-child(20) td:nth-child(7) {background-color: #6BA81E;} table.coverage tr:nth-child(21) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(21) td:nth-child(2) {background-color: #6BA81E;} @@ -72,6 +77,8 @@ Here is a list of all scalar, aggregate, and window functions from Spark, with f table.coverage tr:nth-child(22) td:nth-child(4) {background-color: #6BA81E;} table.coverage tr:nth-child(23) td:nth-child(2) {background-color: #6BA81E;} table.coverage tr:nth-child(23) td:nth-child(3) {background-color: #6BA81E;} + table.coverage tr:nth-child(23) td:nth-child(4) {background-color: #6BA81E;} + table.coverage tr:nth-child(23) td:nth-child(5) {background-color: #6BA81E;} table.coverage tr:nth-child(24) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(24) td:nth-child(2) {background-color: #6BA81E;} table.coverage tr:nth-child(24) td:nth-child(3) {background-color: #6BA81E;} @@ -86,12 +93,15 @@ Here is a list of all scalar, aggregate, and window functions from Spark, with f table.coverage tr:nth-child(26) td:nth-child(7) {background-color: #6BA81E;} table.coverage tr:nth-child(27) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(27) td:nth-child(2) {background-color: #6BA81E;} + table.coverage tr:nth-child(27) td:nth-child(3) {background-color: #6BA81E;} table.coverage tr:nth-child(27) td:nth-child(5) {background-color: #6BA81E;} table.coverage tr:nth-child(28) td:nth-child(3) {background-color: #6BA81E;} table.coverage tr:nth-child(28) td:nth-child(7) {background-color: #6BA81E;} table.coverage tr:nth-child(29) td:nth-child(1) {background-color: #6BA81E;} + table.coverage tr:nth-child(29) td:nth-child(2) {background-color: #6BA81E;} table.coverage tr:nth-child(29) td:nth-child(3) {background-color: #6BA81E;} table.coverage tr:nth-child(29) td:nth-child(4) {background-color: #6BA81E;} + table.coverage tr:nth-child(29) td:nth-child(5) {background-color: #6BA81E;} table.coverage tr:nth-child(29) td:nth-child(7) {background-color: #6BA81E;} table.coverage tr:nth-child(30) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(30) td:nth-child(3) {background-color: #6BA81E;} @@ -106,6 +116,7 @@ Here is a list of all scalar, aggregate, and window functions from Spark, with f table.coverage tr:nth-child(32) td:nth-child(4) {background-color: #6BA81E;} table.coverage tr:nth-child(32) td:nth-child(5) {background-color: #6BA81E;} table.coverage tr:nth-child(32) td:nth-child(7) {background-color: #6BA81E;} + table.coverage tr:nth-child(33) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(33) td:nth-child(3) {background-color: #6BA81E;} table.coverage tr:nth-child(33) td:nth-child(4) {background-color: #6BA81E;} table.coverage tr:nth-child(34) td:nth-child(1) {background-color: #6BA81E;} @@ -150,8 +161,10 @@ Here is a list of all scalar, aggregate, and window functions from Spark, with f table.coverage tr:nth-child(45) td:nth-child(4) {background-color: #6BA81E;} table.coverage tr:nth-child(45) td:nth-child(5) {background-color: #6BA81E;} table.coverage tr:nth-child(45) td:nth-child(7) {background-color: #6BA81E;} + table.coverage tr:nth-child(46) td:nth-child(3) {background-color: #6BA81E;} table.coverage tr:nth-child(47) td:nth-child(2) {background-color: #6BA81E;} table.coverage tr:nth-child(47) td:nth-child(3) {background-color: #6BA81E;} + table.coverage tr:nth-child(47) td:nth-child(4) {background-color: #6BA81E;} table.coverage tr:nth-child(48) td:nth-child(3) {background-color: #6BA81E;} table.coverage tr:nth-child(49) td:nth-child(2) {background-color: #6BA81E;} table.coverage tr:nth-child(49) td:nth-child(3) {background-color: #6BA81E;} @@ -160,6 +173,7 @@ Here is a list of all scalar, aggregate, and window functions from Spark, with f table.coverage tr:nth-child(50) td:nth-child(4) {background-color: #6BA81E;} table.coverage tr:nth-child(50) td:nth-child(5) {background-color: #6BA81E;} table.coverage tr:nth-child(50) td:nth-child(7) {background-color: #6BA81E;} + table.coverage tr:nth-child(51) td:nth-child(2) {background-color: #6BA81E;} table.coverage tr:nth-child(51) td:nth-child(4) {background-color: #6BA81E;} table.coverage tr:nth-child(52) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(52) td:nth-child(3) {background-color: #6BA81E;} @@ -175,9 +189,11 @@ Here is a list of all scalar, aggregate, and window functions from Spark, with f table.coverage tr:nth-child(57) td:nth-child(2) {background-color: #6BA81E;} table.coverage tr:nth-child(58) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(58) td:nth-child(3) {background-color: #6BA81E;} + table.coverage tr:nth-child(58) td:nth-child(4) {background-color: #6BA81E;} table.coverage tr:nth-child(59) td:nth-child(2) {background-color: #6BA81E;} table.coverage tr:nth-child(59) td:nth-child(3) {background-color: #6BA81E;} table.coverage tr:nth-child(59) td:nth-child(4) {background-color: #6BA81E;} + table.coverage tr:nth-child(60) td:nth-child(3) {background-color: #6BA81E;} table.coverage tr:nth-child(60) td:nth-child(4) {background-color: #6BA81E;} table.coverage tr:nth-child(62) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(62) td:nth-child(2) {background-color: #6BA81E;} @@ -198,6 +214,7 @@ Here is a list of all scalar, aggregate, and window functions from Spark, with f table.coverage tr:nth-child(68) td:nth-child(1) {background-color: #6BA81E;} table.coverage tr:nth-child(68) td:nth-child(2) {background-color: #6BA81E;} table.coverage tr:nth-child(68) td:nth-child(4) {background-color: #6BA81E;} + table.coverage tr:nth-child(69) td:nth-child(2) {background-color: #6BA81E;} table.coverage tr:nth-child(69) td:nth-child(4) {background-color: #6BA81E;} @@ -210,7 +227,7 @@ Here is a list of all scalar, aggregate, and window functions from Spark, with f ===================================================================================================================================================================================================================== == ========================================= == ========================================= :spark:func:`abs` count_if inline nvl :spark:func:`sqrt` any cume_dist :spark:func:`acos` count_min_sketch inline_outer nvl2 stack approx_count_distinct :spark:func:`dense_rank` - :spark:func:`acosh` covar_pop input_file_block_length octet_length std approx_percentile first_value + :spark:func:`acosh` covar_pop input_file_block_length octet_length std :spark:func:`approx_percentile` first_value :spark:func:`add_months` covar_samp input_file_block_start or stddev array_agg lag :spark:func:`aggregate` :spark:func:`crc32` input_file_name :spark:func:`overlay` stddev_pop :spark:func:`avg` last_value and cume_dist :spark:func:`instr` parse_url stddev_samp bit_and lead @@ -221,26 +238,26 @@ Here is a list of all scalar, aggregate, and window functions from Spark, with f :spark:func:`array_contains` current_timezone java_method :spark:func:`pmod` :spark:func:`substring` :spark:func:`collect_list` :spark:func:`row_number` :spark:func:`array_distinct` current_user :spark:func:`json_array_length` posexplode :spark:func:`substring_index` :spark:func:`collect_set` :spark:func:`array_except` date :spark:func:`json_object_keys` posexplode_outer sum :spark:func:`corr` - :spark:func:`array_intersect` :spark:func:`date_add` json_tuple position tan count - :spark:func:`array_join` :spark:func:`date_format` kurtosis positive tanh count_if + :spark:func:`array_intersect` :spark:func:`date_add` json_tuple position :spark:func:`tan` count + :spark:func:`array_join` :spark:func:`date_format` kurtosis positive :spark:func:`tanh` count_if :spark:func:`array_max` :spark:func:`date_from_unix_date` lag pow timestamp count_min_sketch :spark:func:`array_min` date_part last :spark:func:`power` :spark:func:`timestamp_micros` covar_pop :spark:func:`array_position` :spark:func:`date_sub` :spark:func:`last_day` printf :spark:func:`timestamp_millis` :spark:func:`covar_samp` - :spark:func:`array_remove` :spark:func:`date_trunc` last_value :spark:func:`quarter` timestamp_seconds every - :spark:func:`array_repeat` :spark:func:`datediff` lcase radians tinyint :spark:func:`first` + :spark:func:`array_remove` :spark:func:`date_trunc` last_value :spark:func:`quarter` :spark:func:`timestamp_seconds` every + :spark:func:`array_repeat` :spark:func:`datediff` lcase :spark:func:`radians` tinyint :spark:func:`first` :spark:func:`array_sort` :spark:func:`day` lead :spark:func:`raise_error` to_csv first_value :spark:func:`array_union` :spark:func:`dayofmonth` :spark:func:`least` :spark:func:`rand` to_date grouping - arrays_overlap :spark:func:`dayofweek` :spark:func:`left` randn to_json grouping_id + arrays_overlap :spark:func:`dayofweek` :spark:func:`left` :spark:func:`randn` :spark:func:`to_json` grouping_id :spark:func:`arrays_zip` :spark:func:`dayofyear` :spark:func:`length` :spark:func:`random` to_timestamp histogram_numeric :spark:func:`ascii` decimal :spark:func:`levenshtein` range :spark:func:`to_unix_timestamp` :spark:func:`kurtosis` :spark:func:`asin` decode :spark:func:`like` rank :spark:func:`to_utc_timestamp` :spark:func:`last` - :spark:func:`asinh` :spark:func:`degrees` ln reflect :spark:func:`transform` last_value + :spark:func:`asinh` :spark:func:`degrees` :spark:func:`ln` reflect :spark:func:`transform` last_value assert_true dense_rank :spark:func:`locate` regexp transform_keys :spark:func:`max` - :spark:func:`atan` div :spark:func:`log` :spark:func:`regexp_extract` transform_values :spark:func:`max_by` + :spark:func:`atan` :spark:func:`div` :spark:func:`log` :spark:func:`regexp_extract` :spark:func:`transform_values` :spark:func:`max_by` :spark:func:`atan2` double :spark:func:`log10` :spark:func:`regexp_extract_all` :spark:func:`translate` mean :spark:func:`atanh` e :spark:func:`log1p` regexp_like :spark:func:`trim` :spark:func:`min` avg :spark:func:`element_at` :spark:func:`log2` :spark:func:`regexp_replace` :spark:func:`trunc` :spark:func:`min_by` - base64 elt :spark:func:`lower` :spark:func:`repeat` try_add percentile + :spark:func:`base64` elt :spark:func:`lower` :spark:func:`repeat` try_add percentile :spark:func:`between` encode :spark:func:`lpad` :spark:func:`replace` try_divide percentile_approx bigint every :spark:func:`ltrim` :spark:func:`reverse` typeof regr_avgx :spark:func:`bin` :spark:func:`exists` :spark:func:`make_date` right ucase regr_avgy @@ -253,21 +270,21 @@ Here is a list of all scalar, aggregate, and window functions from Spark, with f bit_xor :spark:func:`filter` :spark:func:`map_entries` schema_of_csv :spark:func:`unix_timestamp` stddev_pop bool_and :spark:func:`find_in_set` :spark:func:`map_filter` schema_of_json :spark:func:`upper` :spark:func:`stddev_samp` bool_or first :spark:func:`map_from_arrays` :spark:func:`second` :spark:func:`uuid` :spark:func:`sum` - boolean first_value map_from_entries sentences var_pop try_avg - bround :spark:func:`flatten` :spark:func:`map_keys` sequence var_samp try_sum + boolean first_value :spark:func:`map_from_entries` sentences var_pop try_avg + bround :spark:func:`flatten` :spark:func:`map_keys` :spark:func:`sequence` var_samp try_sum btrim float :spark:func:`map_values` session_window variance var_pop cardinality :spark:func:`floor` :spark:func:`map_zip_with` sha version :spark:func:`var_samp` case :spark:func:`forall` max :spark:func:`sha1` :spark:func:`weekday` :spark:func:`variance` - cast format_number max_by :spark:func:`sha2` weekofyear + cast :spark:func:`format_number` max_by :spark:func:`sha2` weekofyear :spark:func:`cbrt` format_string :spark:func:`md5` :spark:func:`shiftleft` when :spark:func:`ceil` from_csv mean :spark:func:`shiftright` :spark:func:`width_bucket` ceiling from_json min shiftrightunsigned window char :spark:func:`from_unixtime` min_by :spark:func:`shuffle` xpath char_length :spark:func:`from_utc_timestamp` :spark:func:`minute` :spark:func:`sign` xpath_boolean character_length :spark:func:`get_json_object` mod signum xpath_double - :spark:func:`chr` getbit :spark:func:`monotonically_increasing_id` sin xpath_float + :spark:func:`chr` getbit :spark:func:`monotonically_increasing_id` :spark:func:`sin` xpath_float coalesce :spark:func:`greatest` :spark:func:`month` :spark:func:`sinh` xpath_int - collect_list grouping months_between :spark:func:`size` xpath_long + collect_list grouping :spark:func:`months_between` :spark:func:`size` xpath_long collect_set grouping_id named_struct skewness xpath_number :spark:func:`concat` :spark:func:`hash` nanvl :spark:func:`slice` xpath_short concat_ws :spark:func:`hex` negative smallint xpath_string @@ -276,5 +293,5 @@ Here is a list of all scalar, aggregate, and window functions from Spark, with f :spark:func:`cos` if now :spark:func:`soundex` :spark:func:`zip_with` :spark:func:`cosh` ifnull nth_value space :spark:func:`cot` :spark:func:`in` ntile :spark:func:`spark_partition_id` - count initcap nullif :spark:func:`split` + count :spark:func:`initcap` nullif :spark:func:`split` ========================================= ========================================= ========================================= ========================================= ========================================= == ========================================= == ========================================= diff --git a/velox/docs/functions/spark/decimal.rst b/velox/docs/functions/spark/decimal.rst index 1afd15d5e88..265d68714dc 100644 --- a/velox/docs/functions/spark/decimal.rst +++ b/velox/docs/functions/spark/decimal.rst @@ -314,6 +314,6 @@ Decimal Special Forms .. spark:function:: make_decimal(x[, nullOnOverflow]) -> decimal - Create ``decimal`` of requsted precision and scale from an unscaled bigint value ``x``. + Create ``decimal`` of requested precision and scale from an unscaled bigint value ``x``. By default, the value of ``nullOnOverflow`` is true, and null will be returned when ``x`` is too large for the result precision. Otherwise, exception will be thrown when ``x`` overflows. diff --git a/velox/docs/functions/spark/map.rst b/velox/docs/functions/spark/map.rst index 09e1bfc0c87..43a330d9d2d 100644 --- a/velox/docs/functions/spark/map.rst +++ b/velox/docs/functions/spark/map.rst @@ -46,10 +46,15 @@ Map Functions .. spark:function:: map_from_arrays(array(K), array(V)) -> map(K,V) - Creates a map with a pair of the given key/value arrays. All elements in keys should not be null. - If key size != value size will throw exception that key and value must have the same length.:: + Creates a map by pairing up the given key and value arrays. Returns NULL if either array is + NULL. Throws if a top-level key is NULL, or if the two arrays have different lengths. Array and + row keys may contain nested NULL values. When a key is repeated, the behavior depends on the + ``throw_exception_on_duplicate_map_keys`` configuration + property: if true, throws; otherwise the last value wins and the key keeps the position of its + first occurrence. :: SELECT map_from_arrays(array(1.0, 3.0), array('2', '4')); -- {1.0 -> 2, 3.0 -> 4} + SELECT map_from_arrays(array(2, 1, 2), array('a', 'b', 'c')); -- {2 -> 'c', 1 -> 'b'} .. spark:function:: map_from_entries(array(struct(K,V))) -> map(K,V) diff --git a/velox/docs/functions/spark/math.rst b/velox/docs/functions/spark/math.rst index e5b74d1781f..e61b09b6959 100644 --- a/velox/docs/functions/spark/math.rst +++ b/velox/docs/functions/spark/math.rst @@ -374,6 +374,25 @@ Mathematical Functions SELECT rand(0); -- 0.7604953758285915 SELECT rand(NULL); -- 0.7604953758285915 +.. spark:function:: randn() -> double + + Returns a random value from the standard normal distribution (mean 0.0 and + standard deviation 1.0). :: + + SELECT randn(); -- -1.1750280342265669 + +.. spark:function:: randn(seed) -> double + + Returns a random value from the standard normal distribution (mean 0.0 and + standard deviation 1.0) using a seed formed by combining user-specified + ``seed`` and the configuration `spark.partition_id`. The framework is + responsible for deterministic partitioning of the data and assigning unique + `spark.partition_id` to each thread (in a deterministic way). + ``seed`` must be constant. NULL ``seed`` is identical to zero ``seed``. :: + + SELECT randn(0); -- 1.6034991609278433 + SELECT randn(NULL); -- 1.6034991609278433 + .. spark:function:: random() -> double An alias for ``rand()``. diff --git a/velox/docs/functions/spark/regexp.rst b/velox/docs/functions/spark/regexp.rst index 087171f3db1..d7c45393c6e 100644 --- a/velox/docs/functions/spark/regexp.rst +++ b/velox/docs/functions/spark/regexp.rst @@ -5,10 +5,10 @@ Regular Expression Functions Regular expression functions use RE2 as the regex engine. RE2 is fast, but supports only a subset of PCRE syntax and in particular does not support backtracking and associated features (e.g. back references). -Java and RE2 regex output can diverage and users should be cautious that +Java and RE2 regex output can diverge and users should be cautious that the patterns they are using perform similarly between RE2 and Java. For example, character class unions, intersections, and differences -``([a[b]], [a&&[b]], [a&&[^b]])`` are intepreted as a single character class +``([a[b]], [a&&[b]], [a&&[^b]])`` are interpreted as a single character class that contain ``[, &, and ^`` rather than union, intersection, or difference of the character classes. diff --git a/velox/docs/monitoring/stats.rst b/velox/docs/monitoring/stats.rst index b40c9ef03f7..e89e869f769 100644 --- a/velox/docs/monitoring/stats.rst +++ b/velox/docs/monitoring/stats.rst @@ -151,6 +151,62 @@ These stats are reported only by TableWriter operator - nanos - The walltime spent on file write data compression. +Nimble Writer +~~~~~~~~~~~~~ +These stats are reported by TableWriter when writing Nimble files. Encoding +CPU time is summed across encoding worker threads and can exceed the +corresponding wall time when parallel encoding is enabled. + +.. list-table:: + :widths: 50 25 50 + :header-rows: 1 + + * - Stats + - Unit + - Description + * - nimble.writtenBytes + - bytes + - Total number of bytes written to the Nimble file. + * - nimble.inputBytes + - bytes + - Uncompressed size of the input written to the Nimble file. + * - nimble.writeCpuNanos + - nanos + - CPU time spent writing encoded stripes through the tablet writer. + * - nimble.writeWallNanos + - nanos + - Wall time spent writing encoded stripes through the tablet writer. + * - nimble.ingestionCpuNanos + - nanos + - CPU time spent ingesting input vectors into field-writer buffers. + * - nimble.ingestionWallNanos + - nanos + - Wall time spent ingesting input vectors into field-writer buffers. + * - nimble.encodingCpuNanos + - nanos + - CPU time spent encoding and compressing streams, summed across all + encoding worker threads. + * - nimble.encodingWallNanos + - nanos + - Wall time spent encoding and compressing streams. + * - nimble.encodingSelectionCpuNanos + - nanos + - CPU time spent selecting encodings. This is a subset of + ``nimble.encodingCpuNanos``. + * - nimble.rowsPerStripe + - + - Distribution of row counts per stripe. The metric count is the number + of stripes written. + * - nimble.chunkSizeBytes + - bytes + - Distribution of encoded chunk sizes. + * - nimble.duplicateStreamCount + - + - Number of streams deduplicated by the tablet writer. + * - nimble.duplicateStreamBytes + - bytes + - Number of encoded bytes deduplicated by the tablet writer. + LookupIndexJoin --------------- These stats are reported only by IndexLookupJoin operator @@ -473,7 +529,12 @@ FileBasedDataSource These stats are reported by the file-based connector data source (Hive connector). Data stream IO stats use the stat names directly (e.g., ``storageReadBytes``). Metadata IO stats (footer, stripe groups, index) use a ``metadata.`` prefix -(e.g., ``metadata.storageReadBytes``, ``metadata.ramReadBytes``). +(e.g., ``metadata.storageReadBytes``, ``metadata.ramReadBytes``). Reader +format-specific stats are prefixed with the file format name +(e.g., ``dwrf.flattenStringDictionaryValues``). Column statistics are also +reported per column using ``column_`` and the column type. For example, +``parquet.pageLoadTimeNanos`` aggregates all Parquet columns, while +``parquet.column_2.BIGINT.pageLoadTimeNanos`` identifies one column. .. list-table:: :widths: 50 25 50 @@ -504,13 +565,6 @@ Metadata IO stats (footer, stripe groups, index) use a ``metadata.`` prefix * - numStripes - - The number of stripes read from the file. - * - flattenStringDictionaryValues - - - - The number of rows returned by the string dictionary reader that were - flattened instead of keeping dictionary encoding. - * - pageLoadTimeNs - - nanos - - The total time spent loading pages. * - numPrefetch - - The number of prefetch operations issued. @@ -547,3 +601,28 @@ Metadata IO stats (footer, stripe groups, index) use a ``metadata.`` prefix coalescing. Measures data locality on disk — smaller gaps indicate co-accessed columns are physically adjacent in the file. Includes min and max per gap. + * - parquet.footerEstimatedBytes + - bytes + - The estimated memory used by the deserialized Parquet footer when + footer memory tracking is enabled. + * - | dwrf.flattenStringDictionaryValues + | dwrf.column_..flattenStringDictionaryValues + - + - The number of rows returned by the DWRF string dictionary reader that + were flattened instead of keeping dictionary encoding. Reported across + all columns and by column. + * - | parquet.pageLoadTimeNanos + | parquet.column_..pageLoadTimeNanos + - nanos + - The time spent loading Parquet pages. Reported across all columns and + by column. + * - | .decompressCPUTimeNanos + | .column_..decompressCPUTimeNanos + - nanos + - The CPU time spent decompressing column data. Reported across all + columns and by column. + * - | .decodeCPUTimeNanos + | .column_..decodeCPUTimeNanos + - nanos + - The CPU time spent decoding column data. Reported across all columns, + and by column. diff --git a/velox/docs/spark_functions.rst b/velox/docs/spark_functions.rst index e6944a702a9..47c7f24db57 100644 --- a/velox/docs/spark_functions.rst +++ b/velox/docs/spark_functions.rst @@ -65,80 +65,89 @@ for :doc:`all ` functions. :widths: auto :class: rows - =========================================== =========================================== =========================================== == =========================================== == =========================================== - Scalar Functions Aggregate Functions Window Functions - ===================================================================================================================================== == =========================================== == =========================================== - :spark:func:`abs` :spark:func:`divide_deny_precision_loss` :spark:func:`not` :spark:func:`avg` :spark:func:`dense_rank` - :spark:func:`acos` :spark:func:`doy` :spark:func:`overlay` :spark:func:`bit_xor` :spark:func:`nth_value` - :spark:func:`acosh` :spark:func:`element_at` :spark:func:`pmod` :spark:func:`bloom_filter_agg` :spark:func:`ntile` - :spark:func:`add` :spark:func:`empty2null` :spark:func:`power` :spark:func:`collect_list` :spark:func:`rank` - :spark:func:`add_deny_precision_loss` :spark:func:`endswith` :spark:func:`quarter` :spark:func:`collect_set` :spark:func:`row_number` - :spark:func:`add_months` :spark:func:`equalnullsafe` :spark:func:`raise_error` :spark:func:`corr` - :spark:func:`aggregate` :spark:func:`equalto` :spark:func:`rand` :spark:func:`covar_samp` - :spark:func:`array` :spark:func:`exists` :spark:func:`random` :spark:func:`first` - :spark:func:`array_append` :spark:func:`exp` :spark:func:`regexp_extract` :spark:func:`first_ignore_null` - :spark:func:`array_compact` :spark:func:`expm1` :spark:func:`regexp_extract_all` :spark:func:`kurtosis` - :spark:func:`array_contains` :spark:func:`factorial` :spark:func:`regexp_replace` :spark:func:`last` - :spark:func:`array_distinct` :spark:func:`filter` :spark:func:`remainder` :spark:func:`last_ignore_null` - :spark:func:`array_except` :spark:func:`find_in_set` :spark:func:`repeat` :spark:func:`max` - :spark:func:`array_insert` :spark:func:`flatten` :spark:func:`replace` :spark:func:`max_by` - :spark:func:`array_intersect` :spark:func:`floor` :spark:func:`reverse` :spark:func:`min` - :spark:func:`array_join` :spark:func:`forall` :spark:func:`rint` :spark:func:`min_by` - :spark:func:`array_max` :spark:func:`from_unixtime` :spark:func:`rlike` :spark:func:`mode` - :spark:func:`array_min` :spark:func:`from_utc_timestamp` :spark:func:`round` :spark:func:`regr_replacement` - :spark:func:`array_position` :spark:func:`get` :spark:func:`rpad` :spark:func:`skewness` - :spark:func:`array_prepend` :spark:func:`get_json_object` :spark:func:`rtrim` :spark:func:`stddev` - :spark:func:`array_remove` :spark:func:`get_timestamp` :spark:func:`sec` :spark:func:`stddev_samp` - :spark:func:`array_repeat` :spark:func:`greaterthan` :spark:func:`second` :spark:func:`sum` - :spark:func:`array_sort` :spark:func:`greaterthanorequal` :spark:func:`sha1` :spark:func:`var_samp` - :spark:func:`array_union` :spark:func:`greatest` :spark:func:`sha2` :spark:func:`variance` - :spark:func:`arrays_zip` :spark:func:`hash` :spark:func:`shiftleft` - :spark:func:`ascii` :spark:func:`hash_with_seed` :spark:func:`shiftright` - :spark:func:`asin` :spark:func:`hex` :spark:func:`shuffle` - :spark:func:`asinh` :spark:func:`hour` :spark:func:`sign` - :spark:func:`atan` :spark:func:`hypot` :spark:func:`sinh` - :spark:func:`atan2` :spark:func:`in` :spark:func:`size` - :spark:func:`atanh` :spark:func:`instr` :spark:func:`slice` - :spark:func:`between` :spark:func:`isnan` :spark:func:`sort_array` - :spark:func:`bin` :spark:func:`isnotnull` :spark:func:`soundex` - :spark:func:`bit_count` :spark:func:`isnull` :spark:func:`spark_partition_id` - :spark:func:`bit_get` :spark:func:`json_array_length` :spark:func:`split` - :spark:func:`bit_length` :spark:func:`json_object_keys` :spark:func:`sqrt` - :spark:func:`bitwise_and` :spark:func:`last_day` :spark:func:`startswith` - :spark:func:`bitwise_not` :spark:func:`least` :spark:func:`str_to_map` - :spark:func:`bitwise_or` :spark:func:`left` :spark:func:`substring` - :spark:func:`bitwise_xor` :spark:func:`length` :spark:func:`substring_index` - :spark:func:`cbrt` :spark:func:`lessthan` :spark:func:`subtract` - :spark:func:`ceil` :spark:func:`lessthanorequal` :spark:func:`subtract_deny_precision_loss` - :spark:func:`checked_add` :spark:func:`levenshtein` :spark:func:`timestamp_micros` - :spark:func:`checked_divide` :spark:func:`like` :spark:func:`timestamp_millis` - :spark:func:`checked_multiply` :spark:func:`locate` :spark:func:`to_unix_timestamp` - :spark:func:`checked_subtract` :spark:func:`log` :spark:func:`to_utc_timestamp` - :spark:func:`chr` :spark:func:`log10` :spark:func:`transform` - :spark:func:`concat` :spark:func:`log1p` :spark:func:`translate` - :spark:func:`contains` :spark:func:`log2` :spark:func:`trim` - :spark:func:`conv` :spark:func:`lower` :spark:func:`trunc` - :spark:func:`cos` :spark:func:`lpad` :spark:func:`unaryminus` - :spark:func:`cosh` :spark:func:`ltrim` :spark:func:`unbase64` - :spark:func:`cot` :spark:func:`luhn_check` :spark:func:`unhex` - :spark:func:`crc32` :spark:func:`make_date` :spark:func:`unix_date` - :spark:func:`csc` :spark:func:`make_timestamp` :spark:func:`unix_micros` - :spark:func:`date_add` :spark:func:`make_ym_interval` :spark:func:`unix_millis` - :spark:func:`date_format` :spark:func:`map` :spark:func:`unix_seconds` - :spark:func:`date_from_unix_date` :spark:func:`map_concat` :spark:func:`unix_timestamp` - :spark:func:`date_sub` :spark:func:`map_entries` :spark:func:`unscaled_value` - :spark:func:`date_trunc` :spark:func:`map_filter` :spark:func:`upper` - :spark:func:`datediff` :spark:func:`map_from_arrays` :spark:func:`url_decode` - :spark:func:`day` :spark:func:`map_keys` :spark:func:`url_encode` - :spark:func:`dayofmonth` :spark:func:`map_values` :spark:func:`uuid` - :spark:func:`dayofweek` :spark:func:`map_zip_with` :spark:func:`varchar_type_write_side_check` - :spark:func:`dayofyear` :spark:func:`mask` :spark:func:`week_of_year` - :spark:func:`decimal_equalto` :spark:func:`md5` :spark:func:`weekday` - :spark:func:`decimal_greaterthan` :spark:func:`might_contain` :spark:func:`width_bucket` - :spark:func:`decimal_greaterthanorequal` :spark:func:`minute` :spark:func:`xxhash64` - :spark:func:`decimal_lessthan` :spark:func:`monotonically_increasing_id` :spark:func:`xxhash64_with_seed` - :spark:func:`decimal_lessthanorequal` :spark:func:`month` :spark:func:`year` - :spark:func:`decimal_notequalto` :spark:func:`multiply` :spark:func:`year_of_week` - :spark:func:`degrees` :spark:func:`multiply_deny_precision_loss` :spark:func:`zip_with` - :spark:func:`divide` :spark:func:`next_day` - =========================================== =========================================== =========================================== == =========================================== == =========================================== + ================================================== ================================================== ================================================== == ================================================== == ================================================== + Scalar Functions Aggregate Functions Window Functions + ========================================================================================================================================================== == ================================================== == ================================================== + :spark:func:`abs` :spark:func:`divide` :spark:func:`radians` :spark:func:`approx_percentile` :spark:func:`dense_rank` + :spark:func:`acos` :spark:func:`divide_deny_precision_loss` :spark:func:`raise_error` :spark:func:`avg` :spark:func:`nth_value` + :spark:func:`acosh` :spark:func:`element_at` :spark:func:`rand` :spark:func:`bit_xor` :spark:func:`ntile` + :spark:func:`add` :spark:func:`empty2null` :spark:func:`randn` :spark:func:`bitmap_construct_agg` :spark:func:`rank` + :spark:func:`add_deny_precision_loss` :spark:func:`endswith` :spark:func:`random` :spark:func:`bitmap_or_agg` :spark:func:`row_number` + :spark:func:`add_months` :spark:func:`equalnullsafe` :spark:func:`randstr` :spark:func:`bloom_filter_agg` + :spark:func:`aggregate` :spark:func:`equalto` :spark:func:`read_side_padding` :spark:func:`collect_list` + :spark:func:`array` :spark:func:`exists` :spark:func:`regexp_extract` :spark:func:`collect_set` + :spark:func:`array_append` :spark:func:`exp` :spark:func:`regexp_extract_all` :spark:func:`corr` + :spark:func:`array_compact` :spark:func:`expm1` :spark:func:`regexp_instr` :spark:func:`covar_samp` + :spark:func:`array_contains` :spark:func:`factorial` :spark:func:`regexp_replace` :spark:func:`first` + :spark:func:`array_distinct` :spark:func:`filter` :spark:func:`remainder` :spark:func:`first_ignore_null` + :spark:func:`array_except` :spark:func:`find_in_set` :spark:func:`repeat` :spark:func:`kurtosis` + :spark:func:`array_insert` :spark:func:`flatten` :spark:func:`replace` :spark:func:`last` + :spark:func:`array_intersect` :spark:func:`floor` :spark:func:`reverse` :spark:func:`last_ignore_null` + :spark:func:`array_join` :spark:func:`forall` :spark:func:`rint` :spark:func:`max` + :spark:func:`array_max` :spark:func:`format_number` :spark:func:`rlike` :spark:func:`max_by` + :spark:func:`array_min` :spark:func:`from_unixtime` :spark:func:`round` :spark:func:`min` + :spark:func:`array_position` :spark:func:`from_utc_timestamp` :spark:func:`rpad` :spark:func:`min_by` + :spark:func:`array_prepend` :spark:func:`get` :spark:func:`rtrim` :spark:func:`mode` + :spark:func:`array_remove` :spark:func:`get_json_object` :spark:func:`sec` :spark:func:`regr_replacement` + :spark:func:`array_repeat` :spark:func:`get_timestamp` :spark:func:`second` :spark:func:`skewness` + :spark:func:`array_sort` :spark:func:`greaterthan` :spark:func:`sequence` :spark:func:`stddev` + :spark:func:`array_sort_desc` :spark:func:`greaterthanorequal` :spark:func:`sha1` :spark:func:`stddev_samp` + :spark:func:`array_union` :spark:func:`greatest` :spark:func:`sha2` :spark:func:`sum` + :spark:func:`arrays_zip` :spark:func:`hash` :spark:func:`shiftleft` :spark:func:`var_samp` + :spark:func:`ascii` :spark:func:`hash_with_seed` :spark:func:`shiftright` :spark:func:`variance` + :spark:func:`asin` :spark:func:`hex` :spark:func:`shuffle` + :spark:func:`asinh` :spark:func:`hour` :spark:func:`sign` + :spark:func:`assert_not_null` :spark:func:`hypot` :spark:func:`sin` + :spark:func:`atan` :spark:func:`initcap` :spark:func:`sinh` + :spark:func:`atan2` :spark:func:`instr` :spark:func:`size` + :spark:func:`atanh` :spark:func:`isnan` :spark:func:`slice` + :spark:func:`base64` :spark:func:`isnotnull` :spark:func:`sort_array` + :spark:func:`between` :spark:func:`isnull` :spark:func:`soundex` + :spark:func:`bin` :spark:func:`json_array_length` :spark:func:`spark_partition_id` + :spark:func:`bit_count` :spark:func:`json_object_keys` :spark:func:`split` + :spark:func:`bit_get` :spark:func:`last_day` :spark:func:`sqrt` + :spark:func:`bit_length` :spark:func:`least` :spark:func:`startswith` + :spark:func:`bitwise_and` :spark:func:`left` :spark:func:`str_to_map` + :spark:func:`bitwise_not` :spark:func:`length` :spark:func:`substring` + :spark:func:`bitwise_or` :spark:func:`lessthan` :spark:func:`substring_index` + :spark:func:`bitwise_xor` :spark:func:`lessthanorequal` :spark:func:`subtract` + :spark:func:`cbrt` :spark:func:`levenshtein` :spark:func:`subtract_deny_precision_loss` + :spark:func:`ceil` :spark:func:`like` :spark:func:`tan` + :spark:func:`char_type_write_side_check` :spark:func:`ln` :spark:func:`tanh` + :spark:func:`checked_add` :spark:func:`locate` :spark:func:`timestamp_micros` + :spark:func:`checked_add_deny_precision_loss` :spark:func:`log` :spark:func:`timestamp_millis` + :spark:func:`checked_div` :spark:func:`log10` :spark:func:`timestamp_seconds` + :spark:func:`checked_divide` :spark:func:`log1p` :spark:func:`timestampadd` + :spark:func:`checked_multiply` :spark:func:`log2` :spark:func:`timestampdiff` + :spark:func:`checked_multiply_deny_precision_loss` :spark:func:`lower` :spark:func:`to_json` + :spark:func:`checked_subtract` :spark:func:`lpad` :spark:func:`to_pretty_string` + :spark:func:`checked_subtract_deny_precision_loss` :spark:func:`ltrim` :spark:func:`to_unix_timestamp` + :spark:func:`chr` :spark:func:`luhn_check` :spark:func:`to_utc_timestamp` + :spark:func:`concat` :spark:func:`make_date` :spark:func:`transform` + :spark:func:`contains` :spark:func:`make_timestamp` :spark:func:`transform_values` + :spark:func:`conv` :spark:func:`make_ym_interval` :spark:func:`translate` + :spark:func:`cos` :spark:func:`map` :spark:func:`trim` + :spark:func:`cosh` :spark:func:`map_concat` :spark:func:`trunc` + :spark:func:`cot` :spark:func:`map_entries` :spark:func:`unaryminus` + :spark:func:`crc32` :spark:func:`map_filter` :spark:func:`unbase64` + :spark:func:`csc` :spark:func:`map_from_arrays` :spark:func:`unhex` + :spark:func:`date_add` :spark:func:`map_from_entries` :spark:func:`unix_date` + :spark:func:`date_format` :spark:func:`map_keys` :spark:func:`unix_micros` + :spark:func:`date_from_unix_date` :spark:func:`map_values` :spark:func:`unix_millis` + :spark:func:`date_sub` :spark:func:`map_zip_with` :spark:func:`unix_seconds` + :spark:func:`date_trunc` :spark:func:`mask` :spark:func:`unix_timestamp` + :spark:func:`datediff` :spark:func:`md5` :spark:func:`unscaled_value` + :spark:func:`day` :spark:func:`might_contain` :spark:func:`upper` + :spark:func:`dayname` :spark:func:`minute` :spark:func:`url_decode` + :spark:func:`dayofmonth` :spark:func:`monotonically_increasing_id` :spark:func:`url_encode` + :spark:func:`dayofweek` :spark:func:`month` :spark:func:`uuid` + :spark:func:`dayofyear` :spark:func:`monthname` :spark:func:`varchar_type_write_side_check` + :spark:func:`decimal_equalto` :spark:func:`months_between` :spark:func:`week_of_year` + :spark:func:`decimal_greaterthan` :spark:func:`multiply` :spark:func:`weekday` + :spark:func:`decimal_greaterthanorequal` :spark:func:`multiply_deny_precision_loss` :spark:func:`width_bucket` + :spark:func:`decimal_lessthan` :spark:func:`next_day` :spark:func:`xxhash64` + :spark:func:`decimal_lessthanorequal` :spark:func:`overlay` :spark:func:`xxhash64_with_seed` + :spark:func:`decimal_notequalto` :spark:func:`pmod` :spark:func:`year` + :spark:func:`degrees` :spark:func:`power` :spark:func:`year_of_week` + :spark:func:`div` :spark:func:`quarter` :spark:func:`zip_with` + ================================================== ================================================== ================================================== == ================================================== == ================================================== diff --git a/velox/duckdb/conversion/DuckParser.cpp b/velox/duckdb/conversion/DuckParser.cpp index 9aa79fbe2c3..586b27dbfbf 100644 --- a/velox/duckdb/conversion/DuckParser.cpp +++ b/velox/duckdb/conversion/DuckParser.cpp @@ -1079,8 +1079,13 @@ core::WindowCallExprPtr buildWindowCallExpr( } auto endType = parseBoundType(windowExpr.end); + // Without ORDER BY every row is a peer, so a RANGE frame ending at the + // current row covers the whole partition. That is the frame a window with no + // frame clause gets. A ROWS frame counts rows rather than peers and ends + // where it says, so read the DuckDB boundary: `parseBoundType` maps both + // spellings of CURRENT ROW to the same bound. if (options.correctWindowFrameDefault && orderByKeys.empty() && - endType == core::WindowCallExpr::BoundType::kCurrentRow) { + windowExpr.end == WindowBoundary::CURRENT_ROW_RANGE) { endType = core::WindowCallExpr::BoundType::kUnboundedFollowing; } diff --git a/velox/duckdb/conversion/DuckParser.h b/velox/duckdb/conversion/DuckParser.h index f8e094109d0..76e3472a9cb 100644 --- a/velox/duckdb/conversion/DuckParser.h +++ b/velox/duckdb/conversion/DuckParser.h @@ -31,8 +31,9 @@ struct ParseOptions { // DuckDB defaults the window frame end bound to CURRENT ROW even when ORDER // BY is absent. The SQL standard requires UNBOUNDED FOLLOWING in that case. - // When true, corrects this default. Cannot distinguish defaulted from - // explicit frames, so an explicit CURRENT ROW may be incorrectly overridden. + // When true, corrects this default. An explicit RANGE frame ending at the + // current row is corrected as well, since unordered it covers the whole + // partition either way. An explicit ROWS frame ends where it says. bool correctWindowFrameDefault = false; /// SQL functions could be registered with different prefixes by the user. diff --git a/velox/duckdb/conversion/tests/DuckParserTest.cpp b/velox/duckdb/conversion/tests/DuckParserTest.cpp index cfd912cec93..ee3646c2345 100644 --- a/velox/duckdb/conversion/tests/DuckParserTest.cpp +++ b/velox/duckdb/conversion/tests/DuckParserTest.cpp @@ -624,6 +624,44 @@ TEST(DuckParserTest, window) { parseWindow("nth_value(x, 3) over ()")); } +TEST(DuckParserTest, correctWindowFrameDefault) { + auto parse = [](const std::string& expr) { + ParseOptions options; + options.correctWindowFrameDefault = true; + return parseWindowExpr(expr, options)->toString(); + }; + + // Without ORDER BY every row is a peer, so a RANGE frame ending at the + // current row covers the whole partition. A window with no frame clause + // gets that frame. + EXPECT_EQ( + "row_number() OVER (PARTITION BY \"a\" " + "RANGE BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)", + parse("row_number() over (partition by a)")); + EXPECT_EQ( + "row_number() OVER (PARTITION BY \"a\" " + "RANGE BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)", + parse( + "row_number() over (partition by a " + "range between unbounded preceding and current row)")); + + // A ROWS frame counts rows rather than peers, so it ends at the current row + // even with nothing ordered. + EXPECT_EQ( + "row_number() OVER (PARTITION BY \"a\" " + "ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)", + parse( + "row_number() over (partition by a " + "rows between unbounded preceding and current row)")); + + // With ORDER BY the peers are the rows that tie, so the default frame ends + // at the current row as written. + EXPECT_EQ( + "row_number() OVER (PARTITION BY \"a\" ORDER BY \"b\" ASC NULLS LAST " + "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)", + parse("row_number() over (partition by a order by b)")); +} + TEST(DuckParserTest, windowWithIntegerConstant) { ParseOptions options; options.parseIntegerAsBigint = false; diff --git a/velox/dwio/CMakeLists.txt b/velox/dwio/CMakeLists.txt index 13a22460e90..240159cd6e1 100644 --- a/velox/dwio/CMakeLists.txt +++ b/velox/dwio/CMakeLists.txt @@ -37,3 +37,7 @@ add_subdirectory(dwrf) add_subdirectory(orc) add_subdirectory(parquet) add_subdirectory(text) + +if(VELOX_ENABLE_NIMBLE) + add_subdirectory(nimble) +endif() diff --git a/velox/dwio/common/ColumnLoader.cpp b/velox/dwio/common/ColumnLoader.cpp index 32d222409da..899c01a4474 100644 --- a/velox/dwio/common/ColumnLoader.cpp +++ b/velox/dwio/common/ColumnLoader.cpp @@ -52,6 +52,9 @@ RowSet read( effectiveRows = selectedRows; } + // Load any deferred input streams for this column before decoding. Done on + // the reader so it runs regardless of which loader subclass is used. + fieldReader->formatData().loadLazyInputStreams(); structReader->advanceFieldReader(fieldReader, offset); fieldReader->scanSpec()->setValueHook(hook); fieldReader->readWithTiming(offset, effectiveRows, incomingNulls); diff --git a/velox/dwio/common/DecoderUtil.h b/velox/dwio/common/DecoderUtil.h index 40151160dad..15355be2795 100644 --- a/velox/dwio/common/DecoderUtil.h +++ b/velox/dwio/common/DecoderUtil.h @@ -17,6 +17,11 @@ #pragma once #include +#include +#include + +#include + #include "velox/common/base/Portability.h" #include "velox/common/memory/RawVector.h" #include "velox/common/process/ProcessBase.h" @@ -145,12 +150,26 @@ inline int32_t firstNullIndex(const uint64_t* nulls, int32_t numRows) { return first; } +// Returns raw bytes because fixed-width page data may be unaligned. +template +const char* +fixedWidthValueBytes(const T* buffer, int32_t row, int32_t rowOffset) { + return reinterpret_cast(buffer) + (row - rowOffset) * sizeof(T); +} + template void scatterDense( const Any* data, const int32_t* indices, int32_t size, T* target) { + if constexpr (std::is_same_v) { + for (auto i = 0; i < size; ++i) { + target[indices[i]] = folly::loadUnaligned(data + i * sizeof(T)); + } + return; + } + auto source = reinterpret_cast(data); if (source >= target && source < target + indices[size - 1]) { for (int32_t i = size - 1; i >= 0; --i) { @@ -309,7 +328,7 @@ void fixedWidthScan( if (isDense(&rows[rowIndex], numRowsInBuffer)) { std::memcpy( rawValues + numValues, - buffer + rows[rowIndex] - rowOffset, + fixedWidthValueBytes(buffer, rows[rowIndex], rowOffset), sizeof(T) * numRowsInBuffer); numValues += numRowsInBuffer; return; @@ -322,23 +341,19 @@ void fixedWidthScan( [&](int32_t rowIndex) { auto firstRow = rows[rowIndex]; if (!hasFilter) { + auto* firstValue = + fixedWidthValueBytes(buffer, firstRow, rowOffset); if (hasHook) { - hook.addValues( - scatterRows + rowIndex, - buffer + firstRow - rowOffset, - kStep); + T values[kStep]; + std::memcpy(values, firstValue, sizeof(T) * kStep); + hook.addValues(scatterRows + rowIndex, values, kStep); } else { if (scatter) { scatterDense( - buffer + firstRow - rowOffset, - scatterRows + rowIndex, - kStep, - rawValues); + firstValue, scatterRows + rowIndex, kStep, rawValues); } else { FOLLY_BUILTIN_MEMCPY( - rawValues + numValues, - buffer + firstRow - rowOffset, - sizeof(T) * kStep); + rawValues + numValues, firstValue, sizeof(T) * kStep); } } numValues += kStep; @@ -525,7 +540,7 @@ bool nonNullRowsFromSparse( template bool useFastPath(Visitor& visitor) { return (!std::is_same_v) && - process::hasAvx2() && Visitor::FilterType::deterministic && + process::hasSimd() && Visitor::FilterType::deterministic && Visitor::kHasBulkPath && (std:: is_same_v || diff --git a/velox/dwio/common/DirectBufferedInput.cpp b/velox/dwio/common/DirectBufferedInput.cpp index ebd64fd43d6..1bbd2c97451 100644 --- a/velox/dwio/common/DirectBufferedInput.cpp +++ b/velox/dwio/common/DirectBufferedInput.cpp @@ -202,7 +202,8 @@ void DirectBufferedInput::readRegion( groupId_.id(), requests, pool_, - options_.loadQuantum()); + options_.loadQuantum(), + options_.directBufferedInputSharedAllocation()); coalescedLoads_.push_back(load); streamToCoalescedLoad_.withWLock([&](auto& loads) { for (auto& request : requests) { @@ -292,20 +293,37 @@ bool duplicateRegion(const LoadRequest& source, const LoadRequest& duplicate) { duplicate.region.length == source.region.length; } -void copyDuplicateRegion( +// Gives 'duplicate' its own copy of the source's bytes, unless it can share the +// source's slice of the shared allocation. Duplicates left uncopied resolve to +// the source's slice in getData(), which finds the first request at the offset. +// Gates on 'useSharedAllocation' rather than the source's 'sharedData' so a 3+ +// duplicate chain still copies correctly. +void maybeCopyDuplicateRegion( const LoadRequest& source, LoadRequest& duplicate, - memory::MemoryPool* pool) { - VELOX_CHECK_EQ(source.loadSize, duplicate.loadSize); - if (source.data.numPages() > 0) { + memory::MemoryPool* pool, + bool useSharedAllocation) { + if (useSharedAllocation && + duplicate.region.length > DirectBufferedInput::kTinySize) { + return; + } + VELOX_CHECK_EQ(source.buffer.requestBytes, duplicate.buffer.requestBytes); + VELOX_CHECK(!source.buffer.empty(), "Duplicate region source is empty"); + VELOX_CHECK( + duplicate.buffer.empty(), "Duplicate region destination is already set"); + if (source.buffer.ownedData.numPages() > 0) { const auto numPages = - memory::AllocationTraits::numPages(duplicate.loadSize); - pool->allocateNonContiguous(numPages, duplicate.data); - memory::Allocation::copy(source.data, duplicate.data, duplicate.loadSize); + memory::AllocationTraits::numPages(duplicate.buffer.requestBytes); + pool->allocateNonContiguous(numPages, duplicate.buffer.ownedData); + memory::Allocation::copy( + source.buffer.ownedData, + duplicate.buffer.ownedData, + duplicate.buffer.requestBytes); } else { VELOX_CHECK( - !source.tinyData.empty(), "Duplicate tiny region source is empty"); - duplicate.tinyData = source.tinyData; + !source.buffer.tinyData.empty(), + "Duplicate tiny region source is empty"); + duplicate.buffer.tinyData = source.buffer.tinyData; } } } // namespace @@ -375,102 +393,169 @@ std::unique_ptr DirectBufferedInput::read( } std::vector DirectCoalescedLoad::loadData(bool prefetch) { - std::vector> buffers; - int64_t lastEnd = requests_[0].region.offset; + const int64_t overAllocatedBytes = computeLoadSizes(); + + // Use the shared allocation only when per-request page-rounding would + // over-allocate more than its minimum run; below that its own floor would + // over-reserve. + constexpr int64_t kMinPaddingBytesToUseSharedAllocation = + static_cast(memory::AllocationPool::kMinPages) * + static_cast(memory::AllocationTraits::kPageSize); + const bool useSharedAllocation = sharedAllocationEnabled_ && + overAllocatedBytes > kMinPaddingBytesToUseSharedAllocation; + int64_t size = 0; int64_t overread = 0; + const auto buffers = buildReadRanges(useSharedAllocation, size, overread); + + uint64_t usecs{0}; + { + MicrosecondWallTimer timer(&usecs); + input_->read(buffers, requests_[0].region.offset, LogType::FILE); + } + + ioStatistics_->read().increment(size + overread); + ioStatistics_->incRawBytesRead(size); + ioStatistics_->incTotalScanTimeNs(static_cast(usecs * 1'000)); + ioStatistics_->storageReadLatencyUs().increment(usecs); + ioStatistics_->incRawOverreadBytes(overread); + if (prefetch) { + ioStatistics_->prefetch().increment(size + overread); + } + + fillDuplicates(useSharedAllocation); + + TestValue::adjust( + "facebook::velox::cache::DirectCoalescedLoad::loadData", this); + return {}; +} +int64_t DirectCoalescedLoad::computeLoadSizes() { + int64_t overAllocatedBytes{0}; for (size_t i = 0; i < requests_.size(); ++i) { auto& request = requests_[i]; const auto& region = request.region; if (i > 0 && duplicateRegion(requests_[i - 1], request)) { - const auto& prev = requests_[i - 1]; - request.loadSize = prev.loadSize; + request.buffer.requestBytes = requests_[i - 1].buffer.requestBytes; continue; } - - if (region.offset > lastEnd) { - buffers.push_back( - folly::Range( - nullptr, - reinterpret_cast( - static_cast(region.offset - lastEnd)))); - overread += buffers.back().size(); - } - if (region.length > DirectBufferedInput::kTinySize) { if (&request != &requests_.back()) { - // Case where request is a little over quantum but is followed by - // another within the max distance. Coalesces and allows reading the - // region of max quantum + max distance in one piece. - request.loadSize = region.length; + // Request is a little over quantum but is followed by another within + // the max distance. Coalesce and read the whole region in one piece. + // This is the only path that does not clamp to 'loadQuantum_', so it is + // the only one where an oversized region could overflow the int32_t + // 'requestBytes'. Fail loudly instead of truncating: a wrapped value + // makes buildReadRanges() read fewer bytes than the region and hand the + // stream silently wrong data. + VELOX_CHECK_LE( + region.length, + static_cast(std::numeric_limits::max()), + "Coalesced region is too large to read in one piece"); + request.buffer.requestBytes = static_cast(region.length); } else { - request.loadSize = std::min(region.length, loadQuantum_); + request.buffer.requestBytes = static_cast( + std::min(region.length, loadQuantum_)); } - const auto numPages = - memory::AllocationTraits::numPages(request.loadSize); - pool_->allocateNonContiguous(numPages, request.data); - appendRanges(request.data, request.loadSize, buffers); + const auto paddedBytes = memory::AllocationTraits::roundUpPageBytes( + request.buffer.requestBytes); + overAllocatedBytes += + static_cast(paddedBytes) - request.buffer.requestBytes; } else { - request.loadSize = region.length; - request.tinyData.resize(region.length); - buffers.push_back(folly::Range(request.tinyData.data(), region.length)); + request.buffer.requestBytes = static_cast(region.length); } - lastEnd = region.offset + request.loadSize; - size += request.loadSize; } + return overAllocatedBytes; +} - uint64_t usecs = 0; - { - MicrosecondWallTimer timer(&usecs); - input_->read(buffers, requests_[0].region.offset, LogType::FILE); - } +std::vector> DirectCoalescedLoad::buildReadRanges( + bool useSharedAllocation, + int64_t& size, + int64_t& overread) { + std::vector> buffers; + uint64_t lastEnd = requests_[0].region.offset; + for (size_t i = 0; i < requests_.size(); ++i) { + auto& request = requests_[i]; + const auto& region = request.region; + if (i > 0 && duplicateRegion(requests_[i - 1], request)) { + // Shares the source's buffer; filled by fillDuplicates() after the read. + continue; + } + if (region.offset > lastEnd) { + buffers.emplace_back( + nullptr, + // NOLINTNEXTLINE(performance-no-int-to-ptr) + reinterpret_cast( + static_cast(region.offset - lastEnd))); + overread += static_cast(buffers.back().size()); + } - ioStatistics_->read().increment(size + overread); - ioStatistics_->incRawBytesRead(size); - ioStatistics_->incTotalScanTimeNs(usecs * 1'000); - ioStatistics_->queryThreadIoLatencyUs().increment(usecs); - ioStatistics_->storageReadLatencyUs().increment(usecs); - ioStatistics_->incRawOverreadBytes(overread); - if (prefetch) { - ioStatistics_->prefetch().increment(size + overread); + if (region.length <= DirectBufferedInput::kTinySize) { + request.buffer.tinyData.resize(region.length); + buffers.emplace_back(request.buffer.tinyData.data(), region.length); + } else if (useSharedAllocation) { + request.buffer.sharedData = + sharedAllocation_.allocateFixed(request.buffer.requestBytes); + buffers.emplace_back( + request.buffer.sharedData, request.buffer.requestBytes); + } else { + const auto numPages = + memory::AllocationTraits::numPages(request.buffer.requestBytes); + pool_->allocateNonContiguous(numPages, request.buffer.ownedData); + appendRanges( + request.buffer.ownedData, request.buffer.requestBytes, buffers); + } + lastEnd = region.offset + request.buffer.requestBytes; + size += request.buffer.requestBytes; } + return buffers; +} + +void DirectCoalescedLoad::fillDuplicates(bool useSharedAllocation) { for (size_t i = 0; i < requests_.size(); ++i) { auto& request = requests_[i]; if (i == 0 || !duplicateRegion(requests_[i - 1], request)) { continue; } - copyDuplicateRegion(requests_[i - 1], request, pool_); + maybeCopyDuplicateRegion( + requests_[i - 1], request, pool_, useSharedAllocation); } - TestValue::adjust( - "facebook::velox::cache::DirectCoalescedLoad::loadData", this); - return {}; } -int32_t DirectCoalescedLoad::getData( - int64_t offset, - memory::Allocation& data, - std::string& tinyData) { +LoadedBuffer DirectCoalescedLoad::getData(uint64_t offset) { + // A failed read may have partially filled the request buffers. Only publish + // them after the complete coalesced load has succeeded. + if (state() != CoalescedLoad::State::kLoaded) { + return {}; + } auto it = std::lower_bound( - requests_.begin(), requests_.end(), offset, [](auto& x, auto offset) { - return x.region.offset < offset; + requests_.begin(), + requests_.end(), + offset, + [](const auto& request, uint64_t targetOffset) { + return request.region.offset < targetOffset; }); - if (it == requests_.cend() || it->region.offset != offset) { - return 0; + if (it == requests_.end() || it->region.offset != offset) { + return {}; } - // Duplicate regions have the same offset. Skip buffers already handed to - // earlier streams so each duplicate stream gets its own copied buffer. + if (it->buffer.sharedData != nullptr) { + // Duplicates backed by the shared allocation share one read-only slice; + // nothing is moved. + return LoadedBuffer{ + .requestBytes = it->buffer.requestBytes, + .sharedData = it->buffer.sharedData, + }; + } + // Skip buffers already handed to earlier duplicate streams. while (it != requests_.end() && it->region.offset == offset && it->bufferConsumed) { ++it; } - if (it == requests_.cend() || it->region.offset != offset) { - return 0; + if (it == requests_.end() || it->region.offset != offset) { + return {}; } - data = std::move(it->data); - tinyData = std::move(it->tinyData); it->bufferConsumed = true; - return it->loadSize; + return std::move(it->buffer); } } // namespace facebook::velox::dwio::common diff --git a/velox/dwio/common/DirectBufferedInput.h b/velox/dwio/common/DirectBufferedInput.h index ac8b88b372d..daa353ef91c 100644 --- a/velox/dwio/common/DirectBufferedInput.h +++ b/velox/dwio/common/DirectBufferedInput.h @@ -24,12 +24,35 @@ #include "velox/common/caching/ScanTracker.h" #include "velox/common/io/IoStatistics.h" #include "velox/common/io/Options.h" +#include "velox/common/memory/AllocationPool.h" #include "velox/dwio/common/BufferedInput.h" #include "velox/dwio/common/CacheInputStream.h" #include "velox/dwio/common/InputStream.h" namespace facebook::velox::dwio::common { +/// Loaded bytes for one region. At most one of the three representations +/// below -- 'sharedData', 'ownedData', 'tinyData' -- is ever set; all three are +/// unset until the region is loaded. +struct LoadedBuffer { + /// Number of request bytes, or 0 if the region was not found. + int32_t requestBytes{0}; + /// Borrowed read-only slice in the load's shared allocation. The consumer + /// keeps the load alive while reading it. + char* sharedData{nullptr}; + /// Owned non-contiguous allocation, moved to the consumer. + memory::Allocation ownedData{}; + /// Owned bytes for a tiny region, moved to the consumer. + std::string tinyData{}; + + /// True when no representation holds bytes, i.e. nothing has been loaded for + /// this region yet. + bool empty() const { + return sharedData == nullptr && ownedData.numPages() == 0 && + tinyData.empty(); + } +}; + struct LoadRequest { LoadRequest() = default; LoadRequest(velox::common::Region& _region, cache::TrackingId _trackingId) @@ -46,14 +69,10 @@ struct LoadRequest { const SeekableInputStream* stream; - /// Buffers to be handed to 'stream' after load. - memory::Allocation data; - std::string tinyData; - /// Number of bytes in 'data/tinyData'. - int32_t loadSize{0}; - // Set after getData() moves 'data/tinyData' to the owning stream. Duplicate - // regions share an offset, so getData() skips consumed buffers to find the - // next duplicate buffer. + /// Loaded bytes for this request; see LoadedBuffer. + LoadedBuffer buffer; + + /// Set once getData() moved 'buffer' out; lets it skip consumed duplicates. bool bufferConsumed{false}; }; @@ -67,13 +86,16 @@ class DirectCoalescedLoad : public cache::CoalescedLoad { uint64_t /* groupId */, const std::vector& requests, memory::MemoryPool* pool, - int32_t loadQuantum) + int32_t loadQuantum, + bool sharedAllocationEnabled) : CoalescedLoad({}, {}), ioStatistics_(ioStatistics), ioStats_(ioStats), input_(std::move(input)), loadQuantum_(loadQuantum), - pool_(pool) { + sharedAllocationEnabled_(sharedAllocationEnabled), + pool_(pool), + sharedAllocation_(pool) { VELOX_DCHECK_NOT_NULL(pool_); VELOX_DCHECK( std::is_sorted( @@ -96,10 +118,9 @@ class DirectCoalescedLoad : public cache::CoalescedLoad { return false; } - /// Returns the buffer for 'region' in either 'data' or 'tinyData'. 'region' - /// must match a region given to DirectBufferedInput::enqueue(). - int32_t - getData(int64_t offset, memory::Allocation& data, std::string& tinyData); + /// Returns the loaded buffer for the request at 'offset'. 'offset' must match + /// a region given to DirectBufferedInput::enqueue(). + LoadedBuffer getData(uint64_t offset); const std::vector& requests() { return requests_; @@ -114,11 +135,33 @@ class DirectCoalescedLoad : public cache::CoalescedLoad { } private: + // Sets each request's 'buffer.requestBytes' and returns the total + // page-padding bytes that per-request allocations would over-allocate. + int64_t computeLoadSizes(); + + // Allocates each non-duplicate request's buffer and returns the file-ordered + // ranges for the coalesced read. Adds read bytes to 'size', gap bytes to + // 'overread'; serves non-tiny buffers from the shared allocation when + // 'useSharedAllocation' is set. + std::vector> + buildReadRanges(bool useSharedAllocation, int64_t& size, int64_t& overread); + + // Gives each duplicate region its buffer after the read: a copy of the source + // for tiny duplicates and for per-request allocations; duplicates backed by + // the shared allocation share the source's slice and are left untouched. + void fillDuplicates(bool useSharedAllocation); + const std::shared_ptr ioStatistics_; const std::shared_ptr ioStats_; const std::shared_ptr input_; const int32_t loadQuantum_; + // Whether the shared allocation may back this load's buffers at all. When + // false, every request keeps its own page-rounded allocation. + const bool sharedAllocationEnabled_; memory::MemoryPool* const pool_; + // Shared allocation backing all non-tiny request buffers; bump-packs them + // into a few allocations, freed as a unit on destruction. + memory::AllocationPool sharedAllocation_; std::vector requests_; }; diff --git a/velox/dwio/common/DirectInputStream.cpp b/velox/dwio/common/DirectInputStream.cpp index b14c989169e..cc4f3a0e911 100644 --- a/velox/dwio/common/DirectInputStream.cpp +++ b/velox/dwio/common/DirectInputStream.cpp @@ -130,20 +130,59 @@ makeRanges(size_t size, memory::Allocation& data, std::string& tinyData) { } } // namespace +void DirectInputStream::LoadedData::set( + LoadedBuffer&& loaded, + std::shared_ptr load) { + // 'owned' must be empty; loadPosition() reaches set() once per stream under + // 'loaded_'. Allocation's move assignment rebinds the runs without + // returning the old ones to the pool, so overwriting a live allocation + // loses its pages while the pool keeps counting them until teardown. + // valid() below does not catch that: it counts live representations, and a + // dropped one still reads as valid. Hard check, not a DCHECK, because the + // failure is a silent leak; not repaired by freeing here, because a second + // set() on one stream means the load lifecycle changed and needs review. + VELOX_CHECK_EQ(owned.numPages(), 0, "set() would drop a live allocation"); + + // Drop any previous slice first. Adopting 'loaded' must not leave a stale + // borrowed pointer live alongside the new representation, since + // loadPosition() dispatches on the shared slice ahead of the owned ones. + resetShared(); + owned = std::move(loaded.ownedData); + tiny = std::move(loaded.tinyData); + if (loaded.sharedData != nullptr) { + sharedPtr = loaded.sharedData; + sharedHolder = std::move(load); + } + VELOX_CHECK(valid(), "Loaded data has multiple live representations"); +} + +bool DirectInputStream::LoadedData::valid() const { + return (static_cast(sharedPtr != nullptr) + + static_cast(owned.numPages() > 0) + + static_cast(!tiny.empty())) <= 1; +} + +void DirectInputStream::LoadedData::resetShared() { + sharedPtr = nullptr; + sharedHolder.reset(); +} + void DirectInputStream::loadSync() { if (region_.length < DirectBufferedInput::kTinySize && - data_.numPages() == 0) { - tinyData_.resize(region_.length); + loadedData_.owned.numPages() == 0) { + loadedData_.tiny.resize(region_.length); } else { const auto numPages = memory::AllocationTraits::numPages(loadedRegion_.length); - if (numPages > data_.numPages()) { - bufferedInput_->pool()->allocateNonContiguous(numPages, data_); + if (numPages > loadedData_.owned.numPages()) { + bufferedInput_->pool()->allocateNonContiguous( + numPages, loadedData_.owned); } } ioStats_->incRawBytesRead(loadedRegion_.length); - auto ranges = makeRanges(loadedRegion_.length, data_, tinyData_); + auto ranges = + makeRanges(loadedRegion_.length, loadedData_.owned, loadedData_.tiny); uint64_t usecs = 0; { MicrosecondWallTimer timer(&usecs); @@ -186,7 +225,9 @@ void DirectInputStream::loadPosition() { waitFuture.wait(); } loadedRegion_.offset = region_.offset; - loadedRegion_.length = load->getData(region_.offset, data_, tinyData_); + auto loaded = load->getData(region_.offset); + loadedRegion_.length = loaded.requestBytes; + loadedData_.set(std::move(loaded), load); } ioStats_->queryThreadIoLatencyUs().increment(loadUs); // DirectCoalescedLoad always reads from remote storage, not SSD. @@ -203,6 +244,9 @@ void DirectInputStream::loadPosition() { region_.offset + offsetInRegion_ < loadedRegion_.offset || region_.offset + offsetInRegion_ >= loadedRegion_.offset + loadedRegion_.length) { + // Outside the loaded range: drop the borrowed slice; loadSync() reloads + // below. + loadedData_.resetShared(); loadedRegion_.offset = region_.offset + offsetInRegion_; loadedRegion_.length = (offsetInRegion_ + loadQuantum_ <= region_.length) ? loadQuantum_ @@ -213,17 +257,26 @@ void DirectInputStream::loadPosition() { loadSync(); } + VELOX_DCHECK( + loadedData_.valid(), + "DirectInputStream has multiple live buffer representations"); + const auto offsetInData = offsetInRegion_ - (loadedRegion_.offset - region_.offset); - if (data_.numPages() == 0) { - run_ = reinterpret_cast(tinyData_.data()); - runSize_ = tinyData_.size(); + if (loadedData_.hasShared()) { + run_ = reinterpret_cast(const_cast(loadedData_.sharedPtr)); + runSize_ = static_cast(loadedRegion_.length); + offsetInRun_ = static_cast(offsetInData); + offsetOfRun_ = 0; + } else if (loadedData_.owned.numPages() == 0) { + run_ = reinterpret_cast(loadedData_.tiny.data()); + runSize_ = static_cast(loadedData_.tiny.size()); offsetInRun_ = offsetInData; offsetOfRun_ = 0; } else { - data_.findRun(offsetInData, &runIndex_, &offsetInRun_); + loadedData_.owned.findRun(offsetInData, &runIndex_, &offsetInRun_); offsetOfRun_ = offsetInData - offsetInRun_; - auto run = data_.runAt(runIndex_); + auto run = loadedData_.owned.runAt(runIndex_); run_ = run.data(); runSize_ = memory::AllocationTraits::pageBytes(run.numPages()); if (offsetOfRun_ + runSize_ > loadedRegion_.length) { diff --git a/velox/dwio/common/DirectInputStream.h b/velox/dwio/common/DirectInputStream.h index b5bc4440066..b4460b06c34 100644 --- a/velox/dwio/common/DirectInputStream.h +++ b/velox/dwio/common/DirectInputStream.h @@ -25,6 +25,7 @@ namespace facebook::velox::dwio::common { class DirectBufferedInput; +struct LoadedBuffer; /// An input stream over possibly coalesced loads. Created by /// DirectBufferedInput. Similar to CacheInputStream but does not use cache. @@ -56,15 +57,15 @@ class DirectInputStream : public SeekableInputStream { memory::Allocation*& data, std::string*& tinyData) { loadedRegion = loadedRegion_; - data = &data_; - tinyData = &tinyData_; + data = &loadedData_.owned; + tinyData = &loadedData_.tiny; } private: - // Ensures that the current position is covered by 'data_'. + // Ensures that the current position is covered by 'loadedData_'. void loadPosition(); - // Synchronously sets 'data_' to cover loadedRegion_'. + // Synchronously sets 'loadedData_' to cover 'loadedRegion_'. void loadSync(); DirectBufferedInput* const bufferedInput_; @@ -80,28 +81,59 @@ class DirectInputStream : public SeekableInputStream { // Maximum number of bytes read from 'input' at a time. const int32_t loadQuantum_; - // The part of 'region_' that is loaded into 'data_'/'tinyData_'. Relative to - // file start. + // The part of 'region_' that is loaded into 'loadedData_'. Relative to file + // start. velox::common::Region loadedRegion_; - // Allocation with loaded data. Has space for region.length or loadQuantum_ - // bytes, whichever is less. - memory::Allocation data_; - - // Contains the data if the range is too small for Allocation. - std::string tinyData_; + // The loaded bytes for 'loadedRegion_', held in exactly one of three + // representations. + struct LoadedData { + // Allocation with loaded data. Has space for region.length or loadQuantum_ + // bytes, whichever is less. + memory::Allocation owned; + + // Contains the data if the range is too small for Allocation. + std::string tiny; + + // Borrowed slice of the load's shared allocation plus a hold on the owning + // load, bundled so the pointer and keep-alive can never desync. Null when + // the bytes are not backed by the shared allocation. + const char* sharedPtr{nullptr}; + std::shared_ptr sharedHolder; + + // Adopts the buffers of a completed coalesced load. 'load' is retained only + // when the bytes are a borrowed slice of that load's shared allocation, so + // the slice cannot outlive its owner. + void set(LoadedBuffer&& loaded, std::shared_ptr load); + + // True when at most one representation holds bytes. The priority dispatch + // in loadPosition() picks the first live one, so a double-set would be + // silently masked rather than caught. + bool valid() const; + + // Drops the borrowed slice. 'owned' is deliberately left in place: an + // Allocation is only freeable through its pool, and loadSync() reuses it + // whenever it is already large enough, so clearing it here would force a + // fresh allocation on every quantum advance. + void resetShared(); + + bool hasShared() const { + return sharedPtr != nullptr; + } + }; + LoadedData loadedData_; // Pointer to start of current run in 'entry->nonContiguousData()' or // 'entry->contiguousData()'. uint8_t* run_{nullptr}; - // Offset of current run from start of 'data_' + // Offset of current run from start of 'loadedData_.owned' uint64_t offsetOfRun_; // Position of stream relative to 'run_'. int offsetInRun_{0}; - // Index of run in 'data_' + // Index of run in 'loadedData_.owned' int runIndex_ = -1; // Number of valid bytes starting at 'run_' diff --git a/velox/dwio/common/FormatData.h b/velox/dwio/common/FormatData.h index cce719d5312..13e149aead3 100644 --- a/velox/dwio/common/FormatData.h +++ b/velox/dwio/common/FormatData.h @@ -36,6 +36,10 @@ class FormatData { return *static_cast(this); } + /// Loads any deferred input streams for this column before it is decoded. + /// The default is a no-op; formats that defer stream loading override it. + virtual void loadLazyInputStreams() {} + /// Reads nulls if the format has nulls separate from the encoded /// data. If there are no nulls, 'nulls' is set to nullptr, else to /// a suitable sized and padded Buffer. 'incomingNulls' may be given @@ -145,7 +149,7 @@ class FormatData { /// Base class for format-specific reader initialization arguments. class FormatParams { public: - FormatParams(memory::MemoryPool& pool, ColumnReaderStatistics& stats) + FormatParams(memory::MemoryPool& pool, SplitStats& stats) : pool_(&pool), stats_(&stats) {} virtual ~FormatParams() = default; @@ -160,13 +164,21 @@ class FormatParams { return *pool_; } - ColumnReaderStatistics& runtimeStatistics() { + /// Returns the runtime statistics for a column, creating them if necessary. + /// @param id Schema node ID identifying the column. + /// @param typeKind Logical column type to record in the statistics. + ColumnRuntimeStats& columnStats(uint32_t id, TypeKind typeKind) { + return stats_->getOrCreateColumnStats(id, typeKind); + } + + /// Returns the runtime statistics accumulator for the current split. + SplitStats& splitStats() { return *stats_; } private: memory::MemoryPool* const pool_; - ColumnReaderStatistics* const stats_; + SplitStats* const stats_; }; } // namespace facebook::velox::dwio::common diff --git a/velox/dwio/common/Reader.h b/velox/dwio/common/Reader.h index 81046d9d1ae..8655c85479a 100644 --- a/velox/dwio/common/Reader.h +++ b/velox/dwio/common/Reader.h @@ -99,7 +99,7 @@ class RowReader { * implementation specific and depends on a format of a file being read. * @param stats stats to update */ - virtual void updateRuntimeStats(RuntimeStatistics& stats) const = 0; + virtual void updateRuntimeStats(RuntimeStats& stats) const = 0; /** * This method should be called whenever filter is modified in a ScanSpec diff --git a/velox/dwio/common/ScanSpec.cpp b/velox/dwio/common/ScanSpec.cpp index ca4d49955c6..2494949b4bf 100644 --- a/velox/dwio/common/ScanSpec.cpp +++ b/velox/dwio/common/ScanSpec.cpp @@ -157,7 +157,7 @@ bool ScanSpec::hasFilter() const { if (hasFilter_.has_value()) { return hasFilter_.value(); } - if (!isConstant() && filter()) { + if (filter()) { hasFilter_ = true; return true; } diff --git a/velox/dwio/common/SelectiveColumnReader.cpp b/velox/dwio/common/SelectiveColumnReader.cpp index 9a2cd25eba7..7dc5cf48715 100644 --- a/velox/dwio/common/SelectiveColumnReader.cpp +++ b/velox/dwio/common/SelectiveColumnReader.cpp @@ -59,10 +59,8 @@ SelectiveColumnReader::SelectiveColumnReader( scanState_.rowsCopy = raw_vector(pool_); scanState_.filterCache = raw_vector(pool_); // Initialize per-column decoding statistics if collection is enabled. - if (params.runtimeStatistics().decodingStatsSet) { - decodingStats_ = params.runtimeStatistics().decodingStatsSet->getOrCreate( - fileType_->id()); - } + auto& stats = params.columnStats(fileType_->id(), fileType_->type()->kind()); + decodingStats_ = stats.decodingStats ? &*stats.decodingStats : nullptr; } void SelectiveColumnReader::readWithTiming( diff --git a/velox/dwio/common/SelectiveColumnReader.h b/velox/dwio/common/SelectiveColumnReader.h index d43f41e55bf..611e235a0ab 100644 --- a/velox/dwio/common/SelectiveColumnReader.h +++ b/velox/dwio/common/SelectiveColumnReader.h @@ -296,15 +296,17 @@ class SelectiveColumnReader { numValues_ = size; } - // The number of passing after filtering. + // The number of result rows after filtering. int32_t numRows() const { return outputRows_.size(); } - // The number of values copied into the results. + // The number of result positions produced so far. This includes null + // positions; it is not the count of decoded non-null values. int32_t numValues() const { return numValues_; } + void setNumRows(vector_size_t size) { outputRows_.resize(size); } @@ -430,7 +432,7 @@ class SelectiveColumnReader { /// is used at read time and is expected to produce the same result. bool useBulkPath() const { auto* filter = scanSpec_->filter(); - return hasBulkPath() && process::hasAvx2() && + return hasBulkPath() && process::hasSimd() && (!filter || (filter->isDeterministic() && (!nullsInReadRange_ || !filter->testNull()))) && diff --git a/velox/dwio/common/SelectiveStructColumnReader.cpp b/velox/dwio/common/SelectiveStructColumnReader.cpp index a91d68223a8..35998a0d7a0 100644 --- a/velox/dwio/common/SelectiveStructColumnReader.cpp +++ b/velox/dwio/common/SelectiveStructColumnReader.cpp @@ -16,21 +16,13 @@ #include "velox/dwio/common/SelectiveStructColumnReader.h" +#include + #include "velox/dwio/common/ColumnLoader.h" namespace facebook::velox::dwio::common { namespace { -bool testFilterOnConstant(const velox::common::ScanSpec& spec) { - if (spec.isConstant() && !spec.constantValue()->isNullAt(0)) { - // Non-null constant is known value during split scheduling and filters on - // them should not be handled at execution level. - return true; - } - // Check filter on missing field. - return !spec.hasFilter() || spec.testNull(); -} - // Recursively makes empty RowVectors for positions in 'children' where the // corresponding child type in 'rowType' is a row. The reader expects RowVector // outputs to be initialized so that the content corresponds to the query schema @@ -338,10 +330,26 @@ void SelectiveStructColumnReaderBase::next( if (hasDeletion_) { fillOutputRowsFromMutation(numValues); numValues = outputRows_.size(); + } else if (useOutputRows()) { + // Nothing on this path populates 'outputRows_', but useOutputRows() is also + // true when the scan spec has a filter, so outputRows() would return an + // empty set while the result carries 'numValues' rows. Callers rely on the + // two agreeing -- the synthesized fields below, and + // RowReader::readWithRowNumber -- so keep the invariant here rather than at + // each use. No row is eliminated on this path: with no child readers the + // only filters are on constant columns, and those are all-or-nothing and + // handled below. + outputRows_.resize(numValues); + std::iota(outputRows_.begin(), outputRows_.end(), 0); } for (const auto& childSpec : scanSpec_->children()) { if (isChildConstant(*childSpec) && !testFilterOnConstant(*childSpec)) { outputRows_.clear(); + // A constant column's filter is not counted by ScanSpec::hasFilter(), so + // useOutputRows() can be false here; clear 'inputRows_' too, otherwise + // outputRows() would fall back to it and report rows that were filtered + // out. + inputRows_ = {}; numValues = 0; break; } @@ -383,6 +391,45 @@ void SelectiveStructColumnReaderBase::next( } } +void SelectiveStructColumnReaderBase::readFlatMapChildren( + int64_t offset, + const RowSet& rows, + const uint64_t* incomingNulls) { + numReads_ = scanSpec_->newRead(); + prepareRead(offset, rows, incomingNulls); + VELOX_DCHECK(!hasDeletion()); + auto activeRows = rows; + const auto* mapNulls = + nullsInReadRange_ ? nullsInReadRange_->as() : nullptr; + if (scanSpec_->filter()) { + const auto kind = scanSpec_->filter()->kind(); + VELOX_CHECK( + kind == velox::common::FilterKind::kIsNull || + kind == velox::common::FilterKind::kIsNotNull); + filterNulls( + rows, kind == velox::common::FilterKind::kIsNull, false); + if (outputRows_.empty()) { + for (auto* child : children_) { + child->addParentNulls(offset, mapNulls, rows); + } + lazyVectorReadOffset_ = offset; + readOffset_ = offset + rows.back() + 1; + return; + } + activeRows = outputRows_; + } + // Separate the loop to be cache friendly. + for (auto* child : children_) { + advanceFieldReader(child, offset); + } + for (auto* child : children_) { + child->readWithTiming(offset, activeRows, mapNulls); + child->addParentNulls(offset, mapNulls, rows); + } + lazyVectorReadOffset_ = offset; + readOffset_ = offset + rows.back() + 1; +} + void SelectiveStructColumnReaderBase::read( int64_t offset, const RowSet& rows, @@ -529,10 +576,16 @@ bool SelectiveStructColumnReaderBase::isChildMissing( TypeKind::MAP // If this is the case it means this is a flat map, // so it can't have "missing" fields. ) && - // Name-based missing-field check applies only to row types, not flat - // maps. + // Name-based missing-field check applies to row types when the mapping + // mode resolves columns by name or by Parquet field ID. In field-ID + // mode getParquetColumnInfo() renames each file column to match the + // requested name, so containsChild() correctly identifies missing ones. + // Channel-based detection is only reliable for kPosition mode, where + // channel i maps directly to file column i. ((fileType_->type()->isRow() && - columnReaderOptions_.columnMappingMode_ == ColumnMappingMode::kName) + (columnReaderOptions_.columnMappingMode_ == ColumnMappingMode::kName || + columnReaderOptions_.columnMappingMode_ == + ColumnMappingMode::kParquetFieldId)) ? !asRowType(fileType_->type())->containsChild(childSpec.fieldName()) : childSpec.channel() >= fileType_->size()); } @@ -692,6 +745,18 @@ void SelectiveStructColumnReaderBase::getValues( resultRow->invalidateContainsLazyNotLoaded(); } +// static +bool SelectiveStructColumnReaderBase::testFilterOnConstant( + const velox::common::ScanSpec& spec) { + if (spec.isConstant() && !spec.constantValue()->isNullAt(0)) { + // Non-null constant is known value during split scheduling and filters on + // them should not be handled at execution level. + return true; + } + // Check filter on missing field. + return !spec.hasFilter() || spec.testNull(); +} + namespace detail { #if XSIMD_WITH_AVX2 diff --git a/velox/dwio/common/SelectiveStructColumnReader.h b/velox/dwio/common/SelectiveStructColumnReader.h index 8fb11eb0782..af7ed943f95 100644 --- a/velox/dwio/common/SelectiveStructColumnReader.h +++ b/velox/dwio/common/SelectiveStructColumnReader.h @@ -105,6 +105,11 @@ class SelectiveStructColumnReaderBase : public SelectiveColumnReader { currentRowNumber_ = value; } + /// Evaluates split-time filtering for a constant field. + /// Returns true for non-null constants, or when no filter is present, or when + /// the filter allows null values. + static bool testFilterOnConstant(const velox::common::ScanSpec& spec); + protected: template friend class SelectiveFlatMapColumnReaderHelper; @@ -133,6 +138,13 @@ class SelectiveStructColumnReaderBase : public SelectiveColumnReader { return hasDeletion_; } + // Reads physical flat-map value streams directly without interpreting the + // logical children in the scan spec. + void readFlatMapChildren( + int64_t offset, + const RowSet& rows, + const uint64_t* incomingNulls); + // Returns true if the file doesn't have this child (in which case it will be // treated as null). bool isChildMissing(const velox::common::ScanSpec& childSpec) const; @@ -314,40 +326,7 @@ void SelectiveFlatMapColumnReaderHelper::read( int64_t offset, RowSet rows, const uint64_t* incomingNulls) { - reader_.numReads_ = reader_.scanSpec_->newRead(); - reader_.prepareRead(offset, rows, incomingNulls); - VELOX_DCHECK(!reader_.hasDeletion()); - auto activeRows = rows; - auto* mapNulls = reader_.nullsInReadRange_ - ? reader_.nullsInReadRange_->as() - : nullptr; - if (reader_.scanSpec_->filter()) { - auto kind = reader_.scanSpec_->filter()->kind(); - VELOX_CHECK( - kind == velox::common::FilterKind::kIsNull || - kind == velox::common::FilterKind::kIsNotNull); - reader_.filterNulls( - rows, kind == velox::common::FilterKind::kIsNull, false); - if (reader_.outputRows_.empty()) { - for (auto* child : reader_.children_) { - child->addParentNulls(offset, mapNulls, rows); - } - reader_.lazyVectorReadOffset_ = offset; - reader_.readOffset_ = offset + rows.back() + 1; - return; - } - activeRows = reader_.outputRows_; - } - // Separate the loop to be cache friendly. - for (auto* child : reader_.children_) { - reader_.advanceFieldReader(child, offset); - } - for (auto* child : reader_.children_) { - child->readWithTiming(offset, activeRows, mapNulls); - child->addParentNulls(offset, mapNulls, rows); - } - reader_.lazyVectorReadOffset_ = offset; - reader_.readOffset_ = offset + rows.back() + 1; + reader_.readFlatMapChildren(offset, rows, incomingNulls); } namespace detail { diff --git a/velox/dwio/common/Statistics.cpp b/velox/dwio/common/Statistics.cpp index cec3f825a2d..f9f61d139e1 100644 --- a/velox/dwio/common/Statistics.cpp +++ b/velox/dwio/common/Statistics.cpp @@ -33,6 +33,16 @@ std::string toStringOr( return std::string{fallback}; } +void mergeRuntimeMetric( + std::string name, + const RuntimeMetric& metric, + std::unordered_map& result) { + auto [it, inserted] = result.emplace(std::move(name), metric); + if (!inserted) { + it->second.merge(metric); + } +} + } // namespace bool KeyInfo::operator==(const KeyInfo& other) const { @@ -127,9 +137,9 @@ std::string TimestampColumnStatistics::toString() const { return folly::to( ColumnStatistics::toString(), ", min: ", - (min_.has_value() ? min_.value().toString() : "unknown"), + toStringOr(min_, kUnknown), ", max: ", - (max_.has_value() ? max_.value().toString() : "unknown")); + toStringOr(max_, kUnknown)); } std::string StringColumnStatistics::toString() const { @@ -164,122 +174,138 @@ void DecodingStats::merge(const DecodingStats& other) { decodeCPUTimeNanos.merge(other.decodeCPUTimeNanos); } -DecodingStats* DecodingStatsSet::getOrCreate( - uint32_t nodeId, - TypeKind typeKind) { - auto locked = map_.wlock(); - auto it = locked->find(nodeId); - if (it == locked->end()) { - it = locked->emplace(nodeId, std::make_unique(typeKind)) - .first; - } - return it->second.get(); +void DecodingStats::toRuntimeMetrics( + std::string_view prefix, + std::unordered_map& result) const { + const auto addCounter = [&](std::string_view name, const auto& counter) { + if (counter.count() == 0) { + return; + } + mergeRuntimeMetric( + fmt::format("{}.{}", prefix, name), + RuntimeMetric{ + saturateCast(counter.sum()), + counter.count(), + saturateCast(counter.min()), + saturateCast(counter.max()), + RuntimeCounter::Unit::kNanos}, + result); + }; + addCounter("decompressCPUTimeNanos", decompressCPUTimeNanos); + addCounter("decodeCPUTimeNanos", decodeCPUTimeNanos); } -void DecodingStatsSet::mergeFrom(const DecodingStatsSet& other) { - auto srcLocked = other.map_.rlock(); - auto dstLocked = map_.wlock(); - for (const auto& [nodeId, srcStats] : *srcLocked) { - auto it = dstLocked->find(nodeId); - if (it == dstLocked->end()) { - it = dstLocked->emplace(nodeId, std::make_unique()).first; - it->second->typeKind = srcStats->typeKind; - } - it->second->merge(*srcStats); +void ColumnRuntimeStats::accumulateStat( + const std::pair& stat, + int64_t value) { + auto [it, inserted] = columnMetrics.try_emplace(stat.first); + if (inserted) { + it->second.unit = stat.second; + } else { + VELOX_CHECK_EQ(it->second.unit, stat.second); } + it->second.addValue(value); } -void DecodingStatsSet::toRuntimeMetrics( - std::unordered_map& result) const { - auto statsLocked = map_.rlock(); - for (const auto& [nodeId, stats] : *statsLocked) { - // Export decompression timing. - const auto& decompressCounter = stats->decompressCPUTimeNanos; - if (decompressCounter.count() > 0) { - result.emplace( - fmt::format( - "column_{}.{}.decompressCPUTimeNanos", - nodeId, - TypeKindName::toName(stats->typeKind)), - RuntimeMetric{ - saturateCast(decompressCounter.sum()), - decompressCounter.count(), - saturateCast(decompressCounter.min()), - saturateCast(decompressCounter.max()), - RuntimeCounter::Unit::kNanos}); - } - // Export decode timing. - const auto& decodeCounter = stats->decodeCPUTimeNanos; - if (decodeCounter.count() > 0) { - result.emplace( - fmt::format( - "column_{}.{}.decodeCPUTimeNanos", - nodeId, - TypeKindName::toName(stats->typeKind)), - RuntimeMetric{ - saturateCast(decodeCounter.sum()), - decodeCounter.count(), - saturateCast(decodeCounter.min()), - saturateCast(decodeCounter.max()), - RuntimeCounter::Unit::kNanos}); - } - } +ColumnRuntimeStats& SplitStats::getOrCreateColumnStats( + uint32_t nodeId, + TypeKind typeKind) { + auto [it, inserted] = columnStats.try_emplace(nodeId, typeKind); + // TODO(#18171): We need to add the following check once we stable column + // ID assignment. For now, we can have different typeKinds for the same + // column ID. + // if (!inserted) { + // VELOX_CHECK_EQ(it->second.typeKind, typeKind); + // } + return it->second; +} + +DecodingStats* SplitStats::decodingStats(uint32_t nodeId) { + const auto it = columnStats.find(nodeId); + return it != columnStats.end() && it->second.decodingStats + ? &*it->second.decodingStats + : nullptr; } -void ColumnReaderStatistics::initColumnStatsCollection( +void SplitStats::initColumnStatsCollection( const TypeWithId& schema, const RowReaderOptions& options) { - if (!options.collectColumnCpuMetrics()) { - return; - } - decodingStatsSet.emplace(); - registerDecodingStatsImpl(schema); + registerColumnStats(schema, options.collectColumnCpuMetrics()); } -void ColumnReaderStatistics::mergeFrom(const ColumnReaderStatistics& other) { - flattenStringDictionaryValues += other.flattenStringDictionaryValues; - pageLoadTimeNs.merge(other.pageLoadTimeNs); - if (other.decodingStatsSet) { - if (!decodingStatsSet) { - decodingStatsSet.emplace(); +void ColumnRuntimeStats::mergeFrom(const ColumnRuntimeStats& other) { + // TODO(#18171): Add the typeKind check once column ID assignment is stable. + // For now, the same column ID can refer to different typeKinds. + // VELOX_CHECK_EQ(typeKind, other.typeKind); + for (const auto& [name, metric] : other.columnMetrics) { + auto [it, inserted] = columnMetrics.emplace(name, metric); + if (!inserted) { + it->second.merge(metric); } - decodingStatsSet->mergeFrom(*other.decodingStatsSet); + } + if (other.decodingStats) { + if (!decodingStats) { + decodingStats.emplace(); + } + decodingStats->merge(*other.decodingStats); } } -void ColumnReaderStatistics::toRuntimeMetrics( +void ColumnRuntimeStats::toRuntimeMetrics( + std::string_view prefix, std::unordered_map& result) const { - if (flattenStringDictionaryValues > 0) { - result.emplace( - "flattenStringDictionaryValues", - RuntimeMetric(flattenStringDictionaryValues)); + for (const auto& [name, metric] : columnMetrics) { + mergeRuntimeMetric(fmt::format("{}.{}", prefix, name), metric, result); } - if (pageLoadTimeNs.sum() > 0) { - result.emplace( - "pageLoadTimeNanos", - RuntimeMetric( - pageLoadTimeNs.sum(), - pageLoadTimeNs.count(), - pageLoadTimeNs.min(), - pageLoadTimeNs.max(), - RuntimeCounter::Unit::kNanos)); + if (decodingStats) { + decodingStats->toRuntimeMetrics(prefix, result); } - if (decodingStatsSet) { - decodingStatsSet->toRuntimeMetrics(result); +} + +void SplitStats::accumulateStat( + const std::pair& stat, + int64_t value) { + auto [it, inserted] = splitMetrics.try_emplace(stat.first); + if (inserted) { + it->second.unit = stat.second; + } else { + VELOX_CHECK_EQ(it->second.unit, stat.second); } + it->second.addValue(value); } -void ColumnReaderStatistics::registerDecodingStatsImpl(const TypeWithId& node) { - decodingStatsSet->getOrCreate(node.id(), node.type()->kind()); +void RuntimeStats::mergeFrom(const SplitStats& split) { + auto& target = formatSpecificStats[split.format]; + for (const auto& [name, metric] : split.splitMetrics) { + auto [it, inserted] = target.emplace(name, metric); + if (!inserted) { + VELOX_CHECK_EQ(it->second.unit, metric.unit); + it->second.merge(metric); + } + } + for (const auto& [nodeId, stats] : split.columnStats) { + auto it = + columnStats[nodeId].try_emplace(split.format, stats.typeKind).first; + it->second.mergeFrom(stats); + } +} + +void SplitStats::registerColumnStats( + const TypeWithId& node, + bool collectDecodingStats) { + auto& stats = getOrCreateColumnStats(node.id(), node.type()->kind()); + if (collectDecodingStats && !stats.decodingStats) { + stats.decodingStats.emplace(); + } for (uint32_t i = 0; i < node.size(); ++i) { if (const auto* child = node.childAt(i).get()) { - registerDecodingStatsImpl(*child); + registerColumnStats(*child, collectDecodingStats); } } } std::unordered_map -RuntimeStatistics::toRuntimeMetricMap() const { +RuntimeStats::toRuntimeMetricMap() const { std::unordered_map result; for (const auto& [name, metric] : unitLoaderStats.stats()) { result.emplace(name, RuntimeMetric(metric.sum, metric.unit)); @@ -317,13 +343,22 @@ RuntimeStatistics::toRuntimeMetricMap() const { if (numStripes > 0) { result.emplace("numStripes", RuntimeMetric(numStripes)); } - if (parquetFooterEstimatedBytes > 0) { - result.emplace( - "parquetFooterEstimatedBytes", - RuntimeMetric( - parquetFooterEstimatedBytes, RuntimeCounter::Unit::kBytes)); + for (const auto& [format, metrics] : formatSpecificStats) { + for (const auto& [name, metric] : metrics) { + result.emplace( + fmt::format("{}.{}", FileFormatName::toName(format), name), metric); + } + } + for (const auto& [nodeId, statsByFormat] : columnStats) { + for (const auto& [format, stats] : statsByFormat) { + const auto formatPrefix = FileFormatName::toName(format); + const auto typeName = TypeKindName::toName(stats.typeKind); + const auto formatAndColumnPrefix = + fmt::format("{}.column_{}.{}", formatPrefix, nodeId, typeName); + stats.toRuntimeMetrics(formatPrefix, result); + stats.toRuntimeMetrics(formatAndColumnPrefix, result); + } } - columnReaderStats.toRuntimeMetrics(result); return result; } } // namespace facebook::velox::dwio::common diff --git a/velox/dwio/common/Statistics.h b/velox/dwio/common/Statistics.h index fc8684a6947..ce28d129d97 100644 --- a/velox/dwio/common/Statistics.h +++ b/velox/dwio/common/Statistics.h @@ -18,11 +18,9 @@ #include #include -#include #include #include #include -#include "velox/common/time/CpuWallTimer.h" #include "velox/common/time/Timer.h" #include "velox/dwio/common/Options.h" #include "velox/dwio/common/TypeWithId.h" @@ -307,9 +305,7 @@ class IntegerColumnStatistics : public virtual ColumnStatistics { std::optional sum_; }; -/** - * Statistics for timestamp columns. - */ +/// Statistics for timestamp columns. class TimestampColumnStatistics : public virtual ColumnStatistics { public: TimestampColumnStatistics( @@ -348,9 +344,7 @@ class TimestampColumnStatistics : public virtual ColumnStatistics { std::optional max_; }; -/** - * Statistics for string columns. - */ +/// Statistics for string columns. class StringColumnStatistics : public virtual ColumnStatistics { public: StringColumnStatistics( @@ -485,71 +479,88 @@ auto withDecompressStats(io::IoCounter* counter, F&& func) /// different types of measurements (decompression, encoding, etc.). /// Can be used by any file format reader (DWRF, Nimble, Parquet, etc.). struct DecodingStats { - explicit DecodingStats(TypeKind type = TypeKind::INVALID) : typeKind(type) {} - - TypeKind typeKind; io::IoCounter decompressCPUTimeNanos; io::IoCounter decodeCPUTimeNanos; /// Merges stats from another DecodingStats instance. void merge(const DecodingStats& other); + + /// Merges non-empty decoding counters into 'result' using 'prefix'. + void toRuntimeMetrics( + std::string_view prefix, + std::unordered_map& result) const; }; -/// Thread-safe collection of per-column decoding statistics keyed by nodeId. -/// Can be used by any file format reader (DWRF, Nimble, Parquet, etc.). -struct DecodingStatsSet { - /// Gets or creates a DecodingStats for a column. Sets typeKind when - /// creating. - DecodingStats* getOrCreate( - uint32_t nodeId, - TypeKind typeKind = TypeKind::INVALID); +/// Collects runtime metrics produced while reading one column. +struct ColumnRuntimeStats { + /// Creates statistics for a column of 'typeKind'. + explicit ColumnRuntimeStats(TypeKind typeKind) : typeKind{typeKind} {} + + // Logical type of this column. + TypeKind typeKind{TypeKind::INVALID}; + + // Format-specific metrics for this column, keyed by metric name. + folly::F14FastMap columnMetrics; - /// Merges all column decoding statistics from another DecodingStatsSet - /// instance. - void mergeFrom(const DecodingStatsSet& other); + // Decoding counters, when collection is enabled. + std::optional decodingStats; - /// Exports per-column metrics into the runtime metrics result map. + /// Adds one sample to a format-specific column metric. + void accumulateStat( + const std::pair& stat, + int64_t value); + + /// Merges all stats from another ColumnRuntimeStats instance. + void mergeFrom(const ColumnRuntimeStats& other); + + /// Merges this column's metrics into 'result' using 'prefix'. void toRuntimeMetrics( + std::string_view prefix, std::unordered_map& result) const; - - private: - folly::Synchronized< - folly::F14FastMap>> - map_; }; -/// Collects runtime metrics produced while reading columns. -struct ColumnReaderStatistics { - // Number of rows returned by string dictionary reader that is flattened - // instead of keeping dictionary encoding. - int64_t flattenStringDictionaryValues{0}; +/// Collects format-specific statistics while processing one file split. +struct SplitStats { + /// Creates an accumulator for a split of 'format'. + explicit SplitStats(FileFormat format) : format{format} { + VELOX_CHECK_NE(format, FileFormat::UNKNOWN); + } + + // File format shared by all metrics collected for this split. + const FileFormat format; - // Total time spent in loading pages, in nanoseconds. - io::IoCounter pageLoadTimeNs; + // Split-level format-specific metrics, keyed by metric name. + folly::F14FastMap splitMetrics; - // Per-column decoding statistics. Only populated when decoding stats - // collection is enabled. - std::optional decodingStatsSet; + // Per-column statistics keyed by schema node ID. + // TODO(#18171): Use a stable column ID rather than schema node ID + // which is not scan-stable. + folly::F14FastMap columnStats; - /// Initializes column stats collection for the given schema if enabled in - /// options. Recursively registers metrics for all columns in the type tree. + /// Returns the statistics for 'nodeId', creating them if necessary. + ColumnRuntimeStats& getOrCreateColumnStats( + uint32_t nodeId, + TypeKind typeKind); + + /// Returns decoding statistics for 'nodeId', or nullptr if unavailable. + DecodingStats* decodingStats(uint32_t nodeId); + + /// Registers every schema node and optionally enables decoding counters. void initColumnStatsCollection( const TypeWithId& schema, const RowReaderOptions& options); - /// Merges all stats from another ColumnReaderStatistics instance. - void mergeFrom(const ColumnReaderStatistics& other); - - /// Exports all metrics into the runtime metrics result map. - void toRuntimeMetrics( - std::unordered_map& result) const; + /// Adds one sample to a split-level format-specific metric. + void accumulateStat( + const std::pair& stat, + int64_t value); private: - void registerDecodingStatsImpl(const TypeWithId& node); + void registerColumnStats(const TypeWithId& node, bool collectDecodingStats); }; /// Aggregates runtime statistics collected while processing a split. -struct RuntimeStatistics { +struct RuntimeStats { // Number of splits skipped based on statistics. int64_t skippedSplits{0}; @@ -577,17 +588,21 @@ struct RuntimeStatistics { // Counts stripes observed in the file. int64_t numStripes{0}; - // Estimated bytes reported to the memory pool for the deserialized - // Parquet file footer, when the parquet reader's footer-memory - // tracking path is engaged. Lets operators compare the estimate - // against actual pool usage. 0 when the reader did not engage - // tracking (e.g. footer below threshold or non-parquet format). - int64_t parquetFooterEstimatedBytes{0}; - // Stores unit-loader runtime metrics. UnitLoaderStats unitLoaderStats; - // Stores reader-side column runtime metrics. - ColumnReaderStatistics columnReaderStats; + + // Split-level format-specific metrics aggregated by file format. + folly::F14FastMap> + formatSpecificStats; + + // Per-column statistics aggregated by schema node ID and file format. + // TODO(#18171): Use a stable column ID rather than schema node ID + // which is not scan-stable. + folly::F14FastMap> + columnStats; + + /// Merges one split's format-specific and per-column statistics. + void mergeFrom(const SplitStats& split); // Exports collected counters as runtime metrics. std::unordered_map toRuntimeMetricMap() const; diff --git a/velox/dwio/common/compression/PagedInputStream.h b/velox/dwio/common/compression/PagedInputStream.h index 7045e98483f..e30d1b7b5d2 100644 --- a/velox/dwio/common/compression/PagedInputStream.h +++ b/velox/dwio/common/compression/PagedInputStream.h @@ -189,8 +189,8 @@ class PagedInputStream : public dwio::common::SeekableInputStream { const std::string streamDebugInfo_; protected: - // Owned by ColumnReaderStatistics. Valid for the lifetime of this stream - // because ColumnReaderStatistics outlives all streams within a DwrfRowReader. + // Owned by ColumnRuntimeStats. Valid for the lifetime of this stream + // because ColumnRuntimeStats outlives all streams within a DwrfRowReader. io::IoCounter* const decompressCounter_{nullptr}; }; diff --git a/velox/dwio/common/tests/BufferedInputTest.cpp b/velox/dwio/common/tests/BufferedInputTest.cpp index c547518b298..14b5ee9c867 100644 --- a/velox/dwio/common/tests/BufferedInputTest.cpp +++ b/velox/dwio/common/tests/BufferedInputTest.cpp @@ -444,7 +444,7 @@ TEST_F(BufferedInputTest, readSorting) { } } -TEST_F(BufferedInputTest, VreadSorting) { +TEST_F(BufferedInputTest, vreadSorting) { std::string content = "aaabbbcccdddeeefffggghhhiiijjjkkklllmmmnnnooopppqqq"; std::vector regions = {{6, 3}, {24, 3}, {3, 3}, {0, 3}, {29, 3}}; @@ -482,7 +482,7 @@ TEST_F(BufferedInputTest, VreadSorting) { } } -TEST_F(BufferedInputTest, VreadSortingWithLabels) { +TEST_F(BufferedInputTest, vreadSortingWithLabels) { std::string content = "aaabbbcccdddeeefffggghhhiiijjjkkklllmmmnnnooopppqqq"; std::vector l = {"a", "b", "c", "d", "e"}; std::vector regions = { @@ -857,7 +857,7 @@ class AdjustedReadPctAccessor : public BufferedInput { // exclude the whole current (not-yet-read) stripe -- not just the last // reference -- so a fully-read column scores ~100% (and is eligible for // prefetch) instead of roughly half. -TEST(BufferedInputAdjustedReadPctTest, ExcludesWholeCurrentStripe) { +TEST(BufferedInputAdjustedReadPctTest, excludesWholeCurrentStripe) { cache::ScanTracker tracker( "test", /*unregisterer=*/nullptr, /*loadQuantum=*/1 << 20); const cache::TrackingId id(1); diff --git a/velox/dwio/common/tests/CMakeLists.txt b/velox/dwio/common/tests/CMakeLists.txt index 9f08f1410c0..6ca6ceea87a 100644 --- a/velox/dwio/common/tests/CMakeLists.txt +++ b/velox/dwio/common/tests/CMakeLists.txt @@ -34,7 +34,7 @@ add_executable( BitConcatenationTest.cpp BitPackDecoderTest.cpp ChainedBufferTests.cpp - ColumnReaderStatisticsTests.cpp + ColumnRuntimeStatsTests.cpp ColumnSelectorTests.cpp DataBufferTests.cpp DecoderUtilTest.cpp diff --git a/velox/dwio/common/tests/ColumnReaderStatisticsTests.cpp b/velox/dwio/common/tests/ColumnReaderStatisticsTests.cpp deleted file mode 100644 index 7531de829e4..00000000000 --- a/velox/dwio/common/tests/ColumnReaderStatisticsTests.cpp +++ /dev/null @@ -1,419 +0,0 @@ -/* - * Copyright (c) Facebook, Inc. and its affiliates. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include -#include -#include -#include "velox/dwio/common/Statistics.h" -#include "velox/type/Type.h" - -using namespace facebook::velox::dwio::common; -using facebook::velox::RuntimeMetric; -using facebook::velox::TypeKind; - -TEST(IoCounterTest, BasicOperations) { - facebook::velox::io::IoCounter counter; - - EXPECT_EQ(counter.sum(), 0); - EXPECT_EQ(counter.count(), 0); - - counter.increment(5'000); - counter.increment(3'000); - - EXPECT_EQ(counter.sum(), 8'000); - EXPECT_EQ(counter.count(), 2); - EXPECT_EQ(counter.min(), 3'000); - EXPECT_EQ(counter.max(), 5'000); -} - -TEST(IoCounterTest, ConcurrentAccess) { - facebook::velox::io::IoCounter counter; - constexpr int kNumThreads = 4; - constexpr int kIterationsPerThread = 1'000; - - std::vector threads; - threads.reserve(kNumThreads); - for (int i = 0; i < kNumThreads; ++i) { - threads.emplace_back([&counter]() { - for (int j = 0; j < kIterationsPerThread; ++j) { - counter.increment(5); - } - }); - } - for (auto& t : threads) { - t.join(); - } - - EXPECT_EQ(counter.sum(), kNumThreads * kIterationsPerThread * 5); - EXPECT_EQ(counter.count(), kNumThreads * kIterationsPerThread); -} - -TEST(DecodingStatsSetTest, GetOrCreate) { - DecodingStatsSet statsSet; - - auto* result = statsSet.getOrCreate(1); - ASSERT_NE(result, nullptr); - - // Returns same instance for same nodeId. - EXPECT_EQ(statsSet.getOrCreate(1), result); - - // Returns different instance for different nodeId. - auto* result2 = statsSet.getOrCreate(2); - EXPECT_NE(result2, result); -} - -TEST(DecodingStatsSetTest, GetOrCreateWithTypeKind) { - DecodingStatsSet statsSet; - - // Pass type when calling getOrCreate. - auto* result = statsSet.getOrCreate(1, TypeKind::BIGINT); - ASSERT_NE(result, nullptr); - result->decompressCPUTimeNanos.increment(1'000); - - // Returns same instance for same nodeId. - auto* result2 = statsSet.getOrCreate(1); - EXPECT_EQ(result2, result); - - // Different nodeId with different type. - auto* result3 = statsSet.getOrCreate(2, TypeKind::VARCHAR); - EXPECT_NE(result3, result); - result3->decompressCPUTimeNanos.increment(2'000); - - // Verify types are used in toRuntimeMetrics. - std::unordered_map metrics; - statsSet.toRuntimeMetrics(metrics); - EXPECT_EQ(metrics["column_1.BIGINT.decompressCPUTimeNanos"].sum, 1'000); - EXPECT_EQ(metrics["column_2.VARCHAR.decompressCPUTimeNanos"].sum, 2'000); -} - -TEST(DecodingStatsSetTest, ToRuntimeMetrics) { - DecodingStatsSet statsSet; - - // Empty stats produces empty result. - std::unordered_map result; - statsSet.toRuntimeMetrics(result); - EXPECT_TRUE(result.empty()); - - // Add timing data with type information. - auto* col1 = statsSet.getOrCreate(1, TypeKind::BIGINT); - col1->decompressCPUTimeNanos.increment(5'000); - col1->decompressCPUTimeNanos.increment(3'000); - - auto* col2 = statsSet.getOrCreate(42, TypeKind::VARCHAR); - col2->decompressCPUTimeNanos.increment(2'000); - - // Create a column with type but no data. - statsSet.getOrCreate(99, TypeKind::DOUBLE); - - result.clear(); - statsSet.toRuntimeMetrics(result); - - // RuntimeMetric has sum/count/min/max, metric name includes type. - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].sum, 8'000); - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].count, 2); - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].min, 3'000); - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].max, 5'000); - EXPECT_EQ(result["column_42.VARCHAR.decompressCPUTimeNanos"].sum, 2'000); - EXPECT_EQ(result["column_42.VARCHAR.decompressCPUTimeNanos"].count, 1); - - // Zero values are not included. - statsSet.getOrCreate(99); - result.clear(); - statsSet.toRuntimeMetrics(result); - EXPECT_EQ(result.count("column_99.DOUBLE.decompressCPUTimeNanos"), 0); -} - -TEST(DecodingStatsSetTest, ToRuntimeMetricsWithInvalidType) { - DecodingStatsSet statsSet; - - // Add timing data without type information (INVALID type). - auto* col1 = statsSet.getOrCreate(1); - col1->decompressCPUTimeNanos.increment(5'000); - - std::unordered_map result; - statsSet.toRuntimeMetrics(result); - - // Should use INVALID as type name. - EXPECT_EQ(result["column_1.INVALID.decompressCPUTimeNanos"].sum, 5'000); -} - -TEST(DecodingStatsSetTest, ToRuntimeMetricsWithDecodeTime) { - DecodingStatsSet statsSet; - - // Add both decompress and decode timing data. - auto* col1 = statsSet.getOrCreate(1, TypeKind::BIGINT); - col1->decompressCPUTimeNanos.increment(5'000); - col1->decodeCPUTimeNanos.increment(10'000); - col1->decodeCPUTimeNanos.increment(8'000); - - auto* col2 = statsSet.getOrCreate(2, TypeKind::VARCHAR); - col2->decodeCPUTimeNanos.increment(3'000); - - std::unordered_map result; - statsSet.toRuntimeMetrics(result); - - // Column 1 has both decompress and decode metrics. - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].sum, 5'000); - EXPECT_EQ(result["column_1.BIGINT.decodeCPUTimeNanos"].sum, 18'000); - EXPECT_EQ(result["column_1.BIGINT.decodeCPUTimeNanos"].count, 2); - EXPECT_EQ(result["column_1.BIGINT.decodeCPUTimeNanos"].min, 8'000); - EXPECT_EQ(result["column_1.BIGINT.decodeCPUTimeNanos"].max, 10'000); - - // Column 2 has only decode metrics. - EXPECT_EQ(result.count("column_2.VARCHAR.decompressCPUTimeNanos"), 0); - EXPECT_EQ(result["column_2.VARCHAR.decodeCPUTimeNanos"].sum, 3'000); - EXPECT_EQ(result["column_2.VARCHAR.decodeCPUTimeNanos"].count, 1); -} - -TEST(RuntimeStatisticsTest, ToRuntimeMetricMap) { - RuntimeStatistics stats; - - // Empty stats produces empty result. - EXPECT_TRUE(stats.toRuntimeMetricMap().empty()); - - // Set various stats. - stats.skippedSplits = 5; - stats.processedSplits = 15; - stats.skippedStrides = 10; - stats.processedStrides = 30; - stats.numStripes = 4; - stats.columnReaderStats.flattenStringDictionaryValues = 1'000; - - // Add per-column stats with type. - stats.columnReaderStats.decodingStatsSet.emplace(); - auto* colStats = stats.columnReaderStats.decodingStatsSet->getOrCreate( - 1, TypeKind::BIGINT); - colStats->decompressCPUTimeNanos.increment(5'000); - colStats->decodeCPUTimeNanos.increment(12'000); - - auto result = stats.toRuntimeMetricMap(); - - EXPECT_EQ(result["skippedSplits"].sum, 5); - EXPECT_EQ(result["processedSplits"].sum, 15); - EXPECT_EQ(result["skippedStrides"].sum, 10); - EXPECT_EQ(result["processedStrides"].sum, 30); - EXPECT_EQ(result["numStripes"].sum, 4); - EXPECT_EQ(result["flattenStringDictionaryValues"].sum, 1'000); - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].sum, 5'000); - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].count, 1); - EXPECT_EQ(result["column_1.BIGINT.decodeCPUTimeNanos"].sum, 12'000); - EXPECT_EQ(result["column_1.BIGINT.decodeCPUTimeNanos"].count, 1); -} - -TEST(DecodingStatsSetConcurrencyTest, ConcurrentGetOrCreate) { - DecodingStatsSet statsSet; - constexpr int kNumThreads = 4; - constexpr int kNumColumns = 10; - - // Pre-populate columns with types before concurrent access. - for (uint32_t colId = 0; colId < kNumColumns; ++colId) { - statsSet.getOrCreate(colId, TypeKind::BIGINT); - } - - std::vector threads; - threads.reserve(kNumThreads); - for (int i = 0; i < kNumThreads; ++i) { - threads.emplace_back([&statsSet]() { - for (uint32_t colId = 0; colId < kNumColumns; ++colId) { - auto* colStats = statsSet.getOrCreate(colId); - colStats->decompressCPUTimeNanos.increment(100); - } - }); - } - for (auto& t : threads) { - t.join(); - } - - std::unordered_map result; - statsSet.toRuntimeMetrics(result); - - for (uint32_t colId = 0; colId < kNumColumns; ++colId) { - auto key = fmt::format("column_{}.BIGINT.decompressCPUTimeNanos", colId); - EXPECT_EQ(result[key].sum, kNumThreads * 100); - EXPECT_EQ(result[key].count, kNumThreads); - } -} - -TEST(IoCounterTest, MergeStats) { - facebook::velox::io::IoCounter counter1; - counter1.increment(5'000); - counter1.increment(3'000); - - facebook::velox::io::IoCounter counter2; - counter2.increment(2'000); - - counter1.merge(counter2); - - EXPECT_EQ(counter1.sum(), 10'000); - EXPECT_EQ(counter1.count(), 3); -} - -TEST(DecodingStatsSetTest, MergeFromWithOverlappingNodeIds) { - DecodingStatsSet src; - auto* srcCol1 = src.getOrCreate(1, TypeKind::BIGINT); - srcCol1->decompressCPUTimeNanos.increment(5'000); - srcCol1->decompressCPUTimeNanos.increment(3'000); - srcCol1->decodeCPUTimeNanos.increment(10'000); - - auto* srcCol2 = src.getOrCreate(2, TypeKind::VARCHAR); - srcCol2->decompressCPUTimeNanos.increment(2'000); - srcCol2->decodeCPUTimeNanos.increment(4'000); - - DecodingStatsSet dst; - auto* dstCol1 = dst.getOrCreate(1, TypeKind::BIGINT); - dstCol1->decompressCPUTimeNanos.increment(1'000); - dstCol1->decodeCPUTimeNanos.increment(6'000); - - dst.mergeFrom(src); - - std::unordered_map result; - dst.toRuntimeMetrics(result); - - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].sum, 9'000); - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].count, 3); - EXPECT_EQ(result["column_1.BIGINT.decodeCPUTimeNanos"].sum, 16'000); - EXPECT_EQ(result["column_1.BIGINT.decodeCPUTimeNanos"].count, 2); - EXPECT_EQ(result["column_2.VARCHAR.decompressCPUTimeNanos"].sum, 2'000); - EXPECT_EQ(result["column_2.VARCHAR.decompressCPUTimeNanos"].count, 1); - EXPECT_EQ(result["column_2.VARCHAR.decodeCPUTimeNanos"].sum, 4'000); - EXPECT_EQ(result["column_2.VARCHAR.decodeCPUTimeNanos"].count, 1); -} - -TEST(DecodingStatsSetTest, MergeFromWithDisjointNodeIds) { - DecodingStatsSet src; - auto* srcCol3 = src.getOrCreate(3, TypeKind::DOUBLE); - srcCol3->decompressCPUTimeNanos.increment(3'000); - - auto* srcCol4 = src.getOrCreate(4, TypeKind::BOOLEAN); - srcCol4->decompressCPUTimeNanos.increment(4'000); - - DecodingStatsSet dst; - auto* dstCol1 = dst.getOrCreate(1, TypeKind::BIGINT); - dstCol1->decompressCPUTimeNanos.increment(1'000); - - auto* dstCol2 = dst.getOrCreate(2, TypeKind::VARCHAR); - dstCol2->decompressCPUTimeNanos.increment(2'000); - - dst.mergeFrom(src); - - std::unordered_map result; - dst.toRuntimeMetrics(result); - - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].sum, 1'000); - EXPECT_EQ(result["column_2.VARCHAR.decompressCPUTimeNanos"].sum, 2'000); - EXPECT_EQ(result["column_3.DOUBLE.decompressCPUTimeNanos"].sum, 3'000); - EXPECT_EQ(result["column_4.BOOLEAN.decompressCPUTimeNanos"].sum, 4'000); -} - -TEST(DecodingStatsSetTest, MergeFromEmpty) { - DecodingStatsSet nonEmpty; - auto* col = nonEmpty.getOrCreate(1, TypeKind::BIGINT); - col->decompressCPUTimeNanos.increment(5'000); - - DecodingStatsSet empty; - - nonEmpty.mergeFrom(empty); - - std::unordered_map result; - nonEmpty.toRuntimeMetrics(result); - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].sum, 5'000); - - DecodingStatsSet empty2; - empty2.mergeFrom(nonEmpty); - - result.clear(); - empty2.toRuntimeMetrics(result); - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].sum, 5'000); -} - -TEST(ColumnReaderStatisticsTest, MergeFromWithDecodingStats) { - ColumnReaderStatistics src; - src.flattenStringDictionaryValues = 100; - src.decodingStatsSet.emplace(); - src.decodingStatsSet->getOrCreate(1, TypeKind::BIGINT) - ->decompressCPUTimeNanos.increment(1'000); - - // Merge into stats without decodingStatsSet - creates and populates it. - ColumnReaderStatistics dst; - dst.flattenStringDictionaryValues = 50; - dst.mergeFrom(src); - - EXPECT_EQ(dst.flattenStringDictionaryValues, 150); - ASSERT_TRUE(dst.decodingStatsSet.has_value()); - - std::unordered_map result; - dst.decodingStatsSet->toRuntimeMetrics(result); - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].sum, 1'000); -} - -TEST(ColumnReaderStatisticsTest, MergeFromBothWithDecodingStats) { - ColumnReaderStatistics src; - src.flattenStringDictionaryValues = 100; - src.decodingStatsSet.emplace(); - src.decodingStatsSet->getOrCreate(1, TypeKind::BIGINT) - ->decompressCPUTimeNanos.increment(1'000); - - ColumnReaderStatistics dst; - dst.flattenStringDictionaryValues = 50; - dst.decodingStatsSet.emplace(); - dst.decodingStatsSet->getOrCreate(1, TypeKind::BIGINT) - ->decompressCPUTimeNanos.increment(2'000); - - dst.mergeFrom(src); - - EXPECT_EQ(dst.flattenStringDictionaryValues, 150); - ASSERT_TRUE(dst.decodingStatsSet.has_value()); - - std::unordered_map result; - dst.decodingStatsSet->toRuntimeMetrics(result); - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].sum, 3'000); -} - -TEST(WithDecompressStatsTest, NonVoidWithCounter) { - facebook::velox::io::IoCounter counter; - int result = withDecompressStats(&counter, [] { return 42; }); - EXPECT_EQ(result, 42); - EXPECT_EQ(counter.count(), 1); -} - -TEST(WithDecompressStatsTest, NullCounter) { - int result = withDecompressStats(nullptr, [] { return 7; }); - EXPECT_EQ(result, 7); - - int sideEffect = 0; - withDecompressStats(nullptr, [&] { sideEffect = 3; }); - EXPECT_EQ(sideEffect, 3); -} - -TEST(ColumnReaderStatisticsTest, MergeFromWithoutDecodingStats) { - ColumnReaderStatistics src; - src.flattenStringDictionaryValues = 100; - - ColumnReaderStatistics dst; - dst.flattenStringDictionaryValues = 50; - dst.decodingStatsSet.emplace(); - dst.decodingStatsSet->getOrCreate(1, TypeKind::BIGINT) - ->decompressCPUTimeNanos.increment(1'000); - - dst.mergeFrom(src); - - EXPECT_EQ(dst.flattenStringDictionaryValues, 150); - ASSERT_TRUE(dst.decodingStatsSet.has_value()); - - std::unordered_map result; - dst.decodingStatsSet->toRuntimeMetrics(result); - EXPECT_EQ(result["column_1.BIGINT.decompressCPUTimeNanos"].sum, 1'000); -} diff --git a/velox/dwio/common/tests/ColumnRuntimeStatsTests.cpp b/velox/dwio/common/tests/ColumnRuntimeStatsTests.cpp new file mode 100644 index 00000000000..993c13cf10b --- /dev/null +++ b/velox/dwio/common/tests/ColumnRuntimeStatsTests.cpp @@ -0,0 +1,306 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include + +#include "velox/dwio/common/Statistics.h" +#include "velox/type/Type.h" + +using namespace facebook::velox::dwio::common; +using facebook::velox::RuntimeMetric; +using facebook::velox::TypeKind; + +namespace { + +constexpr std::string_view kExampleFormatMetricName = "exampleFormatMetric"; + +constexpr std::pair + kExampleFormatMetric = { + kExampleFormatMetricName, + facebook::velox::RuntimeCounter::Unit::kNone}; + +constexpr auto kExampleFormat = FileFormat::PARQUET; +} // namespace + +TEST(IoCounterTest, basicOperations) { + facebook::velox::io::IoCounter counter; + + EXPECT_EQ(counter.sum(), 0); + EXPECT_EQ(counter.count(), 0); + + counter.increment(5'000); + counter.increment(3'000); + + EXPECT_EQ(counter.sum(), 8'000); + EXPECT_EQ(counter.count(), 2); + EXPECT_EQ(counter.min(), 3'000); + EXPECT_EQ(counter.max(), 5'000); +} + +TEST(IoCounterTest, concurrentAccess) { + facebook::velox::io::IoCounter counter; + constexpr int kNumThreads = 4; + constexpr int kIterationsPerThread = 1'000; + + std::vector threads; + threads.reserve(kNumThreads); + for (int i = 0; i < kNumThreads; ++i) { + threads.emplace_back([&counter]() { + for (int j = 0; j < kIterationsPerThread; ++j) { + counter.increment(5); + } + }); + } + for (auto& t : threads) { + t.join(); + } + + EXPECT_EQ(counter.sum(), kNumThreads * kIterationsPerThread * 5); + EXPECT_EQ(counter.count(), kNumThreads * kIterationsPerThread); +} + +TEST(DecodingStatsTest, toRuntimeMetrics) { + DecodingStats stats; + stats.decompressCPUTimeNanos.increment(3'000); + stats.decompressCPUTimeNanos.increment(5'000); + stats.decodeCPUTimeNanos.increment(7'000); + stats.decodeCPUTimeNanos.increment(11'000); + + std::unordered_map result; + stats.toRuntimeMetrics("parquet.column_1.BIGINT", result); + + const auto& decompress = + result.at("parquet.column_1.BIGINT.decompressCPUTimeNanos"); + EXPECT_EQ(decompress.sum, 8'000); + EXPECT_EQ(decompress.count, 2); + EXPECT_EQ(decompress.min, 3'000); + EXPECT_EQ(decompress.max, 5'000); + EXPECT_EQ(decompress.unit, facebook::velox::RuntimeCounter::Unit::kNanos); + + const auto& decode = result.at("parquet.column_1.BIGINT.decodeCPUTimeNanos"); + EXPECT_EQ(decode.sum, 18'000); + EXPECT_EQ(decode.count, 2); + EXPECT_EQ(decode.min, 7'000); + EXPECT_EQ(decode.max, 11'000); + EXPECT_EQ(decode.unit, facebook::velox::RuntimeCounter::Unit::kNanos); +} + +TEST(SplitStatsTest, columnRuntimeStats) { + SplitStats stats{kExampleFormat}; + auto& column1 = stats.getOrCreateColumnStats(1, TypeKind::BIGINT); + EXPECT_EQ(&stats.getOrCreateColumnStats(1, TypeKind::BIGINT), &column1); + auto& column2 = stats.getOrCreateColumnStats(2, TypeKind::VARCHAR); + EXPECT_EQ(&stats.getOrCreateColumnStats(2, TypeKind::VARCHAR), &column2); + ASSERT_EQ(stats.columnStats.size(), 2); + EXPECT_EQ(stats.columnStats.at(1).typeKind, TypeKind::BIGINT); + EXPECT_EQ(stats.columnStats.at(2).typeKind, TypeKind::VARCHAR); +} + +TEST(RuntimeStatsTest, exportWithoutColumnCpuMetrics) { + const auto schema = TypeWithId::create( + facebook::velox::ROW( + {"bigint", "varchar"}, + {facebook::velox::BIGINT(), facebook::velox::VARCHAR()})); + const RowReaderOptions options; + ASSERT_FALSE(options.collectColumnCpuMetrics()); + + SplitStats splitStats{kExampleFormat}; + splitStats.initColumnStatsCollection(*schema, options); + ASSERT_EQ(splitStats.columnStats.size(), 3); + for (const auto& [nodeId, stats] : splitStats.columnStats) { + EXPECT_FALSE(stats.decodingStats.has_value()) << nodeId; + } + + RuntimeStats stats; + stats.mergeFrom(splitStats); + EXPECT_TRUE(stats.toRuntimeMetricMap().empty()); +} + +TEST(RuntimeStatsTest, toRuntimeMetricMap) { + RuntimeStats stats; + SplitStats splitStats{kExampleFormat}; + + // Empty stats produces empty result. + EXPECT_TRUE(stats.toRuntimeMetricMap().empty()); + + // Set various stats. + stats.skippedSplits = 5; + stats.processedSplits = 15; + stats.skippedStrides = 10; + stats.processedStrides = 30; + stats.numStripes = 4; + splitStats.getOrCreateColumnStats(1, TypeKind::BIGINT) + .accumulateStat(kExampleFormatMetric, 1'000); + splitStats.getOrCreateColumnStats(2, TypeKind::VARCHAR) + .accumulateStat(kExampleFormatMetric, 2'000); + + // Add per-column stats with type. + splitStats.getOrCreateColumnStats(1, TypeKind::BIGINT) + .decodingStats.emplace(); + auto* colStats = + &*splitStats.getOrCreateColumnStats(1, TypeKind::BIGINT).decodingStats; + colStats->decompressCPUTimeNanos.increment(5'000); + colStats->decodeCPUTimeNanos.increment(12'000); + splitStats.getOrCreateColumnStats(2, TypeKind::VARCHAR) + .decodingStats.emplace(); + auto* col2Stats = + &*splitStats.getOrCreateColumnStats(2, TypeKind::VARCHAR).decodingStats; + col2Stats->decompressCPUTimeNanos.increment(7'000); + col2Stats->decodeCPUTimeNanos.increment(8'000); + stats.mergeFrom(splitStats); + + auto result = stats.toRuntimeMetricMap(); + + EXPECT_EQ(result["skippedSplits"].sum, 5); + EXPECT_EQ(result["processedSplits"].sum, 15); + EXPECT_EQ(result["skippedStrides"].sum, 10); + EXPECT_EQ(result["processedStrides"].sum, 30); + EXPECT_EQ(result["numStripes"].sum, 4); + const auto prefix = + fmt::format("{}.", FileFormatName::toName(kExampleFormat)); + EXPECT_EQ( + result[prefix + "column_1.BIGINT.decompressCPUTimeNanos"].sum, 5'000); + EXPECT_EQ(result[prefix + "column_1.BIGINT.decompressCPUTimeNanos"].count, 1); + EXPECT_EQ(result[prefix + "column_1.BIGINT.decodeCPUTimeNanos"].sum, 12'000); + EXPECT_EQ(result[prefix + "column_1.BIGINT.decodeCPUTimeNanos"].count, 1); + EXPECT_EQ( + result[prefix + "column_2.VARCHAR.decompressCPUTimeNanos"].sum, 7'000); + EXPECT_EQ( + result[prefix + "column_2.VARCHAR.decompressCPUTimeNanos"].count, 1); + EXPECT_EQ(result[prefix + "column_2.VARCHAR.decodeCPUTimeNanos"].sum, 8'000); + EXPECT_EQ(result[prefix + "column_2.VARCHAR.decodeCPUTimeNanos"].count, 1); + EXPECT_EQ(result[prefix + "decompressCPUTimeNanos"].sum, 12'000); + EXPECT_EQ(result[prefix + "decompressCPUTimeNanos"].count, 2); + EXPECT_EQ(result[prefix + "decompressCPUTimeNanos"].min, 5'000); + EXPECT_EQ(result[prefix + "decompressCPUTimeNanos"].max, 7'000); + EXPECT_EQ(result[prefix + "decodeCPUTimeNanos"].sum, 20'000); + EXPECT_EQ(result[prefix + "decodeCPUTimeNanos"].count, 2); + EXPECT_EQ(result[prefix + "decodeCPUTimeNanos"].min, 8'000); + EXPECT_EQ(result[prefix + "decodeCPUTimeNanos"].max, 12'000); + const auto& column1Metric = result + [prefix + "column_1.BIGINT." + std::string(kExampleFormatMetricName)]; + EXPECT_EQ(column1Metric.sum, 1'000); + EXPECT_EQ(column1Metric.count, 1); + EXPECT_EQ(column1Metric.min, 1'000); + EXPECT_EQ(column1Metric.max, 1'000); + const auto& column2Metric = result + [prefix + "column_2.VARCHAR." + std::string(kExampleFormatMetricName)]; + EXPECT_EQ(column2Metric.sum, 2'000); + EXPECT_EQ(column2Metric.count, 1); + EXPECT_EQ(column2Metric.min, 2'000); + EXPECT_EQ(column2Metric.max, 2'000); + const auto& formatMetric = + result[prefix + std::string(kExampleFormatMetricName)]; + EXPECT_EQ(formatMetric.sum, 3'000); + EXPECT_EQ(formatMetric.count, 2); + EXPECT_EQ(formatMetric.min, 1'000); + EXPECT_EQ(formatMetric.max, 2'000); +} + +TEST(IoCounterTest, mergeStats) { + facebook::velox::io::IoCounter counter1; + counter1.increment(5'000); + counter1.increment(3'000); + + facebook::velox::io::IoCounter counter2; + counter2.increment(2'000); + + counter1.merge(counter2); + + EXPECT_EQ(counter1.sum(), 10'000); + EXPECT_EQ(counter1.count(), 3); +} + +TEST(ColumnRuntimeStatsTest, mergeFromWithDecodingStats) { + ColumnRuntimeStats src{TypeKind::BIGINT}; + src.accumulateStat(kExampleFormatMetric, 100); + src.decodingStats.emplace(); + src.decodingStats->decompressCPUTimeNanos.increment(1'000); + + // Merge into stats without decoding stats creates and populates them. + ColumnRuntimeStats dst{TypeKind::BIGINT}; + dst.accumulateStat(kExampleFormatMetric, 50); + dst.mergeFrom(src); + + ASSERT_NE( + dst.columnMetrics.find(std::string(kExampleFormatMetricName)), + dst.columnMetrics.end()); + EXPECT_EQ( + dst.columnMetrics.at(std::string(kExampleFormatMetricName)).sum, 150); + EXPECT_EQ(dst.typeKind, TypeKind::BIGINT); + ASSERT_TRUE(dst.decodingStats.has_value()); + EXPECT_EQ(dst.decodingStats->decompressCPUTimeNanos.sum(), 1'000); +} + +TEST(ColumnRuntimeStatsTest, mergeFromBothWithDecodingStats) { + ColumnRuntimeStats src{TypeKind::BIGINT}; + src.accumulateStat(kExampleFormatMetric, 100); + src.decodingStats.emplace(); + src.decodingStats->decompressCPUTimeNanos.increment(1'000); + + ColumnRuntimeStats dst{TypeKind::BIGINT}; + dst.accumulateStat(kExampleFormatMetric, 50); + dst.decodingStats.emplace(); + dst.decodingStats->decompressCPUTimeNanos.increment(2'000); + + dst.mergeFrom(src); + + ASSERT_NE( + dst.columnMetrics.find(std::string(kExampleFormatMetricName)), + dst.columnMetrics.end()); + EXPECT_EQ( + dst.columnMetrics.at(std::string(kExampleFormatMetricName)).sum, 150); + ASSERT_TRUE(dst.decodingStats.has_value()); + EXPECT_EQ(dst.decodingStats->decompressCPUTimeNanos.sum(), 3'000); +} + +TEST(WithDecompressStatsTest, nonVoidWithCounter) { + facebook::velox::io::IoCounter counter; + int result = withDecompressStats(&counter, [] { return 42; }); + EXPECT_EQ(result, 42); + EXPECT_EQ(counter.count(), 1); +} + +TEST(WithDecompressStatsTest, nullCounter) { + int result = withDecompressStats(nullptr, [] { return 7; }); + EXPECT_EQ(result, 7); + + int sideEffect = 0; + withDecompressStats(nullptr, [&] { sideEffect = 3; }); + EXPECT_EQ(sideEffect, 3); +} + +TEST(ColumnRuntimeStatsTest, mergeFromWithoutDecodingStats) { + ColumnRuntimeStats src{TypeKind::BIGINT}; + src.accumulateStat(kExampleFormatMetric, 100); + + ColumnRuntimeStats dst{TypeKind::BIGINT}; + dst.accumulateStat(kExampleFormatMetric, 50); + dst.decodingStats.emplace(); + dst.decodingStats->decompressCPUTimeNanos.increment(1'000); + + dst.mergeFrom(src); + + ASSERT_NE( + dst.columnMetrics.find(std::string(kExampleFormatMetricName)), + dst.columnMetrics.end()); + EXPECT_EQ( + dst.columnMetrics.at(std::string(kExampleFormatMetricName)).sum, 150); + ASSERT_TRUE(dst.decodingStats.has_value()); + EXPECT_EQ(dst.decodingStats->decompressCPUTimeNanos.sum(), 1'000); +} diff --git a/velox/dwio/common/tests/DataBufferTests.cpp b/velox/dwio/common/tests/DataBufferTests.cpp index c5e7b816cbc..2e3a7e58ede 100644 --- a/velox/dwio/common/tests/DataBufferTests.cpp +++ b/velox/dwio/common/tests/DataBufferTests.cpp @@ -36,7 +36,7 @@ class DataBufferTest : public testing::Test { const std::shared_ptr pool_ = memoryManager()->addLeafPool(); }; -TEST_F(DataBufferTest, ZeroOut) { +TEST_F(DataBufferTest, zeroOut) { const uint8_t VALUE = 13; DataBuffer buffer(*pool_, 16); for (auto i = 0; i < buffer.size(); i++) { @@ -64,7 +64,7 @@ TEST_F(DataBufferTest, ZeroOut) { } } -TEST_F(DataBufferTest, At) { +TEST_F(DataBufferTest, at) { DataBuffer buffer{*pool_}; for (auto i = 0; i != 15; ++i) { buffer.append(i); @@ -84,7 +84,7 @@ TEST_F(DataBufferTest, At) { } } -TEST_F(DataBufferTest, Reset) { +TEST_F(DataBufferTest, reset) { DataBuffer buffer{*pool_}; buffer.reserve(16); for (auto i = 0; i != 15; ++i) { @@ -138,7 +138,7 @@ TEST_F(DataBufferTest, Reset) { } } -TEST_F(DataBufferTest, Wrap) { +TEST_F(DataBufferTest, wrap) { auto size = 26; auto buffer = velox::AlignedBuffer::allocate(size, pool_.get()); auto raw = buffer->asMutable(); @@ -154,7 +154,7 @@ TEST_F(DataBufferTest, Wrap) { } } -TEST_F(DataBufferTest, Move) { +TEST_F(DataBufferTest, move) { { DataBuffer buffer{*pool_}; buffer.reserve(16); diff --git a/velox/dwio/common/tests/DirectBufferedInputTest.cpp b/velox/dwio/common/tests/DirectBufferedInputTest.cpp index c04169a91f3..62aca5e705c 100644 --- a/velox/dwio/common/tests/DirectBufferedInputTest.cpp +++ b/velox/dwio/common/tests/DirectBufferedInputTest.cpp @@ -349,6 +349,385 @@ TEST_F(DirectBufferedInputTest, duplicateRegionsShareCoalescedRead) { } } +// Shared-allocation path with a duplicate region: each region is read once and +// the duplicate shares the source's slice. +TEST_F(DirectBufferedInputTest, duplicateRegionsShareAllocationRead) { + constexpr int32_t kContentSize = 4 << 20; // 4MB + std::string content; + content.resize(kContentSize); + for (int32_t i = 0; i < kContentSize; ++i) { + content[i] = static_cast(i % 251); + } + + // 32 regions just over a page each over-allocate ~128KB total, above the + // shared-allocation floor. + constexpr int32_t kNumRegions = 32; + constexpr int32_t kRegionSize = 4'097; + + auto readFile = std::make_shared(content); + + io::ReaderOptions readerOptions(pool_.get()); + readerOptions.setDataIoStats(dataIoStats_); + readerOptions.setMetadataIoStats(metadataIoStats_); + readerOptions.setLoadQuantum(1 << 20); + readerOptions.setDirectBufferedInputSharedAllocation(true); + + auto& ids = fileIds(); + StringIdLease fileId(ids, "sharedAllocDup"); + StringIdLease groupId(ids, "sharedAllocDupGroup"); + + DirectBufferedInput input( + readFile, + MetricsLog::voidLog(), + std::move(fileId), + tracker_, + std::move(groupId), + dataIoStats_, + nullptr, + executor_.get(), + readerOptions); + + // Adjacent distinct regions, then a duplicate of region 0. + std::vector> streams; + streams.reserve(kNumRegions); + for (int32_t k = 0; k < kNumRegions; ++k) { + streams.push_back(input.enqueue( + common::Region{static_cast(k) * kRegionSize, kRegionSize}, + nullptr)); + ASSERT_NE(streams.back(), nullptr); + } + auto dupStream = input.enqueue(common::Region{0, kRegionSize}, nullptr); + ASSERT_NE(dupStream, nullptr); + // Second duplicate of region 0: a 3-chain exercising the intermediate + // duplicate (its predecessor is itself a duplicate). + auto dupStream2 = input.enqueue(common::Region{0, kRegionSize}, nullptr); + ASSERT_NE(dupStream2, nullptr); + + input.load(LogType::TEST); + EXPECT_EQ(input.testingCoalescedLoads().size(), 1); + + // Source and its duplicates return the same bytes, read once. + auto src = getNext(*streams[0]); + ASSERT_TRUE(src.has_value()); + EXPECT_EQ(src.value(), content.substr(0, kRegionSize)); + auto dup = getNext(*dupStream); + ASSERT_TRUE(dup.has_value()); + EXPECT_EQ(dup.value(), content.substr(0, kRegionSize)); + auto dup2 = getNext(*dupStream2); + ASSERT_TRUE(dup2.has_value()); + EXPECT_EQ(dup2.value(), content.substr(0, kRegionSize)); + // Only the distinct regions are read; duplicates add none. + EXPECT_EQ(readFile->numReads(), kNumRegions); + + // A non-duplicate shared-allocation stream still reads its own bytes. + auto mid = getNext(*streams[5]); + ASSERT_TRUE(mid.has_value()); + EXPECT_EQ(mid.value(), content.substr(5 * kRegionSize, kRegionSize)); + EXPECT_EQ(readFile->numReads(), kNumRegions); +} + +// A shared-allocation stream whose region exceeds loadQuantum reads the first +// quantum from the shared allocation and the rest into its own buffer, +// reproducing the region intact. +TEST_F(DirectBufferedInputTest, sharedAllocationStreamCrossesQuantum) { + constexpr int32_t kContentSize = 4 << 20; // 4MB + std::string content; + content.resize(kContentSize); + for (int32_t i = 0; i < kContentSize; ++i) { + content[i] = static_cast(i % 251); + } + + // 32 non-tiny regions push overAllocatedBytes over the 64KB floor -> + // useSharedAllocation. + constexpr int32_t kNumSmall = 32; + constexpr int32_t kSmallSize = 4'097; + constexpr uint64_t kLoadQuantum = 1 << 20; // 1MB + const uint64_t bigOffset = static_cast(kNumSmall) * kSmallSize; + constexpr uint64_t kBigSize = (1 << 20) + (512 << 10); // 1.5MB > quantum + + auto readFile = std::make_shared(content); + io::ReaderOptions readerOptions(pool_.get()); + readerOptions.setDataIoStats(dataIoStats_); + readerOptions.setMetadataIoStats(metadataIoStats_); + readerOptions.setLoadQuantum(kLoadQuantum); + readerOptions.setDirectBufferedInputSharedAllocation(true); + + auto& ids = fileIds(); + StringIdLease fileId(ids, "sharedAllocQuantum"); + StringIdLease groupId(ids, "sharedAllocQuantumGroup"); + DirectBufferedInput input( + readFile, + MetricsLog::voidLog(), + std::move(fileId), + tracker_, + std::move(groupId), + dataIoStats_, + nullptr, + executor_.get(), + readerOptions); + + std::vector> streams; + streams.reserve(kNumSmall); + for (int32_t k = 0; k < kNumSmall; ++k) { + streams.push_back(input.enqueue( + common::Region{static_cast(k) * kSmallSize, kSmallSize}, + nullptr)); + } + auto bigStream = input.enqueue(common::Region{bigOffset, kBigSize}, nullptr); + ASSERT_NE(bigStream, nullptr); + + input.load(LogType::TEST); + ASSERT_EQ(input.testingCoalescedLoads().size(), 1); + + // Read the whole region across the quantum boundary. + std::string got; + while (auto chunk = getNext(*bigStream)) { + got += chunk.value(); + } + EXPECT_EQ(got, content.substr(bigOffset, kBigSize)); +} + +// In a shared-allocation load, a tiny request is still served from its own +// 'tinyData' and a tiny duplicate gets its own copy. +TEST_F(DirectBufferedInputTest, sharedAllocationLoadServesTinyAndDuplicates) { + constexpr int32_t kContentSize = 4 << 20; // 4MB + std::string content; + content.resize(kContentSize); + for (int32_t i = 0; i < kContentSize; ++i) { + content[i] = static_cast(i % 251); + } + + constexpr int32_t kNumSmall = 32; + constexpr int32_t kSmallSize = 4'097; // non-tiny, drives useSharedAllocation + constexpr int32_t kTinyLen = 100; // <= kTinySize + const uint64_t tinyOffset = static_cast(kNumSmall) * kSmallSize; + + auto readFile = std::make_shared(content); + io::ReaderOptions readerOptions(pool_.get()); + readerOptions.setDataIoStats(dataIoStats_); + readerOptions.setMetadataIoStats(metadataIoStats_); + readerOptions.setLoadQuantum(1 << 20); + readerOptions.setDirectBufferedInputSharedAllocation(true); + + auto& ids = fileIds(); + StringIdLease fileId(ids, "sharedAllocTiny"); + StringIdLease groupId(ids, "sharedAllocTinyGroup"); + DirectBufferedInput input( + readFile, + MetricsLog::voidLog(), + std::move(fileId), + tracker_, + std::move(groupId), + dataIoStats_, + nullptr, + executor_.get(), + readerOptions); + + std::vector> streams; + streams.reserve(kNumSmall); + for (int32_t k = 0; k < kNumSmall; ++k) { + streams.push_back(input.enqueue( + common::Region{static_cast(k) * kSmallSize, kSmallSize}, + nullptr)); + } + // A tiny region plus a tiny duplicate of it, inside the shared-allocation + // load. + auto tinyStream = + input.enqueue(common::Region{tinyOffset, kTinyLen}, nullptr); + auto tinyDupStream = + input.enqueue(common::Region{tinyOffset, kTinyLen}, nullptr); + ASSERT_NE(tinyStream, nullptr); + ASSERT_NE(tinyDupStream, nullptr); + + input.load(LogType::TEST); + ASSERT_EQ(input.testingCoalescedLoads().size(), 1); + + // Non-tiny request reads from the shared allocation correctly. + auto big = getNext(*streams[7]); + ASSERT_TRUE(big.has_value()); + EXPECT_EQ(big.value(), content.substr(7 * kSmallSize, kSmallSize)); + + // Tiny request and its duplicate both read the correct bytes from 'tinyData'. + auto tiny = getNext(*tinyStream); + ASSERT_TRUE(tiny.has_value()); + EXPECT_EQ(tiny.value(), content.substr(tinyOffset, kTinyLen)); + auto tinyDup = getNext(*tinyDupStream); + ASSERT_TRUE(tinyDup.has_value()); + EXPECT_EQ(tinyDup.value(), content.substr(tinyOffset, kTinyLen)); + + // Each distinct region is read once; the tiny duplicate shares its read. + EXPECT_EQ(readFile->numReads(), kNumSmall + 1); +} + +// With the shared allocation off (the default), a load whose padding would +// otherwise clear the floor still serves every request from its own allocation +// and returns identical bytes. Same 32-region load that +// useSharedAllocationThreshold shows is shared-backed once enabled. +TEST_F(DirectBufferedInputTest, sharedAllocationDisabled) { + constexpr int32_t kContentSize = 4 << 20; // 4MB + std::string content; + content.resize(kContentSize); + for (int32_t i = 0; i < kContentSize; ++i) { + content[i] = static_cast(i % 251); + } + constexpr int32_t kRegionSize = 4'097; + constexpr int32_t kNumRegions = 32; // clears the floor when enabled + + auto readFile = std::make_shared(content); + io::ReaderOptions readerOptions(pool_.get()); + readerOptions.setDataIoStats(dataIoStats_); + readerOptions.setMetadataIoStats(metadataIoStats_); + readerOptions.setLoadQuantum(1 << 20); + // Already the default; set explicitly so the test states its intent. + readerOptions.setDirectBufferedInputSharedAllocation(false); + auto& ids = fileIds(); + StringIdLease fileId(ids, "sharedAllocOff"); + StringIdLease groupId(ids, "sharedAllocOffGroup"); + DirectBufferedInput input( + readFile, + MetricsLog::voidLog(), + std::move(fileId), + tracker_, + std::move(groupId), + dataIoStats_, + nullptr, + executor_.get(), + readerOptions); + std::vector> streams; + streams.reserve(kNumRegions); + for (int32_t k = 0; k < kNumRegions; ++k) { + streams.push_back(input.enqueue( + common::Region{static_cast(k) * kRegionSize, kRegionSize}, + nullptr)); + } + input.load(LogType::TEST); + + // Buffers are allocated on first access, so read every stream before + // inspecting the load. Bytes must be unchanged by the layout. + for (int32_t k = 0; k < kNumRegions; ++k) { + auto bytes = getNext(*streams[k]); + ASSERT_TRUE(bytes.has_value()) << "stream " << k; + EXPECT_EQ( + bytes.value(), + content.substr( + static_cast(k) * kRegionSize, + static_cast(kRegionSize))); + } + + auto* load = dynamic_cast( + input.testingCoalescedLoads()[0].get()); + ASSERT_NE(load, nullptr); + // No request is served from the shared allocation. + for (const auto& request : load->requests()) { + EXPECT_EQ(request.buffer.sharedData, nullptr); + } +} + +// 'useSharedAllocation' engages only when per-request page padding would waste +// more than the shared-allocation floor (kMinPages * kPageSize = 64KB). A 4097B +// region rounds to two pages, wasting ~4095B; ~16 such regions reach the floor. +TEST_F(DirectBufferedInputTest, useSharedAllocationThreshold) { + constexpr int32_t kContentSize = 4 << 20; // 4MB + std::string content; + content.resize(kContentSize); + for (int32_t i = 0; i < kContentSize; ++i) { + content[i] = static_cast(i % 251); + } + constexpr int32_t kRegionSize = 4'097; // non-tiny; ~4095B padding each + + auto firstRequestSharedAllocationBacked = [&](int32_t numRegions) { + auto readFile = std::make_shared(content); + io::ReaderOptions readerOptions(pool_.get()); + readerOptions.setDataIoStats(dataIoStats_); + readerOptions.setMetadataIoStats(metadataIoStats_); + readerOptions.setLoadQuantum(1 << 20); + readerOptions.setDirectBufferedInputSharedAllocation(true); + auto& ids = fileIds(); + StringIdLease fileId(ids, fmt::format("sharedAllocFloor{}", numRegions)); + StringIdLease groupId( + ids, fmt::format("sharedAllocFloorGroup{}", numRegions)); + DirectBufferedInput input( + readFile, + MetricsLog::voidLog(), + std::move(fileId), + tracker_, + std::move(groupId), + dataIoStats_, + nullptr, + executor_.get(), + readerOptions); + std::vector> streams; + streams.reserve(numRegions); + for (int32_t k = 0; k < numRegions; ++k) { + streams.push_back(input.enqueue( + common::Region{static_cast(k) * kRegionSize, kRegionSize}, + nullptr)); + } + input.load(LogType::TEST); + // Buffers are allocated on first access; touch a stream to force it. + EXPECT_TRUE(getNext(*streams[0]).has_value()); + auto* load = dynamic_cast( + input.testingCoalescedLoads()[0].get()); + EXPECT_NE(load, nullptr); + return load->requests()[0].buffer.sharedData != nullptr; + }; + + EXPECT_FALSE(firstRequestSharedAllocationBacked( + 10)); // ~41KB over-allocated -> per-request + EXPECT_TRUE(firstRequestSharedAllocationBacked( + 32)); // ~131KB over-allocated -> shared allocation +} + +// A wide group bump-packs its non-tiny buffers into one shared allocation, +// costing far fewer pool allocations than one per request. +TEST_F(DirectBufferedInputTest, sharedAllocationReducesAllocations) { + constexpr int32_t kContentSize = 4 << 20; // 4MB + std::string content; + content.resize(kContentSize); + for (int32_t i = 0; i < kContentSize; ++i) { + content[i] = static_cast(i % 251); + } + constexpr int32_t kNumRegions = 32; // > shared-allocation floor + constexpr int32_t kRegionSize = 4'097; + + auto readFile = std::make_shared(content); + io::ReaderOptions readerOptions(pool_.get()); + readerOptions.setDataIoStats(dataIoStats_); + readerOptions.setMetadataIoStats(metadataIoStats_); + readerOptions.setLoadQuantum(1 << 20); + readerOptions.setDirectBufferedInputSharedAllocation(true); + auto& ids = fileIds(); + StringIdLease fileId(ids, "sharedAllocs"); + StringIdLease groupId(ids, "sharedAllocsGroup"); + DirectBufferedInput input( + readFile, + MetricsLog::voidLog(), + std::move(fileId), + tracker_, + std::move(groupId), + dataIoStats_, + nullptr, + executor_.get(), + readerOptions); + + std::vector> streams; + streams.reserve(kNumRegions); + for (int32_t k = 0; k < kNumRegions; ++k) { + streams.push_back(input.enqueue( + common::Region{static_cast(k) * kRegionSize, kRegionSize}, + nullptr)); + } + + const auto allocsBefore = pool_->stats().numAllocs; + input.load(LogType::TEST); + for (auto& stream : streams) { + EXPECT_TRUE(getNext(*stream).has_value()); + } + const auto allocsDelta = pool_->stats().numAllocs - allocsBefore; + // 32 requests share one allocation -> far fewer than one allocation each. + EXPECT_LT(allocsDelta, static_cast(kNumRegions)); +} + DEBUG_ONLY_TEST_F(DirectBufferedInputTest, resetInputWithBeforeLoading) { constexpr int32_t kContentSize = 4 << 20; // 4MB std::string content; diff --git a/velox/dwio/common/tests/ExecutorBarrierTest.cpp b/velox/dwio/common/tests/ExecutorBarrierTest.cpp index 1a9533b4a18..1adcfa2637c 100644 --- a/velox/dwio/common/tests/ExecutorBarrierTest.cpp +++ b/velox/dwio/common/tests/ExecutorBarrierTest.cpp @@ -22,7 +22,7 @@ using namespace ::testing; using namespace ::facebook::velox::dwio::common; -TEST(ExecutorBarrierTest, GetNumPriorities) { +TEST(ExecutorBarrierTest, getNumPriorities) { const uint8_t kNumPriorities = 5; auto executor = std::make_shared(10, kNumPriorities); @@ -30,7 +30,7 @@ TEST(ExecutorBarrierTest, GetNumPriorities) { EXPECT_EQ(barrier->getNumPriorities(), kNumPriorities); } -TEST(ExecutorBarrierTest, CanOwn) { +TEST(ExecutorBarrierTest, canOwn) { auto executor = std::make_shared(10); { auto barrier = std::make_shared(executor); @@ -39,7 +39,7 @@ TEST(ExecutorBarrierTest, CanOwn) { EXPECT_EQ(executor.use_count(), 1); } -TEST(ExecutorBarrierTest, CanAwaitMultipleTimes) { +TEST(ExecutorBarrierTest, canAwaitMultipleTimes) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); for (int time = 0, multipleTimes = 10; time < multipleTimes; ++time) { @@ -47,7 +47,7 @@ TEST(ExecutorBarrierTest, CanAwaitMultipleTimes) { } } -TEST(ExecutorBarrierTest, AddCanBeReused) { +TEST(ExecutorBarrierTest, addCanBeReused) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); @@ -66,7 +66,7 @@ TEST(ExecutorBarrierTest, AddCanBeReused) { EXPECT_EQ(count, (2 * kCalls)); } -TEST(ExecutorBarrierTest, AddWithPriorityCanBeReused) { +TEST(ExecutorBarrierTest, addWithPriorityCanBeReused) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); @@ -86,7 +86,7 @@ TEST(ExecutorBarrierTest, AddWithPriorityCanBeReused) { EXPECT_EQ(count, (2 * kCalls)); } -TEST(ExecutorBarrierTest, AddCanBeReusedAfterException) { +TEST(ExecutorBarrierTest, addCanBeReusedAfterException) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); @@ -108,7 +108,7 @@ TEST(ExecutorBarrierTest, AddCanBeReusedAfterException) { EXPECT_EQ(count, (2 * kCalls)); } -TEST(ExecutorBarrierTest, AddWithPriorityCanBeReusedAfterException) { +TEST(ExecutorBarrierTest, addWithPriorityCanBeReusedAfterException) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); @@ -133,7 +133,7 @@ TEST(ExecutorBarrierTest, AddWithPriorityCanBeReusedAfterException) { EXPECT_EQ(count, (2 * kCalls)); } -TEST(ExecutorBarrierTest, Add) { +TEST(ExecutorBarrierTest, add) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); @@ -146,7 +146,7 @@ TEST(ExecutorBarrierTest, Add) { EXPECT_EQ(count, kCalls); } -TEST(ExecutorBarrierTest, AddWithPriority) { +TEST(ExecutorBarrierTest, addWithPriority) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); @@ -160,7 +160,7 @@ TEST(ExecutorBarrierTest, AddWithPriority) { EXPECT_EQ(count, kCalls); } -TEST(ExecutorBarrierTest, AddCanIgnore) { +TEST(ExecutorBarrierTest, addCanIgnore) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); @@ -171,7 +171,7 @@ TEST(ExecutorBarrierTest, AddCanIgnore) { // Discard: barrier->waitAll(); } -TEST(ExecutorBarrierTest, AddWithPriorityCanIgnore) { +TEST(ExecutorBarrierTest, addWithPriorityCanIgnore) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); @@ -182,7 +182,7 @@ TEST(ExecutorBarrierTest, AddWithPriorityCanIgnore) { // Discard: barrier->waitAll(); } -TEST(ExecutorBarrierTest, DestructorDoesntThrow) { +TEST(ExecutorBarrierTest, destructorDoesntThrow) { const int kCalls = 30; std::atomic count{0}; { @@ -201,7 +201,7 @@ TEST(ExecutorBarrierTest, DestructorDoesntThrow) { EXPECT_EQ(count, kCalls); } -TEST(ExecutorBarrierTest, AddException) { +TEST(ExecutorBarrierTest, addException) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); @@ -219,7 +219,7 @@ TEST(ExecutorBarrierTest, AddException) { EXPECT_EQ(count, kCalls); } -TEST(ExecutorBarrierTest, AddWithPriorityException) { +TEST(ExecutorBarrierTest, addWithPriorityException) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); @@ -240,7 +240,7 @@ TEST(ExecutorBarrierTest, AddWithPriorityException) { EXPECT_EQ(count, kCalls); } -TEST(ExecutorBarrierTest, AddNonStdException) { +TEST(ExecutorBarrierTest, addNonStdException) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); @@ -259,7 +259,7 @@ TEST(ExecutorBarrierTest, AddNonStdException) { EXPECT_EQ(count, kCalls); } -TEST(ExecutorBarrierTest, AddWithPriorityNonStdException) { +TEST(ExecutorBarrierTest, addWithPriorityNonStdException) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); @@ -281,7 +281,7 @@ TEST(ExecutorBarrierTest, AddWithPriorityNonStdException) { EXPECT_EQ(count, kCalls); } -TEST(ExecutorBarrierTest, AddExceptions) { +TEST(ExecutorBarrierTest, addExceptions) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); @@ -297,7 +297,7 @@ TEST(ExecutorBarrierTest, AddExceptions) { EXPECT_EQ(count, kCalls); } -TEST(ExecutorBarrierTest, AddWithPriorityExceptions) { +TEST(ExecutorBarrierTest, addWithPriorityExceptions) { auto executor = std::make_shared(10); auto barrier = std::make_shared(executor); diff --git a/velox/dwio/common/tests/MeasureTimeTests.cpp b/velox/dwio/common/tests/MeasureTimeTests.cpp index a3749df4bd6..28725e1aaad 100644 --- a/velox/dwio/common/tests/MeasureTimeTests.cpp +++ b/velox/dwio/common/tests/MeasureTimeTests.cpp @@ -21,18 +21,18 @@ using namespace ::testing; using namespace ::facebook::velox::dwio::common; -TEST(MeasureTimeTests, DoesntCreateMeasureIfNoCallback) { +TEST(MeasureTimeTests, doesntCreateMeasureIfNoCallback) { EXPECT_FALSE(measureTimeIfCallback(nullptr).has_value()); } -TEST(MeasureTimeTests, CreatesMeasureIfCallback) { +TEST(MeasureTimeTests, createsMeasureIfCallback) { auto callback = std::function( [](const auto&) {}); EXPECT_TRUE(measureTimeIfCallback(callback).has_value()); } -TEST(MeasureTimeTests, MeasuresTime) { +TEST(MeasureTimeTests, measuresTime) { bool measured{false}; { auto callback = diff --git a/velox/dwio/common/tests/OnDemandUnitLoaderTests.cpp b/velox/dwio/common/tests/OnDemandUnitLoaderTests.cpp index 245c7d6186d..5c0db914933 100644 --- a/velox/dwio/common/tests/OnDemandUnitLoaderTests.cpp +++ b/velox/dwio/common/tests/OnDemandUnitLoaderTests.cpp @@ -38,31 +38,31 @@ class OnDemandUnitLoaderCommonTests } }; -TEST_F(OnDemandUnitLoaderCommonTests, NoUnitButSkip) { +TEST_F(OnDemandUnitLoaderCommonTests, noUnitButSkip) { testNoUnitButSkip(); } -TEST_F(OnDemandUnitLoaderCommonTests, InitialSkip) { +TEST_F(OnDemandUnitLoaderCommonTests, initialSkip) { testInitialSkip(); } -TEST_F(OnDemandUnitLoaderCommonTests, CanRequestUnitMultipleTimes) { +TEST_F(OnDemandUnitLoaderCommonTests, canRequestUnitMultipleTimes) { testCanRequestUnitMultipleTimes(); } -TEST_F(OnDemandUnitLoaderCommonTests, UnitOutOfRange) { +TEST_F(OnDemandUnitLoaderCommonTests, unitOutOfRange) { testUnitOutOfRange(); } -TEST_F(OnDemandUnitLoaderCommonTests, SeekOutOfRange) { +TEST_F(OnDemandUnitLoaderCommonTests, seekOutOfRange) { testSeekOutOfRange(); } -TEST_F(OnDemandUnitLoaderCommonTests, SeekOutOfRangeReaderError) { +TEST_F(OnDemandUnitLoaderCommonTests, seekOutOfRangeReaderError) { testSeekOutOfRangeReaderError(); } -TEST(OnDemandUnitLoaderTests, LoadsCorrectlyWithReader) { +TEST(OnDemandUnitLoaderTests, loadsCorrectlyWithReader) { size_t blockedOnIoCount = 0; OnDemandUnitLoaderFactory factory([&](auto) { ++blockedOnIoCount; }); ReaderMock readerMock{{10, 20, 30}, {0, 0, 0}, factory, 0}; @@ -99,7 +99,7 @@ TEST(OnDemandUnitLoaderTests, LoadsCorrectlyWithReader) { EXPECT_EQ(blockedOnIoCount, 3); } -TEST(OnDemandUnitLoaderTests, LoadsCorrectlyWithNoCallback) { +TEST(OnDemandUnitLoaderTests, loadsCorrectlyWithNoCallback) { OnDemandUnitLoaderFactory factory(nullptr); ReaderMock readerMock{{10, 20, 30}, {0, 0, 0}, factory, 0}; EXPECT_EQ(readerMock.unitsLoaded(), std::vector({false, false, false})); @@ -127,7 +127,7 @@ TEST(OnDemandUnitLoaderTests, LoadsCorrectlyWithNoCallback) { EXPECT_EQ(readerMock.unitsLoaded(), std::vector({false, false, true})); } -TEST(OnDemandUnitLoaderTests, CanSeek) { +TEST(OnDemandUnitLoaderTests, canSeek) { size_t blockedOnIoCount = 0; OnDemandUnitLoaderFactory factory([&](auto) { ++blockedOnIoCount; }); ReaderMock readerMock{{10, 20, 30}, {0, 0, 0}, factory, 0}; diff --git a/velox/dwio/common/tests/ParallelForTest.cpp b/velox/dwio/common/tests/ParallelForTest.cpp index 023dd6f5278..1c4cbb1df0f 100644 --- a/velox/dwio/common/tests/ParallelForTest.cpp +++ b/velox/dwio/common/tests/ParallelForTest.cpp @@ -99,7 +99,7 @@ void testParallelFor( } // namespace -TEST(ParallelForTest, E2E) { +TEST(ParallelForTest, e2e) { auto inlineExecutor = folly::InlineExecutor::instance(); for (size_t parallelism = 0; parallelism < 25; ++parallelism) { for (size_t begin = 0; begin < 25; ++begin) { @@ -116,7 +116,7 @@ TEST(ParallelForTest, E2E) { } } -TEST(ParallelForTest, E2EParallel) { +TEST(ParallelForTest, e2eParallel) { for (size_t parallelism = 1; parallelism < 2; ++parallelism) { folly::CPUThreadPoolExecutor executor(parallelism); for (size_t begin = 0; begin < 25; ++begin) { @@ -133,7 +133,7 @@ TEST(ParallelForTest, E2EParallel) { } } -TEST(ParallelForTest, CanOwnExecutor) { +TEST(ParallelForTest, canOwnExecutor) { auto executor = std::make_shared(2); const size_t indexInvokedSize = 100; std::unordered_map> indexInvoked; diff --git a/velox/dwio/common/tests/ParallelUnitLoaderTest.cpp b/velox/dwio/common/tests/ParallelUnitLoaderTest.cpp index 690acd9fc11..f6f4b438e89 100644 --- a/velox/dwio/common/tests/ParallelUnitLoaderTest.cpp +++ b/velox/dwio/common/tests/ParallelUnitLoaderTest.cpp @@ -37,31 +37,31 @@ class ParallelUnitLoaderTest std::make_unique(10); }; -TEST_F(ParallelUnitLoaderTest, NoUnitButSkip) { +TEST_F(ParallelUnitLoaderTest, noUnitButSkip) { testNoUnitButSkip(); } -TEST_F(ParallelUnitLoaderTest, InitialSkip) { +TEST_F(ParallelUnitLoaderTest, initialSkip) { testInitialSkip(); } -TEST_F(ParallelUnitLoaderTest, CanRequestUnitMultipleTimes) { +TEST_F(ParallelUnitLoaderTest, canRequestUnitMultipleTimes) { testCanRequestUnitMultipleTimes(); } -TEST_F(ParallelUnitLoaderTest, UnitOutOfRange) { +TEST_F(ParallelUnitLoaderTest, unitOutOfRange) { testUnitOutOfRange(); } -TEST_F(ParallelUnitLoaderTest, SeekOutOfRange) { +TEST_F(ParallelUnitLoaderTest, seekOutOfRange) { testSeekOutOfRange(); } -TEST_F(ParallelUnitLoaderTest, SeekOutOfRangeReaderError) { +TEST_F(ParallelUnitLoaderTest, seekOutOfRangeReaderError) { testSeekOutOfRangeReaderError(); } -TEST_F(ParallelUnitLoaderTest, LoadsCorrectlyWithReader) { +TEST_F(ParallelUnitLoaderTest, loadsCorrectlyWithReader) { auto factory = createFactory(); ReaderMock readerMock{{10, 20, 30}, {0, 0, 0}, factory, 0}; @@ -92,7 +92,7 @@ TEST_F(ParallelUnitLoaderTest, LoadsCorrectlyWithReader) { } // Performance comparison test -TEST_F(ParallelUnitLoaderTest, PerformanceComparison) { +TEST_F(ParallelUnitLoaderTest, performanceComparison) { std::vector rowsPerUnit = {100, 100, 100, 100, 100, 100, 100, 100}; std::vector ioSizes = { 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024}; diff --git a/velox/dwio/common/tests/RangeTests.cpp b/velox/dwio/common/tests/RangeTests.cpp index 1488c95fe2e..f4f7fcdcf89 100644 --- a/velox/dwio/common/tests/RangeTests.cpp +++ b/velox/dwio/common/tests/RangeTests.cpp @@ -50,7 +50,7 @@ TEST(RangeTests, Add) { ASSERT_EQ(ranges.size(), 0); } -TEST(RangeTests, ForEach) { +TEST(RangeTests, forEach) { Ranges ranges; ranges.add(1, 3); ranges.add(6, 9); diff --git a/velox/dwio/common/tests/ReadFileInputStreamTests.cpp b/velox/dwio/common/tests/ReadFileInputStreamTests.cpp index 1b600b6adaf..d12263a5b6e 100644 --- a/velox/dwio/common/tests/ReadFileInputStreamTests.cpp +++ b/velox/dwio/common/tests/ReadFileInputStreamTests.cpp @@ -35,7 +35,7 @@ class ReadFileInputStreamTest : public testing::Test { } }; -TEST_F(ReadFileInputStreamTest, LocalReadFile) { +TEST_F(ReadFileInputStreamTest, localReadFile) { auto tempFile = TempFilePath::create(); const auto& filename = tempFile->getPath(); remove(filename.c_str()); @@ -66,7 +66,7 @@ TEST_F(ReadFileInputStreamTest, LocalReadFile) { remove(filename.c_str()); } -TEST(ReadFileInputStream, SimpleUsage) { +TEST(ReadFileInputStream, simpleUsage) { std::string fileData; { InMemoryWriteFile writeFile(&fileData); @@ -88,7 +88,7 @@ TEST(ReadFileInputStream, SimpleUsage) { ASSERT_EQ(read_value, "aaaaabbbbbccccc"); } -TEST(ReadFileInputStream, VReadIOBufs) { +TEST(ReadFileInputStream, vreadIOBufs) { std::string fileData; { InMemoryWriteFile writeFile(&fileData); diff --git a/velox/dwio/common/tests/ScanSpecTest.cpp b/velox/dwio/common/tests/ScanSpecTest.cpp index 96eb4f10dcb..2ed93315f57 100644 --- a/velox/dwio/common/tests/ScanSpecTest.cpp +++ b/velox/dwio/common/tests/ScanSpecTest.cpp @@ -15,6 +15,7 @@ */ #include "velox/dwio/common/ScanSpec.h" +#include "velox/dwio/common/SelectiveStructColumnReader.h" #include "velox/vector/tests/utils/VectorTestBase.h" #include @@ -105,6 +106,63 @@ TEST_F(ScanSpecTest, setFilterResetsHasFilter) { ASSERT_TRUE(scanSpec.hasFilter()); } +TEST_F(ScanSpecTest, testFilterOnConstant) { + auto test = [&](auto&& setup, bool expected) { + ScanSpec scanSpec(""); + auto* child = scanSpec.addField("c0", 0); + setup(scanSpec, *child); + ASSERT_EQ( + dwio::common::SelectiveStructColumnReaderBase::testFilterOnConstant( + *child), + expected); + }; + + // Non-null constants are accepted regardless of filter kind. + test( + [&](ScanSpec&, ScanSpec& child) { + child.setConstantValue( + BaseVector::createConstant(BIGINT(), 1LL, 1, pool())); + child.setFilter(std::make_shared()); + }, + true); + test( + [&](ScanSpec&, ScanSpec& child) { + child.setConstantValue( + BaseVector::createConstant(BIGINT(), 1LL, 1, pool())); + child.setFilter(std::make_shared()); + }, + true); + + // Null constants are accepted only when the filter can match nulls. + test( + [&](ScanSpec&, ScanSpec& child) { + child.setConstantValue( + BaseVector::createNullConstant(BIGINT(), 1, pool())); + child.setFilter(std::make_shared()); + }, + true); + test( + [&](ScanSpec& scanSpec, ScanSpec& child) { + child.setConstantValue( + BaseVector::createNullConstant(BIGINT(), 1, pool())); + child.setFilter(std::make_shared()); + }, + false); + + // For non-constant specs, there is no filter or the filter accepts nulls. + test([](ScanSpec&, ScanSpec&) {}, true); + test( + [](ScanSpec&, ScanSpec& child) { + child.setFilter(std::make_shared()); + }, + true); + test( + [](ScanSpec&, ScanSpec& child) { + child.setFilter(std::make_shared()); + }, + false); +} + class TypedScanSpecTest : public testing::TestWithParam, public test::VectorTestBase { protected: diff --git a/velox/dwio/common/tests/SelectiveColumnReaderTest.cpp b/velox/dwio/common/tests/SelectiveColumnReaderTest.cpp index d4a237d799c..7be50ec8db0 100644 --- a/velox/dwio/common/tests/SelectiveColumnReaderTest.cpp +++ b/velox/dwio/common/tests/SelectiveColumnReaderTest.cpp @@ -97,7 +97,7 @@ class StubFormatData : public FormatData { // Minimal FormatParams stub that produces a StubFormatData. class StubFormatParams : public FormatParams { public: - StubFormatParams(memory::MemoryPool& pool, ColumnReaderStatistics& stats) + StubFormatParams(memory::MemoryPool& pool, SplitStats& stats) : FormatParams(pool, stats) {} std::unique_ptr toFormatData( @@ -176,7 +176,7 @@ class GetFlatValuesTest : public ::testing::Test { void SetUp() override { pool_ = memory::memoryManager()->addLeafPool("GetFlatValuesTest"); - stats_ = std::make_unique(); + stats_ = std::make_unique(FileFormat::DWRF); params_ = std::make_unique(*pool_, *stats_); scanSpec_ = std::make_unique("test"); scanSpec_->setProjectOut(true); @@ -216,7 +216,7 @@ class GetFlatValuesTest : public ::testing::Test { } std::shared_ptr pool_; - std::unique_ptr stats_; + std::unique_ptr stats_; std::unique_ptr params_; std::unique_ptr scanSpec_; }; diff --git a/velox/dwio/common/tests/UnitLoaderToolsTests.cpp b/velox/dwio/common/tests/UnitLoaderToolsTests.cpp index 0919e42fee2..eeec9d851c3 100644 --- a/velox/dwio/common/tests/UnitLoaderToolsTests.cpp +++ b/velox/dwio/common/tests/UnitLoaderToolsTests.cpp @@ -24,7 +24,7 @@ using namespace ::testing; using namespace ::facebook::velox::dwio::common; using namespace ::facebook::velox::dwio::common::unit_loader_tools; -TEST(UnitLoaderToolsTests, NoCallbacksCreated) { +TEST(UnitLoaderToolsTests, noCallbacksCreated) { std::atomic_size_t callCount = 0; { CallbackOnLastSignal callback([&callCount]() { ++callCount; }); @@ -33,13 +33,13 @@ TEST(UnitLoaderToolsTests, NoCallbacksCreated) { EXPECT_EQ(callCount, 1); } -TEST(UnitLoaderToolsTests, SupportsNullCallbacks) { +TEST(UnitLoaderToolsTests, supportsNullCallbacks) { CallbackOnLastSignal callback(nullptr); auto cb = callback.getCallback(); EXPECT_TRUE(cb == nullptr); } -TEST(UnitLoaderToolsTests, NoExplicitCalls) { +TEST(UnitLoaderToolsTests, noExplicitCalls) { std::atomic_size_t callCount = 0; { CallbackOnLastSignal callback([&callCount]() { ++callCount; }); @@ -62,7 +62,7 @@ TEST(UnitLoaderToolsTests, NoExplicitCalls) { EXPECT_EQ(callCount, 1); } -TEST(UnitLoaderToolsTests, NoExplicitCallsFactoryDeletedFirst) { +TEST(UnitLoaderToolsTests, noExplicitCallsFactoryDeletedFirst) { std::atomic_size_t callCount = 0; { std::function c1, c2; @@ -79,7 +79,7 @@ TEST(UnitLoaderToolsTests, NoExplicitCallsFactoryDeletedFirst) { EXPECT_EQ(callCount, 1); } -TEST(UnitLoaderToolsTests, ExplicitCalls) { +TEST(UnitLoaderToolsTests, explicitCalls) { std::atomic_size_t callCount = 0; { CallbackOnLastSignal callback([&callCount]() { ++callCount; }); @@ -109,7 +109,7 @@ TEST(UnitLoaderToolsTests, ExplicitCalls) { EXPECT_EQ(callCount, 1); } -TEST(UnitLoaderToolsTests, WillOnlyCallbackOnce) { +TEST(UnitLoaderToolsTests, willOnlyCallbackOnce) { std::atomic_size_t callCount = 0; { CallbackOnLastSignal callback([&callCount]() { ++callCount; }); @@ -144,7 +144,7 @@ TEST(UnitLoaderToolsTests, WillOnlyCallbackOnce) { EXPECT_EQ(callCount, 1); } -TEST(UnitLoaderToolsTests, HowMuchToSkip) { +TEST(UnitLoaderToolsTests, howMuchToSkip) { // Helpers auto testSkip = [](uint64_t rowsToSkip, std::vector rowCount) { return howMuchToSkip(rowsToSkip, rowCount.cbegin(), rowCount.cend()); diff --git a/velox/dwio/common/tests/utils/E2EFilterTestBase.cpp b/velox/dwio/common/tests/utils/E2EFilterTestBase.cpp index d5d28c9a854..277782d6783 100644 --- a/velox/dwio/common/tests/utils/E2EFilterTestBase.cpp +++ b/velox/dwio/common/tests/utils/E2EFilterTestBase.cpp @@ -505,7 +505,7 @@ void E2EFilterTestBase::readWithFilter( } OwnershipChecker ownershipChecker; auto rowReader = reader->createRowReader(rowReaderOpts); - runtimeStats_ = dwio::common::RuntimeStatistics(); + runtimeStats_ = dwio::common::RuntimeStats(); auto rowIndex = 0; auto resultBatch = BaseVector::create(rowType_, 1, leafPool_.get()); resetReadBatchSizes(); diff --git a/velox/dwio/common/tests/utils/E2EFilterTestBase.h b/velox/dwio/common/tests/utils/E2EFilterTestBase.h index 1f97b09a6d9..7322b8debd9 100644 --- a/velox/dwio/common/tests/utils/E2EFilterTestBase.h +++ b/velox/dwio/common/tests/utils/E2EFilterTestBase.h @@ -152,7 +152,7 @@ class E2EFilterTestBase : public testing::Test { static bool typeKindSupportsValueHook(TypeKind kind) { return kind != TypeKind::TIMESTAMP && kind != TypeKind::ARRAY && kind != TypeKind::ROW && kind != TypeKind::MAP && - kind != TypeKind::HUGEINT; + kind != TypeKind::HUGEINT && kind != TypeKind::UNKNOWN; } std::vector makeDataset( @@ -428,7 +428,7 @@ class E2EFilterTestBase : public testing::Test { std::shared_ptr rowType_; std::string sinkData_; bool useVInts_ = true; - dwio::common::RuntimeStatistics runtimeStats_; + dwio::common::RuntimeStats runtimeStats_; // Number of calls to flush policy between starting new stripes. int32_t flushEveryNBatches_{10}; int32_t nextReadSizeIndex_{0}; diff --git a/velox/dwio/dwrf/common/CMakeLists.txt b/velox/dwio/dwrf/common/CMakeLists.txt index 311b4ac51ac..b238a310197 100644 --- a/velox/dwio/dwrf/common/CMakeLists.txt +++ b/velox/dwio/dwrf/common/CMakeLists.txt @@ -45,6 +45,7 @@ velox_add_library( FloatingPointDecoder.h IntEncoder.h NextVisitor.h + DwrfRuntimeStats.h RLEv1.h RLEv2.h Statistics.h diff --git a/velox/dwio/dwrf/common/DwrfRuntimeStats.h b/velox/dwio/dwrf/common/DwrfRuntimeStats.h new file mode 100644 index 00000000000..02dd72a723f --- /dev/null +++ b/velox/dwio/dwrf/common/DwrfRuntimeStats.h @@ -0,0 +1,39 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include "velox/common/base/RuntimeMetrics.h" + +namespace facebook::velox::dwrf { + +/// Names and descriptors for DWRF-specific runtime metrics. +struct DwrfRuntimeStats { + /// Count of string dictionary values that were flattened during reading. + inline static constexpr std::string_view kFlattenStringDictionaryValues = + "flattenStringDictionaryValues"; + + /// Describes the flatten-string-dictionary runtime metric. + inline static constexpr std::pair + kFlattenStringDictionaryValuesMetric = { + kFlattenStringDictionaryValues, + RuntimeCounter::Unit::kNone}; +}; + +} // namespace facebook::velox::dwrf diff --git a/velox/dwio/dwrf/reader/DwrfData.h b/velox/dwio/dwrf/reader/DwrfData.h index 04fa978a96e..918806f8d38 100644 --- a/velox/dwio/dwrf/reader/DwrfData.h +++ b/velox/dwio/dwrf/reader/DwrfData.h @@ -127,7 +127,7 @@ class DwrfParams : public dwio::common::FormatParams { explicit DwrfParams( StripeStreams& stripeStreams, const StreamLabels& streamLabels, - dwio::common::ColumnReaderStatistics& stats, + dwio::common::SplitStats& stats, FlatMapContext context = {}) : FormatParams(stripeStreams.getMemoryPool(), stats), streamLabels_(streamLabels), diff --git a/velox/dwio/dwrf/reader/DwrfReader.cpp b/velox/dwio/dwrf/reader/DwrfReader.cpp index 651268818e1..2e6fa297de9 100644 --- a/velox/dwio/dwrf/reader/DwrfReader.cpp +++ b/velox/dwio/dwrf/reader/DwrfReader.cpp @@ -46,7 +46,7 @@ class DwrfUnit : public LoadUnit { DwrfUnit( std::shared_ptr readerBase, const StrideIndexProvider& strideIndexProvider, - std::shared_ptr columnReaderStats, + std::shared_ptr splitStats, uint32_t stripeIndex, std::shared_ptr columnSelector, std::shared_ptr projectedNodes, @@ -55,7 +55,7 @@ class DwrfUnit : public LoadUnit { : stripeReaderBase_{readerBase}, memoryPool_(readerBase->memoryPool().shared_from_this()), strideIndexProvider_{strideIndexProvider}, - columnReaderStats_{std::move(columnReaderStats)}, + splitStats_{std::move(splitStats)}, stripeIndex_{stripeIndex}, columnSelector_{std::move(columnSelector)}, projectedNodes_{std::move(projectedNodes)}, @@ -104,8 +104,7 @@ class DwrfUnit : public LoadUnit { // ColumnReader::next(), where DwrfRowReader is guaranteed to be alive. const StrideIndexProvider& strideIndexProvider_; - const std::shared_ptr - columnReaderStats_; + const std::shared_ptr splitStats_; const uint32_t stripeIndex_; const std::shared_ptr columnSelector_; const std::shared_ptr projectedNodes_; @@ -171,7 +170,7 @@ void DwrfUnit::ensureDecoders() { stripeInfo_.numberOfRows(), strideIndexProvider_, stripeIndex_, - columnReaderStats_.get()); + splitStats_.get()); auto* scanSpec = options_.scanSpec().get(); const auto& fileType = stripeReaderBase_.getReader().schemaWithId(); @@ -187,7 +186,7 @@ void DwrfUnit::ensureDecoders() { fileType, *stripeStreams_, streamLabels, - *columnReaderStats_, + *splitStats_, scanSpec, flatMapContext, /*isRoot=*/true); @@ -270,11 +269,11 @@ DwrfRowReader::DwrfRowReader( reader->schema()))}, decodingTimeCallback_{options_.decodingTimeCallback()}, strideIndex_{0}, - columnReaderStats_( - std::make_shared()), + splitStats_( + std::make_shared( + dwio::common::FileFormat::DWRF)), currentUnit_{nullptr} { - columnReaderStats_->initColumnStatsCollection( - *getReader().schemaWithId(), options_); + splitStats_->initColumnStatsCollection(*getReader().schemaWithId(), options_); const auto& fileFooter = getReader().footer(); const uint32_t numberOfStripes = fileFooter.stripesSize(); currentStripe_ = numberOfStripes; @@ -336,9 +335,8 @@ DwrfRowReader::DwrfRowReader( makeProjectedNodes(*getReader().schemaWithId(), *projectedNodes_); } - // Configure reader options before calling 'getUnitLoader()'. - // Construction is single-threaded, and the unit loader is created only - // after 'columnReaderOptions_' has been initialized. + // Keep this before 'getUnitLoader()': it copies 'columnReaderOptions_' into + // every DwrfUnit, which then uses the copy to build its column readers. columnReaderOptions_ = dwio::common::makeColumnReaderOptions( readerBaseShared()->readerOptions()); unitLoader_ = getUnitLoader(); @@ -366,7 +364,7 @@ std::unique_ptr DwrfRowReader::getUnitLoader() { std::make_unique( /*readerBase=*/readerBaseShared(), /*strideIndexProvider=*/*this, - columnReaderStats_, + splitStats_, stripe, columnSelector_, projectedNodes_, diff --git a/velox/dwio/dwrf/reader/DwrfReader.h b/velox/dwio/dwrf/reader/DwrfReader.h index ee018749e1e..aa46e0b61ff 100644 --- a/velox/dwio/dwrf/reader/DwrfReader.h +++ b/velox/dwio/dwrf/reader/DwrfReader.h @@ -115,13 +115,12 @@ class DwrfRowReader : public StrideIndexProvider, VectorPtr& result, const dwio::common::Mutation* = nullptr) override; - void updateRuntimeStats( - dwio::common::RuntimeStatistics& stats) const override { + void updateRuntimeStats(dwio::common::RuntimeStats& stats) const override { stats.skippedStrides += skippedStrides_; stats.processedStrides += processedStrides_; stats.footerBufferOverread += getReader().footerBufferOverread(); stats.numStripes += stripeCeiling_ - firstStripe_; - stats.columnReaderStats.mergeFrom(*columnReaderStats_); + stats.mergeFrom(*splitStats_); stats.unitLoaderStats.merge(unitLoadStats_); } @@ -230,7 +229,7 @@ class DwrfRowReader : public StrideIndexProvider, // instead of next stripe. bool recomputeStridesToSkip_{false}; - std::shared_ptr columnReaderStats_; + std::shared_ptr splitStats_; std::optional nextRowNumber_; diff --git a/velox/dwio/dwrf/reader/SelectiveDwrfReader.h b/velox/dwio/dwrf/reader/SelectiveDwrfReader.h index 7cc340e4497..19cfe0fca7d 100644 --- a/velox/dwio/dwrf/reader/SelectiveDwrfReader.h +++ b/velox/dwio/dwrf/reader/SelectiveDwrfReader.h @@ -41,7 +41,7 @@ class SelectiveDwrfReader { const std::shared_ptr& fileType, StripeStreams& stripe, const StreamLabels& streamLabels, - dwio::common::ColumnReaderStatistics& stats, + dwio::common::SplitStats& stats, common::ScanSpec* scanSpec, FlatMapContext flatMapContext = {}, bool isRoot = false) { diff --git a/velox/dwio/dwrf/reader/SelectiveFlatMapColumnReader.cpp b/velox/dwio/dwrf/reader/SelectiveFlatMapColumnReader.cpp index 9c90d39f3b4..a40ec780545 100644 --- a/velox/dwio/dwrf/reader/SelectiveFlatMapColumnReader.cpp +++ b/velox/dwio/dwrf/reader/SelectiveFlatMapColumnReader.cpp @@ -116,12 +116,14 @@ std::vector> getKeyNodes( break; } case FlatMapOutput::kFlatMap: - // Remove on filters on keys stream since it doesn't exist (it's common to - // filter out nulls). + // The keys filter is retained (not cleared) so requested keys prune which + // streams are read below. A null-only filter (commonly present on the + // keys stream) passes every key and therefore prunes nothing. The keys + // stream itself is never read; read() ignores this filter to avoid + // mistaking it for a struct-child filter. keysSpec = scanSpec.getOrCreateChild(common::ScanSpec::kMapKeysFieldName); valuesSpec = scanSpec.getOrCreateChild(common::ScanSpec::kMapValuesFieldName); - keysSpec->setFilter(nullptr); VELOX_CHECK(!valuesSpec->hasFilter()); break; } @@ -141,6 +143,10 @@ std::vector> getKeyNodes( auto key = extractKey(keyInfo); common::ScanSpec* childSpec; if (outputType == FlatMapOutput::kFlatMap) { + if (keysSpec->filter() && + !common::applyFilter(*keysSpec->filter(), key.get())) { + return; // Subfield pruning. + } childSpec = scanSpec.getOrCreateChild(toString(key.get())); childSpec->setProjectOut(true); childSpec->setChannel(sequence - 1); @@ -165,7 +171,7 @@ std::vector> getKeyNodes( DwrfParams childParams( stripe, labels, - params.runtimeStatistics(), + params.splitStats(), FlatMapContext{ .sequence = sequence, .inMapDecoder = inMapDecoder.get(), @@ -325,6 +331,11 @@ class SelectiveFlatMapReader } } + void read(int64_t offset, const RowSet& rows, const uint64_t* incomingNulls) + override { + readFlatMapChildren(offset, rows, incomingNulls); + } + const BufferPtr& inMapBuffer(column_index_t childIndex) const override { return children_[childIndex] ->formatData() diff --git a/velox/dwio/dwrf/reader/SelectiveRepeatedColumnReader.cpp b/velox/dwio/dwrf/reader/SelectiveRepeatedColumnReader.cpp index 12bd9c7b2a0..db3cdc1afba 100644 --- a/velox/dwio/dwrf/reader/SelectiveRepeatedColumnReader.cpp +++ b/velox/dwio/dwrf/reader/SelectiveRepeatedColumnReader.cpp @@ -105,7 +105,7 @@ SelectiveListColumnReader::SelectiveListColumnReader( auto childParams = DwrfParams( stripe, params.streamLabels(), - params.runtimeStatistics(), + params.splitStats(), flatMapContextFromEncodingKey(encodingKey)); child_ = SelectiveDwrfReader::build( columnReaderOptions, @@ -137,7 +137,7 @@ void makeMapChildrenReaders( DwrfParams keyParams( stripe, params.streamLabels(), - params.runtimeStatistics(), + params.splitStats(), flatMapContextFromEncodingKey(encodingKey)); keyReader = SelectiveDwrfReader::build( columnReaderOptions, @@ -150,7 +150,7 @@ void makeMapChildrenReaders( DwrfParams elementParams = DwrfParams( stripe, params.streamLabels(), - params.runtimeStatistics(), + params.splitStats(), flatMapContextFromEncodingKey(encodingKey)); elementReader = SelectiveDwrfReader::build( columnReaderOptions, diff --git a/velox/dwio/dwrf/reader/SelectiveStringDictionaryColumnReader.cpp b/velox/dwio/dwrf/reader/SelectiveStringDictionaryColumnReader.cpp index cd0e6b0b15e..95f83b76333 100644 --- a/velox/dwio/dwrf/reader/SelectiveStringDictionaryColumnReader.cpp +++ b/velox/dwio/dwrf/reader/SelectiveStringDictionaryColumnReader.cpp @@ -15,8 +15,10 @@ */ #include "velox/dwio/dwrf/reader/SelectiveStringDictionaryColumnReader.h" + #include "velox/dwio/common/BufferUtil.h" #include "velox/dwio/dwrf/common/DecoderUtil.h" +#include "velox/dwio/dwrf/common/DwrfRuntimeStats.h" namespace facebook::velox::dwrf { @@ -29,7 +31,8 @@ SelectiveStringDictionaryColumnReader::SelectiveStringDictionaryColumnReader( : SelectiveColumnReader(fileType->type(), fileType, params, scanSpec), lastStrideIndex_(-1), provider_(params.stripeStreams().getStrideIndexProvider()), - statistics_(params.runtimeStatistics()) { + statistics_( + params.columnStats(fileType->id(), fileType->type()->kind())) { auto& stripe = params.stripeStreams(); EncodingKey encodingKey{fileType_->id(), params.flatMapContext().sequence}; version_ = convertRleVersion(stripe, encodingKey); @@ -284,7 +287,8 @@ void SelectiveStringDictionaryColumnReader::makeFlat(VectorPtr* result) { numValues_, std::move(values), std::move(stringBuffers)); - statistics_.flattenStringDictionaryValues += numValues_; + statistics_.accumulateStat( + DwrfRuntimeStats::kFlattenStringDictionaryValuesMetric, numValues_); } void SelectiveStringDictionaryColumnReader::getValues( diff --git a/velox/dwio/dwrf/reader/SelectiveStringDictionaryColumnReader.h b/velox/dwio/dwrf/reader/SelectiveStringDictionaryColumnReader.h index d524f3c3f9b..c72d137296d 100644 --- a/velox/dwio/dwrf/reader/SelectiveStringDictionaryColumnReader.h +++ b/velox/dwio/dwrf/reader/SelectiveStringDictionaryColumnReader.h @@ -96,7 +96,7 @@ class SelectiveStringDictionaryColumnReader RleVersion version_; const StrideIndexProvider& provider_; - dwio::common::ColumnReaderStatistics& statistics_; + dwio::common::ColumnRuntimeStats& statistics_; // lazy load the dictionary std::unique_ptr> lengthDecoder_; diff --git a/velox/dwio/dwrf/reader/SelectiveStringDirectColumnReader.cpp b/velox/dwio/dwrf/reader/SelectiveStringDirectColumnReader.cpp index 2a39829acf8..52da9269f4e 100644 --- a/velox/dwio/dwrf/reader/SelectiveStringDirectColumnReader.cpp +++ b/velox/dwio/dwrf/reader/SelectiveStringDirectColumnReader.cpp @@ -430,7 +430,7 @@ void SelectiveStringDirectColumnReader::readWithVisitor( std::is_same_v; auto nulls = nullsInReadRange_ ? nullsInReadRange_->as() : nullptr; - if (process::hasAvx2() && isExtract) { + if (process::hasSimd() && isExtract) { if (nullsInReadRange_) { if (TVisitor::dense) { returnReaderNulls_ = true; diff --git a/velox/dwio/dwrf/reader/SelectiveStructColumnReader.cpp b/velox/dwio/dwrf/reader/SelectiveStructColumnReader.cpp index 0e88d53bebd..e07f6a2a022 100644 --- a/velox/dwio/dwrf/reader/SelectiveStructColumnReader.cpp +++ b/velox/dwio/dwrf/reader/SelectiveStructColumnReader.cpp @@ -82,7 +82,7 @@ SelectiveStructColumnReader::SelectiveStructColumnReader( auto childParams = DwrfParams( stripe, labels, - params.runtimeStatistics(), + params.splitStats(), FlatMapContext{ .sequence = encodingKey.sequence(), .inMapDecoder = nullptr, diff --git a/velox/dwio/dwrf/reader/StripeStream.h b/velox/dwio/dwrf/reader/StripeStream.h index 80473167a2f..92d8058476c 100644 --- a/velox/dwio/dwrf/reader/StripeStream.h +++ b/velox/dwio/dwrf/reader/StripeStream.h @@ -231,7 +231,7 @@ class StripeStreamsImpl : public StripeStreamsBase { int64_t stripeNumberOfRows, const StrideIndexProvider& provider, uint32_t stripeIndex, - dwio::common::ColumnReaderStatistics* columnReaderStats = nullptr) + dwio::common::SplitStats* splitStats = nullptr) : StripeStreamsBase{&readState->readerBase->memoryPool()}, readState_(std::move(readState)), selector_{selector}, @@ -241,7 +241,7 @@ class StripeStreamsImpl : public StripeStreamsBase { stripeNumberOfRows_{stripeNumberOfRows}, provider_(provider), stripeIndex_{stripeIndex}, - columnReaderStats_{columnReaderStats} { + splitStats_{splitStats} { loadStreams(); } @@ -366,10 +366,13 @@ class StripeStreamsImpl : public StripeStreamsBase { } io::IoCounter* getDecompressCounter(uint32_t nodeId) const { - if (!columnReaderStats_ || !columnReaderStats_->decodingStatsSet) { + if (!splitStats_) { + return nullptr; + } + auto* stats = splitStats_->decodingStats(nodeId); + if (!stats) { return nullptr; } - auto* stats = columnReaderStats_->decodingStatsSet->getOrCreate(nodeId); return &stats->decompressCPUTimeNanos; } @@ -385,7 +388,7 @@ class StripeStreamsImpl : public StripeStreamsBase { const int64_t stripeNumberOfRows_; const StrideIndexProvider& provider_; const uint32_t stripeIndex_; - dwio::common::ColumnReaderStatistics* const columnReaderStats_{nullptr}; + dwio::common::SplitStats* const splitStats_{nullptr}; bool readPlanLoaded_{false}; diff --git a/velox/dwio/dwrf/test/ChecksumTests.cpp b/velox/dwio/dwrf/test/ChecksumTests.cpp index e2fcf3c2c51..0416f1303b8 100644 --- a/velox/dwio/dwrf/test/ChecksumTests.cpp +++ b/velox/dwio/dwrf/test/ChecksumTests.cpp @@ -46,7 +46,7 @@ class ChecksumTests : public Test { std::array data; }; -TEST_F(ChecksumTests, Null) { +TEST_F(ChecksumTests, null) { auto checksum = ChecksumFactory::create(proto::ChecksumAlgorithm::NULL_); ASSERT_EQ(checksum, nullptr); } @@ -58,6 +58,6 @@ TEST_F(ChecksumTests, xxHash) { 2625948533963027735); } -TEST_F(ChecksumTests, Crc32) { +TEST_F(ChecksumTests, crc32) { runTest(proto::ChecksumAlgorithm::CRC32, 4133052486, 3074245904); } diff --git a/velox/dwio/dwrf/test/ColumnWriterIndexTest.cpp b/velox/dwio/dwrf/test/ColumnWriterIndexTest.cpp index abf23b57b67..ec47c02ed24 100644 --- a/velox/dwio/dwrf/test/ColumnWriterIndexTest.cpp +++ b/velox/dwio/dwrf/test/ColumnWriterIndexTest.cpp @@ -436,7 +436,7 @@ class TimestampWriterIndexTest : public testing::Test, // is sufficient }; -TEST_F(TimestampWriterIndexTest, TestIndex) { +TEST_F(TimestampWriterIndexTest, testIndex) { // Present Stream has 4. seconds and nanos are Ints with 3 each. // 4+3+3 = 10. There is no backfill. runTest(2, 10, 0, 0); @@ -476,7 +476,7 @@ VELOX_TYPED_TEST_SUITE( IntegerColumnWriterDictionaryEncodingIndexTest, DictionaryTypes); -TYPED_TEST(IntegerColumnWriterDictionaryEncodingIndexTest, WriteAllStreams) { +TYPED_TEST(IntegerColumnWriterDictionaryEncodingIndexTest, writeAllStreams) { // Present stream uses 4 positions. // compressed size, uncompressed size, ByteRLE numLiterals, // Boolean 8-bitsRemaining @@ -492,7 +492,7 @@ TYPED_TEST(IntegerColumnWriterDictionaryEncodingIndexTest, WriteAllStreams) { this->runTest(1, RECORD_POSITION_COUNT, BACKFILL_POSITION_COUNT, 100); } -TYPED_TEST(IntegerColumnWriterDictionaryEncodingIndexTest, OmitInDictStream) { +TYPED_TEST(IntegerColumnWriterDictionaryEncodingIndexTest, omitInDictStream) { // Present stream uses 4 positions. // compressed size, uncompressed size, ByteRLE numLiterals, // Boolean 8-bitsRemaining @@ -533,7 +533,7 @@ class BoolColumnWriterEncodingIndexTest : public testing::Test, } }; -TEST_F(BoolColumnWriterEncodingIndexTest, TestIndex) { +TEST_F(BoolColumnWriterEncodingIndexTest, testIndex) { // Boolean Stream uses one Stream for Presence and Other Stream for data // Each Stream has compressed size, uncompressed size, ByteRLE numLiterals, // Boolean 8-bitsRemaining , so 4+4 =8. It has no back fill streams. @@ -569,7 +569,7 @@ class ByteColumnWriterEncodingIndexTest : public testing::Test, } }; -TEST_F(ByteColumnWriterEncodingIndexTest, TestIndex) { +TEST_F(ByteColumnWriterEncodingIndexTest, testIndex) { // Byte Stream has presence Stream and data stream. // Present Stream has compressed size, uncompressed size, ByteRLE numLiterals, // Boolean 8-bitsRemaining (4) @@ -630,7 +630,7 @@ class BinaryColumnWriterEncodingIndexTest : public testing::Test, } }; -TEST_F(BinaryColumnWriterEncodingIndexTest, TestIndex) { +TEST_F(BinaryColumnWriterEncodingIndexTest, testIndex) { // Binary Stream has present, data and length. // Present Stream has compressed size, uncompressed size, ByteRLE numLiterals, // Boolean 8-bitsRemaining (4) @@ -674,7 +674,7 @@ class FloatColumnWriterEncodingIndexTest : public testing::Test, using FloatTypes = ::testing::Types; VELOX_TYPED_TEST_SUITE(FloatColumnWriterEncodingIndexTest, FloatTypes); -TYPED_TEST(FloatColumnWriterEncodingIndexTest, TestIndex) { +TYPED_TEST(FloatColumnWriterEncodingIndexTest, testIndex) { // Float Stream has present and data. // Present Stream has compressed size, uncompressed size, ByteRLE numLiterals, // Boolean 8-bitsRemaining (4) @@ -868,7 +868,7 @@ class IntegerColumnWriterDirectEncodingIndexTest : public testing::Test { bool abandonDict_; }; -TEST_F(IntegerColumnWriterDirectEncodingIndexTest, ConvertFromDictionary) { +TEST_F(IntegerColumnWriterDirectEncodingIndexTest, convertFromDictionary) { testForAllTypes(1, 6, 0); testForAllTypes(10, 6, 0); // 6 positions: {PRESENT, 4}, {DATA, 2} @@ -876,7 +876,7 @@ TEST_F(IntegerColumnWriterDirectEncodingIndexTest, ConvertFromDictionary) { testForAllTypes(10, 6, 1); } -TEST_F(IntegerColumnWriterDirectEncodingIndexTest, DirectWrites) { +TEST_F(IntegerColumnWriterDirectEncodingIndexTest, directWrites) { // Same 6 positions but recorded at different points in time. testForAllTypes(1, 6, 2); testForAllTypes(10, 6, 100); @@ -889,7 +889,7 @@ class IntegerColumnWriterAbandonDictionaryIndexTest : IntegerColumnWriterDirectEncodingIndexTest{true} {} }; -TEST_F(IntegerColumnWriterAbandonDictionaryIndexTest, AbandonDictionary) { +TEST_F(IntegerColumnWriterAbandonDictionaryIndexTest, abandonDictionary) { testForAllTypes(1, 6, 0, abandonEveryWrite); testForAllTypes(10, 6, 0, abandonEveryWrite); // 6 positions: {PRESENT, 4}, {DATA, 2} @@ -993,7 +993,7 @@ class StringColumnWriterDictionaryEncodingIndexTest : public testing::Test { std::shared_ptr config_; }; -TEST_F(StringColumnWriterDictionaryEncodingIndexTest, WriteAllStreams) { +TEST_F(StringColumnWriterDictionaryEncodingIndexTest, writeAllStreams) { runTest(1, 17, 0); // 17 positions: {PRESENT, 4}, {STRIDE_DICT_DATA, 2}, {STRIDE_DICT_LENGTH, 3}, // {stride dict size, 1}, {IN_DICTIONARY, 4}, {DATA, 3} @@ -1001,7 +1001,7 @@ TEST_F(StringColumnWriterDictionaryEncodingIndexTest, WriteAllStreams) { runTest(1, 17, 100); } -TEST_F(StringColumnWriterDictionaryEncodingIndexTest, OmitInDictStream) { +TEST_F(StringColumnWriterDictionaryEncodingIndexTest, omitInDictStream) { runTest(2, 7, 0); // Writing the batch twice so that all values are in dictionary. // 7 positions: {PRESENT, 4}, {DATA, 3} @@ -1179,7 +1179,7 @@ class StringColumnWriterDirectEncodingIndexTest : public testing::Test { bool abandonDict_; }; -TEST_F(StringColumnWriterDirectEncodingIndexTest, ConvertFromDictionary) { +TEST_F(StringColumnWriterDirectEncodingIndexTest, convertFromDictionary) { runTest(1, 9, 0); runTest(10, 9, 0); // 9 positions: {PRESENT, 4}, {DATA, 2}, {DATA_LENGTH, 3} @@ -1187,7 +1187,7 @@ TEST_F(StringColumnWriterDirectEncodingIndexTest, ConvertFromDictionary) { runTest(10, 9, 1); } -TEST_F(StringColumnWriterDirectEncodingIndexTest, DirectWrites) { +TEST_F(StringColumnWriterDirectEncodingIndexTest, directWrites) { // Same 9 positions but recorded at different points in time. runTest(1, 9, 2); runTest(10, 9, 100); @@ -1200,7 +1200,7 @@ class StringColumnWriterAbandonDictionaryIndexTest : StringColumnWriterDirectEncodingIndexTest{true} {} }; -TEST_F(StringColumnWriterAbandonDictionaryIndexTest, AbandonDictionary) { +TEST_F(StringColumnWriterAbandonDictionaryIndexTest, abandonDictionary) { runTest(1, 9, 0, abandonEveryWrite); runTest(10, 9, 0, abandonEveryWrite); // 9 positions: {PRESENT, 4}, {DATA, 2}, {DATA_LENGTH, 3} @@ -1268,7 +1268,7 @@ class ListColumnWriterEncodingIndexTest : public testing::Test, } }; -TEST_F(ListColumnWriterEncodingIndexTest, TestIndex) { +TEST_F(ListColumnWriterEncodingIndexTest, testIndex) { // List Stream has present and length. // Present Stream has compressed size, uncompressed size, ByteRLE numLiterals, // Boolean 8-bitsRemaining (4) @@ -1343,7 +1343,7 @@ class MapColumnWriterEncodingIndexTest : public testing::Test, } }; -TEST_F(MapColumnWriterEncodingIndexTest, TestIndex) { +TEST_F(MapColumnWriterEncodingIndexTest, testIndex) { // Map Stream has present and length. // Present Stream has compressed size, uncompressed size, ByteRLE numLiterals, // Boolean 8-bitsRemaining (4) @@ -1433,7 +1433,7 @@ class FlatMapColumnWriterEncodingIndexTest : public testing::Test, } }; -TEST_F(FlatMapColumnWriterEncodingIndexTest, TestIndex) { +TEST_F(FlatMapColumnWriterEncodingIndexTest, testIndex) { config_->set(Config::FLATTEN_MAP, true); config_->set(Config::MAP_FLAT_COLS, {0}); @@ -1495,7 +1495,7 @@ class StructColumnWriterEncodingIndexTest : public testing::Test, } }; -TEST_F(StructColumnWriterEncodingIndexTest, TestIndex) { +TEST_F(StructColumnWriterEncodingIndexTest, testIndex) { // Struct Stream has present stream // Present Stream has compressed size, uncompressed size, ByteRLE numLiterals, // Boolean 8-bitsRemaining (4) diff --git a/velox/dwio/dwrf/test/ColumnWriterStatsTests.cpp b/velox/dwio/dwrf/test/ColumnWriterStatsTests.cpp index c9e3c9877da..aa2501596ff 100644 --- a/velox/dwio/dwrf/test/ColumnWriterStatsTests.cpp +++ b/velox/dwio/dwrf/test/ColumnWriterStatsTests.cpp @@ -256,7 +256,7 @@ VectorPtr makeFlatVector( return flatVector; } -TEST_F(ColumnWriterStatsTest, Bool) { +TEST_F(ColumnWriterStatsTest, bool) { auto populateBoolBatch = [](MemoryPool& pool, VectorPtr* vector, size_t size) { BufferPtr nulls = allocateNulls(size, &pool); @@ -282,7 +282,7 @@ TEST_F(ColumnWriterStatsTest, Bool) { verifyTypeStats("struct", populateBoolBatch); } -TEST_F(ColumnWriterStatsTest, TinyInt) { +TEST_F(ColumnWriterStatsTest, tinyInt) { auto populateTinyIntBatch = [](MemoryPool& pool, VectorPtr* vector, size_t size) { BufferPtr nulls = allocateNulls(size, &pool); @@ -308,7 +308,7 @@ TEST_F(ColumnWriterStatsTest, TinyInt) { verifyTypeStats("struct", populateTinyIntBatch); } -TEST_F(ColumnWriterStatsTest, SmallInt) { +TEST_F(ColumnWriterStatsTest, smallInt) { auto populateSmallIntBatch = [](MemoryPool& pool, VectorPtr* vector, size_t size) { BufferPtr nulls = allocateNulls(size, &pool); @@ -336,7 +336,7 @@ TEST_F(ColumnWriterStatsTest, SmallInt) { verifyTypeStats("struct", populateSmallIntBatch); } -TEST_F(ColumnWriterStatsTest, Int) { +TEST_F(ColumnWriterStatsTest, int) { auto populateIntBatch = [](MemoryPool& pool, VectorPtr* vector, size_t size) { BufferPtr nulls = allocateNulls(size, &pool); auto* nullsPtr = nulls->asMutable(); @@ -363,7 +363,7 @@ TEST_F(ColumnWriterStatsTest, Int) { verifyTypeStats("struct", populateIntBatch); } -TEST_F(ColumnWriterStatsTest, Long) { +TEST_F(ColumnWriterStatsTest, long) { auto populateLongBatch = [](MemoryPool& pool, VectorPtr* vector, size_t size) { BufferPtr nulls = allocateNulls(size, &pool); @@ -414,11 +414,11 @@ auto populateFloatBatch = [](MemoryPool& pool, VectorPtr* vector, size_t size) { return std::vector{totalSize, totalSize}; }; -TEST_F(ColumnWriterStatsTest, Float) { +TEST_F(ColumnWriterStatsTest, float) { verifyTypeStats("struct", populateFloatBatch); } -TEST_F(ColumnWriterStatsTest, Double) { +TEST_F(ColumnWriterStatsTest, double) { auto populateDoubleBatch = [](MemoryPool& pool, VectorPtr* vector, size_t size) { BufferPtr nulls = allocateNulls(size, &pool); @@ -446,7 +446,7 @@ TEST_F(ColumnWriterStatsTest, Double) { verifyTypeStats("struct", populateDoubleBatch); } -TEST_F(ColumnWriterStatsTest, String) { +TEST_F(ColumnWriterStatsTest, string) { auto populateStringBatch = [](MemoryPool& pool, VectorPtr* vector, size_t size) { std::mt19937 gen{}; @@ -466,7 +466,7 @@ TEST_F(ColumnWriterStatsTest, String) { verifyTypeStats("struct", populateStringBatch); } -TEST_F(ColumnWriterStatsTest, Binary) { +TEST_F(ColumnWriterStatsTest, binary) { auto populateBinaryBatch = [](MemoryPool& pool, VectorPtr* vector, size_t size) { std::mt19937 gen{}; @@ -486,7 +486,7 @@ TEST_F(ColumnWriterStatsTest, Binary) { verifyTypeStats("struct", populateBinaryBatch); } -TEST_F(ColumnWriterStatsTest, Timestamp) { +TEST_F(ColumnWriterStatsTest, timestamp) { auto populateTimestampBatch = [](MemoryPool& pool, VectorPtr* vector, size_t size) { @@ -515,7 +515,7 @@ TEST_F(ColumnWriterStatsTest, Timestamp) { verifyTypeStats("struct", populateTimestampBatch); } -TEST_F(ColumnWriterStatsTest, List) { +TEST_F(ColumnWriterStatsTest, list) { auto populateListBatch = [](MemoryPool& pool, VectorPtr* vector, size_t size) { @@ -561,7 +561,7 @@ TEST_F(ColumnWriterStatsTest, List) { verifyTypeStats("struct>", populateListBatch); } -TEST_F(ColumnWriterStatsTest, Map) { +TEST_F(ColumnWriterStatsTest, map) { auto populateMapBatch = [](MemoryPool& pool, VectorPtr* vector, size_t size) { BufferPtr nulls = allocateNulls(size, &pool); auto* nullsPtr = nulls->asMutable(); @@ -616,7 +616,7 @@ TEST_F(ColumnWriterStatsTest, Map) { "struct>", populateMapBatch, FLAT_MAP_COL_ID); } -TEST_F(ColumnWriterStatsTest, Struct) { +TEST_F(ColumnWriterStatsTest, struct) { auto populateStructBatch = [](MemoryPool& pool, VectorPtr* vector, size_t size) { BufferPtr nulls = allocateNulls(size, &pool); diff --git a/velox/dwio/dwrf/test/ColumnWriterTest.cpp b/velox/dwio/dwrf/test/ColumnWriterTest.cpp index b70333559a3..260828da5a4 100644 --- a/velox/dwio/dwrf/test/ColumnWriterTest.cpp +++ b/velox/dwio/dwrf/test/ColumnWriterTest.cpp @@ -446,7 +446,7 @@ TEST_F(ColumnWriterTest, StringDictionaryEncodingEnabledConfig) { EXPECT_FALSE(writer->useDictionaryEncoding()); } -TEST_F(ColumnWriterTest, TestBooleanWriter) { +TEST_F(ColumnWriterTest, testBooleanWriter) { std::vector> data; for (auto i = 0; i < ITERATIONS; ++i) { bool value = (bool)(Random::rand32() & 1); @@ -458,7 +458,7 @@ TEST_F(ColumnWriterTest, TestBooleanWriter) { testDataTypeWriter(BOOLEAN(), data, 3); } -TEST_F(ColumnWriterTest, TestNullBooleanWriter) { +TEST_F(ColumnWriterTest, testNullBooleanWriter) { std::vector> data; for (auto i = 0; i < ITERATIONS; ++i) { data.emplace_back(); @@ -503,7 +503,7 @@ TEST_F(ColumnWriterTest, testDecimalWriter) { testDataTypeWriter(DECIMAL(38, 4), longValues, /*sequence=*/0, format); } -TEST_F(ColumnWriterTest, TestTimestampEpochWriter) { +TEST_F(ColumnWriterTest, testTimestampEpochWriter) { std::vector> data; // This value will be corrupted. verified in verifyValue method. data.emplace_back(Timestamp(-1, 1)); @@ -517,7 +517,7 @@ TEST_F(ColumnWriterTest, TestTimestampEpochWriter) { testDataTypeWriter(TIMESTAMP(), data); } -TEST_F(ColumnWriterTest, TestTimestampWriter) { +TEST_F(ColumnWriterTest, testTimestampWriter) { std::vector> data; for (int64_t i = 0; i < ITERATIONS; ++i) { Timestamp ts(i, i); @@ -529,7 +529,7 @@ TEST_F(ColumnWriterTest, TestTimestampWriter) { testDataTypeWriter(TIMESTAMP(), data, 6); } -TEST_F(ColumnWriterTest, TestTimestampBoundaryValuesWriter) { +TEST_F(ColumnWriterTest, testTimestampBoundaryValuesWriter) { std::vector> data; for (int64_t i = 0; i < ITERATIONS; ++i) { if (i & 1) { @@ -544,7 +544,7 @@ TEST_F(ColumnWriterTest, TestTimestampBoundaryValuesWriter) { testDataTypeWriter(TIMESTAMP(), data); } -TEST_F(ColumnWriterTest, TestTimestampMixedWriter) { +TEST_F(ColumnWriterTest, testTimestampMixedWriter) { std::vector> data; for (int64_t i = 0; i < ITERATIONS; ++i) { int64_t seconds = Random::rand64(Timestamp::kMaxSeconds); @@ -569,7 +569,7 @@ void verifyInvalidTimestamp(int64_t seconds, int64_t nanos) { testDataTypeWriter(TIMESTAMP(), data), exception::LoggedException); } -TEST_F(ColumnWriterTest, TestTimestampNullWriter) { +TEST_F(ColumnWriterTest, testTimestampNullWriter) { std::vector> data; for (int64_t i = 0; i < ITERATIONS; ++i) { data.emplace_back(); @@ -577,7 +577,7 @@ TEST_F(ColumnWriterTest, TestTimestampNullWriter) { testDataTypeWriter(TIMESTAMP(), data); } -TEST_F(ColumnWriterTest, TestBooleanMixedWriter) { +TEST_F(ColumnWriterTest, testBooleanMixedWriter) { std::vector> data; for (auto i = 0; i < ITERATIONS; ++i) { bool value = (bool)(Random::rand32() & 1); @@ -587,7 +587,7 @@ TEST_F(ColumnWriterTest, TestBooleanMixedWriter) { testDataTypeWriter(BOOLEAN(), data); } -TEST_F(ColumnWriterTest, TestAllBytesWriter) { +TEST_F(ColumnWriterTest, testAllBytesWriter) { std::vector> data; for (int16_t i = INT8_MIN; i <= INT8_MAX; ++i) { data.emplace_back(i); @@ -598,7 +598,7 @@ TEST_F(ColumnWriterTest, TestAllBytesWriter) { testDataTypeWriter(TINYINT(), data); } -TEST_F(ColumnWriterTest, TestRepeatedValuesByteWriter) { +TEST_F(ColumnWriterTest, testRepeatedValuesByteWriter) { std::vector> data; for (auto i = 0; i < ITERATIONS; ++i) { data.emplace_back(INT8_MIN); @@ -606,7 +606,7 @@ TEST_F(ColumnWriterTest, TestRepeatedValuesByteWriter) { testDataTypeWriter(TINYINT(), data); } -TEST_F(ColumnWriterTest, TestOnlyNullByteWriter) { +TEST_F(ColumnWriterTest, testOnlyNullByteWriter) { std::vector> data; for (auto i = 0; i <= ITERATIONS; ++i) { data.emplace_back(); @@ -614,7 +614,7 @@ TEST_F(ColumnWriterTest, TestOnlyNullByteWriter) { testDataTypeWriter(TINYINT(), data); } -TEST_F(ColumnWriterTest, TestByteNullAndExtremeValueMixed) { +TEST_F(ColumnWriterTest, testByteNullAndExtremeValueMixed) { std::vector> data; for (auto i = 0; i < ITERATIONS; ++i) { data.emplace_back(INT8_MIN); @@ -637,7 +637,7 @@ void generateSampleData(std::vector>& data) { } } -TEST_F(ColumnWriterTest, TestByteWriter) { +TEST_F(ColumnWriterTest, testByteWriter) { std::vector> data; generateSampleData(data); testDataTypeWriter(TINYINT(), data); @@ -646,7 +646,7 @@ TEST_F(ColumnWriterTest, TestByteWriter) { testDataTypeWriter(TINYINT(), data, 5); } -TEST_F(ColumnWriterTest, TestShortWriter) { +TEST_F(ColumnWriterTest, testShortWriter) { std::vector> data; generateSampleData(data); testDataTypeWriter(SMALLINT(), data); @@ -655,7 +655,7 @@ TEST_F(ColumnWriterTest, TestShortWriter) { testDataTypeWriter(SMALLINT(), data, 23); } -TEST_F(ColumnWriterTest, TestIntWriter) { +TEST_F(ColumnWriterTest, testIntWriter) { std::vector> data; generateSampleData(data); testDataTypeWriter(INTEGER(), data); @@ -664,7 +664,7 @@ TEST_F(ColumnWriterTest, TestIntWriter) { testDataTypeWriter(INTEGER(), data, 1); } -TEST_F(ColumnWriterTest, TestLongWriter) { +TEST_F(ColumnWriterTest, testLongWriter) { std::vector> data; generateSampleData(data); testDataTypeWriter(BIGINT(), data); @@ -673,7 +673,7 @@ TEST_F(ColumnWriterTest, TestLongWriter) { testDataTypeWriter(BIGINT(), data, 42); } -TEST_F(ColumnWriterTest, TestBinaryWriter) { +TEST_F(ColumnWriterTest, testBinaryWriter) { std::vector> data; const size_t size = 100; for (size_t i = 0; i < size; ++i) { @@ -690,7 +690,7 @@ TEST_F(ColumnWriterTest, TestBinaryWriter) { testDataTypeWriter(VARBINARY(), data, 42); } -TEST_F(ColumnWriterTest, TestBinaryWriterAllNulls) { +TEST_F(ColumnWriterTest, testBinaryWriterAllNulls) { std::vector> data{100}; testDataTypeWriter(VARBINARY(), data); } @@ -1292,7 +1292,7 @@ void testMapWriterRowImpl() { testMapWriterRow(*pool, batches, true, true); } -TEST_F(ColumnWriterTest, TestMapWriterNestedRow) { +TEST_F(ColumnWriterTest, testMapWriterNestedRow) { testMapWriterRowImpl(); testMapWriterRowImpl>(); testMapWriterRowImpl>(); @@ -1314,7 +1314,7 @@ TEST_F(ColumnWriterTest, TestMapWriterNestedRow) { // reserve(MAP_FLAT_MAX_KEYS) capacity, reallocating it mid-loop and dangling // the structKeys_ StringViews still being iterated -> heap-use-after-free (same // F14 dangling-key crash family as D96817300). Reproduces under ASAN. -TEST_F(ColumnWriterTest, FlatMapStructKeysDanglingStringView) { +TEST_F(ColumnWriterTest, flatMapStructKeysDanglingStringView) { const auto rowType = ROW({{"c0", MAP(VARCHAR(), BIGINT())}}); const auto writerSchema = TypeWithId::create(rowType); const auto mapColumn = writerSchema->childAt(0); @@ -1422,7 +1422,7 @@ void testMapWriterNumericKeyUseFlatMap(bool useFlatMap) { testMapWriterNumericKey(useFlatMap, MapWriterInputType::kFlatMap); } -TEST_F(ColumnWriterTest, TestMapWriterFloatKey) { +TEST_F(ColumnWriterTest, testMapWriterFloatKey) { testMapWriterNumericKey(/* useFlatMap */ false); EXPECT_THROW( @@ -1444,14 +1444,14 @@ TEST_F(ColumnWriterTest, TestMapWriterFloatKey) { exception::LoggedException); } -TEST_F(ColumnWriterTest, TestMapWriterInt64Key) { +TEST_F(ColumnWriterTest, testMapWriterInt64Key) { testMapWriterNumericKey(/* useFlatMap */ false); testMapWriterNumericKey(/* useFlatMap */ true); testMapWriterNumericKeyUseStruct(/* useFlatMap */ true); testMapWriterNumericKeyUseFlatMap(/* useFlatMap */ true); } -TEST_F(ColumnWriterTest, TestMapWriterDuplicatedInt64Key) { +TEST_F(ColumnWriterTest, testMapWriterDuplicatedInt64Key) { using T = int64_t; using b = MapBuilder; @@ -1464,28 +1464,28 @@ TEST_F(ColumnWriterTest, TestMapWriterDuplicatedInt64Key) { "Duplicated key in map: 5"); } -TEST_F(ColumnWriterTest, TestMapWriterInt32Key) { +TEST_F(ColumnWriterTest, testMapWriterInt32Key) { testMapWriterNumericKey(/* useFlatMap */ false); testMapWriterNumericKey(/* useFlatMap */ true); testMapWriterNumericKeyUseStruct(/* useFlatMap */ true); testMapWriterNumericKeyUseFlatMap(/* useFlatMap */ true); } -TEST_F(ColumnWriterTest, TestMapWriterInt16Key) { +TEST_F(ColumnWriterTest, testMapWriterInt16Key) { testMapWriterNumericKey(/* useFlatMap */ false); testMapWriterNumericKey(/* useFlatMap */ true); testMapWriterNumericKeyUseStruct(/* useFlatMap */ true); testMapWriterNumericKeyUseFlatMap(/* useFlatMap */ true); } -TEST_F(ColumnWriterTest, TestMapWriterInt8Key) { +TEST_F(ColumnWriterTest, testMapWriterInt8Key) { testMapWriterNumericKey(/* useFlatMap */ false); testMapWriterNumericKey(/* useFlatMap */ true); testMapWriterNumericKeyUseStruct(/* useFlatMap */ true); testMapWriterNumericKeyUseFlatMap(/* useFlatMap */ true); } -TEST_F(ColumnWriterTest, TestMapWriterStringKey) { +TEST_F(ColumnWriterTest, testMapWriterStringKey) { using keyType = StringView; using valueType = StringView; using b = MapBuilder; @@ -1580,7 +1580,7 @@ void testFlatMapWriter( // // With the bug: crashes with SIGABRT in F14Table::rehashImpl. // With the fix: passes (keys are properly owned). -TEST_F(ColumnWriterTest, TestFlatMapDanglingStringViewKeyOnRehash) { +TEST_F(ColumnWriterTest, testFlatMapDanglingStringViewKeyOnRehash) { const auto rowType = CppToType>>::create(); const auto writerSchema = TypeWithId::create(rowType); const auto writerDataTypeWithId = writerSchema->childAt(0); @@ -1651,7 +1651,7 @@ TEST_F(ColumnWriterTest, TestFlatMapDanglingStringViewKeyOnRehash) { writer->createIndexEntry(); } -TEST_F(ColumnWriterTest, TestFlatMapKeyNotInAllBatches) { +TEST_F(ColumnWriterTest, testFlatMapKeyNotInAllBatches) { VectorMaker maker(pool_.get()); // Test the case where not all keys appear in all batches. const std::vector batches{ @@ -1665,7 +1665,7 @@ TEST_F(ColumnWriterTest, TestFlatMapKeyNotInAllBatches) { testFlatMapWriter(batches, pool_.get()); } -TEST_F(ColumnWriterTest, TesFlatMapDuplicatedKey) { +TEST_F(ColumnWriterTest, testFlatMapDuplicatedKey) { const size_t size = 3; const BufferPtr inMaps = AlignedBuffer::allocate(size, pool_.get()); bits::fillBits(inMaps->asMutable(), 1, size, pool_.get()); @@ -1685,7 +1685,7 @@ TEST_F(ColumnWriterTest, TesFlatMapDuplicatedKey) { testFlatMapWriter({batch}, pool_.get()), "Duplicated key in map: 2"); } -TEST_F(ColumnWriterTest, TestMapWriterDuplicatedStringKey) { +TEST_F(ColumnWriterTest, testMapWriterDuplicatedStringKey) { using keyType = StringView; using valueType = StringView; using b = MapBuilder; @@ -1700,7 +1700,7 @@ TEST_F(ColumnWriterTest, TestMapWriterDuplicatedStringKey) { "Duplicated key in map: 2"); } -TEST_F(ColumnWriterTest, TestMapWriterDifferentNumericKeyValue) { +TEST_F(ColumnWriterTest, testMapWriterDifferentNumericKeyValue) { using keyType = float; using valueType = int32_t; using b = MapBuilder; @@ -1714,7 +1714,7 @@ TEST_F(ColumnWriterTest, TestMapWriterDifferentNumericKeyValue) { testMapWriter(*pool_, batch, /* useFlatMap */ false); } -TEST_F(ColumnWriterTest, TestMapWriterDifferentKeyValue) { +TEST_F(ColumnWriterTest, testMapWriterDifferentKeyValue) { using keyType = float; using valueType = StringView; using b = MapBuilder; @@ -1728,7 +1728,7 @@ TEST_F(ColumnWriterTest, TestMapWriterDifferentKeyValue) { testMapWriter(*pool_, batch, /* useFlatMap */ false); } -TEST_F(ColumnWriterTest, TestMapWriterMixedBatchTypeHandling) { +TEST_F(ColumnWriterTest, testMapWriterMixedBatchTypeHandling) { using keyType = int32_t; using valueType = int32_t; using b = MapBuilder; @@ -1761,7 +1761,7 @@ TEST_F(ColumnWriterTest, TestMapWriterMixedBatchTypeHandling) { ""); } -TEST_F(ColumnWriterTest, TestMapWriterBinaryKey) { +TEST_F(ColumnWriterTest, testMapWriterBinaryKey) { using keyType = StringView; using valueType = int32_t; using b = MapBuilder; @@ -1799,7 +1799,7 @@ void testMapWriterImpl() { testMapWriter(*pool, batch, /* useFlatMap */ true); } -TEST_F(ColumnWriterTest, TestMapWriterNestedMap) { +TEST_F(ColumnWriterTest, testMapWriterNestedMap) { testMapWriterImpl(); testMapWriterImpl>(); testMapWriterImpl>(); @@ -1813,7 +1813,7 @@ TEST_F(ColumnWriterTest, TestMapWriterNestedMap) { testMapWriterImpl>(); } -TEST_F(ColumnWriterTest, TestMapWriterDifferentStripeBatches) { +TEST_F(ColumnWriterTest, testMapWriterDifferentStripeBatches) { using keyType = int32_t; using valueType = int32_t; using b = MapBuilder; @@ -1847,7 +1847,7 @@ TEST_F(ColumnWriterTest, TestMapWriterDifferentStripeBatches) { false); } -TEST_F(ColumnWriterTest, TestMapWriterNullValues) { +TEST_F(ColumnWriterTest, testMapWriterNullValues) { using keyType = int32_t; using valueType = int32_t; using b = MapBuilder; @@ -1862,7 +1862,7 @@ TEST_F(ColumnWriterTest, TestMapWriterNullValues) { testMapWriter(*pool_, batch, /* useFlatMap */ true); } -TEST_F(ColumnWriterTest, TestMapWriterNullRows) { +TEST_F(ColumnWriterTest, testMapWriterNullRows) { using keyType = int32_t; using valueType = int32_t; using b = MapBuilder; @@ -1880,7 +1880,7 @@ TEST_F(ColumnWriterTest, TestMapWriterNullRows) { testMapWriter(*pool_, batch, /* useFlatMap */ true); } -TEST_F(ColumnWriterTest, TestMapWriterDuplicateKeys) { +TEST_F(ColumnWriterTest, testMapWriterDuplicateKeys) { using keyType = int32_t; using valueType = int32_t; using b = MapBuilder; @@ -1900,7 +1900,7 @@ TEST_F(ColumnWriterTest, TestMapWriterDuplicateKeys) { exception::LoggedException); } -TEST_F(ColumnWriterTest, TestMapWriterBigBatch) { +TEST_F(ColumnWriterTest, testMapWriterBigBatch) { using keyType = int32_t; using valueType = float; using b = MapBuilder; @@ -1936,7 +1936,7 @@ TEST_F(ColumnWriterTest, TestMapWriterBigBatch) { /* useFlatMap */ true); } -TEST_F(ColumnWriterTest, TestMapWriterUnalignedKeyValueCount) { +TEST_F(ColumnWriterTest, testMapWriterUnalignedKeyValueCount) { VectorMaker maker(pool_.get()); auto keys = maker.flatVector(11, folly::identity); auto values = maker.flatVector(12, folly::identity); @@ -1978,7 +1978,7 @@ TEST_F(ColumnWriterTest, TestMapWriterUnalignedKeyValueCount) { (testMapWriter(*pool_, batch, true)), ""); } -TEST_F(ColumnWriterTest, TestStructKeysConfigSerializationDeserialization) { +TEST_F(ColumnWriterTest, testStructKeysConfigSerializationDeserialization) { const std::vector> columns{ {"1.45", "hi, you;", "29102819", "1e-4"}, {"291", "world"}, @@ -2072,7 +2072,7 @@ void testMapWriterStats(const std::shared_ptr type) { } } -TEST_F(ColumnWriterTest, TestMapWriterCompareStatsBinaryKey) { +TEST_F(ColumnWriterTest, testMapWriterCompareStatsBinaryKey) { using keyType = Varbinary; // We create a complex map with complex value structure to test that value // aggregation work well in flat maps @@ -2082,7 +2082,7 @@ TEST_F(ColumnWriterTest, TestMapWriterCompareStatsBinaryKey) { testMapWriterStats(type); } -TEST_F(ColumnWriterTest, TestMapWriterCompareStatsStringKey) { +TEST_F(ColumnWriterTest, testMapWriterCompareStatsStringKey) { using keyType = std::string; // We create a complex map with complex value structure to test that value // aggregation work well in flat maps @@ -2092,7 +2092,7 @@ TEST_F(ColumnWriterTest, TestMapWriterCompareStatsStringKey) { testMapWriterStats(type); } -TEST_F(ColumnWriterTest, TestMapWriterCompareStatsInt8Key) { +TEST_F(ColumnWriterTest, testMapWriterCompareStatsInt8Key) { using keyType = int8_t; // We create a complex map with complex value structure to test that value @@ -2103,7 +2103,7 @@ TEST_F(ColumnWriterTest, TestMapWriterCompareStatsInt8Key) { testMapWriterStats(type); } -TEST_F(ColumnWriterTest, TestMapWriterCompareStatsInt16Key) { +TEST_F(ColumnWriterTest, testMapWriterCompareStatsInt16Key) { using keyType = int16_t; // We create a complex map with complex value structure to test that value @@ -2114,7 +2114,7 @@ TEST_F(ColumnWriterTest, TestMapWriterCompareStatsInt16Key) { testMapWriterStats(type); } -TEST_F(ColumnWriterTest, TestMapWriterCompareStatsInt32Key) { +TEST_F(ColumnWriterTest, testMapWriterCompareStatsInt32Key) { using keyType = int32_t; // We create a complex map with complex value structure to test that value @@ -2125,7 +2125,7 @@ TEST_F(ColumnWriterTest, TestMapWriterCompareStatsInt32Key) { testMapWriterStats(type); } -TEST_F(ColumnWriterTest, TestMapWriterCompareStatsInt64Key) { +TEST_F(ColumnWriterTest, testMapWriterCompareStatsInt64Key) { using keyType = int64_t; // We create a complex map with complex value structure to test that value @@ -2143,11 +2143,11 @@ void testFractionalWrite(const TypePtr& t) { testDataTypeWriter(t, data); } -TEST_F(ColumnWriterTest, TestFloatWriter) { +TEST_F(ColumnWriterTest, testFloatWriter) { testFractionalWrite(REAL()); } -TEST_F(ColumnWriterTest, TestDoubleWriter) { +TEST_F(ColumnWriterTest, testDoubleWriter) { testFractionalWrite(DOUBLE()); } @@ -2161,11 +2161,11 @@ void testFractionalInfinityWrite(const TypePtr& t) { testDataTypeWriter(t, data); } -TEST_F(ColumnWriterTest, TestFloatInfinityWriter) { +TEST_F(ColumnWriterTest, testFloatInfinityWriter) { testFractionalInfinityWrite(REAL()); } -TEST_F(ColumnWriterTest, TestDoubleInfinityWriter) { +TEST_F(ColumnWriterTest, testDoubleInfinityWriter) { testFractionalInfinityWrite(DOUBLE()); } @@ -2179,11 +2179,11 @@ void testFractionalNegativeInfinityWrite(const TypePtr& t) { testDataTypeWriter(t, data); } -TEST_F(ColumnWriterTest, TestFloatNegativeInfinityWriter) { +TEST_F(ColumnWriterTest, testFloatNegativeInfinityWriter) { testFractionalNegativeInfinityWrite(REAL()); } -TEST_F(ColumnWriterTest, TestDoubleNegativeInfinityWriter) { +TEST_F(ColumnWriterTest, testDoubleNegativeInfinityWriter) { testFractionalNegativeInfinityWrite(DOUBLE()); } @@ -2197,11 +2197,11 @@ void testFractionalNaNWrite(const TypePtr& t) { testDataTypeWriter(t, data); } -TEST_F(ColumnWriterTest, TestFloatNanWriter) { +TEST_F(ColumnWriterTest, testFloatNanWriter) { testFractionalNaNWrite(REAL()); } -TEST_F(ColumnWriterTest, TestDoubleNanWriter) { +TEST_F(ColumnWriterTest, testDoubleNanWriter) { testFractionalNaNWrite(DOUBLE()); } @@ -2214,11 +2214,11 @@ void testFractionalNullWrite(const TypePtr& t) { testDataTypeWriter(t, data); } -TEST_F(ColumnWriterTest, TestFloatAllNullWriter) { +TEST_F(ColumnWriterTest, testFloatAllNullWriter) { testFractionalNullWrite(REAL()); } -TEST_F(ColumnWriterTest, TestDoubleAllNullWriter) { +TEST_F(ColumnWriterTest, testDoubleAllNullWriter) { testFractionalNullWrite(DOUBLE()); } @@ -2242,11 +2242,11 @@ void testFractionalMixedWrite(const TypePtr& t) { testDataTypeWriter(t, data); } -TEST_F(ColumnWriterTest, TestFloatMixedWriter) { +TEST_F(ColumnWriterTest, testFloatMixedWriter) { testFractionalMixedWrite(REAL()); } -TEST_F(ColumnWriterTest, TestDoubleMixedWriter) { +TEST_F(ColumnWriterTest, testDoubleMixedWriter) { testFractionalMixedWrite(DOUBLE()); } @@ -2747,7 +2747,7 @@ struct IntegerColumnWriterDirectEncodingUniversalTestCase flushCount} {} }; -TEST_F(ColumnWriterTest, IntegerTypeDictionaryEncodingWrites) { +TEST_F(ColumnWriterTest, integerTypeDictionaryEncodingWrites) { struct TestCase : public IntegerColumnWriterDictionaryEncodingUniversalTestCase { TestCase( @@ -2787,7 +2787,7 @@ TEST_F(ColumnWriterTest, IntegerTypeDictionaryEncodingWrites) { } } -TEST_F(ColumnWriterTest, IntegerTypeDictionaryEncodingWritesWithNulls) { +TEST_F(ColumnWriterTest, integerTypeDictionaryEncodingWritesWithNulls) { struct DictionaryEncodingTestCase : public IntegerColumnWriterDictionaryEncodingUniversalTestCase { DictionaryEncodingTestCase( @@ -2856,7 +2856,7 @@ TEST_F(ColumnWriterTest, IntegerTypeDictionaryEncodingWritesWithNulls) { } } -TEST_F(ColumnWriterTest, IntegerTypeDictionaryEncodingHugeWrites) { +TEST_F(ColumnWriterTest, integerTypeDictionaryEncodingHugeWrites) { struct TestCase : public IntegerColumnWriterDictionaryEncodingUniversalTestCase { TestCase( @@ -2900,7 +2900,7 @@ TEST_F(ColumnWriterTest, IntegerTypeDictionaryEncodingHugeWrites) { } // Split test to avoid sandcastle timeouts. -TEST_F(ColumnWriterTest, IntegerTypeDictionaryEncodingHugeRepeatedWrites) { +TEST_F(ColumnWriterTest, integerTypeDictionaryEncodingHugeRepeatedWrites) { struct TestCase : public IntegerColumnWriterDictionaryEncodingUniversalTestCase { TestCase( @@ -2937,7 +2937,7 @@ TEST_F(ColumnWriterTest, IntegerTypeDictionaryEncodingHugeRepeatedWrites) { } } -TEST_F(ColumnWriterTest, IntegerTypeDirectEncodingWrites) { +TEST_F(ColumnWriterTest, integerTypeDirectEncodingWrites) { struct TestCase : public IntegerColumnWriterDirectEncodingUniversalTestCase { TestCase( size_t size, @@ -2971,7 +2971,7 @@ TEST_F(ColumnWriterTest, IntegerTypeDirectEncodingWrites) { } } -TEST_F(ColumnWriterTest, IntegerTypeDirectEncodingWritesWithNulls) { +TEST_F(ColumnWriterTest, integerTypeDirectEncodingWritesWithNulls) { struct TestCase : public IntegerColumnWriterDirectEncodingUniversalTestCase { TestCase( size_t size, @@ -3007,7 +3007,7 @@ TEST_F(ColumnWriterTest, IntegerTypeDirectEncodingWritesWithNulls) { } } -TEST_F(ColumnWriterTest, IntegerTypeDirectEncodingHugeWrites) { +TEST_F(ColumnWriterTest, integerTypeDirectEncodingHugeWrites) { struct TestCase : public IntegerColumnWriterDirectEncodingUniversalTestCase { TestCase( size_t size, @@ -3037,7 +3037,7 @@ TEST_F(ColumnWriterTest, IntegerTypeDirectEncodingHugeWrites) { } // Split test to avoid sandcastle timeouts. -TEST_F(ColumnWriterTest, IntegerTypeDirectEncodingHugeRepeatedWrites) { +TEST_F(ColumnWriterTest, integerTypeDirectEncodingHugeRepeatedWrites) { struct TestCase : public IntegerColumnWriterDirectEncodingUniversalTestCase { TestCase( size_t size, @@ -3070,7 +3070,7 @@ TEST_F(ColumnWriterTest, IntegerTypeDirectEncodingHugeRepeatedWrites) { } } -TEST_F(ColumnWriterTest, IntegerTypeDictionaryWriteThreshold) { +TEST_F(ColumnWriterTest, integerTypeDictionaryWriteThreshold) { struct DictionaryEncodingTestCase : public IntegerColumnWriterDictionaryEncodingUniversalTestCase { DictionaryEncodingTestCase( @@ -3161,7 +3161,7 @@ TEST_F(ColumnWriterTest, IntegerTypeDictionaryWriteThreshold) { } } -TEST_F(ColumnWriterTest, IntegerColumnWriterAbandonDictionaries) { +TEST_F(ColumnWriterTest, integerColumnWriterAbandonDictionaries) { struct TestCase : public IntegerColumnWriterUniversalTestCase { TestCase( size_t size, @@ -3293,7 +3293,7 @@ TEST_F(ColumnWriterTest, IntegerColumnWriterAbandonDictionaries) { } } -TEST_F(ColumnWriterTest, IntegerColumnWriterAbandonDictionariesWithNulls) { +TEST_F(ColumnWriterTest, integerColumnWriterAbandonDictionariesWithNulls) { struct TestCase : public IntegerColumnWriterUniversalTestCase { TestCase( size_t size, @@ -3425,7 +3425,7 @@ TEST_F(ColumnWriterTest, IntegerColumnWriterAbandonDictionariesWithNulls) { } } -TEST_F(ColumnWriterTest, IntegerColumnWriterAbandonLowValueDictionaries) { +TEST_F(ColumnWriterTest, integerColumnWriterAbandonLowValueDictionaries) { struct TestCase : public IntegerColumnWriterUniversalTestCase { TestCase( size_t size, @@ -3630,7 +3630,7 @@ void testIntegerDictionaryEncodableWriterConstructor() { } } -TEST_F(ColumnWriterTest, IntegerDictionaryDictionaryEncodableWriterCtor) { +TEST_F(ColumnWriterTest, integerDictionaryDictionaryEncodableWriterCtor) { testIntegerDictionaryEncodableWriterConstructor(); testIntegerDictionaryEncodableWriterConstructor(); testIntegerDictionaryEncodableWriterConstructor(); @@ -3879,7 +3879,7 @@ struct StringDirectEncodingTestCase : public StringColumnWriterTestCase { flushCount} {} }; -TEST_F(ColumnWriterTest, StringDictionaryEncodingWrite) { +TEST_F(ColumnWriterTest, stringDictionaryEncodingWrite) { struct TestCase : public StringDictionaryEncodingTestCase { explicit TestCase( size_t size, @@ -3953,7 +3953,7 @@ bool genNulls_ForStride2( return strideIndex == 2; } -TEST_F(ColumnWriterTest, StrideStringWithSomeDataNotInDictionary) { +TEST_F(ColumnWriterTest, strideStringWithSomeDataNotInDictionary) { struct TestCase : public StringDictionaryEncodingTestCase { explicit TestCase( size_t size, @@ -3985,7 +3985,7 @@ TEST_F(ColumnWriterTest, StrideStringWithSomeDataNotInDictionary) { } } -TEST_F(ColumnWriterTest, StringDictionaryEncodingWritesWithNulls) { +TEST_F(ColumnWriterTest, stringDictionaryEncodingWritesWithNulls) { struct DictionaryEncodingTestCase : public StringDictionaryEncodingTestCase { DictionaryEncodingTestCase( size_t size, @@ -4059,7 +4059,7 @@ TEST_F(ColumnWriterTest, StringDictionaryEncodingWritesWithNulls) { } } -TEST_F(ColumnWriterTest, StringDirectEncodingWrites) { +TEST_F(ColumnWriterTest, stringDirectEncodingWrites) { struct TestCase : public StringDirectEncodingTestCase { TestCase( size_t size, @@ -4091,7 +4091,7 @@ TEST_F(ColumnWriterTest, StringDirectEncodingWrites) { } } -TEST_F(ColumnWriterTest, StringDirectEncodingWritesWithNulls) { +TEST_F(ColumnWriterTest, stringDirectEncodingWritesWithNulls) { struct TestCase : public StringDirectEncodingTestCase { TestCase( size_t size, @@ -4127,7 +4127,7 @@ TEST_F(ColumnWriterTest, StringDirectEncodingWritesWithNulls) { } } -TEST_F(ColumnWriterTest, StringColumnWriterAbandonDictionaries) { +TEST_F(ColumnWriterTest, stringColumnWriterAbandonDictionaries) { struct TestCase : public StringColumnWriterTestCase { TestCase( size_t size, @@ -4261,7 +4261,7 @@ TEST_F(ColumnWriterTest, StringColumnWriterAbandonDictionaries) { } // TODO: how about all nulls? -TEST_F(ColumnWriterTest, StringColumnWriterAbandonDictionariesWithNulls) { +TEST_F(ColumnWriterTest, stringColumnWriterAbandonDictionariesWithNulls) { struct TestCase : public StringColumnWriterTestCase { TestCase( size_t size, @@ -4394,7 +4394,7 @@ TEST_F(ColumnWriterTest, StringColumnWriterAbandonDictionariesWithNulls) { } } -TEST_F(ColumnWriterTest, StringColumnWriterAbandonLowValueDictionaries) { +TEST_F(ColumnWriterTest, stringColumnWriterAbandonLowValueDictionaries) { struct TestCase : public StringColumnWriterTestCase { TestCase( size_t size, @@ -4576,7 +4576,7 @@ TEST_F(ColumnWriterTest, StringColumnWriterAbandonLowValueDictionaries) { } } -TEST_F(ColumnWriterTest, IntDictWriterDirectValueOverflow) { +TEST_F(ColumnWriterTest, intDictWriterDirectValueOverflow) { auto config = std::make_shared(); WriterContext context{ config, @@ -4619,7 +4619,7 @@ TEST_F(ColumnWriterTest, IntDictWriterDirectValueOverflow) { } } -TEST_F(ColumnWriterTest, ShortDictWriterDictValueOverflow) { +TEST_F(ColumnWriterTest, shortDictWriterDictValueOverflow) { auto config = std::make_shared(); WriterContext context{config, memory::memoryManager()->addRootPool()}; context.initBuffer(); @@ -4665,7 +4665,7 @@ TEST_F(ColumnWriterTest, ShortDictWriterDictValueOverflow) { } } -TEST_F(ColumnWriterTest, RemovePresentStream) { +TEST_F(ColumnWriterTest, removePresentStream) { auto config = std::make_shared(); std::vector> data; @@ -4696,7 +4696,7 @@ TEST_F(ColumnWriterTest, RemovePresentStream) { ASSERT_EQ(streams.getStream(si, {}, false), nullptr); } -TEST_F(ColumnWriterTest, ColumnIdInStream) { +TEST_F(ColumnWriterTest, columnIdInStream) { auto config = std::make_shared(); std::vector> data; @@ -4915,7 +4915,7 @@ void testDictionary( .runTest(valueAt, [](int) { return false; }); } -TEST_F(ColumnWriterTest, ColumnWriterDictionarySimple) { +TEST_F(ColumnWriterTest, columnWriterDictionarySimple) { testDictionary(TIMESTAMP(), randomNulls(11), [](vector_size_t i) { return Timestamp(i * 5, i * 2); }); diff --git a/velox/dwio/dwrf/test/CommonTests.cpp b/velox/dwio/dwrf/test/CommonTests.cpp index 03b69182c49..2aa4e981014 100644 --- a/velox/dwio/dwrf/test/CommonTests.cpp +++ b/velox/dwio/dwrf/test/CommonTests.cpp @@ -49,7 +49,7 @@ class CommonTest : public ::testing::Test { TEST_F( CommonTest, - DwrfStreamIdentifierGetFormat_WithStreamTypes_GetCorrectFormat) { + dwrfStreamIdentifierGetFormatWithStreamTypesGetCorrectFormat) { EXPECT_EQ( DwrfStreamIdentifier(createDwrfStream({}, {}, {}, {})).format(), DwrfFormat::kDwrf); @@ -58,7 +58,7 @@ TEST_F( DwrfStreamIdentifier(createOrcStream({}, {})).format(), DwrfFormat::kOrc); } -TEST_F(CommonTest, DwrfStreamIdentifier_WithDwrfStream_GetCorrectInfo) { +TEST_F(CommonTest, dwrfStreamIdentifierWithDwrfStreamGetCorrectInfo) { auto streamId = DwrfStreamIdentifier( createDwrfStream(proto::Stream::Kind::Stream_Kind_DATA, 1, 2, 3)); @@ -69,7 +69,7 @@ TEST_F(CommonTest, DwrfStreamIdentifier_WithDwrfStream_GetCorrectInfo) { EXPECT_EQ(streamId.column(), 3); } -TEST_F(CommonTest, DwrfStreamIdentifier_WithOrcStream_GetCorrectInfo) { +TEST_F(CommonTest, dwrfStreamIdentifierWithOrcStreamGetCorrectInfo) { auto streamId = DwrfStreamIdentifier( createOrcStream(proto::orc::Stream::Kind::Stream_Kind_DATA, 1)); @@ -83,7 +83,7 @@ TEST_F(CommonTest, DwrfStreamIdentifier_WithOrcStream_GetCorrectInfo) { TEST_F( CommonTest, - DwrfStreamIdentifierGetKind_WithAllStreamKinds_GetCorrectConversion) { + dwrfStreamIdentifierGetKindWithAllStreamKindsGetCorrectConversion) { // DWRF for (auto [dwrfStreamKind, veloxStreamKind] : std::vector>{ @@ -133,7 +133,7 @@ TEST_F( TEST_F( CommonTest, - EncodingKeyGetKindFor_WithAllStreamKinds_GetCorrectConversion) { + encodingKeyGetKindForWithAllStreamKindsGetCorrectConversion) { // DWRF for (auto [dwrfStreamKind, veloxStreamKind] : std::vector>{ diff --git a/velox/dwio/dwrf/test/ConfigTests.cpp b/velox/dwio/dwrf/test/ConfigTests.cpp index 62120e9d079..aa3033c48ca 100644 --- a/velox/dwio/dwrf/test/ConfigTests.cpp +++ b/velox/dwio/dwrf/test/ConfigTests.cpp @@ -46,7 +46,7 @@ std::shared_ptr createWriterConfig( return Writer::makeWriterConfig(options, *dwrfOptions); } -TEST(ConfigTests, Set) { +TEST(ConfigTests, set) { Config config; // set auto val = folly::Random::rand32(); @@ -63,7 +63,7 @@ TEST(ConfigTests, Set) { EXPECT_TRUE(config.get(Config::CREATE_INDEX)); } -TEST(ConfigTests, EnumConfig) { +TEST(ConfigTests, enumConfig) { Config config; config.set(Config::COMPRESSION, CompressionKind::CompressionKind_ZLIB); EXPECT_EQ( @@ -73,14 +73,14 @@ TEST(ConfigTests, EnumConfig) { config.get(Config::COMPRESSION), CompressionKind::CompressionKind_NONE); } -TEST(ConfigTests, UInt32Config) { +TEST(ConfigTests, uint32Config) { Config config; auto val = folly::Random::rand32(); config.set(Config::ROW_INDEX_STRIDE, val); EXPECT_EQ(config.get(Config::ROW_INDEX_STRIDE), val); } -TEST(ConfigTests, BoolConfig) { +TEST(ConfigTests, boolConfig) { Config config; config.set(Config::CREATE_INDEX, false); EXPECT_FALSE(config.get(Config::CREATE_INDEX)); @@ -162,7 +162,7 @@ TEST(ConfigTests, writerOptionsOverrideSession) { } class ConfigTests : public ::testing::TestWithParam {}; -TEST_P(ConfigTests, FlatMapCols) { +TEST_P(ConfigTests, flatMapCols) { const auto& params = GetParam(); std::map inputConfigMap{ {"orc.map.flat.cols", params.inputCols}}; diff --git a/velox/dwio/dwrf/test/DataBufferHolderTests.cpp b/velox/dwio/dwrf/test/DataBufferHolderTests.cpp index d4b5f6e1238..ce12bee4c99 100644 --- a/velox/dwio/dwrf/test/DataBufferHolderTests.cpp +++ b/velox/dwio/dwrf/test/DataBufferHolderTests.cpp @@ -31,7 +31,7 @@ class DataBufferHolderTest : public testing::Test { std::shared_ptr pool_{memoryManager()->addLeafPool()}; }; -TEST_F(DataBufferHolderTest, InputCheck) { +TEST_F(DataBufferHolderTest, inputCheck) { VELOX_ASSERT_THROW((DataBufferHolder{*pool_, 0}), ""); VELOX_ASSERT_THROW((DataBufferHolder{*pool_, 1024, 2048}), ""); VELOX_ASSERT_THROW((DataBufferHolder{*pool_, 1024, 1024, 1.1f}), ""); @@ -47,7 +47,7 @@ TEST_F(DataBufferHolderTest, InputCheck) { } } -TEST_F(DataBufferHolderTest, TakeAndGetBuffer) { +TEST_F(DataBufferHolderTest, takeAndGetBuffer) { MemorySink sink{1024, {.pool = pool_.get()}}; DataBufferHolder holder{*pool_, 1024, 0, 2.0f, &sink}; DataBuffer buffer{*pool_, 512}; @@ -69,7 +69,7 @@ TEST_F(DataBufferHolderTest, TakeAndGetBuffer) { ASSERT_EQ(holder.getBuffers().size(), 0); } -TEST_F(DataBufferHolderTest, TruncateBufferHolder) { +TEST_F(DataBufferHolderTest, truncateBufferHolder) { DataBufferHolder holder{*pool_, 1024}; constexpr size_t BUF_SIZE = 10; DataBuffer buffer{*pool_, BUF_SIZE}; @@ -93,7 +93,7 @@ TEST_F(DataBufferHolderTest, TruncateBufferHolder) { } } -TEST_F(DataBufferHolderTest, TakeAndGetBufferNoOutput) { +TEST_F(DataBufferHolderTest, takeAndGetBufferNoOutput) { DataBufferHolder holder{*pool_, 1024}; DataBuffer buffer{*pool_, 512}; std::memset(buffer.data(), 'a', 512); @@ -122,7 +122,7 @@ TEST_F(DataBufferHolderTest, TakeAndGetBufferNoOutput) { } } -TEST_F(DataBufferHolderTest, Reset) { +TEST_F(DataBufferHolderTest, reset) { DataBufferHolder holder{*pool_, 1024}; DataBuffer buffer{*pool_, 512}; std::memset(buffer.data(), 'a', 512); @@ -135,7 +135,7 @@ TEST_F(DataBufferHolderTest, Reset) { ASSERT_FALSE(holder.isSuppressed()); } -TEST_F(DataBufferHolderTest, TryResize) { +TEST_F(DataBufferHolderTest, tryResize) { DataBufferHolder holder{*pool_, 1024, 128}; auto runTest = [&](uint64_t size, @@ -248,7 +248,7 @@ TEST_F(DataBufferHolderTest, TryResize) { 1024 + headerSize); } -TEST_F(DataBufferHolderTest, TestGrowRatio) { +TEST_F(DataBufferHolderTest, testGrowRatio) { DataBufferHolder holder{*pool_, 1024, 16, 4.0f}; DataBuffer buffer{*pool_, 16}; ASSERT_TRUE(holder.tryResize(buffer, 0, 1)); diff --git a/velox/dwio/dwrf/test/DecompressionTest.cpp b/velox/dwio/dwrf/test/DecompressionTest.cpp index 7a4d6b28377..e5c848e63e9 100644 --- a/velox/dwio/dwrf/test/DecompressionTest.cpp +++ b/velox/dwio/dwrf/test/DecompressionTest.cpp @@ -953,18 +953,18 @@ class TestSeek : public ::testing::Test { std::shared_ptr pool_ = memory::memoryManager()->addLeafPool(); }; -TEST_F(TestSeek, Zlib) { +TEST_F(TestSeek, zlib) { auto codec = zlib::getCodec( zlib::Options(zlib::Options::Format::RAW), COMPRESSION_LEVEL_DEFAULT); runTest(*codec, CompressionKind_ZLIB); } -TEST_F(TestSeek, Zstd) { +TEST_F(TestSeek, zstd) { auto codec = getCodec(CodecType::ZSTD); runTest(*codec, CompressionKind_ZSTD); } -TEST_F(TestSeek, Snappy) { +TEST_F(TestSeek, snappy) { auto codec = getCodec(CodecType::SNAPPY); runTest(*codec, CompressionKind_SNAPPY); } diff --git a/velox/dwio/dwrf/test/DecryptionTests.cpp b/velox/dwio/dwrf/test/DecryptionTests.cpp index 7cff7a4acb2..28f0aef2485 100644 --- a/velox/dwio/dwrf/test/DecryptionTests.cpp +++ b/velox/dwio/dwrf/test/DecryptionTests.cpp @@ -28,7 +28,7 @@ using namespace facebook::velox::dwrf; using namespace facebook::velox::dwrf::encryption; using namespace facebook::velox::type::fbhive; -TEST(Decryption, NotEncrypted) { +TEST(Decryption, notEncrypted) { HiveTypeParser parser; auto type = parser.parse("struct"); proto::Footer footer; @@ -39,7 +39,7 @@ TEST(Decryption, NotEncrypted) { ASSERT_FALSE(handler->isEncrypted()); } -TEST(Decryption, NoKeyProvider) { +TEST(Decryption, noKeyProvider) { HiveTypeParser parser; auto type = parser.parse("struct"); proto::Footer footer; @@ -51,7 +51,7 @@ TEST(Decryption, NoKeyProvider) { DecryptionHandler::create(footer, &factory), exception::LoggedException); } -TEST(Decryption, EmptyGroup) { +TEST(Decryption, emptyGroup) { HiveTypeParser parser; auto type = parser.parse("struct"); proto::Footer footer; @@ -64,7 +64,7 @@ TEST(Decryption, EmptyGroup) { DecryptionHandler::create(footer, &factory), exception::LoggedException); } -TEST(Decryption, EmptyNodes) { +TEST(Decryption, emptyNodes) { HiveTypeParser parser; auto type = parser.parse("struct"); proto::Footer footer; @@ -79,7 +79,7 @@ TEST(Decryption, EmptyNodes) { DecryptionHandler::create(footer, &factory), exception::LoggedException); } -TEST(Decryption, StatsMismatch) { +TEST(Decryption, statsMismatch) { HiveTypeParser parser; auto type = parser.parse("struct"); proto::Footer footer; @@ -97,7 +97,7 @@ TEST(Decryption, StatsMismatch) { DecryptionHandler::create(footer, &factory), exception::LoggedException); } -TEST(Decryption, KeyExistenceMismatch) { +TEST(Decryption, keyExistenceMismatch) { HiveTypeParser parser; auto type = parser.parse("struct"); proto::Footer footer; @@ -118,7 +118,7 @@ TEST(Decryption, KeyExistenceMismatch) { DecryptionHandler::create(footer, &factory), exception::LoggedException); } -TEST(Decryption, ReuseStripeKey) { +TEST(Decryption, reuseStripeKey) { HiveTypeParser parser; auto type = parser.parse("struct"); proto::Footer footer; @@ -138,7 +138,7 @@ TEST(Decryption, ReuseStripeKey) { ASSERT_EQ(td.getKey(), "foobar"); } -TEST(Decryption, StripeKeyMismatch) { +TEST(Decryption, stripeKeyMismatch) { HiveTypeParser parser; auto type = parser.parse("struct"); proto::Footer footer; @@ -157,7 +157,7 @@ TEST(Decryption, StripeKeyMismatch) { DecryptionHandler::create(footer, &factory), exception::LoggedException); } -TEST(Decryption, Basic) { +TEST(Decryption, basic) { HiveTypeParser parser; auto type = parser.parse("struct"); proto::Footer footer; @@ -187,7 +187,7 @@ TEST(Decryption, Basic) { } } -TEST(Decryption, NestedType) { +TEST(Decryption, nestedType) { HiveTypeParser parser; auto type = parser.parse( "struct>,c:struct,d:array>"); @@ -228,7 +228,7 @@ TEST(Decryption, NestedType) { } } -TEST(Decryption, RootNode) { +TEST(Decryption, rootNode) { HiveTypeParser parser; auto type = parser.parse("struct"); proto::Footer footer; @@ -245,7 +245,7 @@ TEST(Decryption, RootNode) { ASSERT_EQ(handler->getEncryptionGroupCount(), 1); } -TEST(Decryption, GroupOverlap) { +TEST(Decryption, groupOverlap) { HiveTypeParser parser; auto type = parser.parse("struct>"); proto::Footer footer; diff --git a/velox/dwio/dwrf/test/DirectBufferedInputTest.cpp b/velox/dwio/dwrf/test/DirectBufferedInputTest.cpp index 8d540d34c31..9ba0e8e632b 100644 --- a/velox/dwio/dwrf/test/DirectBufferedInputTest.cpp +++ b/velox/dwio/dwrf/test/DirectBufferedInputTest.cpp @@ -18,13 +18,19 @@ #include #include #include +#include +#include +#include +#include "velox/common/file/LocalFile.h" #include "velox/common/io/IoStatistics.h" #include "velox/common/memory/MmapAllocator.h" +#include "velox/common/testutil/TempFilePath.h" #include "velox/dwio/common/Options.h" #include "velox/dwio/dwrf/common/Common.h" #include "velox/dwio/dwrf/test/TestReadFile.h" #include +#include using namespace facebook::velox; using namespace facebook::velox::dwio; @@ -42,6 +48,42 @@ struct TestRegion { bool read = true; }; +class QueuedExecutor final : public folly::Executor { + public: + void add(folly::Func func) override { + functions_.push_back(std::move(func)); + } + + size_t size() const { + return functions_.size(); + } + + std::exception_ptr runNext() { + VELOX_CHECK(!functions_.empty()); + auto func = std::move(functions_.front()); + functions_.pop_front(); + try { + func(); + return nullptr; + } catch (...) { + return std::current_exception(); + } + } + + private: + std::deque functions_; +}; + +std::string readAll(SeekableInputStream& stream) { + std::string result; + const void* buffer; + int32_t size; + while (stream.Next(&buffer, &size)) { + result.append(static_cast(buffer), size); + } + return result; +} + class DirectBufferedInputTest : public testing::Test { protected: static constexpr int32_t kLoadQuantum = 8 << 20; @@ -226,3 +268,55 @@ TEST_F(DirectBufferedInputTest, ioStatsLifeTimeTest) { t.join(); } } + +TEST_F(DirectBufferedInputTest, cancelledLocalLoadFallsBackToSyncRead) { + constexpr int32_t kRegionSize = 1024; + std::string content(4 * kRegionSize, 0); + for (size_t i = 0; i < content.size(); ++i) { + content[i] = 'a' + (i % 26); + } + + auto tempFile = facebook::velox::common::testutil::TempFilePath::create(); + tempFile->append(content); + auto localFile = std::make_shared(tempFile->getPath()); + QueuedExecutor executor; + auto input = std::make_unique( + localFile, + dwio::common::MetricsLog::voidLog(), + StringIdLease{}, + tracker_, + StringIdLease{}, + ioStatistics_, + ioStats_, + &executor, + *opts_); + + auto first = input->enqueue({0, kRegionSize}, nullptr); + auto second = input->enqueue({kRegionSize, kRegionSize}, nullptr); + input->load(LogType::FILE); + ASSERT_EQ(executor.size(), 1); + + // Keep the size cached by LocalReadFile, but truncate the backing file so + // preadv fills the first region and half of the second before returning a + // short read. ReadFileInputStream turns the short read into an exception and + // the coalesced load becomes cancelled. + ASSERT_EQ( + ::truncate(tempFile->getPath().c_str(), kRegionSize + kRegionSize / 2), + 0); + ASSERT_NE(executor.runNext(), nullptr); + + // Restore the file so the foreground synchronous fallback can succeed. + { + std::ofstream output( + tempFile->getPath(), std::ios::binary | std::ios::trunc); + ASSERT_TRUE(output.good()); + output.write(content.data(), content.size()); + output.close(); + ASSERT_TRUE(output.good()); + } + + EXPECT_EQ(readAll(*first), content.substr(0, kRegionSize)); + EXPECT_EQ(readAll(*second), content.substr(kRegionSize, kRegionSize)); + EXPECT_EQ(ioStatistics_->read().count(), 2); + EXPECT_EQ(ioStatistics_->read().sum(), 2 * kRegionSize); +} diff --git a/velox/dwio/dwrf/test/E2EReaderTest.cpp b/velox/dwio/dwrf/test/E2EReaderTest.cpp index 05aaa4e84ef..c9a23e0fa19 100644 --- a/velox/dwio/dwrf/test/E2EReaderTest.cpp +++ b/velox/dwio/dwrf/test/E2EReaderTest.cpp @@ -112,7 +112,7 @@ class E2EReaderTest : public testing::TestWithParam { }; } // namespace -TEST_P(E2EReaderTest, SharedDictionaryFlatmapReadAsStruct) { +TEST_P(E2EReaderTest, sharedDictionaryFlatmapReadAsStruct) { const size_t batchCount = 10; size_t size = 1; auto pool = memory::memoryManager()->addLeafPool(); diff --git a/velox/dwio/dwrf/test/E2EWriterTest.cpp b/velox/dwio/dwrf/test/E2EWriterTest.cpp index 8c3ca302d94..0d740ca0e8e 100644 --- a/velox/dwio/dwrf/test/E2EWriterTest.cpp +++ b/velox/dwio/dwrf/test/E2EWriterTest.cpp @@ -322,7 +322,7 @@ RowTypePtr readWithFieldIds( return reader->rowType(); } -TEST_F(E2EWriterTest, FieldIdMappingRenameReorderDrop) { +TEST_F(E2EWriterTest, fieldIdMappingRenameReorderDrop) { // File node ids: 0=root, 1=a, 2=b, 3=c. auto fileSchema = ROW({"a", "b", "c"}, {INTEGER(), BIGINT(), VARCHAR()}); // Requested: c renamed to c2 (id 3, reordered first), a kept (id 1), b @@ -358,7 +358,7 @@ TEST_F(E2EWriterTest, rejectWrongFormatSpecificOptions) { dwrf::Writer(std::move(sink), options), "DwrfWriterOptions"); } -TEST_F(E2EWriterTest, FieldIdMappingDropReaddSameName) { +TEST_F(E2EWriterTest, fieldIdMappingDropReaddSameName) { // File has column c with field id 5; the table dropped it and re-added a new // column c with field id 9. The stale file column must NOT bind to the new c. auto fileSchema = ROW({"c"}, {INTEGER()}); @@ -375,7 +375,7 @@ TEST_F(E2EWriterTest, FieldIdMappingDropReaddSameName) { EXPECT_EQ(rowType->nameOf(0), "$dwrf_unmatched_1"); } -TEST_F(E2EWriterTest, FieldIdMappingNestedStruct) { +TEST_F(E2EWriterTest, fieldIdMappingNestedStruct) { // File node ids: 0=root, 1=s, 2=s.x, 3=s.y. auto fileSchema = ROW({"s"}, {ROW({"x", "y"}, {INTEGER(), INTEGER()})}); // Requested struct reorders children and renames y->y2; ids: s=10, x=11, @@ -405,7 +405,7 @@ TEST_F(E2EWriterTest, FieldIdMappingNestedStruct) { // buck test velox/dwio/dwrf/test:velox_dwrf_e2e_writer_tests -- // DISABLED_TestFileCreation // --run-disabled -TEST_F(E2EWriterTest, DISABLED_TestFileCreation) { +TEST_F(E2EWriterTest, DISABLED_testFileCreation) { const size_t batchCount = 4; const size_t batchSize = 200; @@ -466,7 +466,7 @@ VectorPtr createRowVector( /*nullCount=*/0); } -TEST_F(E2EWriterTest, E2E) { +TEST_F(E2EWriterTest, e2e) { const size_t batchCount = 4; // Start with a size larger than stride to cover splitting into // strides. Continue with smaller size for faster test. @@ -510,7 +510,7 @@ TEST_F(E2EWriterTest, E2E) { } // Disabled because test is failing in continuous runs T193531984. -TEST_F(E2EWriterTest, DISABLED_DisableLinearHeuristics) { +TEST_F(E2EWriterTest, DISABLED_disableLinearHeuristics) { const size_t batchCount = 100; size_t batchSize = 3000; @@ -556,7 +556,7 @@ TEST_F(E2EWriterTest, DISABLED_DisableLinearHeuristics) { // Beside writing larger files, this test also uses regular maps only. // Disabled because test is failing in continuous runs T193531984. -TEST_F(E2EWriterTest, DISABLED_DisableLinearHeuristicsLargeAnalytics) { +TEST_F(E2EWriterTest, DISABLED_disableLinearHeuristicsLargeAnalytics) { const size_t batchCount = 500; size_t batchSize = 3000; @@ -598,7 +598,7 @@ TEST_F(E2EWriterTest, DISABLED_DisableLinearHeuristicsLargeAnalytics) { dwrf::E2EWriterTestUtil::testWriter(*leafPool_, type, batches, 8, 8, config); } -TEST_F(E2EWriterTest, FlatMapDictionaryEncoding) { +TEST_F(E2EWriterTest, flatMapDictionaryEncoding) { const size_t batchCount = 4; // Start with a size larger than stride to cover splitting into // strides. Continue with smaller size for faster test. @@ -635,7 +635,7 @@ TEST_F(E2EWriterTest, FlatMapDictionaryEncoding) { dwrf::E2EWriterTestUtil::testWriter(*pool, type, batches, 1, 1, config); } -TEST_F(E2EWriterTest, MaxFlatMapKeys) { +TEST_F(E2EWriterTest, maxFlatMapKeys) { using keyType = int32_t; using valueType = int32_t; using b = MapBuilder; @@ -667,7 +667,7 @@ TEST_F(E2EWriterTest, MaxFlatMapKeys) { config); } -TEST_F(E2EWriterTest, PresentStreamIsSuppressedOnFlatMap) { +TEST_F(E2EWriterTest, presentStreamIsSuppressedOnFlatMap) { using keyType = int32_t; using valueType = int64_t; using b = MapBuilder; @@ -717,7 +717,7 @@ TEST_F(E2EWriterTest, PresentStreamIsSuppressedOnFlatMap) { } } -TEST_F(E2EWriterTest, TooManyFlatMapKeys) { +TEST_F(E2EWriterTest, tooManyFlatMapKeys) { using keyType = int32_t; using valueType = int32_t; using b = MapBuilder; @@ -751,7 +751,7 @@ TEST_F(E2EWriterTest, TooManyFlatMapKeys) { ""); } -TEST_F(E2EWriterTest, FlatMapBackfill) { +TEST_F(E2EWriterTest, flatMapBackfill) { auto pool = memory::memoryManager()->addLeafPool(); using keyType = int32_t; @@ -860,7 +860,7 @@ void testFlatMapWithNulls( dwrf::E2EWriterTestUtil::simpleFlushPolicyFactory(false)); } -TEST_F(E2EWriterTest, FlatMapWithNulls) { +TEST_F(E2EWriterTest, flatMapWithNulls) { testFlatMapWithNulls( /*firstRowNotNull=*/false, /*enableFlatmapDictionaryEncoding=*/false); testFlatMapWithNulls( @@ -871,7 +871,7 @@ TEST_F(E2EWriterTest, FlatMapWithNulls) { /*firstRowNotNull=*/true, /*enableFlatmapDictionaryEncoding=*/true); } -TEST_F(E2EWriterTest, FlatMapWithNullsSharedDict) { +TEST_F(E2EWriterTest, flatMapWithNullsSharedDict) { testFlatMapWithNulls( /*firstRowNotNull=*/false, /*enableFlatmapDictionaryEncoding=*/true, @@ -882,7 +882,7 @@ TEST_F(E2EWriterTest, FlatMapWithNullsSharedDict) { /*shareDictionary=*/true); } -TEST_F(E2EWriterTest, FlatMapEmpty) { +TEST_F(E2EWriterTest, flatMapEmpty) { auto pool = memory::memoryManager()->addLeafPool(); using keyType = int32_t; @@ -923,7 +923,7 @@ TEST_F(E2EWriterTest, FlatMapEmpty) { dwrf::E2EWriterTestUtil::simpleFlushPolicyFactory(false)); } -TEST_F(E2EWriterTest, FlatMapConfigSingleColumn) { +TEST_F(E2EWriterTest, flatMapConfigSingleColumn) { HiveTypeParser parser; auto type = parser.parse( "struct<" @@ -934,7 +934,7 @@ TEST_F(E2EWriterTest, FlatMapConfigSingleColumn) { testFlatMapConfig(type, {}, {}); } -TEST_F(E2EWriterTest, FlatMapConfigMixedTypes) { +TEST_F(E2EWriterTest, flatMapConfigMixedTypes) { HiveTypeParser parser; auto type = parser.parse( "struct<" @@ -946,7 +946,7 @@ TEST_F(E2EWriterTest, FlatMapConfigMixedTypes) { testFlatMapConfig(type, {}, {}); } -TEST_F(E2EWriterTest, FlatMapConfigNestedMap) { +TEST_F(E2EWriterTest, flatMapConfigNestedMap) { HiveTypeParser parser; auto type = parser.parse( "struct<" @@ -958,7 +958,7 @@ TEST_F(E2EWriterTest, FlatMapConfigNestedMap) { testFlatMapConfig(type, {}, {}); } -TEST_F(E2EWriterTest, FlatMapConfigMixedMaps) { +TEST_F(E2EWriterTest, flatMapConfigMixedMaps) { HiveTypeParser parser; auto type = parser.parse( "struct<" @@ -972,7 +972,7 @@ TEST_F(E2EWriterTest, FlatMapConfigMixedMaps) { testFlatMapConfig(type, {}, {}); } -TEST_F(E2EWriterTest, FlatMapConfigNotMapColumn) { +TEST_F(E2EWriterTest, flatMapConfigNotMapColumn) { HiveTypeParser parser; auto type = parser.parse( "struct<" @@ -1023,7 +1023,7 @@ TEST_F(E2EWriterTest, mapStatsMultiStrides) { testFlatMapFileStats(type, {0, 1, 2, 3, 4, 5}, /*strideSize=*/1000); } -TEST_F(E2EWriterTest, PartialStride) { +TEST_F(E2EWriterTest, partialStride) { auto type = ROW({"bool_val"}, {INTEGER()}); size_t batchSize = 1'000; @@ -1084,7 +1084,7 @@ TEST_F(E2EWriterTest, PartialStride) { ASSERT_EQ(true, reader->columnStatistics(1)->hasNull().value()); } -TEST_F(E2EWriterTest, OversizeRows) { +TEST_F(E2EWriterTest, oversizeRows) { auto pool = facebook::velox::memory::memoryManager()->addLeafPool(); HiveTypeParser parser; @@ -1122,7 +1122,7 @@ TEST_F(E2EWriterTest, OversizeRows) { false); } -TEST_F(E2EWriterTest, OversizeBatches) { +TEST_F(E2EWriterTest, oversizeBatches) { auto pool = facebook::velox::memory::memoryManager()->addLeafPool(); HiveTypeParser parser; @@ -1170,7 +1170,7 @@ TEST_F(E2EWriterTest, OversizeBatches) { false); } -TEST_F(E2EWriterTest, OverflowLengthIncrements) { +TEST_F(E2EWriterTest, overflowLengthIncrements) { auto pool = facebook::velox::memory::memoryManager()->addLeafPool(); HiveTypeParser parser; @@ -1326,7 +1326,7 @@ class E2EEncryptionTest : public E2EWriterTest { std::vector batches_; }; -TEST_F(E2EEncryptionTest, EncryptRoot) { +TEST_F(E2EEncryptionTest, encryptRoot) { auto spec = std::make_shared(EncryptionProvider::Unknown); spec->withRootEncryptionProperties( @@ -1377,7 +1377,7 @@ TEST_F(E2EEncryptionTest, EncryptRoot) { validateFileContent(*reader); } -TEST_F(E2EEncryptionTest, EncryptSelectedFields) { +TEST_F(E2EEncryptionTest, encryptSelectedFields) { auto spec = std::make_shared(EncryptionProvider::Unknown); spec->withEncryptedField( @@ -1461,7 +1461,7 @@ TEST_F(E2EEncryptionTest, EncryptSelectedFields) { validateFileContent(*reader); } -TEST_F(E2EEncryptionTest, EncryptEmptyFile) { +TEST_F(E2EEncryptionTest, encryptEmptyFile) { auto spec = std::make_shared(EncryptionProvider::Unknown); spec->withEncryptedField( @@ -1476,7 +1476,7 @@ TEST_F(E2EEncryptionTest, EncryptEmptyFile) { ASSERT_FALSE(handler.isEncrypted()); } -TEST_F(E2EEncryptionTest, ReadWithoutKey) { +TEST_F(E2EEncryptionTest, readWithoutKey) { auto spec = std::make_shared(EncryptionProvider::Unknown); spec->withEncryptedField( diff --git a/velox/dwio/dwrf/test/EncodingManagerTests.cpp b/velox/dwio/dwrf/test/EncodingManagerTests.cpp index e53c276e29c..c1a142d6d53 100644 --- a/velox/dwio/dwrf/test/EncodingManagerTests.cpp +++ b/velox/dwio/dwrf/test/EncodingManagerTests.cpp @@ -279,7 +279,7 @@ void testEncodingIter( } } // namespace -TEST(TestEncodingManager, EncodingIter) { +TEST(TestEncodingManager, encodingIter) { testEncodingIter({{1, 0}}, {}); testEncodingIter({}, {{{1, 0}}}); testEncodingIter({{1, 0}}, {{{2, 1}, {2, 3}}}); diff --git a/velox/dwio/dwrf/test/EncryptionTests.cpp b/velox/dwio/dwrf/test/EncryptionTests.cpp index c823c1ef1d1..f1f5e83bb35 100644 --- a/velox/dwio/dwrf/test/EncryptionTests.cpp +++ b/velox/dwio/dwrf/test/EncryptionTests.cpp @@ -27,7 +27,7 @@ using namespace facebook::velox::dwio::common::encryption::test; using namespace facebook::velox::dwrf::encryption; using namespace facebook::velox::type::fbhive; -TEST(Encryption, NotEncrypted) { +TEST(Encryption, notEncrypted) { HiveTypeParser parser; auto type = parser.parse("struct"); EncryptionSpecification spec{EncryptionProvider::Unknown}; @@ -36,7 +36,7 @@ TEST(Encryption, NotEncrypted) { ASSERT_FALSE(handler->isEncrypted()); } -TEST(Encryption, RootThenField) { +TEST(Encryption, rootThenField) { HiveTypeParser parser; auto type = parser.parse("struct"); EncryptionSpecification spec{EncryptionProvider::Unknown}; @@ -49,7 +49,7 @@ TEST(Encryption, RootThenField) { exception::LoggedException); } -TEST(Encryption, FieldThenRoot) { +TEST(Encryption, fieldThenRoot) { HiveTypeParser parser; auto type = parser.parse("struct"); EncryptionSpecification spec{EncryptionProvider::Unknown}; @@ -62,7 +62,7 @@ TEST(Encryption, FieldThenRoot) { exception::LoggedException); } -TEST(Encryption, Root) { +TEST(Encryption, root) { HiveTypeParser parser; auto type = parser.parse("struct"); EncryptionSpecification spec{EncryptionProvider::Unknown}; @@ -77,7 +77,7 @@ TEST(Encryption, Root) { } } -TEST(Encryption, InvalidField) { +TEST(Encryption, invalidField) { HiveTypeParser parser; auto type = parser.parse("struct"); EncryptionSpecification spec{EncryptionProvider::Unknown}; @@ -90,7 +90,7 @@ TEST(Encryption, InvalidField) { exception::LoggedException); } -TEST(Encryption, InvalidProperties) { +TEST(Encryption, invalidProperties) { HiveTypeParser parser; auto type = parser.parse("struct"); EncryptionSpecification spec{EncryptionProvider::Unknown}; @@ -101,7 +101,7 @@ TEST(Encryption, InvalidProperties) { exception::LoggedException); } -TEST(Encryption, DifferentEncrypter) { +TEST(Encryption, differentEncrypter) { HiveTypeParser parser; auto type = parser.parse("struct"); EncryptionSpecification spec{EncryptionProvider::Unknown}; @@ -116,7 +116,7 @@ TEST(Encryption, DifferentEncrypter) { ASSERT_EQ(handler->getEncryptionGroupCount(), 2); } -TEST(Encryption, SameEncrypter) { +TEST(Encryption, sameEncrypter) { HiveTypeParser parser; auto type = parser.parse("struct"); EncryptionSpecification spec{EncryptionProvider::Unknown}; @@ -132,7 +132,7 @@ TEST(Encryption, SameEncrypter) { ASSERT_EQ(handler->getEncryptionGroupCount(), 1); } -TEST(Encryption, SameEncrypter2) { +TEST(Encryption, sameEncrypter2) { HiveTypeParser parser; auto type = parser.parse("struct"); EncryptionSpecification spec{EncryptionProvider::Unknown}; @@ -147,7 +147,7 @@ TEST(Encryption, SameEncrypter2) { ASSERT_EQ(handler->getEncryptionGroupCount(), 1); } -TEST(Encryption, DupeField) { +TEST(Encryption, dupeField) { HiveTypeParser parser; auto type = parser.parse("struct"); EncryptionSpecification spec{EncryptionProvider::Unknown}; @@ -164,7 +164,7 @@ TEST(Encryption, DupeField) { exception::LoggedException); } -TEST(Encryption, Basic) { +TEST(Encryption, basic) { HiveTypeParser parser; auto type = parser.parse("struct"); EncryptionSpecification spec{EncryptionProvider::Unknown}; @@ -191,7 +191,7 @@ TEST(Encryption, Basic) { } } -TEST(Encryption, NestedType) { +TEST(Encryption, nestedType) { HiveTypeParser parser; auto type = parser.parse( "struct>,c:struct,d:array>"); diff --git a/velox/dwio/dwrf/test/FileMetadataTest.cpp b/velox/dwio/dwrf/test/FileMetadataTest.cpp index 404762a7a31..45f2024f8fe 100644 --- a/velox/dwio/dwrf/test/FileMetadataTest.cpp +++ b/velox/dwio/dwrf/test/FileMetadataTest.cpp @@ -37,14 +37,12 @@ class FileMetadataTest : public ::testing::Test { StripeFooterWrapper(orcStripeFooter_); }; -TEST_F(FileMetadataTest, StripeFooterWrapper_GetFormat_CorrectFormat) { +TEST_F(FileMetadataTest, stripeFooterWrapperGetFormatCorrectFormat) { EXPECT_EQ(dwrfStripeFooterWrapper_.format(), DwrfFormat::kDwrf); EXPECT_EQ(orcStripeFooterWrapper_.format(), DwrfFormat::kOrc); } -TEST_F( - FileMetadataTest, - StripeFooterWrapper_GetStripeFooter_ReturnsProtoTypes) { +TEST_F(FileMetadataTest, stripeFooterWrapperGetStripeFooterReturnsProtoTypes) { EXPECT_EQ( &dwrfStripeFooterWrapper_.getStripeFooterDwrf(), dwrfStripeFooter_.get()); EXPECT_ANY_THROW(dwrfStripeFooterWrapper_.getStripeFooterOrc()); @@ -56,7 +54,7 @@ TEST_F( TEST_F( FileMetadataTest, - StripeFooterWrapper_GetStreamData_ReturnsCorrispondingProtoInfo) { + stripeFooterWrapperGetStreamDataReturnsCorrespondingProtoInfo) { // 2 streams with incrementing length for validation dwrfStripeFooter_->add_streams()->set_length(1); dwrfStripeFooter_->add_streams()->set_length(2); @@ -92,7 +90,7 @@ TEST_F( TEST_F( FileMetadataTest, - StripeFooterWrapper_GetEncodings_ReturnsCorrispondingProtoInfo) { + stripeFooterWrapperGetEncodingsReturnsCorrespondingProtoInfo) { // 2 encoding with incrementing node for validation dwrfStripeFooter_->add_encoding()->set_kind( proto::ColumnEncoding_Kind::ColumnEncoding_Kind_DICTIONARY); @@ -143,7 +141,7 @@ TEST_F( TEST_F( FileMetadataTest, - StripeFooterWrapper_GetEncryptionGroups_ReturnsCorrispondingProtoInfo) { + stripeFooterWrapperGetEncryptionGroupsReturnsCorrespondingProtoInfo) { // 2 encryption groups with incrementing group for validation dwrfStripeFooter_->add_encryptiongroups()->append("test_encryption_group_1"); dwrfStripeFooter_->add_encryptiongroups()->append("test_encryption_group_2"); diff --git a/velox/dwio/dwrf/test/FlushPolicyTest.cpp b/velox/dwio/dwrf/test/FlushPolicyTest.cpp index b045326664e..cd06801899d 100644 --- a/velox/dwio/dwrf/test/FlushPolicyTest.cpp +++ b/velox/dwio/dwrf/test/FlushPolicyTest.cpp @@ -32,7 +32,7 @@ class DefaultFlushPolicyTest : public testing::Test { } }; -TEST_F(DefaultFlushPolicyTest, StripeProgressTest) { +TEST_F(DefaultFlushPolicyTest, stripeProgressTest) { struct TestCase { const uint64_t stripeSizeThreshold; const int64_t stripeSize; @@ -57,7 +57,7 @@ TEST_F(DefaultFlushPolicyTest, StripeProgressTest) { } } -TEST_F(DefaultFlushPolicyTest, AdditionalCriteriaTest) { +TEST_F(DefaultFlushPolicyTest, additionalCriteriaTest) { struct TestCase { const bool flushStripe; const bool overMemoryBudget; @@ -134,7 +134,7 @@ TEST_F(DefaultFlushPolicyTest, AdditionalCriteriaTest) { } } -TEST_F(DefaultFlushPolicyTest, EarlyDictionaryEvaluation) { +TEST_F(DefaultFlushPolicyTest, earlyDictionaryEvaluation) { // Test the precedence of decisions. struct TestCase { dwio::common::StripeProgress stripeProgress; @@ -237,7 +237,7 @@ TEST_F(DefaultFlushPolicyTest, EarlyDictionaryEvaluation) { FlushDecision::SKIP); } -TEST_F(DefaultFlushPolicyTest, EmptyFile) { +TEST_F(DefaultFlushPolicyTest, emptyFile) { // Empty vector creation succeeds. RowsPerStripeFlushPolicy policy({}); @@ -246,7 +246,7 @@ TEST_F(DefaultFlushPolicyTest, EmptyFile) { exception::LoggedException); } -TEST_F(DefaultFlushPolicyTest, InvalidCases) { +TEST_F(DefaultFlushPolicyTest, invalidCases) { // Vector with 0 rows, throws ASSERT_THROW( RowsPerStripeFlushPolicy policy({5, 7, 0, 10}), @@ -254,7 +254,7 @@ TEST_F(DefaultFlushPolicyTest, InvalidCases) { } // RowsPerStripeFlushPolicy has no dictionary flush criteria. -TEST_F(DefaultFlushPolicyTest, DictionaryCriteriaTest) { +TEST_F(DefaultFlushPolicyTest, dictionaryCriteriaTest) { auto config = std::make_shared(); WriterContext context{ config, memory::memoryManager()->addRootPool("DictionaryCriteriaTest")}; @@ -294,7 +294,7 @@ TEST_F(DefaultFlushPolicyTest, DictionaryCriteriaTest) { context)); } -TEST_F(DefaultFlushPolicyTest, FlushTest) { +TEST_F(DefaultFlushPolicyTest, flushTest) { RowsPerStripeFlushPolicy policy({5, 7, 12}); ASSERT_FALSE(policy.shouldFlush( diff --git a/velox/dwio/dwrf/test/IndexBuilderTests.cpp b/velox/dwio/dwrf/test/IndexBuilderTests.cpp index ee5914e236a..97d33a18634 100644 --- a/velox/dwio/dwrf/test/IndexBuilderTests.cpp +++ b/velox/dwio/dwrf/test/IndexBuilderTests.cpp @@ -41,14 +41,14 @@ class IndexBuilderTest : public testing::Test { StatisticsBuilderOptions options_{16}; }; -TEST_F(IndexBuilderTest, Constructor) { +TEST_F(IndexBuilderTest, constructor) { IndexBuilder builder{nullptr}; EXPECT_EQ(1, builder.getEntrySize()); // Ensure a clean start. EXPECT_EQ(0, getEntry(builder, 0).positionsSize()); } -TEST_F(IndexBuilderTest, AddEntry) { +TEST_F(IndexBuilderTest, addEntry) { IndexBuilder builder{nullptr}; ASSERT_EQ(1, builder.getEntrySize()); builder.add(0uL); @@ -70,7 +70,7 @@ TEST_F(IndexBuilderTest, AddEntry) { } } -TEST_F(IndexBuilderTest, Add) { +TEST_F(IndexBuilderTest, add) { IndexBuilder builder{nullptr}; builder.add(0uL); EXPECT_THAT(getPositions(builder, 0), ElementsAreArray({0uL})); @@ -90,7 +90,7 @@ TEST_F(IndexBuilderTest, Add) { EXPECT_THAT(getPositions(builder, 2), ElementsAreArray({144uL})); } -TEST_F(IndexBuilderTest, Backfill) { +TEST_F(IndexBuilderTest, backfill) { IndexBuilder builder{nullptr}; StatisticsBuilder sb{options_}; builder.addEntry(sb); @@ -124,7 +124,7 @@ TEST_F(IndexBuilderTest, Backfill) { ASSERT_EQ(0, getEntry(builder, 6).positionsSize()); } -TEST_F(IndexBuilderTest, RemovePresentStreamPositions) { +TEST_F(IndexBuilderTest, removePresentStreamPositions) { IndexBuilder builder{nullptr}; auto indexCount = 3; StatisticsBuilder sb{options_}; diff --git a/velox/dwio/dwrf/test/LayoutPlannerTests.cpp b/velox/dwio/dwrf/test/LayoutPlannerTests.cpp index 32cf1a5062b..3b7dac598c9 100644 --- a/velox/dwio/dwrf/test/LayoutPlannerTests.cpp +++ b/velox/dwio/dwrf/test/LayoutPlannerTests.cpp @@ -44,7 +44,7 @@ class LayoutPlannerTest : public testing::Test { }; } // namespace -TEST_F(LayoutPlannerTest, CreateNodeToColumnIdMapping) { +TEST_F(LayoutPlannerTest, createNodeToColumnIdMapping) { testCreateNodeToColumnIdMapping(ROW({BOOLEAN()}), {{0, 0}, {1, 0}}); testCreateNodeToColumnIdMapping( ROW( @@ -96,7 +96,7 @@ TEST_F(LayoutPlannerTest, CreateNodeToColumnIdMapping) { {17, 5}}); } -TEST_F(LayoutPlannerTest, Basic) { +TEST_F(LayoutPlannerTest, basic) { auto config = std::make_shared(); config->set( Config::COMPRESSION, common::CompressionKind::CompressionKind_NONE); diff --git a/velox/dwio/dwrf/test/PhysicalSizeAggregatorTest.cpp b/velox/dwio/dwrf/test/PhysicalSizeAggregatorTest.cpp index 91e5bbf366d..8c1f91cf47f 100644 --- a/velox/dwio/dwrf/test/PhysicalSizeAggregatorTest.cpp +++ b/velox/dwio/dwrf/test/PhysicalSizeAggregatorTest.cpp @@ -21,7 +21,7 @@ using namespace ::testing; namespace facebook::velox::dwrf { -TEST(PhysicalSizeAggregatorTest, UpdateLeaf) { +TEST(PhysicalSizeAggregatorTest, updateLeaf) { auto leafOne = std::make_unique(nullptr); auto leafTwo = std::make_unique(nullptr); ASSERT_EQ(0, leafOne->getResult()); @@ -43,7 +43,7 @@ TEST(PhysicalSizeAggregatorTest, UpdateLeaf) { EXPECT_EQ(2, leafTwo->getResult()); } -TEST(PhysicalSizeAggregatorTest, UpdateParent) { +TEST(PhysicalSizeAggregatorTest, updateParent) { auto parent = std::make_unique(nullptr); auto childOne = std::make_unique(parent.get()); auto childTwo = std::make_unique(parent.get()); diff --git a/velox/dwio/dwrf/test/RatioTrackerTest.cpp b/velox/dwio/dwrf/test/RatioTrackerTest.cpp index 4f4bddf0662..359a8697b52 100644 --- a/velox/dwio/dwrf/test/RatioTrackerTest.cpp +++ b/velox/dwio/dwrf/test/RatioTrackerTest.cpp @@ -23,7 +23,7 @@ using namespace ::testing; namespace facebook::velox::dwrf { -TEST(RatioTrackerTest, BasicTests) { +TEST(RatioTrackerTest, basicTests) { struct TestCase { explicit TestCase( std::shared_ptr tracker, @@ -70,7 +70,7 @@ TEST(RatioTrackerTest, BasicTests) { } } -TEST(RatioTrackerTest, EmptyInputTests) { +TEST(RatioTrackerTest, emptyInputTests) { struct TestCase { explicit TestCase( std::shared_ptr tracker, diff --git a/velox/dwio/dwrf/test/ReaderBaseTests.cpp b/velox/dwio/dwrf/test/ReaderBaseTests.cpp index 08eb3897c41..c281699407f 100644 --- a/velox/dwio/dwrf/test/ReaderBaseTests.cpp +++ b/velox/dwio/dwrf/test/ReaderBaseTests.cpp @@ -237,7 +237,7 @@ class ReaderBaseTest : public Test { } }; -TEST_F(ReaderBaseTest, InvalidPostScriptThrows) { +TEST_F(ReaderBaseTest, invalidPostScriptThrows) { VELOX_ASSERT_THROW( createCorruptedFileReader(1'000'000, 0), "Corrupted file, footer size is invalid"); diff --git a/velox/dwio/dwrf/test/ReaderTest.cpp b/velox/dwio/dwrf/test/ReaderTest.cpp index f28edf1550d..e5c082b835f 100644 --- a/velox/dwio/dwrf/test/ReaderTest.cpp +++ b/velox/dwio/dwrf/test/ReaderTest.cpp @@ -29,6 +29,7 @@ #include "velox/dwio/common/FileSink.h" #include "velox/dwio/common/tests/utils/BatchMaker.h" #include "velox/dwio/dwrf/common/Common.h" +#include "velox/dwio/dwrf/common/DwrfRuntimeStats.h" #include "velox/dwio/dwrf/reader/DwrfReader.h" #include "velox/dwio/dwrf/test/OrcTest.h" #include "velox/dwio/dwrf/test/utils/E2EWriterTestUtil.h" @@ -1687,7 +1688,7 @@ TEST_F(TestReader, fileColumnNamesReadAsLowerCaseComplexStruct) { EXPECT_EQ(col0_1_1_0_0->childByName("ccint3"), col0_1_1_0_0_0); } -TEST_F(TestReader, TestStripeSizeCallback) { +TEST_F(TestReader, testStripeSizeCallback) { dwio::common::ReaderOptions readerOpts{pool()}; readerOpts.setDataIoStats(dataIoStats_); readerOpts.setMetadataIoStats(metadataIoStats_); @@ -1717,7 +1718,7 @@ TEST_F(TestReader, TestStripeSizeCallback) { EXPECT_EQ(numCalls, 1); } -TEST_F(TestReader, TestStripeSizeCallbackLimitsOneStripe) { +TEST_F(TestReader, testStripeSizeCallbackLimitsOneStripe) { dwio::common::ReaderOptions readerOpts{pool()}; readerOpts.setDataIoStats(dataIoStats_); readerOpts.setMetadataIoStats(metadataIoStats_); @@ -1748,7 +1749,7 @@ TEST_F(TestReader, TestStripeSizeCallbackLimitsOneStripe) { EXPECT_EQ(numCalls, 1); } -TEST_F(TestReader, TestStripeSizeCallbackLimitsTwoStripe) { +TEST_F(TestReader, testStripeSizeCallbackLimitsTwoStripe) { dwio::common::ReaderOptions readerOpts{pool()}; readerOpts.setDataIoStats(dataIoStats_); readerOpts.setMetadataIoStats(metadataIoStats_); @@ -2453,6 +2454,46 @@ createWriterReader( return std::make_pair(std::move(writer), std::move(reader)); } +struct FlatMapReaderWithKeyFilter { + std::unique_ptr writer; + std::unique_ptr reader; + std::unique_ptr rowReader; + RowTypePtr schema; + VectorPtr batch; +}; + +FlatMapReaderWithKeyFilter createFlatMapReaderWithKeyFilter( + const std::vector& inputs, + std::shared_ptr keyFilter, + memory::MemoryPool* pool, + const std::shared_ptr& dataIoStats, + const std::shared_ptr& metadataIoStats) { + auto config = std::make_shared(); + config->set(dwrf::Config::FLATTEN_MAP, true); + config->set(dwrf::Config::MAP_FLAT_COLS, {0}); + + auto [writer, reader] = + createWriterReader(inputs, pool, dataIoStats, metadataIoStats, config); + auto schema = asRowType(inputs.front()->type()); + auto scanSpec = std::make_shared(""); + scanSpec->addAllChildFields(*schema); + scanSpec->childByName("c0") + ->childByName(common::ScanSpec::kMapKeysFieldName) + ->setFilter(std::move(keyFilter)); + + RowReaderOptions rowReaderOptions; + rowReaderOptions.setScanSpec(scanSpec); + rowReaderOptions.setPreserveFlatMapsInMemory(true); + auto rowReader = reader->createRowReader(rowReaderOptions); + auto batch = BaseVector::create(schema, 0, pool); + return { + std::move(writer), + std::move(reader), + std::move(rowReader), + std::move(schema), + std::move(batch)}; +} + } // namespace TEST_F(TestReader, setRowNumberColumnInfo) { @@ -2830,6 +2871,189 @@ TEST_F(TestReader, readFlatMapsAsFlatMaps) { {{0, 12}, {1, 13}, {2, 14}, {3, 15}}})); } +TEST_F(TestReader, readFlatMapsAsFlatMapsWithKeyFilter) { + // Reading a flat map with preserveFlatMapsInMemory=true must honor a + // requested-key filter and project only the selected keys into the output + // FlatMapVector. + auto flatMap = makeFlatMapVector({ + {{0, 0}, {1, 1}, {2, 2}, {3, 3}}, + {{0, 4}, {1, 5}, {2, 6}, {3, 7}}, + {{0, 8}, {1, 9}, {2, 10}, {3, 11}}, + }); + auto input = makeRowVector({flatMap->toMapVector()}); + auto result = createFlatMapReaderWithKeyFilter( + {input}, + common::createBigintValues({1, 2}, false), + pool(), + dataIoStats_, + metadataIoStats_); + + ASSERT_EQ( + result.rowReader->next(flatMap->size(), result.batch), flatMap->size()); + auto rowVector = result.batch->as(); + auto resultFlatMap = + rowVector->childAt(0)->loadedVector()->as(); + ASSERT_TRUE(resultFlatMap); + + // Only the selected keys are projected out. + auto distinctKeys = + resultFlatMap->distinctKeys()->as>(); + std::unordered_set keySet; + for (auto i = 0; i < distinctKeys->size(); ++i) { + keySet.insert(distinctKeys->valueAt(i)); + } + EXPECT_EQ(keySet, (std::unordered_set{1, 2})); + + auto expected = makeFlatMapVector({ + {{1, 1}, {2, 2}}, + {{1, 5}, {2, 6}}, + {{1, 9}, {2, 10}}, + }); + assertEqualVectors(expected->toMapVector(), resultFlatMap->toMapVector()); +} + +TEST_F(TestReader, readFlatMapsAsFlatMapsWithStringKeyFilter) { + // Key pruning in the preserving path must also work for string keys. + auto flatMap = makeFlatMapVector({ + {{"a", 1}, {"b", 2}, {"c", 3}}, + {{"a", 4}, {"b", 5}, {"c", 6}}, + }); + auto input = makeRowVector({flatMap->toMapVector()}); + auto result = createFlatMapReaderWithKeyFilter( + {input}, + std::make_unique( + std::vector{"a", "c"}, false), + pool(), + dataIoStats_, + metadataIoStats_); + + ASSERT_EQ( + result.rowReader->next(flatMap->size(), result.batch), flatMap->size()); + auto resultFlatMap = result.batch->as() + ->childAt(0) + ->loadedVector() + ->as(); + ASSERT_TRUE(resultFlatMap); + + auto expected = makeFlatMapVector({ + {{"a", 1}, {"c", 3}}, + {{"a", 4}, {"c", 6}}, + }); + assertEqualVectors(expected->toMapVector(), resultFlatMap->toMapVector()); +} + +TEST_F(TestReader, readFlatMapsAsFlatMapsKeyFilterExcludesAllKeys) { + // A key filter that matches no key in the stripe yields N empty maps, not + // zero rows. + auto flatMap = makeFlatMapVector({ + {{0, 0}, {1, 1}}, + {{0, 2}, {1, 3}}, + }); + auto input = makeRowVector({flatMap->toMapVector()}); + auto result = createFlatMapReaderWithKeyFilter( + {input}, + common::createBigintValues({99}, false), + pool(), + dataIoStats_, + metadataIoStats_); + + ASSERT_EQ( + result.rowReader->next(flatMap->size(), result.batch), flatMap->size()); + auto resultFlatMap = result.batch->as() + ->childAt(0) + ->loadedVector() + ->as(); + ASSERT_TRUE(resultFlatMap); + EXPECT_EQ(resultFlatMap->distinctKeys()->size(), 0); + + auto resultMaps = resultFlatMap->toMapVector(); + ASSERT_EQ(resultMaps->size(), 2); + for (vector_size_t row = 0; row < resultMaps->size(); ++row) { + EXPECT_FALSE(resultMaps->isNullAt(row)); + EXPECT_EQ(resultMaps->sizeAt(row), 0); + } +} + +TEST_F(TestReader, readFlatMapsAsFlatMapsWithKeyFilterAndNullMaps) { + // Key pruning must compose with null maps: null rows stay null and non-null + // rows are pruned to the requested keys. + auto flatMap = makeNullableFlatMapVector({ + {{{0, 0}, {1, 1}, {2, 2}, {3, 3}}}, + {std::nullopt}, + {{{0, 4}, {1, 5}, {2, 6}, {3, 7}}}, + {std::nullopt}, + }); + auto input = makeRowVector({flatMap->toMapVector()}); + auto result = createFlatMapReaderWithKeyFilter( + {input}, + common::createBigintValues({1, 2, 3}, false), + pool(), + dataIoStats_, + metadataIoStats_); + + ASSERT_EQ( + result.rowReader->next(flatMap->size(), result.batch), flatMap->size()); + auto resultFlatMap = result.batch->as() + ->childAt(0) + ->loadedVector() + ->as(); + ASSERT_TRUE(resultFlatMap); + + auto expected = makeNullableFlatMapVector({ + {{{1, 1}, {2, 2}, {3, 3}}}, + {std::nullopt}, + {{{1, 5}, {2, 6}, {3, 7}}}, + {std::nullopt}, + }); + assertEqualVectors(expected->toMapVector(), resultFlatMap->toMapVector()); +} + +TEST_F(TestReader, readFlatMapsAsFlatMapsMultiStripeWithKeyFilter) { + // The key filter must prune every stripe, not just the first. The ScanSpec + // is shared across stripes, so a fix that consumes/clears the filter would + // silently stop pruning after stripe 1. + auto stripe1 = makeRowVector({makeMapVector({ + {{0, 10}, {1, 11}, {2, 12}, {3, 13}, {4, 14}}, + {{0, 20}, {1, 21}, {2, 22}, {3, 23}, {4, 24}}, + })}); + auto stripe2 = makeRowVector({makeMapVector({ + {{0, 30}, {1, 31}, {2, 32}}, + {{0, 40}, {1, 41}, {2, 42}}, + {{0, 50}, {1, 51}, {2, 52}}, + })}); + + // simpleFlushPolicyFactory(true) produces one stripe per batch. + auto result = createFlatMapReaderWithKeyFilter( + {stripe1, stripe2}, + common::createBigintValues({1, 2}, false), + pool(), + dataIoStats_, + metadataIoStats_); + ASSERT_EQ(result.reader->getNumberOfStripes(), 2); + + uint64_t totalRows = 0; + // Read one row at a time to exercise repeated lazy materialization at + // non-zero reader offsets as well as the stripe transition. + while (result.rowReader->next(1, result.batch) > 0) { + auto resultMaps = result.batch->as() + ->childAt(0) + ->loadedVector() + ->as() + ->toMapVector(); + auto* resultKeys = resultMaps->mapKeys()->as>(); + for (vector_size_t row = 0; row < resultMaps->size(); ++row) { + std::unordered_set keySet; + const auto offset = resultMaps->offsetAt(row); + for (vector_size_t i = 0; i < resultMaps->sizeAt(row); ++i) { + keySet.insert(resultKeys->valueAt(offset + i)); + } + EXPECT_EQ(keySet, (std::unordered_set{1, 2})); + } + totalRows += result.batch->size(); + } + EXPECT_EQ(totalRows, 5); // 2 rows from stripe 1 + 3 rows from stripe 2. +} + // Regression test: reading a multi-stripe flatmap file with // preserveFlatMapsInMemory=true used to crash when stripes had different // key sets. The ScanSpec accumulated stale children across stripes, causing @@ -2978,9 +3202,13 @@ TEST_F(TestReader, readStringDictionaryAsFlat) { ASSERT_EQ(c0->encoding(), VectorEncoding::Simple::DICTIONARY); ASSERT_TRUE(c0->valueVector()->isFlatEncoding()); ASSERT_EQ(c0->valueVector()->size(), dictionary.size()); - dwio::common::RuntimeStatistics stats; + dwio::common::RuntimeStats stats; rowReader->updateRuntimeStats(stats); - ASSERT_EQ(stats.columnReaderStats.flattenStringDictionaryValues, 0); + const auto metricName = + std::string(DwrfRuntimeStats::kFlattenStringDictionaryValues); + ASSERT_FALSE(stats.columnStats.at(1) + .at(FileFormat::DWRF) + .columnMetrics.contains(metricName)); spec->childByName("c0")->setFilter( std::make_unique( std::vector{"aaaaaaaaaaaaaaaaaaaa"}, false)); @@ -2989,9 +3217,17 @@ TEST_F(TestReader, readStringDictionaryAsFlat) { ASSERT_EQ(rowReader->next(20, actual), 20); ASSERT_EQ(actual->size(), 1); ASSERT_TRUE(actual->as()->childAt(0)->isFlatEncoding()); - stats = {}; + stats = dwio::common::RuntimeStats(); rowReader->updateRuntimeStats(stats); - ASSERT_EQ(stats.columnReaderStats.flattenStringDictionaryValues, 1); + ASSERT_TRUE(stats.columnStats.at(1) + .at(FileFormat::DWRF) + .columnMetrics.contains(metricName)); + ASSERT_EQ( + stats.columnStats.at(1) + .at(FileFormat::DWRF) + .columnMetrics.at(metricName) + .sum, + 1); } // A primitive subfield is missing in file, and result is not reused. diff --git a/velox/dwio/dwrf/test/StreamLabelsTests.cpp b/velox/dwio/dwrf/test/StreamLabelsTests.cpp index ff7c27f0c43..a29b061e79b 100644 --- a/velox/dwio/dwrf/test/StreamLabelsTests.cpp +++ b/velox/dwio/dwrf/test/StreamLabelsTests.cpp @@ -30,7 +30,7 @@ class StreamLabelsTest : public testing::Test { } }; -TEST_F(StreamLabelsTest, E2E) { +TEST_F(StreamLabelsTest, e2e) { auto pool = facebook::velox::memory::memoryManager()->addLeafPool(); AllocationPool allocationPool(pool.get()); StreamLabels root(allocationPool); diff --git a/velox/dwio/dwrf/test/StripeReaderBaseTests.cpp b/velox/dwio/dwrf/test/StripeReaderBaseTests.cpp index 09ab5503f19..bed6c5829a2 100644 --- a/velox/dwio/dwrf/test/StripeReaderBaseTests.cpp +++ b/velox/dwio/dwrf/test/StripeReaderBaseTests.cpp @@ -118,22 +118,22 @@ class StripeLoadKeysTest : public Test { } // namespace facebook::velox::dwrf -TEST_F(StripeLoadKeysTest, FirstStripeHasKey) { +TEST_F(StripeLoadKeysTest, firstStripeHasKey) { runTest(0); ASSERT_EQ(enc_->getKey(), "stripe0"); } -TEST_F(StripeLoadKeysTest, SecondStripeNoKey) { +TEST_F(StripeLoadKeysTest, secondStripeNoKey) { runTest(1); ASSERT_EQ(enc_->getKey(), "stripe0"); } -TEST_F(StripeLoadKeysTest, ThirdStripeHasKey) { +TEST_F(StripeLoadKeysTest, thirdStripeHasKey) { runTest(2); ASSERT_EQ(enc_->getKey(), "stripe2"); } -TEST_F(StripeLoadKeysTest, KeyMismatch) { +TEST_F(StripeLoadKeysTest, keyMismatch) { try { static_cast(runTest(3)); FAIL() << "Expected an exception"; diff --git a/velox/dwio/dwrf/test/StripeStreamTest.cpp b/velox/dwio/dwrf/test/StripeStreamTest.cpp index 72bebf35a0d..7e5985ec252 100644 --- a/velox/dwio/dwrf/test/StripeStreamTest.cpp +++ b/velox/dwio/dwrf/test/StripeStreamTest.cpp @@ -898,10 +898,10 @@ TEST_F(StripeStreamTest, shareDictionary) { char nonSharedDictBuffer[1024]; size_t nonSharedDictBufferSize = writeRange(nonSharedDictBuffer, 0, 100); EXPECT_CALL(ss, getStreamProxy(1, 0, proto::Stream_Kind_DICTIONARY_DATA, _)) - .WillOnce(InvokeWithoutArgs([&]() { + .WillOnce([&]() { return new SeekableArrayInputStream( nonSharedDictBuffer, nonSharedDictBufferSize); - })); + }); auto sharedDictionaryEncoding2_2 = genColumnEncoding(2, 2, proto::ColumnEncoding_Kind_DICTIONARY, 100); EXPECT_CALL(ss, getEncodingProxy(2, 2)) @@ -929,10 +929,10 @@ TEST_F(StripeStreamTest, shareDictionary) { char sharedDictBuffer[2048]; size_t sharedDictBufferSize = writeRange(sharedDictBuffer, 100, 200); EXPECT_CALL(ss, getStreamProxy(2, 0, proto::Stream_Kind_DICTIONARY_DATA, _)) - .WillRepeatedly(InvokeWithoutArgs([&]() { + .WillRepeatedly([&]() { return new SeekableArrayInputStream( sharedDictBuffer, sharedDictBufferSize); - })); + }); EXPECT_CALL( ss, getStreamProxy(2, Not(0), proto::Stream_Kind_DICTIONARY_DATA, _)) .WillRepeatedly(Return(nullptr)); diff --git a/velox/dwio/dwrf/test/TestBinaryStreamReader.cpp b/velox/dwio/dwrf/test/TestBinaryStreamReader.cpp index 074e81b4d74..b9a95b5e11b 100644 --- a/velox/dwio/dwrf/test/TestBinaryStreamReader.cpp +++ b/velox/dwio/dwrf/test/TestBinaryStreamReader.cpp @@ -137,7 +137,7 @@ TEST_F(BinaryStreamReaderTest, columnIdsEmpty) { "At least one column expected to be read"); } -TEST_F(BinaryStreamReaderTest, EmptyFile) { +TEST_F(BinaryStreamReaderTest, emptyFile) { auto pool = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); constexpr uint32_t STRIDE_LEN = 100; @@ -203,7 +203,7 @@ void verifyStream( } } -TEST_F(BinaryStreamReaderTest, BasicFlow) { +TEST_F(BinaryStreamReaderTest, basicFlow) { auto type = HiveTypeParser().parse("struct"); auto config = std::make_shared(); diff --git a/velox/dwio/dwrf/test/TestBufferedOutputStream.cpp b/velox/dwio/dwrf/test/TestBufferedOutputStream.cpp index 8dd41d8dfb4..651486abd29 100644 --- a/velox/dwio/dwrf/test/TestBufferedOutputStream.cpp +++ b/velox/dwio/dwrf/test/TestBufferedOutputStream.cpp @@ -246,7 +246,7 @@ class AppendOnlyBufferedStreamTest : public testing::Test { std::shared_ptr pool_ = memoryManager()->addLeafPool(); }; -TEST_F(AppendOnlyBufferedStreamTest, Basic) { +TEST_F(AppendOnlyBufferedStreamTest, basic) { MemorySink memSink(1024, {.pool = pool_.get()}); uint64_t block = 10; DataBufferHolder holder{*pool_, block, 0, DEFAULT_PAGE_GROW_RATIO, &memSink}; diff --git a/velox/dwio/dwrf/test/TestByteRLEEncoder.cpp b/velox/dwio/dwrf/test/TestByteRLEEncoder.cpp index 9377206804c..46e92f26080 100644 --- a/velox/dwio/dwrf/test/TestByteRLEEncoder.cpp +++ b/velox/dwio/dwrf/test/TestByteRLEEncoder.cpp @@ -117,7 +117,7 @@ class ByteRleEncoderTest : public testing::Test { } }; -TEST_F(ByteRleEncoderTest, random_chars) { +TEST_F(ByteRleEncoderTest, randomChars) { auto pool = memory::memoryManager()->addLeafPool(); MemorySink memSink(DEFAULT_MEM_STREAM_SIZE, {.pool = pool.get()}); @@ -137,7 +137,7 @@ TEST_F(ByteRleEncoderTest, random_chars) { delete[] data; } -TEST_F(ByteRleEncoderTest, random_chars_with_null) { +TEST_F(ByteRleEncoderTest, randomCharsWithNull) { auto pool = memory::memoryManager()->addLeafPool(); MemorySink memSink(DEFAULT_MEM_STREAM_SIZE, {.pool = pool.get()}); @@ -166,7 +166,7 @@ class BooleanRleEncoderTest : public testing::Test { } }; -TEST_F(BooleanRleEncoderTest, random_bits_not_aligned) { +TEST_F(BooleanRleEncoderTest, randomBitsNotAligned) { auto pool = memory::memoryManager()->addLeafPool(); MemorySink memSink(DEFAULT_MEM_STREAM_SIZE, {.pool = pool.get()}); @@ -186,7 +186,7 @@ TEST_F(BooleanRleEncoderTest, random_bits_not_aligned) { delete[] data; } -TEST_F(BooleanRleEncoderTest, random_bits_aligned) { +TEST_F(BooleanRleEncoderTest, randomBitsAligned) { auto pool = memory::memoryManager()->addLeafPool(); MemorySink memSink(DEFAULT_MEM_STREAM_SIZE, {.pool = pool.get()}); @@ -206,7 +206,7 @@ TEST_F(BooleanRleEncoderTest, random_bits_aligned) { delete[] data; } -TEST_F(BooleanRleEncoderTest, random_bits_aligned_with_null) { +TEST_F(BooleanRleEncoderTest, randomBitsAlignedWithNull) { auto pool = memory::memoryManager()->addLeafPool(); MemorySink memSink(DEFAULT_MEM_STREAM_SIZE, {.pool = pool.get()}); diff --git a/velox/dwio/dwrf/test/TestColumnReader.cpp b/velox/dwio/dwrf/test/TestColumnReader.cpp index 1ead2ca74b1..53d069939e4 100644 --- a/velox/dwio/dwrf/test/TestColumnReader.cpp +++ b/velox/dwio/dwrf/test/TestColumnReader.cpp @@ -152,7 +152,7 @@ class ColumnReaderTestBase { fileTypeWithId, streams_, labels_, - columnReaderStatistics_, + splitStatistics_, scanSpec, FlatMapContext{}); selectiveColumnReader_->setIsTopLevel(); @@ -233,7 +233,7 @@ class ColumnReaderTestBase { private: std::unique_ptr scanSpec_; - ColumnReaderStatistics columnReaderStatistics_; + SplitStats splitStatistics_{FileFormat::DWRF}; }; struct StringReaderTestParams { diff --git a/velox/dwio/dwrf/test/TestDictionaryEncodingUtils.cpp b/velox/dwio/dwrf/test/TestDictionaryEncodingUtils.cpp index 0fb09b43fa4..231a08c01bd 100644 --- a/velox/dwio/dwrf/test/TestDictionaryEncodingUtils.cpp +++ b/velox/dwio/dwrf/test/TestDictionaryEncodingUtils.cpp @@ -32,7 +32,7 @@ class DictionaryEncodingUtilsTest : public testing::Test { } }; -TEST_F(DictionaryEncodingUtilsTest, StringGetSortedIndexLookupTable) { +TEST_F(DictionaryEncodingUtilsTest, stringGetSortedIndexLookupTable) { struct TestCase { explicit TestCase( bool sort, @@ -146,7 +146,7 @@ TEST_F(DictionaryEncodingUtilsTest, StringGetSortedIndexLookupTable) { } } -TEST_F(DictionaryEncodingUtilsTest, StringStrideDictOptimization) { +TEST_F(DictionaryEncodingUtilsTest, stringStrideDictOptimization) { constexpr size_t kStrideSize{10}; struct TestCase { explicit TestCase( diff --git a/velox/dwio/dwrf/test/TestEncodingSelector.cpp b/velox/dwio/dwrf/test/TestEncodingSelector.cpp index 7599a4f11a1..ccf5dd1e333 100644 --- a/velox/dwio/dwrf/test/TestEncodingSelector.cpp +++ b/velox/dwio/dwrf/test/TestEncodingSelector.cpp @@ -28,7 +28,7 @@ class EntropyEncodingSelectorTests : public testing::Test { } }; -TEST_F(EntropyEncodingSelectorTests, Ctor) { +TEST_F(EntropyEncodingSelectorTests, ctor) { auto pool = memoryManager()->addLeafPool(); float slightlyOver = 1.0f + std::numeric_limits::epsilon() * 2; float slightlyUnder = -std::numeric_limits::epsilon(); @@ -94,7 +94,7 @@ class EntropyEncodingSelectorTest { } }; -TEST_F(EntropyEncodingSelectorTests, NoHeuristic) { +TEST_F(EntropyEncodingSelectorTests, noHeuristic) { class TestCase : public EntropyEncodingSelectorTest { public: explicit TestCase( @@ -137,7 +137,7 @@ TEST_F(EntropyEncodingSelectorTests, NoHeuristic) { } } -TEST_F(EntropyEncodingSelectorTests, NoSampling) { +TEST_F(EntropyEncodingSelectorTests, noSampling) { class TestCase : public EntropyEncodingSelectorTest { public: explicit TestCase( @@ -187,7 +187,7 @@ std::string alphabeticRoundRobin(size_t index, size_t size) { return std::string(size % (index + 1), element); } -TEST_F(EntropyEncodingSelectorTests, Sampling) { +TEST_F(EntropyEncodingSelectorTests, sampling) { class TestCase : public EntropyEncodingSelectorTest { public: explicit TestCase( diff --git a/velox/dwio/dwrf/test/TestIntegerDictionaryEncoder.cpp b/velox/dwio/dwrf/test/TestIntegerDictionaryEncoder.cpp index dbb50bbe86e..e5543f45a38 100644 --- a/velox/dwio/dwrf/test/TestIntegerDictionaryEncoder.cpp +++ b/velox/dwio/dwrf/test/TestIntegerDictionaryEncoder.cpp @@ -33,7 +33,7 @@ class TestIntegerDictionaryEncoder : public ::testing::Test { } }; -TEST_F(TestIntegerDictionaryEncoder, AddKey) { +TEST_F(TestIntegerDictionaryEncoder, addKey) { struct TestCase { explicit TestCase( const std::vector& addKeySequence, @@ -99,7 +99,7 @@ TEST_F(TestIntegerDictionaryEncoder, GetCount) { } } -TEST_F(TestIntegerDictionaryEncoder, GetTotalCount) { +TEST_F(TestIntegerDictionaryEncoder, getTotalCount) { struct TestCase { explicit TestCase( const std::vector& addKeySequence, @@ -207,7 +207,7 @@ TEST_F(TestIntegerDictionaryEncoder, Clear) { } } -TEST_F(TestIntegerDictionaryEncoder, RepeatedFlush) { +TEST_F(TestIntegerDictionaryEncoder, repeatedFlush) { auto pool = memoryManager()->addLeafPool(); IntegerDictionaryEncoder intDictEncoder{*pool, *pool}; std::vector keys{0, 1, 4, 9, 16, 25, 9, 1}; @@ -233,7 +233,7 @@ TEST_F(TestIntegerDictionaryEncoder, RepeatedFlush) { EXPECT_ANY_THROW(intDictEncoder.getLookupTable()); } -TEST_F(TestIntegerDictionaryEncoder, Limit) { +TEST_F(TestIntegerDictionaryEncoder, limit) { auto pool = memoryManager()->addLeafPool(); IntegerDictionaryEncoder intDictEncoder{*pool, *pool}; for (size_t iter = 0; iter < 2; ++iter) { @@ -310,13 +310,13 @@ void testGetSortedIndexLookupTable() { } } -TEST_F(TestIntegerDictionaryEncoder, GetSortedIndexLookupTable) { +TEST_F(TestIntegerDictionaryEncoder, getSortedIndexLookupTable) { testGetSortedIndexLookupTable(); testGetSortedIndexLookupTable(); testGetSortedIndexLookupTable(); } -TEST_F(TestIntegerDictionaryEncoder, ShortIntegerDictionary) { +TEST_F(TestIntegerDictionaryEncoder, shortIntegerDictionary) { // DictionaryEncoding lookupTable can contain the index into dictionary // or the actual value. For short integer, index can be [0,2^16-1] // and the values can be from [-2^15, 2^15-1]. Integer writers always @@ -494,7 +494,7 @@ void testInfrequentKeyOptimization() { } } -TEST_F(TestIntegerDictionaryEncoder, InfrequentKeyOptimization) { +TEST_F(TestIntegerDictionaryEncoder, infrequentKeyOptimization) { testInfrequentKeyOptimization(); testInfrequentKeyOptimization(); testInfrequentKeyOptimization(); diff --git a/velox/dwio/dwrf/test/TestStringDictionaryEncoder.cpp b/velox/dwio/dwrf/test/TestStringDictionaryEncoder.cpp index f8c33d91ff6..a5a579ed7c0 100644 --- a/velox/dwio/dwrf/test/TestStringDictionaryEncoder.cpp +++ b/velox/dwio/dwrf/test/TestStringDictionaryEncoder.cpp @@ -30,7 +30,7 @@ class TestStringDictionaryEncoder : public ::testing::Test { } }; -TEST_F(TestStringDictionaryEncoder, AddKey) { +TEST_F(TestStringDictionaryEncoder, addKey) { struct TestCase { explicit TestCase( const std::vector& addKeySequence, @@ -238,7 +238,7 @@ TEST_F(TestStringDictionaryEncoder, Clear) { EXPECT_LT(pool->usedBytes(), peakMemory); } -TEST_F(TestStringDictionaryEncoder, MemBenchmark) { +TEST_F(TestStringDictionaryEncoder, memBenchmark) { auto pool = memory::memoryManager()->addLeafPool(); StringDictionaryEncoder stringDictEncoder{*pool, *pool}; std::string baseString{"jjkkll"}; @@ -249,7 +249,7 @@ TEST_F(TestStringDictionaryEncoder, MemBenchmark) { LOG(INFO) << "Total memory bytes: " << pool->usedBytes(); } -TEST_F(TestStringDictionaryEncoder, Limit) { +TEST_F(TestStringDictionaryEncoder, limit) { auto pool = memory::memoryManager()->addLeafPool(); StringDictionaryEncoder encoder{*pool, *pool}; encoder.addKey(std::string_view{"abc"}, 0); diff --git a/velox/dwio/dwrf/test/TestStripeDictionaryCache.cpp b/velox/dwio/dwrf/test/TestStripeDictionaryCache.cpp index b4ca2e3c5d2..f89673e4c6c 100644 --- a/velox/dwio/dwrf/test/TestStripeDictionaryCache.cpp +++ b/velox/dwio/dwrf/test/TestStripeDictionaryCache.cpp @@ -100,7 +100,7 @@ TEST_F(StripeDictionaryCacheTest, RegisterDictionary) { } } -TEST_F(StripeDictionaryCacheTest, GetDictionaryBuffer) { +TEST_F(StripeDictionaryCacheTest, getDictionaryBuffer) { { StripeDictionaryCache cache{pool_.get()}; diff --git a/velox/dwio/dwrf/test/WriterContextTest.cpp b/velox/dwio/dwrf/test/WriterContextTest.cpp index cd67deddb3d..8db85c96197 100644 --- a/velox/dwio/dwrf/test/WriterContextTest.cpp +++ b/velox/dwio/dwrf/test/WriterContextTest.cpp @@ -117,7 +117,7 @@ TEST_F(WriterContextTest, RemoveIntDictionaryEncoderForNode) { EXPECT_EQ(0, context.dictEncoders_.size()); } -TEST_F(WriterContextTest, BuildPhysicalSizeAggregators) { +TEST_F(WriterContextTest, buildPhysicalSizeAggregators) { auto config = std::make_shared(); WriterContext context{ config, diff --git a/velox/dwio/dwrf/test/WriterExtendedTests.cpp b/velox/dwio/dwrf/test/WriterExtendedTests.cpp index c066d7cc1b2..ea0701172c4 100644 --- a/velox/dwio/dwrf/test/WriterExtendedTests.cpp +++ b/velox/dwio/dwrf/test/WriterExtendedTests.cpp @@ -96,7 +96,7 @@ class E2EWriterTest : public testing::Test { } }; -TEST_F(E2EWriterTest, FlushPolicySimpleEncoding) { +TEST_F(E2EWriterTest, flushPolicySimpleEncoding) { const size_t batchCount = 200; const size_t batchSize = 1000; auto pool = facebook::velox::memory::memoryManager()->addLeafPool(); @@ -152,7 +152,7 @@ TEST_F(E2EWriterTest, FlushPolicySimpleEncoding) { // Many streams are not yet allocated prior to the first flush, hence first // flush is delayed if we rely on stream usage to estimate stripe size. -TEST_F(E2EWriterTest, FlushPolicyDictionaryEncoding) { +TEST_F(E2EWriterTest, flushPolicyDictionaryEncoding) { const size_t batchCount = 500; const size_t batchSize = 1000; auto pool = facebook::velox::memory::memoryManager()->addLeafPool(); @@ -277,7 +277,7 @@ TEST_F(E2EWriterTest, FlushPolicyDictionaryEncoding) { } // stream usage seems to have a delta that is close to compression block size? -TEST_F(E2EWriterTest, FlushPolicyNestedTypes) { +TEST_F(E2EWriterTest, flushPolicyNestedTypes) { const size_t batchCount = 10; const size_t batchSize = 1000; auto pool = facebook::velox::memory::memoryManager()->addLeafPool(); @@ -326,7 +326,7 @@ TEST_F(E2EWriterTest, FlushPolicyNestedTypes) { } // Flat map has 1.5 orders of magnitude inflated stream memory usage. -TEST_F(E2EWriterTest, FlushPolicyFlatMap) { +TEST_F(E2EWriterTest, flushPolicyFlatMap) { const size_t batchCount = 10; const size_t batchSize = 500; auto pool = facebook::velox::memory::memoryManager()->addLeafPool(); @@ -423,7 +423,7 @@ TEST_F(E2EWriterTest, FlushPolicyFlatMap) { } } -TEST_F(E2EWriterTest, MemoryPoolBasedFlushPolicySimpleEncoding) { +TEST_F(E2EWriterTest, memoryPoolBasedFlushPolicySimpleEncoding) { const size_t batchCount = 2000; const size_t batchSize = 5000; auto pool = facebook::velox::memory::memoryManager()->addLeafPool(); @@ -481,7 +481,7 @@ TEST_F(E2EWriterTest, MemoryPoolBasedFlushPolicySimpleEncoding) { // Many streams are not yet allocated prior to the first flush, hence first // flush is delayed if we rely on stream usage to estimate stripe size. -TEST_F(E2EWriterTest, MemoryPoolBasedFlushPolicyDictionaryEncoding) { +TEST_F(E2EWriterTest, memoryPoolBasedFlushPolicyDictionaryEncoding) { const size_t batchCount = 1000; const size_t batchSize = 2000; auto pool = facebook::velox::memory::memoryManager()->addLeafPool(); @@ -618,7 +618,7 @@ TEST_F(E2EWriterTest, MemoryPoolBasedFlushPolicyDictionaryEncoding) { } // stream usage seems to have a delta that is close to compression block size? -TEST_F(E2EWriterTest, MemoryPoolBasedFlushPolicyNestedTypes) { +TEST_F(E2EWriterTest, memoryPoolBasedFlushPolicyNestedTypes) { const size_t batchCount = 100; const size_t batchSize = 1000; auto pool = facebook::velox::memory::memoryManager()->addLeafPool(); @@ -676,7 +676,7 @@ TEST_F(E2EWriterTest, MemoryPoolBasedFlushPolicyNestedTypes) { // Flat map has 1.5 orders of magnitude inflated stream memory usage. // Disabled because test is failing in continuous runs T193531984. -TEST_F(E2EWriterTest, DISABLED_MemoryPoolBasedFlushPolicyFlatMap) { +TEST_F(E2EWriterTest, DISABLED_memoryPoolBasedFlushPolicyFlatMap) { const size_t batchCount = 500; const size_t batchSize = 500; auto pool = facebook::velox::memory::memoryManager()->addLeafPool(); diff --git a/velox/dwio/dwrf/test/WriterFlushTest.cpp b/velox/dwio/dwrf/test/WriterFlushTest.cpp index 749e9b1184c..ab2f7d1bbfc 100644 --- a/velox/dwio/dwrf/test/WriterFlushTest.cpp +++ b/velox/dwio/dwrf/test/WriterFlushTest.cpp @@ -651,7 +651,7 @@ TEST_F(TestWriterFlush, CheckAgainstMemoryBudget) { } // Tests the number of stripes produced based on random results. -TEST_F(TestWriterFlush, MemoryBasedFlushRandom) { +TEST_F(TestWriterFlush, memoryBasedFlushRandom) { struct TestCase { TestCase( uint32_t seed, diff --git a/velox/dwio/dwrf/test/WriterSinkTest.cpp b/velox/dwio/dwrf/test/WriterSinkTest.cpp index ddb9384ae8a..da6da92b6e1 100644 --- a/velox/dwio/dwrf/test/WriterSinkTest.cpp +++ b/velox/dwio/dwrf/test/WriterSinkTest.cpp @@ -50,7 +50,7 @@ class WriterSinkTest : public Test { std::array data; }; -TEST_F(WriterSinkTest, Checksum) { +TEST_F(WriterSinkTest, checksum) { auto pool = memoryManager()->addLeafPool(); MemorySink out{1024 + 3, {.pool = pool.get()}}; Config config; @@ -93,7 +93,7 @@ TEST_F(WriterSinkTest, Checksum) { ASSERT_EQ(sink.size(), out.size() - offset); } -TEST_F(WriterSinkTest, NoChecksum) { +TEST_F(WriterSinkTest, noChecksum) { auto pool = memoryManager()->addLeafPool(); MemorySink out{1024 + 3, {.pool = pool.get()}}; Config config; @@ -110,7 +110,7 @@ TEST_F(WriterSinkTest, NoChecksum) { checkOutput(out, offset); } -TEST_F(WriterSinkTest, NoCache) { +TEST_F(WriterSinkTest, noCache) { auto pool = memoryManager()->addLeafPool(); MemorySink out{1024, {.pool = pool.get()}}; Config config; @@ -135,7 +135,7 @@ TEST_F(WriterSinkTest, NoCache) { ASSERT_EQ(out.size() - offset, 512); } -TEST_F(WriterSinkTest, CacheIndex) { +TEST_F(WriterSinkTest, cacheIndex) { auto pool = memoryManager()->addLeafPool(); MemorySink out{1024, {.pool = pool.get()}}; Config config; @@ -178,7 +178,7 @@ TEST_F(WriterSinkTest, CacheIndex) { std::string(out.data() + offset, 128)); } -TEST_F(WriterSinkTest, CacheFooter) { +TEST_F(WriterSinkTest, cacheFooter) { auto pool = memoryManager()->addLeafPool(); MemorySink out{1024, {.pool = pool.get()}}; Config config; @@ -221,7 +221,7 @@ TEST_F(WriterSinkTest, CacheFooter) { std::string(out.data() + offset + 384, 128)); } -TEST_F(WriterSinkTest, CacheBothEmptyIndex) { +TEST_F(WriterSinkTest, cacheBothEmptyIndex) { auto pool = memoryManager()->addLeafPool(); MemorySink out{1024, {.pool = pool.get()}}; Config config; @@ -250,7 +250,7 @@ TEST_F(WriterSinkTest, CacheBothEmptyIndex) { ASSERT_EQ(sink.getCacheSize(), 128); } -TEST_F(WriterSinkTest, CacheBoth) { +TEST_F(WriterSinkTest, cacheBoth) { auto pool = memoryManager()->addLeafPool(); MemorySink out{1024, {.pool = pool.get()}}; Config config; @@ -299,7 +299,7 @@ TEST_F(WriterSinkTest, CacheBoth) { std::string(out.data() + offset + 512, 128)); } -TEST_F(WriterSinkTest, CacheExceedsLimit) { +TEST_F(WriterSinkTest, cacheExceedsLimit) { auto pool = memoryManager()->addLeafPool(); MemorySink out{1024, {.pool = pool.get()}}; Config config; @@ -427,7 +427,7 @@ TEST_F(WriterSinkTest, CacheExceedsLimit) { } } -TEST_F(WriterSinkTest, CacheLarge) { +TEST_F(WriterSinkTest, cacheLarge) { auto pool = memoryManager()->addLeafPool(); MemorySink out{10 * 1024 * 1024 + 3, {.pool = pool.get()}}; Config config; @@ -454,7 +454,7 @@ TEST_F(WriterSinkTest, CacheLarge) { ASSERT_EQ(out.size() - offset, total * 2); } -TEST_F(WriterSinkTest, SetModeOutOfOrder) { +TEST_F(WriterSinkTest, setModeOutOfOrder) { auto pool = memoryManager()->addLeafPool(); MemorySink out{1024, {.pool = pool.get()}}; Config config; diff --git a/velox/dwio/dwrf/test/WriterTest.cpp b/velox/dwio/dwrf/test/WriterTest.cpp index b3e9ce2114f..ae2c69292c1 100644 --- a/velox/dwio/dwrf/test/WriterTest.cpp +++ b/velox/dwio/dwrf/test/WriterTest.cpp @@ -205,7 +205,7 @@ TEST_P(AllWriterCompressionTest, compression) { : compressionKind_); } -TEST_F(WriterTest, SchemaAttributesStamped) { +TEST_F(WriterTest, schemaAttributesStamped) { // Attributes set on the writer must be stamped into the footer per node id // and survive a write/read round-trip. Node ids: 0=root, 1=a, 2=b, 3=c. auto config = std::make_shared(); @@ -240,7 +240,7 @@ TEST_F(WriterTest, SchemaAttributesStamped) { std::make_pair(std::string("iceberg.id"), std::string("30"))); } -TEST_P(SupportedCompressionTest, WriteFooter) { +TEST_P(SupportedCompressionTest, writeFooter) { auto config = std::make_shared(); config->set(Config::COMPRESSION, supportedCompressionKind_); auto& writer = createWriter(config); @@ -333,7 +333,7 @@ TEST_P(SupportedCompressionTest, WriteFooter) { } } -TEST_P(SupportedCompressionTest, AddStripeInfo) { +TEST_P(SupportedCompressionTest, addStripeInfo) { auto config = std::make_shared(); config->set(Config::COMPRESSION, supportedCompressionKind_); auto& writer = createWriter(config); @@ -355,7 +355,7 @@ TEST_P(SupportedCompressionTest, AddStripeInfo) { writer.close(); } -TEST_P(SupportedCompressionTest, NoChecksum) { +TEST_P(SupportedCompressionTest, noChecksum) { auto config = std::make_shared(); config->set(Config::CHECKSUM_ALGORITHM, proto::ChecksumAlgorithm::NULL_); config->set(Config::COMPRESSION, supportedCompressionKind_); @@ -387,7 +387,7 @@ TEST_P(SupportedCompressionTest, NoChecksum) { ASSERT_EQ(footer.checksumAlgorithm(), proto::ChecksumAlgorithm::NULL_); } -TEST_P(SupportedCompressionTest, NoCache) { +TEST_P(SupportedCompressionTest, noCache) { auto config = std::make_shared(); config->set(Config::STRIPE_CACHE_MODE, StripeCacheMode::NA); config->set(Config::COMPRESSION, supportedCompressionKind_); @@ -425,7 +425,7 @@ TEST_P(SupportedCompressionTest, NoCache) { ASSERT_EQ(reader->metadataCache(), nullptr); } -TEST_P(SupportedCompressionTest, ValidateStreamSizeConfigDisabled) { +TEST_P(SupportedCompressionTest, validateStreamSizeConfigDisabled) { auto config = std::make_shared(); config->set(Config::STREAM_SIZE_ABOVE_THRESHOLD_CHECK_ENABLED, false); config->set(Config::COMPRESSION, supportedCompressionKind_); @@ -435,7 +435,7 @@ TEST_P(SupportedCompressionTest, ValidateStreamSizeConfigDisabled) { writer.close(); } -TEST_P(SupportedCompressionTest, ValidateStreamSizeConfigEnabled) { +TEST_P(SupportedCompressionTest, validateStreamSizeConfigEnabled) { auto config = std::make_shared(); ASSERT_TRUE(config->get(Config::STREAM_SIZE_ABOVE_THRESHOLD_CHECK_ENABLED)); config->set(Config::COMPRESSION, supportedCompressionKind_); @@ -486,7 +486,7 @@ void abandonWriterWithoutClosing() { // guard.dismiss(); } -TEST_F(WriterTest, DoNotCrashDbgModeOnAbort) { +TEST_F(WriterTest, doNotCrashDbgModeOnAbort) { EXPECT_THROW(abandonWriterWithoutClosing(), std::runtime_error); } @@ -514,7 +514,7 @@ class MockFileSink : public dwio::common::FileSink { #pragma GCC diagnostic pop }; -TEST_F(WriterTest, FlushWriterSinkUponClose) { +TEST_F(WriterTest, flushWriterSinkUponClose) { auto config = std::make_shared(); auto pool = memory::memoryManager()->addRootPool("FlushWriterSinkUponClose"); auto sink = std::make_unique(); diff --git a/velox/dwio/dwrf/utils/test/BitIteratorTests.cpp b/velox/dwio/dwrf/utils/test/BitIteratorTests.cpp index e596f8eb03c..d714cda2402 100644 --- a/velox/dwio/dwrf/utils/test/BitIteratorTests.cpp +++ b/velox/dwio/dwrf/utils/test/BitIteratorTests.cpp @@ -24,7 +24,7 @@ using namespace ::testing; namespace facebook::velox::dwrf::utils { -TEST(BulkBitIterator, Basic) { +TEST(BulkBitIterator, basic) { std::vector charBuffer1{ std::numeric_limits::max(), 102, 40, -120}; std::vector charBuffer2{125, 85, 42, std::numeric_limits::min()}; diff --git a/velox/dwio/dwrf/utils/test/BufferedWriterTest.cpp b/velox/dwio/dwrf/utils/test/BufferedWriterTest.cpp index 8260f44dda2..fbb900819a7 100644 --- a/velox/dwio/dwrf/utils/test/BufferedWriterTest.cpp +++ b/velox/dwio/dwrf/utils/test/BufferedWriterTest.cpp @@ -44,7 +44,7 @@ class BufferedWriterTest : public testing::TestWithParam { const std::shared_ptr pool_; }; -TEST_P(BufferedWriterTest, Basic) { +TEST_P(BufferedWriterTest, basic) { const int bufferSize = 1024; struct { diff --git a/velox/dwio/dwrf/utils/test/ProtoUtilsTests.cpp b/velox/dwio/dwrf/utils/test/ProtoUtilsTests.cpp index 877a27adce0..b63059fd747 100644 --- a/velox/dwio/dwrf/utils/test/ProtoUtilsTests.cpp +++ b/velox/dwio/dwrf/utils/test/ProtoUtilsTests.cpp @@ -22,7 +22,7 @@ using namespace facebook::velox::dwrf; using namespace facebook::velox::type::fbhive; -TEST(ProtoUtilsTests, AllTypes) { +TEST(ProtoUtilsTests, allTypes) { std::vector types{ "struct", "struct,b:array>>>>"}; @@ -41,7 +41,7 @@ TEST(ProtoUtilsTests, AllTypes) { } } -TEST(ProtoUtilsTests, Projection) { +TEST(ProtoUtilsTests, projection) { HiveTypeParser parser; auto schema = parser.parse( "struct>"); @@ -56,7 +56,7 @@ TEST(ProtoUtilsTests, Projection) { EXPECT_EQ("struct>", res); } -TEST(ProtoUtilsTests, AttributesRoundTrip) { +TEST(ProtoUtilsTests, attributesRoundTrip) { // iceberg.id stamped on a subset of nodes must survive footer serialization // and come back keyed by the same pre-order node id, leaving the schema // intact. Node ids: 0=root, 1=a, 2=b, 3=c, 4=c.x, 5=c.y. @@ -94,7 +94,7 @@ TEST(ProtoUtilsTests, AttributesRoundTrip) { "struct>"); } -TEST(ProtoUtilsTests, AttributesRoundTripOrc) { +TEST(ProtoUtilsTests, attributesRoundTripOrc) { // The same iceberg.id round-trip must work for ORC footers: DWRF/ORC Iceberg // reads resolve columns by field id from these attributes, and Iceberg // manifest-tags DWRF files as ORC. Node ids: 0=root, 1=a, 2=b, 3=c, 4=c.x, @@ -130,7 +130,7 @@ TEST(ProtoUtilsTests, AttributesRoundTripOrc) { EXPECT_EQ(ProtoUtils::readAttributes(FooterWrapper(&parsed)), expected); } -TEST(ProtoUtilsTests, AttributesAbsentByDefault) { +TEST(ProtoUtilsTests, attributesAbsentByDefault) { // A type written without an attribute provider -- the existing path for every // DWRF file today -- yields an empty attribute map. HiveTypeParser parser; diff --git a/velox/dwio/dwrf/utils/test/TypeAttributesTest.cpp b/velox/dwio/dwrf/utils/test/TypeAttributesTest.cpp index c54723e0b83..6162c1b73aa 100644 --- a/velox/dwio/dwrf/utils/test/TypeAttributesTest.cpp +++ b/velox/dwio/dwrf/utils/test/TypeAttributesTest.cpp @@ -47,7 +47,7 @@ proto::Type buildTypeWithAttributes( } // namespace -TEST(TypeAttributesTest, AttributesAbsentByDefault) { +TEST(TypeAttributesTest, attributesAbsentByDefault) { // A Type built without attributes -- the existing wire format for every // DWRF file written today -- must round-trip and surface as an empty // attributes list. Protects the no-op upgrade path for existing files. @@ -63,7 +63,7 @@ TEST(TypeAttributesTest, AttributesAbsentByDefault) { EXPECT_EQ(parsed.attributes_size(), 0); } -TEST(TypeAttributesTest, AttributesRoundTripIcebergKeys) { +TEST(TypeAttributesTest, attributesRoundTripIcebergKeys) { // All Iceberg ORC-spec attribute keys must survive the proto round-trip // with their string values, exactly mirroring the Apache ORC attribute // convention. @@ -91,7 +91,7 @@ TEST(TypeAttributesTest, AttributesRoundTripIcebergKeys) { } } -TEST(TypeAttributesTest, LegacyBufferIsForwardCompatible) { +TEST(TypeAttributesTest, legacyBufferIsForwardCompatible) { // A buffer produced before this proto change has no attributes field set. // The new proto schema must parse it cleanly with an empty attributes // list, leaving the other fields intact. Forward-compat invariant for diff --git a/velox/dwio/nimble/CMakeLists.txt b/velox/dwio/nimble/CMakeLists.txt new file mode 100644 index 00000000000..6caa39f4019 --- /dev/null +++ b/velox/dwio/nimble/CMakeLists.txt @@ -0,0 +1,85 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Nimble is reached only through velox/dwio/CMakeLists.txt when +# VELOX_ENABLE_NIMBLE=ON. FlatBuffers, OpenZL and the vendored FSST are all +# resolved by the top-level Velox CMakeLists before this point, so this file +# only configures what is specific to Nimble. + +# FsstEncoding.h includes , and DuckDB vendors a header of the same name. +# Resolved BUNDLED, DuckDB contributes its copy to Velox targets through the link +# interface, which CMake emits after a target's own include directories. Adding +# the vendored copy at directory scope therefore puts it first for every target +# in this subtree, including those that only pull FsstEncoding.h in transitively +# and never link fsst. Relies on the fsst target exporting this path as a plain +# -I: while it was SYSTEM, CMake de-duplicated the two and kept the -isystem, +# which is searched after every -I. +include_directories(BEFORE "${CMAKE_SOURCE_DIR}/velox/external/fsst") + +# Nimble code expects an upper case suffix to the generated file. +set(FLATBUFFERS_FLATC_SCHEMA_EXTRA_ARGS "--filename-suffix" "Generated") + +# Velox only resolves Abseil for its own test targets, but Nimble's field reader +# needs absl::flat_hash_map unconditionally. +if(NOT TARGET absl::flat_hash_map) + velox_set_source(absl) + velox_resolve_dependency(absl) +endif() + +# Nimble's OSS build has no access to the Meta-internal compressor. This is read +# by widely included headers (compression/CompressionPolicy.h, +# encodings/legacy/EncodingSelectionPolicy.h), so every translation unit that +# sees those headers must agree on it. +add_compile_definitions(DISABLE_META_INTERNAL_COMPRESSOR=1) + +option(NIMBLE_ENABLE_EXPERIMENTAL_ENCODINGS "Enable experimental encodings" OFF) +if(NIMBLE_ENABLE_EXPERIMENTAL_ENCODINGS) + add_compile_definitions(NIMBLE_ENABLE_EXPERIMENTAL_ENCODINGS) +endif() + +if(VELOX_MONO_LIBRARY AND TARGET velox) + # add_compile_definitions() above sets a directory property, which only + # reaches targets created in this directory tree. Under VELOX_MONO_LIBRARY + # Nimble's sources are folded into the single `velox` target, created + # elsewhere, so the definitions have to be applied to it explicitly. + target_compile_definitions(velox PUBLIC DISABLE_META_INTERNAL_COMPRESSOR=1) + if(NIMBLE_ENABLE_EXPERIMENTAL_ENCODINGS) + target_compile_definitions(velox PUBLIC NIMBLE_ENABLE_EXPERIMENTAL_ENCODINGS) + endif() +endif() + +# Nimble's tests follow Velox's testing switch. The separate variable is kept so +# a build can opt into Nimble's tests alone. +if(NOT DEFINED NIMBLE_BUILD_TESTING) + set(NIMBLE_BUILD_TESTING ${VELOX_BUILD_TESTING}) +endif() + +# Ordered by dependency: common is the root, and the tablet, index and velox +# layers reference each other's targets at declaration time. +if(VELOX_ENABLE_BENCHMARKS) + add_subdirectory(benchmarks) +endif() +add_subdirectory(common) +add_subdirectory(compression) +add_subdirectory(encodings) +add_subdirectory(index) +add_subdirectory(tablet) +add_subdirectory(velox) +add_subdirectory(writer) +add_subdirectory(serializer) +add_subdirectory(tools) + +if(NIMBLE_BUILD_TESTING) + add_subdirectory(fuzzer) +endif() diff --git a/velox/dwio/nimble/benchmarks/BenchmarkSuite.h b/velox/dwio/nimble/benchmarks/BenchmarkSuite.h new file mode 100644 index 00000000000..6fb8eda58fc --- /dev/null +++ b/velox/dwio/nimble/benchmarks/BenchmarkSuite.h @@ -0,0 +1,311 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// BenchmarkSuite — runs each benchmark in its own BenchmarkingState so +// start/end hooks can fire per-benchmark. Uses a standalone +// folly::detail::BenchmarkingState per benchmark, enabling std::function +// callbacks that fire exactly once per benchmark function (start and end). + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace facebook::nimble { + +class BenchmarkSuite { + public: + using Clock = std::chrono::high_resolution_clock; + using StartHook = std::function; + using EndHook = std::function< + void(const std::string& name, const folly::detail::BenchmarkResult&)>; + + void setOnStart(StartHook hook) { + onStart_ = std::move(hook); + } + void setOnEnd(EndHook hook) { + onEnd_ = std::move(hook); + } + + void setTitle(std::string title) { + title_ = std::move(title); + } + + void setJsonOutputPath(std::string path) { + jsonOutputPath_ = std::move(path); + } + + void addBenchmark(std::string name, std::function fn) { + entries_.push_back(Entry{std::move(name), std::move(fn), /*sep=*/false}); + } + + void addSeparator() { + entries_.push_back(Entry{"", nullptr, /*sep=*/true}); + } + + void run() { + std::vector allResults; + + // Print table header. + printHeader(); + + bool pendingSeparator = false; + + for (size_t idx = 0; idx < entries_.size(); ++idx) { + auto& entry = entries_[idx]; + + if (entry.isSeparator) { + pendingSeparator = true; + continue; + } + + // Create a fresh BenchmarkingState for this benchmark. + folly::detail::BenchmarkingState state; + + // Register baseline benchmarks (required by the measurement engine). + auto baselineName = std::string("baseline"); + auto suspenderBaselineName = std::string("suspender_baseline"); + + state.addBenchmark( + "BenchmarkSuite", std::move(baselineName), [](unsigned) { + // Minimal baseline — matches folly's global baseline. +#ifdef _MSC_VER + _ReadWriteBarrier(); +#else + asm volatile(""); +#endif + return 1u; + }); + state.addBenchmark( + "BenchmarkSuite", std::move(suspenderBaselineName), [](unsigned) { + folly::BenchmarkSuspender sus; + return 1u; + }); + + // Register the actual benchmark. + auto fn = entry.fn; // copy the function for the lambda capture + state.addBenchmark( + "BenchmarkSuite", entry.name, [fn](unsigned iters) -> unsigned { + fn(iters); + return iters; + }); + + // Fire start hook. + if (onStart_) { + onStart_(entry.name); + } + + // Run measurement. + auto results = state.runBenchmarksWithResults(); + + // Find the result for our benchmark (skip baselines). + folly::detail::BenchmarkResult bmResult; + bmResult.name = entry.name; + bmResult.file = "BenchmarkSuite"; + bmResult.timeInNs = 0; + for (auto& r : results) { + if (r.name == entry.name) { + bmResult = r; + break; + } + } + // Print separator if one was pending before this benchmark. + if (pendingSeparator) { + printSeparator('-'); + pendingSeparator = false; + } else if (!allResults.empty()) { + printSeparator('.'); + } + + // Print this result immediately. + printResult(bmResult); + + // Fire end hook. + if (onEnd_) { + onEnd_(entry.name, bmResult); + } + + allResults.push_back(bmResult); + } + + // Print table footer. + printSeparator('='); + + // Write JSON output if a path was provided. + if (!jsonOutputPath_.empty()) { + writeJsonResults(allResults, jsonOutputPath_); + } + } + + private: + struct Entry { + std::string name; + std::function fn; + bool isSeparator; + }; + + // Human-readable formatting helpers matching folly's output. + struct ScaleInfo { + double boundary; + const char* suffix; + }; + + static std::string + humanReadable(double n, unsigned int decimals, const ScaleInfo* scales) { + if (std::isinf(n) || std::isnan(n)) { + return folly::to(n); + } + const double absValue = fabs(n); + const ScaleInfo* scale = scales; + while (absValue < scale[0].boundary && scale[1].suffix != nullptr) { + ++scale; + } + const double scaledValue = n / scale->boundary; + return fmt::format("{:.{}f}{}", scaledValue, decimals, scale->suffix); + } + + static std::string readableTime(double n, unsigned int decimals) { + static const ScaleInfo kTimeSuffixes[] = { + {365.25 * 24 * 3600, "years"}, + {24 * 3600, "days"}, + {3600, "hr"}, + {60, "min"}, + {1, "s"}, + {1E-3, "ms"}, + {1E-6, "us"}, + {1E-9, "ns"}, + {1E-12, "ps"}, + {1E-15, "fs"}, + {0, nullptr}, + }; + return humanReadable(n, decimals, kTimeSuffixes); + } + + static std::string metricReadable(double n, unsigned int decimals) { + static const ScaleInfo kMetricSuffixes[] = { + {1E24, "Y"}, + {1E21, "Z"}, + {1E18, "X"}, + {1E15, "P"}, + {1E12, "T"}, + {1E9, "G"}, + {1E6, "M"}, + {1E3, "K"}, + {1, ""}, + {1E-3, "m"}, + {1E-6, "u"}, + {1E-9, "n"}, + {1E-12, "p"}, + {1E-15, "f"}, + {1E-18, "a"}, + {1E-21, "z"}, + {1E-24, "y"}, + {0, nullptr}, + }; + return humanReadable(n, decimals, kMetricSuffixes); + } + + static constexpr unsigned int kColumns = 76; + static constexpr std::string_view kHeader = "time/iter iters/s"; + + // ANSI escape helpers. + static constexpr const char* kBold = "\033[1m"; + static constexpr const char* kBoldBlue = "\033[1;34m"; + static constexpr const char* kBoldCyan = "\033[1;36m"; + static constexpr const char* kBoldGreen = "\033[1;32m"; + static constexpr const char* kBoldYellow = "\033[1;33m"; + static constexpr const char* kReset = "\033[0m"; + + static void printSeparator(char pad) { + printf("%s%s%s\n", kBoldBlue, std::string(kColumns, pad).c_str(), kReset); + } + + void printHeader() const { + printSeparator('='); + static const std::string kDefaultTitle = "BenchmarkSuite"; + const std::string& title = title_.empty() ? kDefaultTitle : title_; + size_t padLen = (kColumns > title.size() + kHeader.size()) + ? (kColumns - title.size() - kHeader.size()) + : 1; + printf( + "%s%s%*s%s%.*s%s\n", + kBold, + title.c_str(), + static_cast(padLen), + "", + kBoldYellow, + static_cast(kHeader.size()), + kHeader.data(), + kReset); + printSeparator('='); + } + + static void printResult(const folly::detail::BenchmarkResult& r) { + const double nsPerIter = r.timeInNs; + const double secPerIter = nsPerIter / 1E9; + const double itersPerSec = (secPerIter == 0) + ? std::numeric_limits::infinity() + : (1.0 / secPerIter); + + std::string name = r.name; + size_t nameWidth = kColumns - kHeader.size(); + name.resize(nameWidth, ' '); + + printf( + "%s%*s%s%s%9.9s %s%8.8s%s\n", + kBoldGreen, + static_cast(name.size()), + name.c_str(), + kReset, + kBoldCyan, + readableTime(secPerIter, 2).c_str(), + kBoldYellow, + metricReadable(itersPerSec, 2).c_str(), + kReset); + } + + static void writeJsonResults( + const std::vector& results, + const std::string& outputPath) { + folly::dynamic d; + folly::benchmarkResultsToDynamic(results, d); + auto jsonStr = folly::toPrettyJson(d); + if (folly::writeFile(jsonStr, outputPath.c_str())) { + LOG(INFO) << "JSON results written to " << outputPath; + } else { + LOG(ERROR) << "Failed to write JSON results to " << outputPath; + } + } + + std::vector entries_; + StartHook onStart_; + EndHook onEnd_; + std::string title_; + std::string jsonOutputPath_; +}; + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/benchmarks/CMakeLists.txt b/velox/dwio/nimble/benchmarks/CMakeLists.txt new file mode 100644 index 00000000000..ef909eb5a4b --- /dev/null +++ b/velox/dwio/nimble/benchmarks/CMakeLists.txt @@ -0,0 +1,18 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Header-only helper shared by the per-directory benchmark binaries. It has no +# sources of its own; the target exists so the header is owned by the build +# graph rather than showing up as untracked. +velox_add_library(velox_dwio_nimble_benchmark_suite INTERFACE HEADERS BenchmarkSuite.h) diff --git a/velox/dwio/nimble/common/BitEncoder.h b/velox/dwio/nimble/common/BitEncoder.h new file mode 100644 index 00000000000..97b540241f1 --- /dev/null +++ b/velox/dwio/nimble/common/BitEncoder.h @@ -0,0 +1,88 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +namespace facebook::nimble { + +// Class to put bits into, and read bits from, a buffer. +class BitEncoder { + public: + // It's up to the user to not write/read beyond the provided buffer. + // Note that you probably want to ensure the memory is zero'd out. + explicit BitEncoder(char* buffer) + : buffer_(buffer), + writeWord_(reinterpret_cast(buffer)), + readWord_(reinterpret_cast(buffer)) {} + // We don't prevent you from using the putBits call after using this + // constructor, but you really shouldn't. + explicit BitEncoder(const char* buffer) + : BitEncoder(const_cast(buffer)) {} + + // Places |value| occupying |numBits| into the stream. Behavior + // is undefined if |value| is >= 2^numBits. + void putBits(uint64_t value, int numBits) { + *writeWord_ |= value << writeOffset_; + const int nextOffset = writeOffset_ + numBits; + if (nextOffset >= 64) { + ++writeWord_; + const int spilloverBits = nextOffset - 64; + if (numBits == 64 && spilloverBits == 0) { + return; + } + *writeWord_ |= value >> (numBits - spilloverBits); + writeOffset_ = spilloverBits; + } else { + writeOffset_ = nextOffset; + } + } + + uint64_t bitsWritten() { + return ((reinterpret_cast(writeWord_) - buffer_) << 3) + + writeOffset_; + } + + uint64_t getBits(int numBits) { + const int nextOffset = readOffset_ + numBits; + if (nextOffset >= 64) { + const uint64_t lowBits = *readWord_ >> readOffset_; + ++readWord_; + const int spilloverBits = nextOffset - 64; + readOffset_ = spilloverBits; + if (spilloverBits == 0) { + return lowBits; + } + return lowBits | + (*readWord_ & ((1ULL << spilloverBits) - 1ULL)) + << (numBits - spilloverBits); + } else { + const uint64_t result = + (*readWord_ >> readOffset_) & ((1ULL << numBits) - 1ULL); + readOffset_ = nextOffset; + return result; + } + } + + private: + char* buffer_; + uint64_t* writeWord_; + const uint64_t* readWord_; + int writeOffset_ = 0; + int readOffset_ = 0; +}; + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/Buffer.cpp b/velox/dwio/nimble/common/Buffer.cpp new file mode 100644 index 00000000000..371327b303e --- /dev/null +++ b/velox/dwio/nimble/common/Buffer.cpp @@ -0,0 +1,119 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/dwio/nimble/common/Buffer.h" +#include "velox/dwio/nimble/common/Exceptions.h" + +namespace facebook::nimble { + +char* Buffer::reserve(uint64_t bytes) { + std::scoped_lock l(mutex_); + if (reserveEnd_ + bytes <= chunkEnd_) { + pos_ = reserveEnd_; + reserveEnd_ += bytes; + } else if (!tryAdvanceToNextChunk(bytes)) { + addChunk(bytes); + } + return pos_; +} + +std::vector Buffer::transferBuffers() { + std::scoped_lock l(mutex_); + auto transferred = std::move(chunks_); + chunkIndex_ = 0; + chunkEnd_ = nullptr; + pos_ = nullptr; + reserveEnd_ = nullptr; + return transferred; +} + +void Buffer::reset() { + std::scoped_lock l(mutex_); + chunkIndex_ = 0; + pos_ = chunks_.front()->asMutable(); + chunkEnd_ = pos_ + chunks_.front()->capacity(); + reserveEnd_ = pos_; +} + +bool Buffer::tryAdvanceToNextChunk(uint64_t bytes) { + while (chunkIndex_ + 1 < chunks_.size()) { + ++chunkIndex_; + auto& chunk = chunks_[chunkIndex_]; + if (chunk->capacity() >= bytes) { + pos_ = chunk->asMutable(); + chunkEnd_ = pos_ + chunk->capacity(); + reserveEnd_ = pos_ + bytes; + return true; + } + } + return false; +} + +void Buffer::addChunk(uint64_t bytes) { + const uint64_t chunkSize = std::max(bytes, kMinChunkSize); + auto bufferPtr = velox::AlignedBuffer::allocateExact(chunkSize, pool_); + pos_ = bufferPtr->asMutable(); + chunkEnd_ = pos_ + chunkSize; + reserveEnd_ = pos_ + bytes; + chunks_.emplace_back(std::move(bufferPtr)); + chunkIndex_ = chunks_.size() - 1; +} + +EncodingBufferPool::EncodingBufferPool( + MemoryPool* pool, + uint32_t maxCachedBuffers) + : pool_{pool}, maxCachedBuffers_{maxCachedBuffers} { + NIMBLE_CHECK_NOT_NULL(pool_, "Memory pool cannot be null"); +} + +std::unique_ptr EncodingBufferPool::acquire() { + if (!buffers_.empty()) { + auto buffer = std::move(buffers_.back()); + buffers_.pop_back(); + // Keep acquire() as the handoff point that guarantees a clean scratch + // buffer, even though release() also resets before caching. + buffer->reset(); + return buffer; + } + + return std::make_unique(*pool_); +} + +void EncodingBufferPool::release(std::unique_ptr buffer) { + NIMBLE_CHECK_NOT_NULL(buffer, "Buffer cannot be null"); + + buffer->reset(); + if (buffers_.size() < maxCachedBuffers_) { + buffers_.emplace_back(std::move(buffer)); + } +} + +ScopedEncodingBuffer::ScopedEncodingBuffer( + MemoryPool* memoryPool, + EncodingBufferPool* bufferPool) + : bufferPool_{bufferPool} { + NIMBLE_CHECK_NOT_NULL(memoryPool, "Memory pool cannot be null"); + buffer_ = bufferPool_ != nullptr ? bufferPool_->acquire() + : std::make_unique(*memoryPool); +} + +ScopedEncodingBuffer::~ScopedEncodingBuffer() { + if (bufferPool_ != nullptr) { + bufferPool_->release(std::move(buffer_)); + } +} + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/Buffer.h b/velox/dwio/nimble/common/Buffer.h new file mode 100644 index 00000000000..7f690d837a8 --- /dev/null +++ b/velox/dwio/nimble/common/Buffer.h @@ -0,0 +1,163 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "velox/buffer/Buffer.h" +#include "velox/common/memory/Memory.h" + +#include +#include +#include +#include +#include + +// Basic memory buffer interface, aka arena. +// +// Standard usage: +// char* pos = buffer->reserve(100); +// write data to [pos, pos + 100) +// pos = buffer->Reserve(27); +// write data to [pos, pos + 27) +// and so on + +namespace facebook::nimble { + +/// Internally manages memory in chunks. Releases memory only upon destruction. +/// +/// NOTE: This class is not thread-safe. External locking is required. +class Buffer { + using MemoryPool = facebook::velox::memory::MemoryPool; + + public: + explicit Buffer(MemoryPool& pool, uint64_t initialChunkSize = kMinChunkSize) + : pool_{&pool} { + addChunk(initialChunkSize); + reserveEnd_ = pos_; + } + + /// Returns a pointer to a block of memory of size bytes that can be written + /// to, and guarantees for the lifetime of *this that that region will remain + /// valid. Does NOT guarantee that the region is initially 0'd. + char* reserve(uint64_t bytes); + + /// Copies |data| into the chunk, returning a view to the copied data. + std::string_view writeString(std::string_view data) { + char* pos = reserve(data.size()); + // @lint-ignore CLANGSECURITY facebook-security-vulnerable-memcpy + std::memcpy(pos, data.data(), data.size()); + return {pos, data.size()}; + } + + MemoryPool& getMemoryPool() { + return *pool_; + } + + std::string_view takeOwnership(velox::BufferPtr&& bufferPtr) { + std::string_view chunk{bufferPtr->as(), bufferPtr->size()}; + chunks_.push_back(std::move(bufferPtr)); + return chunk; + } + + /// Transfers ownership of allocated chunks. + std::vector transferBuffers(); + + /// Resets write pointers to the beginning of the first chunk. + /// Keeps all allocated chunks around for reuse. Previously returned + /// pointers/string_views are invalidated. + void reset(); + + /// Returns the number of allocated chunks. Only for testing. + uint32_t testingChunkCount() const { + return chunks_.size(); + } + + /// Returns the current chunk index. Only for testing. + uint32_t testingCurrentChunkIndex() const { + return chunkIndex_; + } + + private: + static constexpr uint64_t kMinChunkSize = 1LL << 20; + + // Tries to advance to the next existing chunk that can fit 'bytes'. + // Returns true if a suitable chunk was found, false if a new allocation + // is needed. Must be called under mutex_. + bool tryAdvanceToNextChunk(uint64_t bytes); + + void addChunk(uint64_t bytes); + + // --- Const members --- + MemoryPool* const pool_; + + // --- Mutable members (protected by mutex_) --- + + // NOTE: this is temporary fix, to quickly enable parallel access to the + // buffer class. In the near future, we are going to templatize this class to + // produce a concurrent and a non-concurrent variants, and change the call + // sites to use each variant when needed. + std::mutex mutex_; + char* chunkEnd_; + char* pos_; + char* reserveEnd_; + uint32_t chunkIndex_{0}; + std::vector chunks_; +}; + +/// Reuses Nimble encoding scratch buffers across nested encode calls. +/// Not thread-safe: each writer thread or encode task should use its own pool. +class EncodingBufferPool { + private: + using MemoryPool = facebook::velox::memory::MemoryPool; + + public: + static constexpr uint32_t kDefaultMaxCachedBuffers = 8; + + explicit EncodingBufferPool( + MemoryPool* pool, + uint32_t maxCachedBuffers = kDefaultMaxCachedBuffers); + + std::unique_ptr acquire(); + + void release(std::unique_ptr buffer); + + private: + MemoryPool* const pool_; + const uint32_t maxCachedBuffers_; + + std::vector> buffers_; +}; + +class ScopedEncodingBuffer { + using MemoryPool = facebook::velox::memory::MemoryPool; + + public: + ScopedEncodingBuffer(MemoryPool* memoryPool, EncodingBufferPool* bufferPool); + + ~ScopedEncodingBuffer(); + + ScopedEncodingBuffer(const ScopedEncodingBuffer&) = delete; + ScopedEncodingBuffer& operator=(const ScopedEncodingBuffer&) = delete; + + Buffer& get() { + return *buffer_; + } + + private: + EncodingBufferPool* const bufferPool_; + std::unique_ptr buffer_; +}; + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/CMakeLists.txt b/velox/dwio/nimble/common/CMakeLists.txt new file mode 100644 index 00000000000..73cc1b37aab --- /dev/null +++ b/velox/dwio/nimble/common/CMakeLists.txt @@ -0,0 +1,48 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +if(NIMBLE_BUILD_TESTING) + add_subdirectory(tests) +endif() + +add_library( + nimble_common + Buffer.cpp + Checksum.cpp + DataTypeDispatch.h + Exceptions.cpp + FeatureGate.cpp + FixedBitArray.cpp + MetricsLogger.cpp + NimbleException.cpp + Types.cpp + Varint.cpp + Zigzag.h + Vector.h + Varint.h + Types.h + StatsUtil.h + NimbleException.h + MetricsLogger.h + FixedBitArray.h + FeatureGate.h + Exceptions.h + ExceptionHelper.h + Constants.h + ChunkHeader.h + Checksum.h + Buffer.h + BitEncoder.h +) + +target_link_libraries(nimble_common velox_common_base velox_memory velox_exception Folly::folly) diff --git a/velox/dwio/nimble/common/Checksum.cpp b/velox/dwio/nimble/common/Checksum.cpp new file mode 100644 index 00000000000..08ff6cb1e52 --- /dev/null +++ b/velox/dwio/nimble/common/Checksum.cpp @@ -0,0 +1,72 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "velox/dwio/nimble/common/Checksum.h" +#include "velox/dwio/nimble/common/Exceptions.h" + +#define XXH_INLINE_ALL +#include + +namespace facebook::nimble { + +namespace { +class Xxh3_64Checksum : public Checksum { + public: + Xxh3_64Checksum() : state_{XXH3_createState()} { + NIMBLE_DCHECK_NOT_NULL(state_, "Failed to initialize Xxh3_64Checksum."); + reset(); + } + + ~Xxh3_64Checksum() override { + XXH3_freeState(state_); + } + + void update(std::string_view data) override { + const auto result = XXH3_64bits_update(state_, data.data(), data.size()); + NIMBLE_CHECK(result != XXH_ERROR, "XXH3_64bits_update error."); + } + + uint64_t getChecksum(bool reset) override { + auto ret = static_cast(XXH3_64bits_digest(state_)); + if (UNLIKELY(reset)) { + this->reset(); + } + return ret; + } + + ChecksumType getType() const override { + return ChecksumType::XXH3_64; + } + + private: + XXH3_state_t* state_; + + void reset() { + const auto result = XXH3_64bits_reset(state_); + NIMBLE_CHECK(result != XXH_ERROR, "XXH3_64bits_reset error."); + } +}; +} // namespace + +std::unique_ptr ChecksumFactory::create(ChecksumType type) { + switch (type) { + case ChecksumType::XXH3_64: + return std::make_unique(); + default: + NIMBLE_UNSUPPORTED("Unsupported checksum type: {}", toString(type)); + } +} + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/Checksum.h b/velox/dwio/nimble/common/Checksum.h new file mode 100644 index 00000000000..49e4522a8e3 --- /dev/null +++ b/velox/dwio/nimble/common/Checksum.h @@ -0,0 +1,38 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "velox/dwio/nimble/common/Types.h" + +#include +#include + +namespace facebook::nimble { + +class Checksum { + public: + virtual ~Checksum() = default; + virtual void update(std::string_view data) = 0; + virtual uint64_t getChecksum(bool reset = false) = 0; + virtual ChecksumType getType() const = 0; +}; + +class ChecksumFactory { + public: + static std::unique_ptr create(ChecksumType type); +}; + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/ChunkHeader.h b/velox/dwio/nimble/common/ChunkHeader.h new file mode 100644 index 00000000000..f15cd4ba199 --- /dev/null +++ b/velox/dwio/nimble/common/ChunkHeader.h @@ -0,0 +1,49 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +#include "velox/dwio/nimble/common/Types.h" +#include "velox/dwio/nimble/encodings/common/EncodingPrimitives.h" + +namespace facebook::nimble { + +/// Chunk header: 4 bytes for compressed chunk length + 1 byte for compression +/// type. This constant must remain 5 to maintain backward compatibility with +/// existing Nimble files. +constexpr int kChunkHeaderSize = 5; + +struct ChunkHeader { + uint32_t length; + CompressionType compressionType; +}; + +/// Reads a chunk header from 'pos', advancing 'pos' by kChunkHeaderSize bytes. +inline ChunkHeader readChunkHeader(const char*& pos) { + return { + encoding::readUint32(pos), + static_cast(encoding::readChar(pos))}; +} + +/// Writes a chunk header to 'pos', advancing 'pos' by kChunkHeaderSize bytes. +inline void +writeChunkHeader(uint32_t length, CompressionType compressionType, char*& pos) { + encoding::writeUint32(length, pos); + encoding::write(compressionType, pos); +} + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/Constants.h b/velox/dwio/nimble/common/Constants.h new file mode 100644 index 00000000000..2c4f241a6b0 --- /dev/null +++ b/velox/dwio/nimble/common/Constants.h @@ -0,0 +1,53 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +namespace facebook::nimble { +/// WARNING: These values have been derived experimentally. + +/// Attempt to compress a data stream only if the data size is equal or greater +/// than this threshold. +constexpr uint32_t kMetaInternalMinCompressionSize{40}; +constexpr uint32_t kZstdMinCompressionSize{25}; +constexpr uint32_t kLz4MinCompressionSize{12}; +constexpr uint32_t kOpenZLMinCompressionSize{40}; + +/// Default options for the ChunkFlushPolicy. +/// Threshold to trigger chunking to relieve memory pressure +constexpr uint64_t kChunkingWriterMemoryHighThreshold{2ULL << 30}; // 2 GB +/// Threshold below which chunking stops. +constexpr uint64_t kChunkingWriterMemoryLowThreshold{1ULL << 30}; // 1 GB +/// Target size for encoded stripes. +constexpr uint64_t kChunkingWriterTargetStripeStorageSize{100 << 20}; // 100MB +/// Expected ratio of raw to encoded data. +constexpr double kChunkingWriterEstimatedCompressionFactor{3.7}; +/// When flushing data streams into chunks, streams with raw data size smaller +/// than this threshold will not be flushed. +/// Note: this threshold is ignored when it is time to flush a stripe. +constexpr uint64_t kChunkingWriterMinChunkSize{512 << 10}; // 512KB +/// When flushing data streams into chunks, streams with raw data size larger +/// than this threshold will be broken down into multiple smaller chunks. Each +/// chunk will be at most this size. +constexpr uint64_t kChunkingWriterMaxChunkSize{20 << 20}; // 20MB +/// Used in place of kChunkingWriterMaxChunkSize for tables with large schemas. +constexpr uint64_t kChunkingWriterWideSchemaMaxChunkSize{2 << 20}; // 2MB + +/// Default block size for BlockBitPacking encoding and its statistics. +constexpr uint16_t kBlockBitPackingBlockSize{1024}; + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/DataTypeDispatch.h b/velox/dwio/nimble/common/DataTypeDispatch.h new file mode 100644 index 00000000000..061228ed400 --- /dev/null +++ b/velox/dwio/nimble/common/DataTypeDispatch.h @@ -0,0 +1,329 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +#include "velox/dwio/nimble/common/Exceptions.h" +#include "velox/dwio/nimble/common/Types.h" + +namespace facebook::nimble { + +#define NIMBLE_RETURN_BY_DATA_TYPE_OR(dataType, Type, expression, ...) \ + switch (dataType) { \ + case ::facebook::nimble::DataType::Int8: { \ + using Type = int8_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint8: { \ + using Type = uint8_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Int16: { \ + using Type = int16_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint16: { \ + using Type = uint16_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Int32: { \ + using Type = int32_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint32: { \ + using Type = uint32_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Int64: { \ + using Type = int64_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint64: { \ + using Type = uint64_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Float: { \ + using Type = float; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Double: { \ + using Type = double; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Bool: { \ + using Type = bool; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::String: { \ + using Type = std::string_view; \ + return (expression); \ + } \ + default: { \ + __VA_ARGS__; \ + } \ + } + +#define NIMBLE_RETURN_BY_DATA_TYPE(dataType, Type, expression) \ + NIMBLE_RETURN_BY_DATA_TYPE_OR( \ + dataType, \ + Type, \ + expression, \ + NIMBLE_UNREACHABLE("Unsupported data type {}.", dataType)) + +#define NIMBLE_TRY_RETURN_BY_DATA_TYPE(dataType, Type, expression) \ + NIMBLE_RETURN_BY_DATA_TYPE_OR(dataType, Type, expression, return {}) + +#define NIMBLE_RETURN_BY_VARINT_DATA_TYPE_OR(dataType, Type, expression, ...) \ + switch (dataType) { \ + case ::facebook::nimble::DataType::Int32: { \ + using Type = int32_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint32: { \ + using Type = uint32_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Int64: { \ + using Type = int64_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint64: { \ + using Type = uint64_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Float: { \ + using Type = float; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Double: { \ + using Type = double; \ + return (expression); \ + } \ + default: { \ + __VA_ARGS__; \ + } \ + } + +#define NIMBLE_RETURN_BY_VARINT_DATA_TYPE(dataType, Type, expression) \ + NIMBLE_RETURN_BY_VARINT_DATA_TYPE_OR( \ + dataType, \ + Type, \ + expression, \ + NIMBLE_UNREACHABLE("Unsupported varint data type {}.", dataType)) + +#define NIMBLE_TRY_RETURN_BY_VARINT_DATA_TYPE(dataType, Type, expression) \ + NIMBLE_RETURN_BY_VARINT_DATA_TYPE_OR(dataType, Type, expression, return {}) + +#define NIMBLE_RETURN_BY_NON_BOOL_DATA_TYPE_OR( \ + dataType, Type, expression, ...) \ + switch (dataType) { \ + case ::facebook::nimble::DataType::Int8: { \ + using Type = int8_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint8: { \ + using Type = uint8_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Int16: { \ + using Type = int16_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint16: { \ + using Type = uint16_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Int32: { \ + using Type = int32_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint32: { \ + using Type = uint32_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Int64: { \ + using Type = int64_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint64: { \ + using Type = uint64_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Float: { \ + using Type = float; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Double: { \ + using Type = double; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::String: { \ + using Type = std::string_view; \ + return (expression); \ + } \ + default: { \ + __VA_ARGS__; \ + } \ + } + +#define NIMBLE_RETURN_BY_NON_BOOL_DATA_TYPE(dataType, Type, expression) \ + NIMBLE_RETURN_BY_NON_BOOL_DATA_TYPE_OR( \ + dataType, \ + Type, \ + expression, \ + NIMBLE_UNREACHABLE("Unsupported non-bool data type {}.", dataType)) + +#define NIMBLE_TRY_RETURN_BY_NON_BOOL_DATA_TYPE(dataType, Type, expression) \ + NIMBLE_RETURN_BY_NON_BOOL_DATA_TYPE_OR(dataType, Type, expression, return {}) + +#define NIMBLE_RETURN_BY_NUMERIC_DATA_TYPE_OR(dataType, Type, expression, ...) \ + switch (dataType) { \ + case ::facebook::nimble::DataType::Int8: { \ + using Type = int8_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint8: { \ + using Type = uint8_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Int16: { \ + using Type = int16_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint16: { \ + using Type = uint16_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Int32: { \ + using Type = int32_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint32: { \ + using Type = uint32_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Int64: { \ + using Type = int64_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint64: { \ + using Type = uint64_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Float: { \ + using Type = float; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Double: { \ + using Type = double; \ + return (expression); \ + } \ + default: { \ + __VA_ARGS__; \ + } \ + } + +#define NIMBLE_RETURN_BY_NUMERIC_DATA_TYPE(dataType, Type, expression) \ + NIMBLE_RETURN_BY_NUMERIC_DATA_TYPE_OR( \ + dataType, \ + Type, \ + expression, \ + NIMBLE_UNREACHABLE("Unsupported numeric data type {}.", dataType)) + +#define NIMBLE_TRY_RETURN_BY_NUMERIC_DATA_TYPE(dataType, Type, expression) \ + NIMBLE_RETURN_BY_NUMERIC_DATA_TYPE_OR(dataType, Type, expression, return {}) + +#define NIMBLE_RETURN_BY_FLOATING_POINT_DATA_TYPE_OR( \ + dataType, Type, expression, ...) \ + switch (dataType) { \ + case ::facebook::nimble::DataType::Float: { \ + using Type = float; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Double: { \ + using Type = double; \ + return (expression); \ + } \ + default: { \ + __VA_ARGS__; \ + } \ + } + +#define NIMBLE_RETURN_BY_FLOATING_POINT_DATA_TYPE(dataType, Type, expression) \ + NIMBLE_RETURN_BY_FLOATING_POINT_DATA_TYPE_OR( \ + dataType, \ + Type, \ + expression, \ + NIMBLE_UNREACHABLE( \ + "Unsupported floating point data type {}.", dataType)) + +#define NIMBLE_TRY_RETURN_BY_FLOATING_POINT_DATA_TYPE( \ + dataType, Type, expression) \ + NIMBLE_RETURN_BY_FLOATING_POINT_DATA_TYPE_OR( \ + dataType, Type, expression, return {}) + +#define NIMBLE_RETURN_BY_INTEGER_DATA_TYPE_OR(dataType, Type, expression, ...) \ + switch (dataType) { \ + case ::facebook::nimble::DataType::Int8: { \ + using Type = int8_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint8: { \ + using Type = uint8_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Int16: { \ + using Type = int16_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint16: { \ + using Type = uint16_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Int32: { \ + using Type = int32_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint32: { \ + using Type = uint32_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Int64: { \ + using Type = int64_t; \ + return (expression); \ + } \ + case ::facebook::nimble::DataType::Uint64: { \ + using Type = uint64_t; \ + return (expression); \ + } \ + default: { \ + __VA_ARGS__; \ + } \ + } + +#define NIMBLE_RETURN_BY_INTEGER_DATA_TYPE(dataType, Type, expression) \ + NIMBLE_RETURN_BY_INTEGER_DATA_TYPE_OR( \ + dataType, \ + Type, \ + expression, \ + NIMBLE_UNREACHABLE("Unsupported integer data type {}.", dataType)) + +#define NIMBLE_TRY_RETURN_BY_INTEGER_DATA_TYPE(dataType, Type, expression) \ + NIMBLE_RETURN_BY_INTEGER_DATA_TYPE_OR(dataType, Type, expression, return {}) + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/ExceptionHelper.h b/velox/dwio/nimble/common/ExceptionHelper.h new file mode 100644 index 00000000000..12d917a73db --- /dev/null +++ b/velox/dwio/nimble/common/ExceptionHelper.h @@ -0,0 +1,56 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +namespace facebook::nimble { + +struct CompileTimeEmptyString { + CompileTimeEmptyString() = default; + constexpr operator const char*() const { + return ""; + } + constexpr operator std::string_view() const { + return {}; + } + operator std::string() const { + return {}; + } +}; + +// When there is no message passed, we can statically detect this case +// and avoid passing even a single unnecessary argument pointer, +// minimizing size and thus maximizing eligibility for inlining. +inline CompileTimeEmptyString errorMessage() { + return {}; +} + +inline const char* errorMessage(const char* s) { + return s; +} + +inline std::string errorMessage(const std::string& str) { + return str; +} + +template +std::string errorMessage(fmt::string_view fmt, const Args&... args) { + return fmt::vformat(fmt, fmt::make_format_args(args...)); +} + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/Exceptions.cpp b/velox/dwio/nimble/common/Exceptions.cpp new file mode 100644 index 00000000000..a635fe6b958 --- /dev/null +++ b/velox/dwio/nimble/common/Exceptions.cpp @@ -0,0 +1,27 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/dwio/nimble/common/Exceptions.h" + +namespace facebook::nimble { + +// Explicit template instantiations to reduce binary bloat +// Note: NimbleExternalError is not included here because it has a different +// constructor signature (requires externalSource parameter) +NIMBLE_DEFINE_CHECK_FAIL_TEMPLATES(NimbleUserError); +NIMBLE_DEFINE_CHECK_FAIL_TEMPLATES(NimbleInternalError); + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/Exceptions.h b/velox/dwio/nimble/common/Exceptions.h new file mode 100644 index 00000000000..4f4cc9aab23 --- /dev/null +++ b/velox/dwio/nimble/common/Exceptions.h @@ -0,0 +1,360 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +#include "folly/FixedString.h" +#include "folly/Preprocessor.h" +#include "velox/dwio/nimble/common/ExceptionHelper.h" +#include "velox/dwio/nimble/common/NimbleException.h" + +// Standard errors used throughout the codebase. + +namespace facebook::nimble { + +namespace detail { + +// Struct containing the arguments needed to throw a Nimble exception +struct NimbleCheckFailArgs { + const char* file; + size_t line; + const char* function; + const char* expression; + std::string_view errorCode; + bool isRetryable; +}; + +// Out-of-line nimbleCheckFail implementations to reduce binary bloat +template +[[noreturn]] void nimbleCheckFail(const NimbleCheckFailArgs& args, Msg msg) { + throw Exception( + args.file, + args.line, + args.function, + args.expression, + msg, + args.errorCode, + args.isRetryable); +} + +// NimbleCheckFailStringType helps us pass by reference to +// NimbleCheckFail exactly when the string type is std::string. +template +struct NimbleCheckFailStringType; + +template <> +struct NimbleCheckFailStringType { + using type = CompileTimeEmptyString; +}; + +template <> +struct NimbleCheckFailStringType { + using type = const char*; +}; + +template <> +struct NimbleCheckFailStringType { + using type = const std::string&; +}; + +// Declare explicit instantiations of nimbleCheckFail for the given +// exceptionType. Just the signatures go in this macro, not the definitions. +#define NIMBLE_DECLARE_CHECK_FAIL_TEMPLATES(exceptionType) \ + namespace detail { \ + extern template void nimbleCheckFail( \ + const NimbleCheckFailArgs&, \ + const char*); \ + extern template void nimbleCheckFail( \ + const NimbleCheckFailArgs&, \ + const std::string&); \ + extern template void nimbleCheckFail( \ + const NimbleCheckFailArgs&, \ + CompileTimeEmptyString); \ + } + +// Define explicit instantiations of nimbleCheckFail for the given +// exceptionType. The actual template instantiations go in this macro. +#define NIMBLE_DEFINE_CHECK_FAIL_TEMPLATES(exceptionType) \ + namespace detail { \ + template void nimbleCheckFail( \ + const NimbleCheckFailArgs&, \ + const char*); \ + template void nimbleCheckFail( \ + const NimbleCheckFailArgs&, \ + const std::string&); \ + template void nimbleCheckFail( \ + const NimbleCheckFailArgs&, \ + CompileTimeEmptyString); \ + } +} // namespace detail + +// Declare template instantiations for common exception types +NIMBLE_DECLARE_CHECK_FAIL_TEMPLATES(::facebook::nimble::NimbleUserError); +NIMBLE_DECLARE_CHECK_FAIL_TEMPLATES(::facebook::nimble::NimbleInternalError); + +// Base throw implementation - Velox-style with NimbleCheckFailArgs +#define _NIMBLE_THROW_IMPL(exception, exprStr, errorCode, retryable, ...) \ + do { \ + /* GCC 9.2.1 doesn't accept this code with constexpr. */ \ + static const ::facebook::nimble::detail::NimbleCheckFailArgs \ + nimbleCheckFailArgs = { \ + __FILE__, __LINE__, __FUNCTION__, exprStr, errorCode, retryable}; \ + auto message = ::facebook::nimble::errorMessage(__VA_ARGS__); \ + ::facebook::nimble::detail::nimbleCheckFail< \ + exception, \ + typename ::facebook::nimble::detail::NimbleCheckFailStringType< \ + decltype(message)>::type>(nimbleCheckFailArgs, message); \ + } while (0) + +#define NIMBLE_RAISE_USER_ERROR(expression, code, retryable, ...) \ + _NIMBLE_THROW_IMPL( \ + ::facebook::nimble::NimbleUserError, \ + expression, \ + code, \ + retryable, \ + ##__VA_ARGS__) + +#define NIMBLE_RAISE_INTERNAL_ERROR(expression, code, retryable, ...) \ + _NIMBLE_THROW_IMPL( \ + ::facebook::nimble::NimbleInternalError, \ + expression, \ + code, \ + retryable, \ + ##__VA_ARGS__) + +// Check internal preconditions and invariants - Velox-style implementation +#define _NIMBLE_CHECK_AND_THROW_IMPL( \ + exprStr, expr, exception, errorCode, retryable, ...) \ + if (UNLIKELY(!(expr))) { \ + _NIMBLE_THROW_IMPL( \ + exception, exprStr, errorCode, retryable, ##__VA_ARGS__); \ + } + +#define _NIMBLE_CHECK_IMPL(expr, exprStr, ...) \ + _NIMBLE_CHECK_AND_THROW_IMPL( \ + exprStr, \ + expr, \ + ::facebook::nimble::NimbleInternalError, \ + ::facebook::nimble::error_code::InvalidArgument, \ + false, \ + ##__VA_ARGS__) + +#define NIMBLE_CHECK(expr, ...) _NIMBLE_CHECK_IMPL(expr, #expr, ##__VA_ARGS__) + +#define NIMBLE_USER_CHECK(expr, ...) \ + _NIMBLE_USER_CHECK_IMPL(expr, #expr, ##__VA_ARGS__) + +// Verify an expected file format conditions. +// the file is corrupted (e.g. passed magic number and version verification, but +// got unexpected format). This will trigger a user error. +#define _NIMBLE_CHECK_FILE_IMPL(condition, conditionString, ...) \ + if (UNLIKELY(!(condition))) { \ + NIMBLE_RAISE_USER_ERROR( \ + conditionString, \ + ::facebook::nimble::error_code::CorruptedFile, \ + /* retryable */ false, \ + __VA_ARGS__); \ + } + +#define NIMBLE_CHECK_FILE(condition, ...) \ + _NIMBLE_CHECK_FILE_IMPL(condition, #condition, ##__VA_ARGS__) + +// Should be raised when we don't expect to hit a code path, but we did. This +// means a bug in Nimble. +#define NIMBLE_UNREACHABLE(...) \ + NIMBLE_RAISE_INTERNAL_ERROR( \ + "", \ + ::facebook::nimble::error_code::UnreachableCode, \ + /* retryable */ false, \ + __VA_ARGS__); + +// Should be raised in places where we still didn't implement the required +// functionality, but intend to do so in the future. This raises an internal +// error to indicate users needs this functionality, but we don't provide it. +#define NIMBLE_NOT_IMPLEMENTED(...) \ + NIMBLE_RAISE_INTERNAL_ERROR( \ + "", \ + ::facebook::nimble::error_code::NotImplemented, \ + /* retryable */ false, \ + __VA_ARGS__); + +// Should be raised in places where we don't support a functionality, and have +// no intention to support it in the future. This raises a user error, as the +// user should not expect this functionality to exist in the first place. +#define NIMBLE_UNSUPPORTED(...) \ + NIMBLE_RAISE_USER_ERROR( \ + "", \ + ::facebook::nimble::error_code::NotSupported, \ + /* retryable */ false, \ + __VA_ARGS__); + +// Incompatible Encoding errors are used in Nimble's encoding optimization, to +// indicate that an attempted encoding is incompatible with the data and should +// be avoided. +#define NIMBLE_INCOMPATIBLE_ENCODING(...) \ + NIMBLE_RAISE_USER_ERROR( \ + "", \ + ::facebook::nimble::error_code::IncompatibleEncoding, \ + /* retryable */ false, \ + __VA_ARGS__); + +// Should be used in "catch all" exception handlers, where we can't classify the +// error correctly. These errors mean that we are missing error classification. +#define NIMBLE_UNKNOWN(...) \ + NIMBLE_RAISE_INTERNAL_ERROR( \ + "", \ + ::facebook::nimble::error_code::Unknown, \ + /* retryable */ true, \ + __VA_ARGS__); + +// Comparison macros - Velox-style +#define _NIMBLE_CHECK_OP_WITH_USER_FMT_HELPER( \ + implmacro, expr1, expr2, op, user_fmt, ...) \ + implmacro( \ + (expr1)op(expr2), \ + #expr1 " " #op " " #expr2, \ + "({} vs. {}) " user_fmt, \ + expr1, \ + expr2, \ + ##__VA_ARGS__) + +#define _NIMBLE_CHECK_OP_HELPER(implmacro, expr1, expr2, op, ...) \ + do { \ + if constexpr (FOLLY_PP_DETAIL_NARGS(__VA_ARGS__) > 0) { \ + _NIMBLE_CHECK_OP_WITH_USER_FMT_HELPER( \ + implmacro, expr1, expr2, op, __VA_ARGS__); \ + } else { \ + implmacro( \ + (expr1)op(expr2), \ + #expr1 " " #op " " #expr2, \ + "({} vs. {})", \ + expr1, \ + expr2); \ + } \ + } while (0) + +#define _NIMBLE_CHECK_OP(expr1, expr2, op, ...) \ + _NIMBLE_CHECK_OP_HELPER(_NIMBLE_CHECK_IMPL, expr1, expr2, op, ##__VA_ARGS__) + +#define _NIMBLE_CHECK_FILE_OP(expr1, expr2, op, ...) \ + _NIMBLE_CHECK_OP_HELPER( \ + _NIMBLE_CHECK_FILE_IMPL, expr1, expr2, op, ##__VA_ARGS__) + +#define _NIMBLE_USER_CHECK_IMPL(expr, exprStr, ...) \ + _NIMBLE_CHECK_AND_THROW_IMPL( \ + exprStr, \ + expr, \ + ::facebook::nimble::NimbleUserError, \ + ::facebook::nimble::error_code::InvalidArgument, \ + /* retryable */ false, \ + ##__VA_ARGS__) + +#define _NIMBLE_USER_CHECK_OP(expr1, expr2, op, ...) \ + _NIMBLE_CHECK_OP_HELPER( \ + _NIMBLE_USER_CHECK_IMPL, expr1, expr2, op, ##__VA_ARGS__) + +// Comparison check macros - internal errors +#define NIMBLE_CHECK_GT(e1, e2, ...) _NIMBLE_CHECK_OP(e1, e2, >, ##__VA_ARGS__) +#define NIMBLE_CHECK_GE(e1, e2, ...) _NIMBLE_CHECK_OP(e1, e2, >=, ##__VA_ARGS__) +#define NIMBLE_CHECK_LT(e1, e2, ...) _NIMBLE_CHECK_OP(e1, e2, <, ##__VA_ARGS__) +#define NIMBLE_CHECK_LE(e1, e2, ...) _NIMBLE_CHECK_OP(e1, e2, <=, ##__VA_ARGS__) +#define NIMBLE_CHECK_EQ(e1, e2, ...) _NIMBLE_CHECK_OP(e1, e2, ==, ##__VA_ARGS__) +#define NIMBLE_CHECK_NE(e1, e2, ...) _NIMBLE_CHECK_OP(e1, e2, !=, ##__VA_ARGS__) + +// Comparison check macros - corrupted file errors +#define NIMBLE_CHECK_FILE_GT(e1, e2, ...) \ + _NIMBLE_CHECK_FILE_OP(e1, e2, >, ##__VA_ARGS__) +#define NIMBLE_CHECK_FILE_GE(e1, e2, ...) \ + _NIMBLE_CHECK_FILE_OP(e1, e2, >=, ##__VA_ARGS__) +#define NIMBLE_CHECK_FILE_LT(e1, e2, ...) \ + _NIMBLE_CHECK_FILE_OP(e1, e2, <, ##__VA_ARGS__) +#define NIMBLE_CHECK_FILE_LE(e1, e2, ...) \ + _NIMBLE_CHECK_FILE_OP(e1, e2, <=, ##__VA_ARGS__) +#define NIMBLE_CHECK_FILE_EQ(e1, e2, ...) \ + _NIMBLE_CHECK_FILE_OP(e1, e2, ==, ##__VA_ARGS__) +#define NIMBLE_CHECK_FILE_NE(e1, e2, ...) \ + _NIMBLE_CHECK_FILE_OP(e1, e2, !=, ##__VA_ARGS__) + +// Comparison check macros - user errors +#define NIMBLE_USER_CHECK_GT(e1, e2, ...) \ + _NIMBLE_USER_CHECK_OP(e1, e2, >, ##__VA_ARGS__) +#define NIMBLE_USER_CHECK_GE(e1, e2, ...) \ + _NIMBLE_USER_CHECK_OP(e1, e2, >=, ##__VA_ARGS__) +#define NIMBLE_USER_CHECK_LT(e1, e2, ...) \ + _NIMBLE_USER_CHECK_OP(e1, e2, <, ##__VA_ARGS__) +#define NIMBLE_USER_CHECK_LE(e1, e2, ...) \ + _NIMBLE_USER_CHECK_OP(e1, e2, <=, ##__VA_ARGS__) +#define NIMBLE_USER_CHECK_EQ(e1, e2, ...) \ + _NIMBLE_USER_CHECK_OP(e1, e2, ==, ##__VA_ARGS__) +#define NIMBLE_USER_CHECK_NE(e1, e2, ...) \ + _NIMBLE_USER_CHECK_OP(e1, e2, !=, ##__VA_ARGS__) + +// Null pointer checks +#define NIMBLE_CHECK_NULL(e, ...) NIMBLE_CHECK((e) == nullptr, ##__VA_ARGS__) +#define NIMBLE_CHECK_NOT_NULL(e, ...) \ + NIMBLE_CHECK((e) != nullptr, ##__VA_ARGS__) +#define NIMBLE_CHECK_FILE_NOT_NULL(e, ...) \ + NIMBLE_CHECK_FILE((e) != nullptr, ##__VA_ARGS__) +#define NIMBLE_USER_CHECK_NULL(e, ...) \ + NIMBLE_USER_CHECK((e) == nullptr, ##__VA_ARGS__) +#define NIMBLE_USER_CHECK_NOT_NULL(e, ...) \ + NIMBLE_USER_CHECK((e) != nullptr, ##__VA_ARGS__) + +// Failure macros without conditions +#define NIMBLE_FAIL(...) \ + NIMBLE_RAISE_INTERNAL_ERROR( \ + "", \ + ::facebook::nimble::error_code::InvalidState, \ + /* retryable */ false, \ + __VA_ARGS__) + +#define NIMBLE_USER_FAIL(...) \ + NIMBLE_RAISE_USER_ERROR( \ + "", \ + ::facebook::nimble::error_code::InvalidArgument, \ + /* retryable */ false, \ + __VA_ARGS__) + +// Debug variants +#ifndef NDEBUG +#define NIMBLE_DCHECK(condition, ...) NIMBLE_CHECK(condition, ##__VA_ARGS__) +#define NIMBLE_DCHECK_GT(e1, e2, ...) NIMBLE_CHECK_GT(e1, e2, ##__VA_ARGS__) +#define NIMBLE_DCHECK_GE(e1, e2, ...) NIMBLE_CHECK_GE(e1, e2, ##__VA_ARGS__) +#define NIMBLE_DCHECK_LT(e1, e2, ...) NIMBLE_CHECK_LT(e1, e2, ##__VA_ARGS__) +#define NIMBLE_DCHECK_LE(e1, e2, ...) NIMBLE_CHECK_LE(e1, e2, ##__VA_ARGS__) +#define NIMBLE_DCHECK_EQ(e1, e2, ...) NIMBLE_CHECK_EQ(e1, e2, ##__VA_ARGS__) +#define NIMBLE_DCHECK_NE(e1, e2, ...) NIMBLE_CHECK_NE(e1, e2, ##__VA_ARGS__) +#define NIMBLE_DCHECK_NULL(e, ...) NIMBLE_CHECK_NULL(e, ##__VA_ARGS__) +#define NIMBLE_DCHECK_NOT_NULL(e, ...) NIMBLE_CHECK_NOT_NULL(e, ##__VA_ARGS__) + +#define NIMBLE_DEBUG_ONLY +#else +#define NIMBLE_DCHECK(condition, ...) NIMBLE_CHECK(true, "") +#define NIMBLE_DCHECK_GT(e1, e2, ...) NIMBLE_CHECK(true, "") +#define NIMBLE_DCHECK_GE(e1, e2, ...) NIMBLE_CHECK(true, "") +#define NIMBLE_DCHECK_LT(e1, e2, ...) NIMBLE_CHECK(true, "") +#define NIMBLE_DCHECK_LE(e1, e2, ...) NIMBLE_CHECK(true, "") +#define NIMBLE_DCHECK_EQ(e1, e2, ...) NIMBLE_CHECK(true, "") +#define NIMBLE_DCHECK_NE(e1, e2, ...) NIMBLE_CHECK(true, "") +#define NIMBLE_DCHECK_NULL(e, ...) NIMBLE_CHECK(true, "") +#define NIMBLE_DCHECK_NOT_NULL(e, ...) NIMBLE_CHECK(true, "") + +#define NIMBLE_DEBUG_ONLY [[maybe_unused]] +#endif + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/FeatureGate.cpp b/velox/dwio/nimble/common/FeatureGate.cpp new file mode 100644 index 00000000000..0c63cd86402 --- /dev/null +++ b/velox/dwio/nimble/common/FeatureGate.cpp @@ -0,0 +1,44 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "velox/dwio/nimble/common/FeatureGate.h" + +#include + +namespace facebook::nimble { + +namespace { +// Process-wide gate, defaulting to the base no-op. Held by shared_ptr so that +// featureGate() can hand out an owning pointer that outlives a concurrent +// re-registration. +folly::Synchronized>& gateStorage() { + static auto* storage = new folly::Synchronized>( + std::make_shared()); + return *storage; +} +} // namespace + +void registerFeatureGate(std::shared_ptr gate) { + if (gate == nullptr) { + gate = std::make_shared(); + } + *gateStorage().wlock() = std::move(gate); +} + +std::shared_ptr featureGate() { + return *gateStorage().rlock(); +} + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/FeatureGate.h b/velox/dwio/nimble/common/FeatureGate.h new file mode 100644 index 00000000000..fa23829367c --- /dev/null +++ b/velox/dwio/nimble/common/FeatureGate.h @@ -0,0 +1,67 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +namespace facebook::nimble { + +/// Runtime enablement gate for optional writer features (e.g. rollout +/// killswitches). The base implementation applies no runtime override: +/// enabled() returns the caller-provided default, so OSS builds honor the +/// static writer config alone. Internal builds install a dynamic-config-backed +/// implementation via registerFeatureGate() to gate gradual rollouts. +class FeatureGate { + public: + /// Stable identifiers for the runtime-gated writer features. OSS-neutral: no + /// dynamic-config (e.g. JustKnobs) names leak into open source. The internal + /// FeatureGate implementation maps each identifier to its backing knob. + class FeatureSet { + public: + static constexpr std::string_view kChunkedEncoding = "chunked_encoding"; + static constexpr std::string_view kStreamDeduplication = + "stream_deduplication"; + static constexpr std::string_view kDisableSharedStringBuffers = + "disable_shared_string_buffers"; + }; + + virtual ~FeatureGate() = default; + + /// Resolves whether `feature` is enabled at runtime, given the caller's + /// requested `defaultValue`. With no gate installed (the OSS default), + /// returns `defaultValue` unchanged. A registered gate may consult dynamic + /// config to hold a feature back during rollout (return false) or force it on + /// (return true), independent of `defaultValue`. + virtual bool enabled(std::string_view feature, bool defaultValue) const { + return defaultValue; + } +}; + +/// Installs the process-wide FeatureGate, replacing any previously registered +/// one. Intended to be called once at startup by internal (non-OSS) code with a +/// dynamic-config-backed gate. Passing nullptr restores the default no-op gate. +/// Thread-safe. +void registerFeatureGate(std::shared_ptr gate); + +/// Returns the process-wide FeatureGate: the one installed via +/// registerFeatureGate(), or a default no-op gate when none has been registered +/// (the OSS case). The returned pointer is never null and owns the gate, so it +/// stays valid even if a different gate is registered concurrently. +/// Thread-safe. +std::shared_ptr featureGate(); + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/FixedBitArray.cpp b/velox/dwio/nimble/common/FixedBitArray.cpp new file mode 100644 index 00000000000..436e2eb76be --- /dev/null +++ b/velox/dwio/nimble/common/FixedBitArray.cpp @@ -0,0 +1,971 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "velox/dwio/nimble/common/FixedBitArray.h" + +#include + +#include + +namespace facebook::nimble { + +// Warning: do not change this function or lots of horrible data corruption +// will probably happen. +uint64_t FixedBitArray::bufferSize(uint64_t elementCount, int bitWidth) { + // We may read or write up to 7 bytes beyond the last theoretically needed + // byte, as we access whole machine words. + constexpr int kSlopSize = 7; + return velox::bits::nbytes(elementCount * bitWidth) + kSlopSize; +} + +FixedBitArray::FixedBitArray(char* buffer, int bitWidth) + : buffer_(buffer), bitWidth_(bitWidth) { + DCHECK(bitWidth_ >= 0); + DCHECK(bitWidth_ <= 64); + // mask_ has the first bitWidth bits set to 1, rest 0. + mask_ = bitWidth == 64 ? (~0ULL) : ((1ULL << bitWidth_) - 1); +} + +uint64_t FixedBitArray::get(uint64_t index) const { + const uint64_t bits = index * bitWidth_; + const uint64_t offset = bits >> 3; + const uint64_t remainder = bits & 7; + const uint64_t word = *reinterpret_cast(buffer_ + offset); + // For widths > 58 bits, the value may overflow into the next word. + if (bitWidth_ > 58) { + const int overflow = bitWidth_ + remainder - 64; + if (overflow > 0) { + const uint64_t nextWord = + *reinterpret_cast(buffer_ + offset + 8); + return ((word >> remainder) | (nextWord << (bitWidth_ - overflow))) & + mask_; + } + } + return (word >> remainder) & mask_; +} + +uint32_t FixedBitArray::get32(uint64_t index) const { + const uint64_t bits = index * bitWidth_; + const uint64_t offset = bits >> 3; + const uint64_t remainder = bits & 7; + const uint64_t word = *reinterpret_cast(buffer_ + offset); + // Don't have to worry about overflow here since bitWidth_ <= 32. + return (word >> remainder) & mask_; +} + +void FixedBitArray::set(uint64_t index, uint64_t value) { + const uint64_t bits = index * bitWidth_; + const uint64_t offset = bits >> 3; + const uint64_t remainder = bits & 7; + uint64_t& word = *reinterpret_cast(buffer_ + offset); + // For widths > 58 bits, the value may overflow into the next word. + if (bitWidth_ > 58) { + const int overflow = bitWidth_ + remainder - 64; + if (overflow > 0) { + uint64_t& nextWord = *reinterpret_cast(buffer_ + offset + 8); + nextWord |= value >> (bitWidth_ - overflow); + } + } + word |= value << remainder; +} + +void FixedBitArray::set32(uint64_t index, uint32_t value) { + const uint64_t bits = index * bitWidth_; + const uint64_t offset = bits >> 3; + const uint64_t remainder = bits & 7; + uint64_t& word = *reinterpret_cast(buffer_ + offset); + // Don't have to worry about overflow here since bitWidth_ <= 32. + word |= static_cast(value) << remainder; +} + +void FixedBitArray::zeroAndSet(uint64_t index, uint64_t value) { + const uint64_t bits = index * bitWidth_; + const uint64_t offset = bits >> 3; + const uint64_t remainder = bits & 7; + uint64_t& word = *reinterpret_cast(buffer_ + offset); + word &= ~(mask_ << remainder); + // For widths > 58 bits, the value may overflow into the next word. + if (bitWidth_ > 58) { + const int overflow = bitWidth_ + remainder - 64; + if (overflow > 0) { + uint64_t& nextWord = *reinterpret_cast(buffer_ + offset + 8); + nextWord &= 0xFFFFFFFFFFFFFFFF << overflow; + nextWord |= value >> (bitWidth_ - overflow); + } + } + word |= value << remainder; +} + +namespace { + +// T is the output data types, namely uint32_t or uint64_t. +template +void bulkGet32Loop( + uint64_t& word, + const uint64_t** nextWord, + T** values, + T baseline) { + constexpr uint64_t kMask = (1ULL << bitWidth) - 1ULL; + // Some bits for the next value may still be present in word. + constexpr uint64_t spillover = (loopPosition * bitWidth) % 64 == 0 + ? 0 + : ((loopPosition + 1) * bitWidth) % 64; + // Note that this isn't a real branch as its on a constexpr. + if constexpr (spillover > 0) { + uint64_t remainder = word; + word = **nextWord; + ++(*nextWord); + if constexpr (withBaseline) { + **values = + ((remainder | word << (bitWidth - spillover)) & kMask) + baseline; + } else { + **values = (remainder | word << (bitWidth - spillover)) & kMask; + } + word >>= spillover; + ++(*values); + } else { + word = **nextWord; + ++(*nextWord); + } + // How many remaining values are in this word? + constexpr int valueCount = (64 - spillover) / bitWidth; + for (int i = 0; i < valueCount; ++i) { + if constexpr (withBaseline) { + **values = (word & kMask) + baseline; + } else { + **values = word & kMask; + } + ++(*values); + word >>= bitWidth; + } + constexpr int nextLoopPosition = loopPosition + valueCount + (spillover > 0); + bulkGet32Loop( + word, nextWord, values, baseline); +} + +// Unfortunately we cannot partially specialize the template for the +// terminal case of loopPosition = 64 so we must explicitly specify them. +#define BULK_GET32_LOOP_TERMINAL_CASE(bitWidth) \ + template <> \ + void bulkGet32Loop( \ + uint64_t& word, \ + const uint64_t** nextWord, \ + uint32_t** values, \ + uint32_t baseline) {} \ + template <> \ + void bulkGet32Loop( \ + uint64_t& word, \ + const uint64_t** nextWord, \ + uint64_t** values, \ + uint64_t baseline) {} \ + template <> \ + void bulkGet32Loop( \ + uint64_t& word, \ + const uint64_t** nextWord, \ + uint32_t** values, \ + uint32_t baseline) {} \ + template <> \ + void bulkGet32Loop( \ + uint64_t& word, \ + const uint64_t** nextWord, \ + uint64_t** values, \ + uint64_t baseline) {} + +BULK_GET32_LOOP_TERMINAL_CASE(1) +BULK_GET32_LOOP_TERMINAL_CASE(2) +BULK_GET32_LOOP_TERMINAL_CASE(3) +BULK_GET32_LOOP_TERMINAL_CASE(4) +BULK_GET32_LOOP_TERMINAL_CASE(5) +BULK_GET32_LOOP_TERMINAL_CASE(6) +BULK_GET32_LOOP_TERMINAL_CASE(7) +BULK_GET32_LOOP_TERMINAL_CASE(8) +BULK_GET32_LOOP_TERMINAL_CASE(9) +BULK_GET32_LOOP_TERMINAL_CASE(10) +BULK_GET32_LOOP_TERMINAL_CASE(11) +BULK_GET32_LOOP_TERMINAL_CASE(12) +BULK_GET32_LOOP_TERMINAL_CASE(13) +BULK_GET32_LOOP_TERMINAL_CASE(14) +BULK_GET32_LOOP_TERMINAL_CASE(15) +BULK_GET32_LOOP_TERMINAL_CASE(16) +BULK_GET32_LOOP_TERMINAL_CASE(17) +BULK_GET32_LOOP_TERMINAL_CASE(18) +BULK_GET32_LOOP_TERMINAL_CASE(19) +BULK_GET32_LOOP_TERMINAL_CASE(20) +BULK_GET32_LOOP_TERMINAL_CASE(21) +BULK_GET32_LOOP_TERMINAL_CASE(22) +BULK_GET32_LOOP_TERMINAL_CASE(23) +BULK_GET32_LOOP_TERMINAL_CASE(24) +BULK_GET32_LOOP_TERMINAL_CASE(25) +BULK_GET32_LOOP_TERMINAL_CASE(26) +BULK_GET32_LOOP_TERMINAL_CASE(27) +BULK_GET32_LOOP_TERMINAL_CASE(28) +BULK_GET32_LOOP_TERMINAL_CASE(29) +BULK_GET32_LOOP_TERMINAL_CASE(30) +BULK_GET32_LOOP_TERMINAL_CASE(31) +BULK_GET32_LOOP_TERMINAL_CASE(32) + +#undef BULK_GET32_LOOP_TERMINAL_CASE + +template +void bulkGet32Internal( + const FixedBitArray& fixedBitArray, + const char* buffer, + uint64_t start, + uint64_t length, + T* values, + T baseline) { + // Every 64 elements we know the slot will end on a word boundary. + // (It might end on a boundary before that, which should technically + // let us be more efficient if we used that instead because we can use + // the non-bulk API less, but the benefit is pretty small so we do the + // simple thing for now). + // + // We first use the non-bulk method to align ourselves to a 64-element + // boundary, then dispatch to the appropriate bit-width-specific code + // for the 64-element loops, then finish with the non-bulk method. + const uint64_t alignedStart = velox::bits::divRoundUp(start, 64) << 6; + if (start + length < alignedStart) { + for (uint64_t i = start; i < start + length; ++i) { + if constexpr (withBaseline) { + // TODO: An alternative would be to have a separate (constexpr) loop + // adding baselines at the end and hopefully the compiler will vectorize + // it (with -mavx2?). But it might be slower due to the extra loop + // condition checks. Need to benchmark it. + *values = fixedBitArray.get32(i) + baseline; + } else { + *values = fixedBitArray.get32(i); + } + ++values; + } + return; + } + for (uint64_t i = start; i < alignedStart; ++i) { + if constexpr (withBaseline) { + *values = fixedBitArray.get32(i) + baseline; + } else { + *values = fixedBitArray.get32(i); + } + ++values; + } + const uint64_t loopCount = (length - (alignedStart - start)) >> 6; + switch (fixedBitArray.bitWidth()) { +#define BULK_GET32_SWITCH_CASE(bitWidth) \ + case bitWidth: { \ + const uint64_t* nextWord = reinterpret_cast( \ + buffer + ((alignedStart * bitWidth) >> 3)); \ + uint64_t word; \ + for (uint64_t i = 0; i < loopCount; ++i) { \ + bulkGet32Loop( \ + word, &nextWord, &values, baseline); \ + } \ + break; \ + } + + BULK_GET32_SWITCH_CASE(1) + BULK_GET32_SWITCH_CASE(2) + BULK_GET32_SWITCH_CASE(3) + BULK_GET32_SWITCH_CASE(4) + BULK_GET32_SWITCH_CASE(5) + BULK_GET32_SWITCH_CASE(6) + BULK_GET32_SWITCH_CASE(7) + BULK_GET32_SWITCH_CASE(8) + BULK_GET32_SWITCH_CASE(9) + BULK_GET32_SWITCH_CASE(10) + BULK_GET32_SWITCH_CASE(11) + BULK_GET32_SWITCH_CASE(12) + BULK_GET32_SWITCH_CASE(13) + BULK_GET32_SWITCH_CASE(14) + BULK_GET32_SWITCH_CASE(15) + BULK_GET32_SWITCH_CASE(16) + BULK_GET32_SWITCH_CASE(17) + BULK_GET32_SWITCH_CASE(18) + BULK_GET32_SWITCH_CASE(19) + BULK_GET32_SWITCH_CASE(20) + BULK_GET32_SWITCH_CASE(21) + BULK_GET32_SWITCH_CASE(22) + BULK_GET32_SWITCH_CASE(23) + BULK_GET32_SWITCH_CASE(24) + BULK_GET32_SWITCH_CASE(25) + BULK_GET32_SWITCH_CASE(26) + BULK_GET32_SWITCH_CASE(27) + BULK_GET32_SWITCH_CASE(28) + BULK_GET32_SWITCH_CASE(29) + BULK_GET32_SWITCH_CASE(30) + BULK_GET32_SWITCH_CASE(31) + BULK_GET32_SWITCH_CASE(32) + +#undef BULK_GET32_SWITCH_CASE + + default: + LOG(FATAL) << "bit width must lie in [1, 32], got: " + << fixedBitArray.bitWidth(); + } + const uint64_t remainderStart = alignedStart + (loopCount << 6); + const uint64_t remainderEnd = start + length; + for (uint64_t i = remainderStart; i < remainderEnd; ++i) { + if constexpr (withBaseline) { + *values = fixedBitArray.get32(i) + baseline; + } else { + *values = fixedBitArray.get32(i); + } + ++values; + } + return; +} + +template +inline uint64_t loadByteAlignedResidual(const char* next) { + static_assert(byteWidth >= 4 && byteWidth <= 8); + if constexpr (byteWidth == 4) { + return folly::loadUnaligned(next); + } else if constexpr (byteWidth == 8) { + return folly::loadUnaligned(next); + } else { + // 5/6/7 bytes: a single unaligned 8-byte load (FixedBitArray buffers are + // required to be sized with bufferSize(), which reserves 7 bytes of slop, + // so reading past the last value is safe) masked to the byte width. + // Faster than stitching the value from a uint32 plus narrow loads. + constexpr uint64_t kMask = (uint64_t{1} << (byteWidth * 8)) - 1; + return folly::loadUnaligned(next) & kMask; + } +} + +template +inline void bulkGetByteAlignedWithBaseline( + const char* buffer, + uint64_t start, + uint64_t length, + uint64_t* values, + uint64_t baseline) { + const char* next = buffer + start * byteWidth; + uint64_t* nextValue = values; + for (uint64_t i = 0; i < length; ++i) { + *nextValue++ = loadByteAlignedResidual(next) + baseline; + next += byteWidth; + } +} + +template +inline void bulkSetByteAlignedWithBaseline( + char* buffer, + uint64_t start, + uint64_t length, + const uint64_t* values, + uint64_t baseline) { + static_assert(byteWidth >= 4 && byteWidth <= 8); + char* next = buffer + start * byteWidth; + const uint64_t* nextValue = values; + for (uint64_t i = 0; i < length; ++i) { + const uint64_t residual{*nextValue++ - baseline}; + if constexpr (byteWidth == 4) { + *reinterpret_cast(next) = static_cast(residual); + } else if constexpr (byteWidth == 5) { + *reinterpret_cast(next) = static_cast(residual); + next[4] = static_cast(residual >> 32); + } else if constexpr (byteWidth == 6) { + *reinterpret_cast(next) = static_cast(residual); + *reinterpret_cast(next + 4) = + static_cast(residual >> 32); + } else if constexpr (byteWidth == 7) { + *reinterpret_cast(next) = static_cast(residual); + *reinterpret_cast(next + 4) = + static_cast(residual >> 32); + next[6] = static_cast(residual >> 48); + } else { + static_assert(byteWidth == 8); + *reinterpret_cast(next) = residual; + } + next += byteWidth; + } +} + +} // namespace + +void FixedBitArray::bulkGet32(uint64_t start, uint64_t length, uint32_t* values) + const { + bulkGet32Internal(*this, buffer_, start, length, values, 0); +} + +void FixedBitArray::bulkGet32Into64( + uint64_t start, + uint64_t length, + uint64_t* values) const { + bulkGet32Internal(*this, buffer_, start, length, values, 0); +} + +void FixedBitArray::bulkGetWithBaseline32( + uint64_t start, + uint64_t length, + uint32_t* values, + uint32_t baseline) const { + bulkGet32Internal( + *this, buffer_, start, length, values, baseline); +} + +void FixedBitArray::bulkGetWithBaseline32Into64( + uint64_t start, + uint64_t length, + uint64_t* values, + uint64_t baseline) const { + bulkGet32Internal( + *this, buffer_, start, length, values, baseline); +} + +void FixedBitArray::bulkGet64WithBaseline( + uint64_t start, + uint64_t length, + uint64_t* values, + uint64_t baseline) const { + const int bitWidth = bitWidth_; + if (bitWidth < 32) { + // Delegate to the optimized template-unrolled 32-bit path. + bulkGetWithBaseline32Into64(start, length, values, baseline); + return; + } + + switch (bitWidth) { + case 32: { + bulkGetByteAlignedWithBaseline<4>( + buffer_, start, length, values, baseline); + return; + } + case 40: { + bulkGetByteAlignedWithBaseline<5>( + buffer_, start, length, values, baseline); + return; + } + case 48: { + bulkGetByteAlignedWithBaseline<6>( + buffer_, start, length, values, baseline); + return; + } + case 56: { + bulkGetByteAlignedWithBaseline<7>( + buffer_, start, length, values, baseline); + return; + } + case 64: { + bulkGetByteAlignedWithBaseline<8>( + buffer_, start, length, values, baseline); + return; + } + default: + break; + } + + // Hoist members to prevent reload on every iteration -- the compiler + // cannot prove that writes to values[] don't alias this->buffer_ etc. + const char* const buffer = buffer_; + const uint64_t mask = mask_; + // Absolute bit offset of the next value from the beginning of the packed + // buffer. + uint64_t bitsOffset = start * static_cast(bitWidth); + // A single 64-bit load is safe when bitWidth + bitsRemainder <= 64. For + // bitWidth 58, bitsRemainder is always even because bitsOffset advances by + // 58 bits, so the maximum possible bitsRemainder is 6, not 7. + if (bitWidth <= 58) { + for (uint64_t i = 0; i < length; ++i) { + const uint64_t byteOffset = bitsOffset >> 3; + const uint64_t bitsRemainder = bitsOffset & 7; + const uint64_t word = + *reinterpret_cast(buffer + byteOffset); + values[i] = ((word >> bitsRemainder) & mask) + baseline; + bitsOffset += bitWidth; + } + return; + } + + for (uint64_t i = 0; i < length; ++i) { + const uint64_t byteOffset = bitsOffset >> 3; + const uint64_t bitsRemainder = bitsOffset & 7; + const uint64_t word = + *reinterpret_cast(buffer + byteOffset); + const int overflow = bitWidth + static_cast(bitsRemainder) - 64; + if (overflow > 0) { + const uint64_t nextWord = + *reinterpret_cast(buffer + byteOffset + 8); + values[i] = + (((word >> bitsRemainder) | (nextWord << (bitWidth - overflow))) & + mask) + + baseline; + } else { + values[i] = ((word >> bitsRemainder) & mask) + baseline; + } + bitsOffset += bitWidth; + } +} + +namespace { + +template < + int bitWidth, + int loopPosition, + bool withBaseline, + typename InputT = uint32_t> +void bulkSet32Loop( + uint64_t** nextWord, + const InputT** values, + InputT baseline) { + // Some bits for the next value may need to put be in the current word. + constexpr int spillover = (loopPosition * bitWidth) % 64 == 0 + ? 0 + : ((loopPosition + 1) * bitWidth) % 64; + // Note that this isn't a real branch as its on a constexpr. + if constexpr (spillover > 0) { + if constexpr (withBaseline) { + **nextWord |= static_cast(**values - baseline) + << (64 - bitWidth + spillover); + ++(*nextWord); + **nextWord |= + static_cast(**values - baseline) >> (bitWidth - spillover); + ++(*values); + } else { + **nextWord |= static_cast(**values) + << (64 - bitWidth + spillover); + ++(*nextWord); + **nextWord |= static_cast(**values) >> (bitWidth - spillover); + ++(*values); + } + } else { + ++(*nextWord); + } + // How many remaining values are in this word? + constexpr int valueCount = (64 - spillover) / bitWidth; + int offset = spillover; + for (int i = 0; i < valueCount; ++i) { + if constexpr (withBaseline) { + **nextWord |= static_cast(**values - baseline) << offset; + } else { + **nextWord |= static_cast(**values) << offset; + } + offset += bitWidth; + ++(*values); + } + constexpr int nextLoopPosition = loopPosition + valueCount + (spillover > 0); + bulkSet32Loop( + nextWord, values, baseline); +} + +// Unfortunately we cannot partially specialize the template for the +// terminal case of loopPosition = 64 so we must explicitly specify them. +// Only the (uint32_t, false/true) and (uint64_t, true) cases are +// instantiated; bulkSet64WithBaseline always uses withBaseline=true, so +// the (uint64_t, false) case is intentionally omitted to avoid unused +// function lints. +#define BULK_SET32_LOOP_TERMINAL_CASE(bitWidth) \ + template <> \ + void bulkSet32Loop( \ + uint64_t** nextWord, const uint32_t** values, uint32_t baseline) {} \ + template <> \ + void bulkSet32Loop( \ + uint64_t** nextWord, const uint32_t** values, uint32_t baseline) {} \ + template <> \ + void bulkSet32Loop( \ + uint64_t** nextWord, const uint64_t** values, uint64_t baseline) {} + +BULK_SET32_LOOP_TERMINAL_CASE(1) +BULK_SET32_LOOP_TERMINAL_CASE(2) +BULK_SET32_LOOP_TERMINAL_CASE(3) +BULK_SET32_LOOP_TERMINAL_CASE(4) +BULK_SET32_LOOP_TERMINAL_CASE(5) +BULK_SET32_LOOP_TERMINAL_CASE(6) +BULK_SET32_LOOP_TERMINAL_CASE(7) +BULK_SET32_LOOP_TERMINAL_CASE(8) +BULK_SET32_LOOP_TERMINAL_CASE(9) +BULK_SET32_LOOP_TERMINAL_CASE(10) +BULK_SET32_LOOP_TERMINAL_CASE(11) +BULK_SET32_LOOP_TERMINAL_CASE(12) +BULK_SET32_LOOP_TERMINAL_CASE(13) +BULK_SET32_LOOP_TERMINAL_CASE(14) +BULK_SET32_LOOP_TERMINAL_CASE(15) +BULK_SET32_LOOP_TERMINAL_CASE(16) +BULK_SET32_LOOP_TERMINAL_CASE(17) +BULK_SET32_LOOP_TERMINAL_CASE(18) +BULK_SET32_LOOP_TERMINAL_CASE(19) +BULK_SET32_LOOP_TERMINAL_CASE(20) +BULK_SET32_LOOP_TERMINAL_CASE(21) +BULK_SET32_LOOP_TERMINAL_CASE(22) +BULK_SET32_LOOP_TERMINAL_CASE(23) +BULK_SET32_LOOP_TERMINAL_CASE(24) +BULK_SET32_LOOP_TERMINAL_CASE(25) +BULK_SET32_LOOP_TERMINAL_CASE(26) +BULK_SET32_LOOP_TERMINAL_CASE(27) +BULK_SET32_LOOP_TERMINAL_CASE(28) +BULK_SET32_LOOP_TERMINAL_CASE(29) +BULK_SET32_LOOP_TERMINAL_CASE(30) +BULK_SET32_LOOP_TERMINAL_CASE(31) +BULK_SET32_LOOP_TERMINAL_CASE(32) + +#undef BULK_SET32_LOOP_TERMINAL_CASE + +template +void bulkSetInternal32( + FixedBitArray& fixedBitArray, + char* buffer, + uint64_t start, + uint64_t length, + const InputT* values, + InputT baseline) { + // Same general logic as BulkGet32. See the comments there. + switch (fixedBitArray.bitWidth()) { +#define BULK_SET32_SWITCH_CASE(bitWidth) \ + case bitWidth: { \ + const uint64_t alignedStart = velox::bits::divRoundUp(start, 64) << 6; \ + if (start + length < alignedStart) { \ + for (uint64_t i = start; i < start + length; ++i) { \ + if constexpr (withBaseline) { \ + fixedBitArray.set32(i, static_cast(*values - baseline)); \ + } else { \ + fixedBitArray.set32(i, static_cast(*values)); \ + } \ + ++values; \ + } \ + return; \ + } \ + for (uint64_t i = start; i < alignedStart; ++i) { \ + if constexpr (withBaseline) { \ + fixedBitArray.set32(i, static_cast(*values - baseline)); \ + } else { \ + fixedBitArray.set32(i, static_cast(*values)); \ + } \ + ++values; \ + } \ + const uint64_t loopCount = (length - (alignedStart - start)) >> 6; \ + uint64_t* nextWord = reinterpret_cast( \ + buffer + ((alignedStart * bitWidth) >> 3) - 8); \ + for (uint64_t i = 0; i < loopCount; ++i) { \ + bulkSet32Loop( \ + &nextWord, &values, baseline); \ + } \ + const uint64_t remainderStart = alignedStart + (loopCount << 6); \ + const uint64_t remainderEnd = start + length; \ + for (uint64_t i = remainderStart; i < remainderEnd; ++i) { \ + if constexpr (withBaseline) { \ + fixedBitArray.set32(i, static_cast(*values - baseline)); \ + } else { \ + fixedBitArray.set32(i, static_cast(*values)); \ + } \ + ++values; \ + } \ + return; \ + } + + BULK_SET32_SWITCH_CASE(1) + BULK_SET32_SWITCH_CASE(2) + BULK_SET32_SWITCH_CASE(3) + BULK_SET32_SWITCH_CASE(4) + BULK_SET32_SWITCH_CASE(5) + BULK_SET32_SWITCH_CASE(6) + BULK_SET32_SWITCH_CASE(7) + BULK_SET32_SWITCH_CASE(8) + BULK_SET32_SWITCH_CASE(9) + BULK_SET32_SWITCH_CASE(10) + BULK_SET32_SWITCH_CASE(11) + BULK_SET32_SWITCH_CASE(12) + BULK_SET32_SWITCH_CASE(13) + BULK_SET32_SWITCH_CASE(14) + BULK_SET32_SWITCH_CASE(15) + BULK_SET32_SWITCH_CASE(16) + BULK_SET32_SWITCH_CASE(17) + BULK_SET32_SWITCH_CASE(18) + BULK_SET32_SWITCH_CASE(19) + BULK_SET32_SWITCH_CASE(20) + BULK_SET32_SWITCH_CASE(21) + BULK_SET32_SWITCH_CASE(22) + BULK_SET32_SWITCH_CASE(23) + BULK_SET32_SWITCH_CASE(24) + BULK_SET32_SWITCH_CASE(25) + BULK_SET32_SWITCH_CASE(26) + BULK_SET32_SWITCH_CASE(27) + BULK_SET32_SWITCH_CASE(28) + BULK_SET32_SWITCH_CASE(29) + BULK_SET32_SWITCH_CASE(30) + BULK_SET32_SWITCH_CASE(31) + BULK_SET32_SWITCH_CASE(32) + +#undef BULK_SET32_SWITCH_CASE + + default: + LOG(FATAL) << "bit width must lie in [1, 32], got: " + << fixedBitArray.bitWidth(); + } +} + +} // namespace + +void FixedBitArray::bulkSet32( + uint64_t start, + uint64_t length, + const uint32_t* values) { + bulkSetInternal32(*this, buffer_, start, length, values, 0); +} + +void FixedBitArray::bulkSet32WithBaseline( + uint64_t start, + uint64_t length, + const uint32_t* values, + uint32_t baseline) { + bulkSetInternal32(*this, buffer_, start, length, values, baseline); +} + +void FixedBitArray::bulkSet64WithBaseline( + uint64_t start, + uint64_t length, + const uint64_t* values, + uint64_t baseline) { + const int bitWidth = bitWidth_; + if (bitWidth < 32) { + bulkSetInternal32( + *this, buffer_, start, length, values, baseline); + return; + } + + switch (bitWidth) { + case 32: { + bulkSetByteAlignedWithBaseline<4>( + buffer_, start, length, values, baseline); + return; + } + case 40: { + bulkSetByteAlignedWithBaseline<5>( + buffer_, start, length, values, baseline); + return; + } + case 48: { + bulkSetByteAlignedWithBaseline<6>( + buffer_, start, length, values, baseline); + return; + } + case 56: { + bulkSetByteAlignedWithBaseline<7>( + buffer_, start, length, values, baseline); + return; + } + case 64: { + if (baseline == 0) { + std::memcpy( + buffer_ + start * sizeof(uint64_t), + values, + length * sizeof(uint64_t)); + return; + } + bulkSetByteAlignedWithBaseline<8>( + buffer_, start, length, values, baseline); + return; + } + default: + break; + } + + char* const buffer = buffer_; + // Absolute bit offset of the next value from the beginning of the packed + // buffer. + uint64_t bitsOffset = start * static_cast(bitWidth); + if (bitWidth <= 58) { + const uint64_t* nextValue = values; + for (uint64_t i = 0; i < length; ++i) { + const uint64_t residualValue = *nextValue++ - baseline; + const uint64_t byteOffset = bitsOffset >> 3; + const uint64_t bitsRemainder = bitsOffset & 7; + uint64_t& word = *reinterpret_cast(buffer + byteOffset); + word |= residualValue << bitsRemainder; + bitsOffset += bitWidth; + } + return; + } + + // For wide values, the residual can straddle two 64-bit words when the + // starting bit offset is not byte/word aligned. Write the low bits into the + // current word and spill the remaining high bits into the next word. + const uint64_t* nextValue = values; + for (uint64_t i = 0; i < length; ++i) { + const uint64_t residualValue = *nextValue++ - baseline; + const uint64_t byteOffset = bitsOffset >> 3; + const uint64_t bitsRemainder = bitsOffset & 7; + uint64_t& word = *reinterpret_cast(buffer + byteOffset); + word |= residualValue << bitsRemainder; + const int overflow = bitWidth + static_cast(bitsRemainder) - 64; + if (overflow > 0) { + uint64_t& nextWord = + *reinterpret_cast(buffer + byteOffset + 8); + nextWord |= residualValue >> (bitWidth - overflow); + } + bitsOffset += bitWidth; + } +} + +namespace { + +template +void equals32Loop( + const uint32_t value, + const uint64_t equalsMask, + uint64_t word, + uint64_t** nextWord, + uint64_t* outputWord) { + constexpr uint64_t kMask = (1ULL << bitWidth) - 1ULL; + constexpr uint64_t spillover = (loopPosition * bitWidth) % 64 == 0 + ? 0 + : ((loopPosition + 1) * bitWidth) % 64; + if (spillover > 0) { + uint64_t remainder = word; + word = **nextWord; + ++(*nextWord); + if (((remainder | word << (bitWidth - spillover)) & kMask) == value) { + *outputWord |= (1ULL << loopPosition); + } + word >>= spillover; + } else { + word = **nextWord; + ++(*nextWord); + } + constexpr int valueCount = (64 - spillover) / bitWidth; + uint64_t maskedWord = word ^ equalsMask; + constexpr int offset = loopPosition + (spillover > 0); + for (int i = 0; i < valueCount; ++i) { + *outputWord |= static_cast((maskedWord & kMask) == 0) + << (offset + i); + maskedWord >>= bitWidth; + } + constexpr int nextWordShift = valueCount * bitWidth; + const uint64_t nextSpillWord = + nextWordShift == 64 ? 0 : (word >> nextWordShift); + constexpr int nextLoopPosition = offset + valueCount; + equals32Loop( + value, equalsMask, nextSpillWord, nextWord, outputWord); +} + +#define EQUALS32_TERMINAL_CASE(bitWidth) \ + template <> \ + void equals32Loop( \ + uint32_t value, \ + uint64_t equalsMask, \ + uint64_t word, \ + uint64_t** nextWord, \ + uint64_t* outputWord) {} + +EQUALS32_TERMINAL_CASE(1) +EQUALS32_TERMINAL_CASE(2) +EQUALS32_TERMINAL_CASE(3) +EQUALS32_TERMINAL_CASE(4) +EQUALS32_TERMINAL_CASE(5) +EQUALS32_TERMINAL_CASE(6) +EQUALS32_TERMINAL_CASE(7) +EQUALS32_TERMINAL_CASE(8) +EQUALS32_TERMINAL_CASE(9) +EQUALS32_TERMINAL_CASE(10) +EQUALS32_TERMINAL_CASE(11) +EQUALS32_TERMINAL_CASE(12) +EQUALS32_TERMINAL_CASE(13) +EQUALS32_TERMINAL_CASE(14) +EQUALS32_TERMINAL_CASE(15) +EQUALS32_TERMINAL_CASE(16) +EQUALS32_TERMINAL_CASE(17) +EQUALS32_TERMINAL_CASE(18) +EQUALS32_TERMINAL_CASE(19) +EQUALS32_TERMINAL_CASE(20) +EQUALS32_TERMINAL_CASE(21) +EQUALS32_TERMINAL_CASE(22) +EQUALS32_TERMINAL_CASE(23) +EQUALS32_TERMINAL_CASE(24) +EQUALS32_TERMINAL_CASE(25) +EQUALS32_TERMINAL_CASE(26) +EQUALS32_TERMINAL_CASE(27) +EQUALS32_TERMINAL_CASE(28) +EQUALS32_TERMINAL_CASE(29) +EQUALS32_TERMINAL_CASE(30) +EQUALS32_TERMINAL_CASE(31) +EQUALS32_TERMINAL_CASE(32) + +#undef EQUALS32_TERMINAL_CASE + +} // namespace + +void FixedBitArray::equals32( + uint64_t start, + uint64_t length, + uint32_t value, + char* bitVector) const { + // Per the header comment we require that start be a multiple of 64. + CHECK_EQ(start & 63, 0); + // First build the equality mask we'll use during the loop. + uint64_t equalsMask = 0; + FixedBitArray maskFixedBitArray((char*)&equalsMask, bitWidth_); + const int maskSlots = 64 / bitWidth_; + for (int i = 0; i < maskSlots; ++i) { + maskFixedBitArray.set(i, value); + } + const uint64_t loopCount = length >> 6; + uint64_t* nextWord = + reinterpret_cast(buffer_ + ((start * bitWidth_) >> 3)); + uint64_t* outputWord = reinterpret_cast(bitVector); + switch (bitWidth_) { +#define EQUALS32_SWITCH_CASE(bitWidth) \ + case bitWidth: { \ + for (uint64_t i = 0; i < loopCount; ++i) { \ + equals32Loop(value, equalsMask, 0, &nextWord, outputWord); \ + ++outputWord; \ + } \ + break; \ + } + + EQUALS32_SWITCH_CASE(1) + EQUALS32_SWITCH_CASE(2) + EQUALS32_SWITCH_CASE(3) + EQUALS32_SWITCH_CASE(4) + EQUALS32_SWITCH_CASE(5) + EQUALS32_SWITCH_CASE(6) + EQUALS32_SWITCH_CASE(7) + EQUALS32_SWITCH_CASE(8) + EQUALS32_SWITCH_CASE(9) + EQUALS32_SWITCH_CASE(10) + EQUALS32_SWITCH_CASE(11) + EQUALS32_SWITCH_CASE(12) + EQUALS32_SWITCH_CASE(13) + EQUALS32_SWITCH_CASE(14) + EQUALS32_SWITCH_CASE(15) + EQUALS32_SWITCH_CASE(16) + EQUALS32_SWITCH_CASE(17) + EQUALS32_SWITCH_CASE(18) + EQUALS32_SWITCH_CASE(19) + EQUALS32_SWITCH_CASE(20) + EQUALS32_SWITCH_CASE(21) + EQUALS32_SWITCH_CASE(22) + EQUALS32_SWITCH_CASE(23) + EQUALS32_SWITCH_CASE(24) + EQUALS32_SWITCH_CASE(25) + EQUALS32_SWITCH_CASE(26) + EQUALS32_SWITCH_CASE(27) + EQUALS32_SWITCH_CASE(28) + EQUALS32_SWITCH_CASE(29) + EQUALS32_SWITCH_CASE(30) + EQUALS32_SWITCH_CASE(31) + EQUALS32_SWITCH_CASE(32) + +#undef EQUALS32_SWITCH_CASE + } + const uint64_t remainderStart = start + (loopCount >> 6); + const uint64_t remainderEnd = start + length; + // Hrm actually in the case we are talking about here the final piece + // could be a written as a single word itself. + for (uint64_t i = remainderStart; i < remainderEnd; ++i) { + if (get32(i) == value) { + velox::bits::setBit(reinterpret_cast(bitVector), i - start); + } + } + return; +} + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/FixedBitArray.h b/velox/dwio/nimble/common/FixedBitArray.h new file mode 100644 index 00000000000..a245745e257 --- /dev/null +++ b/velox/dwio/nimble/common/FixedBitArray.h @@ -0,0 +1,256 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +#include "velox/common/base/BitUtil.h" +#include "velox/dwio/nimble/common/Types.h" + +/// Packs integers into a fixed number (1-64) of bits and retrieves them. +/// The 'bulk' API provided is particularly efficient at setting/getting +/// ranges of values with small bit widths. +/// +/// Typical usage: +/// std::vector data = ... +/// int bitsRequired = BitsRequired(*maxElement(data.begin(), data.end()); +/// auto buffer = std::make_unique( +/// FixedBitArray::BufferSize(data.size(), bitsRequired)); +/// FixedBitArray fixedBitArray(buffer.get(), bitsRequired); +/// fixedBitArray.bulkSet32(0, data.size(), data.data()); +/// +/// And then you'd store the buffer somewhere, and later read the values back, +/// recovering the info stored in data. Note that the FBA does not own +/// any data. + +namespace facebook::nimble { + +class FixedBitArray { + public: + /// Computes the buffer size needed to hold |elementCount| values with the + /// specified |bitWidth|. Note that we allocate a few bytes of 'slop' space + /// beyond what is mathematically required to speed things up. + static uint64_t bufferSize(uint64_t elementCount, int bitWidth); + + /// Not legal to use; included so we can default construct on the stack. + FixedBitArray() = default; + + /// Creates a fixed bit array stored at buffer whose elements can lie within + /// the range [0, 2^|bitWidth|). The |buffer| must already be preallocated to + /// the appropriate size, as given by BufferSize. + FixedBitArray(char* buffer, int bitWidth); + + /// Convenience constructor for reading read-only data. Note that you are NOT + /// prevented by the code from using the non-const calls on a FBA constructed + /// via this method, but to do so is undefined behavior. + FixedBitArray(std::string_view buffer, int bitWidth) + : FixedBitArray(const_cast(buffer.data()), bitWidth) {} + + /// Convenience constructor for reading read-only data. Calling non-const + /// methods on instances constructed from this is undefined behavior. + FixedBitArray(const char* buffer, int bitWidth) + : FixedBitArray(const_cast(buffer), bitWidth) {} + + /// Sets the |index|'th slot to the given |value|. If value + /// is >= 2^|bitWidth| behavior is undefined. + /// + /// IMPORTANT NOTE: this does NOT function as normal assignment. + /// We do NOT first 0 out the slot, i.e. we use or semantics. + /// If you need to first 0 it out, use ZeroAndSet. + void set(uint64_t index, uint64_t value); + + /// Zeroes a slot and then sets it (like normal assignment). + void zeroAndSet(uint64_t index, uint64_t value); + + /// Gets the |index|'th value. + uint64_t get(uint64_t index) const; + + /// Versions of set/get that only work with bitWidth <= 32, but are slightly + /// faster. Same assignment caveat applies to set32 as to set. + void set32(uint64_t index, uint32_t value); + + uint32_t get32(uint64_t index) const; + + /// Retrieves a contiguous subarray from slots [start, start + length). + /// Considerably faster than looping a get call. Only callable when + /// bit width <= 32. + void bulkGet32(uint64_t start, uint64_t length, uint32_t* values) const; + + /// Same as above, but outputs into 8-bytes values. Still requires that bit + /// width <= 32. + void bulkGet32Into64(uint64_t start, uint64_t length, uint64_t* values) const; + + /// Same as above. Add baseline to every value. + void bulkGetWithBaseline32Into64( + uint64_t start, + uint64_t length, + uint64_t* values, + uint64_t baseline) const; + + /// Unified bulk get for any unsigned integral element type, adding baseline + /// to each value. Mirror of bulkSetWithBaseline: 4-byte outputs reuse + /// bulkGetWithBaseline32 and 8-byte outputs reuse bulkGet64WithBaseline (both + /// zero-copy). Narrow 1- and 2-byte outputs are read into fixed-size stack + /// chunks via the fast bulkGet32 path and narrowed down -- faster than + /// per-element get, with no heap allocation. + template + void bulkGetWithBaseline( + uint64_t start, + uint64_t length, + T* values, + uint64_t baseline) const { + static_assert( + isUnsignedIntegralType(), + "bulkGetWithBaseline requires an unsigned integral type."); + if constexpr (isFourByteIntegralType()) { + bulkGetWithBaseline32( + start, + length, + reinterpret_cast(values), + static_cast(baseline)); + } else if constexpr (isEightByteIntegralType()) { + bulkGet64WithBaseline( + start, length, reinterpret_cast(values), baseline); + } else { + // Narrow (1- or 2-byte) types cannot be reinterpreted as a wider array in + // place, so read into fixed-size stack chunks via the fast bulkGet32 path + // and narrow down -- no heap allocation regardless of length. + constexpr uint64_t kWidenChunk = 1024; + uint32_t widened[kWidenChunk]; + const auto baseline32 = static_cast(baseline); + for (uint64_t offset = 0; offset < length; offset += kWidenChunk) { + const uint64_t chunk = + length - offset < kWidenChunk ? length - offset : kWidenChunk; + bulkGetWithBaseline32(start + offset, chunk, widened, baseline32); + for (uint64_t i = 0; i < chunk; ++i) { + values[offset + i] = static_cast(widened[i]); + } + } + } + } + + /// Sets a contiguous subarray of slots from [start, start + length). + /// Considerably faster than looping set calls. Only callable when bitWidth + /// <= 32. Same semantics as set -- see the warning there. + void bulkSet32(uint64_t start, uint64_t length, const uint32_t* values); + + /// Unified bulk set for any unsigned integral element type, subtracting + /// baseline from each value before packing. Dispatches at compile time to the + /// optimal path: 4-byte values reuse bulkSet32WithBaseline and 8-byte values + /// reuse bulkSet64WithBaseline (both zero-copy). Narrow 1- and 2-byte values + /// are widened in fixed-size stack chunks onto the bulkSet32 path -- faster + /// than per-element packing, with no heap allocation. Requires the same + /// zeroed-buffer precondition as set. + template + void bulkSetWithBaseline( + uint64_t start, + uint64_t length, + const T* values, + uint64_t baseline) { + static_assert( + isUnsignedIntegralType(), + "bulkSetWithBaseline requires an unsigned integral type."); + if constexpr (isFourByteIntegralType()) { + bulkSet32WithBaseline( + start, + length, + reinterpret_cast(values), + static_cast(baseline)); + } else if constexpr (isEightByteIntegralType()) { + bulkSet64WithBaseline( + start, length, reinterpret_cast(values), baseline); + } else { + // Narrow (1- or 2-byte) types cannot be reinterpreted as a wider array in + // place, so widen in fixed-size stack chunks and use the fast bulkSet32 + // path -- no heap allocation regardless of length. + constexpr uint64_t kWidenChunk = 1024; + uint32_t widened[kWidenChunk]; + const auto baseline32 = static_cast(baseline); + for (uint64_t offset = 0; offset < length; offset += kWidenChunk) { + const uint64_t chunk = + length - offset < kWidenChunk ? length - offset : kWidenChunk; + for (uint64_t i = 0; i < chunk; ++i) { + widened[i] = static_cast(values[offset + i]); + } + bulkSet32WithBaseline(start + offset, chunk, widened, baseline32); + } + } + } + + /// Fills in the bit-packed |bitVector| with whether each value in + /// [start, start + length) equals |value|. Should be faster than + /// doing bulkGet32 and then doing your own equality check + bitpack. + /// + /// bitVector must point to a buffer of at least size BufferSize(length, 1); + /// + /// REQUIRES: start is a multiple of 64. If you need to run equals on + /// a range where that isn't true, you'll have to manually do the part + /// up to the first multiple of 64 yourself + adjust the results accordingly. + void equals32( + uint64_t start, + uint64_t length, + uint32_t value, + char* bitVector) const; + + int bitWidth() const { + return bitWidth_; + } + + private: + // Width-specific bulk-set/get implementations. Callers use the public + // bulkSetWithBaseline / bulkGetWithBaseline, which dispatch to these by + // element type. + + // Only callable when bitWidth <= 32; subtracts baseline from each value. + // Same zeroed-buffer precondition as set. + void bulkSet32WithBaseline( + uint64_t start, + uint64_t length, + const uint32_t* values, + uint32_t baseline); + + // Packs 64-bit values, subtracting baseline; supports bit widths up to 64. + // For bitWidth < 32 delegates to bulkSet32; 32/40/48/56/64 use direct stores; + // up to 58 use a single 64-bit OR per value; 59-63 handle cross-word + // overflow. Same zeroed-buffer precondition as set. + void bulkSet64WithBaseline( + uint64_t start, + uint64_t length, + const uint64_t* values, + uint64_t baseline); + + // Only callable when bitWidth <= 32; adds baseline to every value. + void bulkGetWithBaseline32( + uint64_t start, + uint64_t length, + uint32_t* values, + uint32_t baseline) const; + + // Reads 64-bit values, adding baseline; supports bit widths up to 64. For + // bitWidth < 32 delegates to bulkGet32; 32/40/48/56/64 use direct loads; up + // to 58 use a single 64-bit load per value; 59-63 handle cross-word overflow. + void bulkGet64WithBaseline( + uint64_t start, + uint64_t length, + uint64_t* values, + uint64_t baseline) const; + + char* buffer_; + int bitWidth_; + uint64_t mask_; +}; + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/MetricsLogger.cpp b/velox/dwio/nimble/common/MetricsLogger.cpp new file mode 100644 index 00000000000..b8756450290 --- /dev/null +++ b/velox/dwio/nimble/common/MetricsLogger.cpp @@ -0,0 +1,54 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "velox/dwio/nimble/common/MetricsLogger.h" + +namespace facebook::nimble { + +folly::dynamic StripeLoadMetrics::serialize() const { + folly::dynamic obj = folly::dynamic::object; + obj["stripeIndex"] = stripeIndex; + obj["rowsInStripe"] = rowsInStripe; + obj["streamCount"] = streamCount; + obj["totalStreamSize"] = totalStreamSize; + return obj; +} + +folly::dynamic StripeFlushMetrics::serialize() const { + folly::dynamic obj = folly::dynamic::object; + obj["inputSize"] = inputSize; + obj["rowCount"] = rowCount; + obj["stripeSize"] = stripeSize; + obj["trackedMemory"] = trackedMemory; + return obj; +} + +folly::dynamic FileCloseMetrics::serialize() const { + folly::dynamic obj = folly::dynamic::object; + obj["rowCount"] = rowCount; + obj["inputSize"] = inputSize; + obj["stripeCount"] = stripeCount; + obj["fileSize"] = fileSize; + obj["encodingCpuNs"] = encodingCpuNs; + obj["encodingWallNs"] = encodingWallNs; + return obj; +} + +LoggingScope::Context& LoggingScope::Context::get() { + thread_local static Context ctx; + return ctx; +} + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/MetricsLogger.h b/velox/dwio/nimble/common/MetricsLogger.h new file mode 100644 index 00000000000..a5ed8822bbc --- /dev/null +++ b/velox/dwio/nimble/common/MetricsLogger.h @@ -0,0 +1,120 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +namespace facebook::nimble { + +struct StripeLoadMetrics { + uint32_t stripeIndex; + uint32_t rowsInStripe; + uint32_t streamCount{0}; + uint32_t totalStreamSize{0}; + // TODO: add IO sizes. + + // TODO: add encoding summary + + size_t cpuUsec; + size_t wallTimeUsec; + + folly::dynamic serialize() const; +}; + +// Might be good to capture via some kind of run stats struct. +// We can then adapt the run stats to file writer run stats. +struct StripeFlushMetrics { + // Stripe shape summary. + uint64_t inputSize; + uint64_t rowCount; + uint64_t stripeSize; + + // We would add flush policy states here when wired up in the future. + // uint64_t inputChunkSize_; + + // TODO: add some type of encoding summary in a follow-up diff. + + // Memory footprint + uint64_t trackedMemory; + // uint64_t residentMemory; + + // Perf stats. + uint64_t flushCpuUsec{}; + uint64_t flushWallTimeUsec{}; + // Add IOStatistics when we have finished WS api consolidations. + + folly::dynamic serialize() const; +}; + +struct FileCloseMetrics { + uint64_t rowCount; + uint64_t inputSize{}; + uint64_t stripeCount; + uint64_t fileSize; + + // Perf stats. + uint64_t encodingCpuNs; + uint64_t encodingWallNs; + // Add IOStatistics when we have finished WS api consolidations. + + folly::dynamic serialize() const; +}; + +enum class LogOperation { + Write, + Flush, + Close, + StripeLoad, + CompressionContext, +}; + +class MetricsLogger { + public: + virtual ~MetricsLogger() = default; + + virtual void logException( + LogOperation /* operation */, + const std::string& /* errorMessage */) const {} + + virtual void logStripeLoad(const StripeLoadMetrics& /* metrics */) const {} + virtual void logStripeFlush(const StripeFlushMetrics& /* metrics */) const {} + virtual void logFileClose(const FileCloseMetrics& /* metrics */) const {} + virtual void logCompressionContext(const std::string&) const {} +}; + +class LoggingScope { + public: + explicit LoggingScope(const MetricsLogger& logger) { + Context::get().logger = &logger; + } + + ~LoggingScope() { + Context::get().logger = nullptr; + } + + static const MetricsLogger* getLogger() { + return Context::get().logger; + } + + private: + struct Context { + const MetricsLogger* logger; + + static Context& get(); + }; +}; + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/NimbleException.cpp b/velox/dwio/nimble/common/NimbleException.cpp new file mode 100644 index 00000000000..297e7fb429b --- /dev/null +++ b/velox/dwio/nimble/common/NimbleException.cpp @@ -0,0 +1,196 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/dwio/nimble/common/NimbleException.h" + +#include +#include +#include + +namespace facebook::nimble { + +namespace { + +std::string captureContextMessage( + velox::VeloxException::Type veloxExceptionType) { + auto* context = &velox::getExceptionContext(); + std::string contextMessage = context->message(veloxExceptionType); + while (context->parent) { + context = context->parent; + if (!context->isEssential) { + continue; + } + const auto message = context->message(veloxExceptionType); + if (message.empty()) { + continue; + } + if (!contextMessage.empty()) { + contextMessage += ' '; + } + contextMessage += message; + } + return contextMessage; +} + +} // namespace + +NimbleException::NimbleException( + std::string_view exceptionName, + const char* fileName, + size_t fileLine, + const char* functionName, + std::string_view failingExpression, + std::string_view errorMessage, + std::string_view errorCode, + bool retryable, + velox::VeloxException::Type veloxExceptionType) + : exceptionName_{std::move(exceptionName)}, + fileName_{fileName}, + fileLine_{fileLine}, + functionName_{functionName}, + failingExpression_{std::move(failingExpression)}, + errorMessage_{std::move(errorMessage)}, + errorCode_{std::move(errorCode)}, + retryable_{retryable}, + context_{captureContextMessage(veloxExceptionType)} { + captureStackTraceFrames(); +} + +const char* NimbleException::what() const noexcept { + try { + folly::call_once(once_, [&] { finalizeMessage(); }); + return finalizedMessage_.c_str(); + } catch (...) { + return ""; + } +} + +void NimbleException::captureStackTraceFrames() { + try { + constexpr size_t skipFrames = 2; + constexpr size_t maxFrames = 200; + uintptr_t addresses[maxFrames]; + ssize_t n = folly::symbolizer::getStackTrace(addresses, maxFrames); + + if (n < skipFrames) { + return; + } + + exceptionFrames_.assign(addresses + skipFrames, addresses + n); + } catch (const std::exception& ex) { + LOG(WARNING) << "Unable to capture stack trace: " << ex.what(); + } catch (...) { + LOG(WARNING) << "Unable to capture stack trace."; + } +} + +void NimbleException::finalizeMessage() const { + finalizedMessage_ += exceptionName_; + finalizedMessage_ += "\nError Source: "; + finalizedMessage_ += errorSource(); + finalizedMessage_ += "\nError Code: "; + finalizedMessage_ += errorCode_; + if (!errorMessage_.empty()) { + finalizedMessage_ += "\nError Message: "; + finalizedMessage_ += errorMessage_; + } + finalizedMessage_ += "\n"; + appendMessage(finalizedMessage_); + finalizedMessage_ += "Retryable: "; + finalizedMessage_ += retryable_ ? "True" : "False"; + finalizedMessage_ += "\nLocation: "; + finalizedMessage_ += functionName_; + finalizedMessage_ += "@"; + finalizedMessage_ += fileName_; + finalizedMessage_ += ":"; + finalizedMessage_ += folly::to(fileLine_); + + if (!failingExpression_.empty()) { + finalizedMessage_ += "\nExpression: "; + finalizedMessage_ += failingExpression_; + } + + if (!context_.empty()) { + finalizedMessage_ += "\nContext: "; + finalizedMessage_ += context_; + } + + if (LIKELY(!exceptionFrames_.empty())) { + std::vector symbolizedFrames; + symbolizedFrames.resize(exceptionFrames_.size()); + + folly::symbolizer::Symbolizer symbolizer{ + folly::symbolizer::LocationInfoMode::FULL}; + symbolizer.symbolize( + exceptionFrames_.data(), + symbolizedFrames.data(), + symbolizedFrames.size()); + + folly::symbolizer::StringSymbolizePrinter printer{ + folly::symbolizer::StringSymbolizePrinter::COLOR}; + printer.println(symbolizedFrames.data(), symbolizedFrames.size()); + + finalizedMessage_ += "\nStack Trace:\n"; + finalizedMessage_ += printer.str(); + } +} + +NimbleUserError::NimbleUserError( + const char* fileName, + size_t fileLine, + const char* functionName, + std::string_view failingExpression, + std::string_view errorMessage, + std::string_view errorCode, + bool retryable) + : NimbleException( + "NimbleUserError", + fileName, + fileLine, + functionName, + failingExpression, + errorMessage, + errorCode, + retryable, + velox::VeloxException::Type::kUser) {} + +const std::string_view NimbleUserError::errorSource() const { + return "USER"; +} + +NimbleInternalError::NimbleInternalError( + const char* fileName, + size_t fileLine, + const char* functionName, + std::string_view failingExpression, + std::string_view errorMessage, + std::string_view errorCode, + bool retryable) + : NimbleException( + "NimbleInternalError", + fileName, + fileLine, + functionName, + failingExpression, + errorMessage, + errorCode, + retryable, + velox::VeloxException::Type::kSystem) {} + +const std::string_view NimbleInternalError::errorSource() const { + return "INTERNAL"; +} +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/NimbleException.h b/velox/dwio/nimble/common/NimbleException.h new file mode 100644 index 00000000000..1e7f9fc754e --- /dev/null +++ b/velox/dwio/nimble/common/NimbleException.h @@ -0,0 +1,192 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "velox/common/base/VeloxException.h" + +namespace facebook::nimble { + +namespace error_source { +using namespace folly::string_literals; + +// Errors where the root cause of the problem is either because of bad input +// or an unsupported pattern of use are classified with source USER. +inline constexpr auto User = "USER"_fs; + +// Errors where the root cause of the problem is an unexpected internal state in +// the system. +inline constexpr auto Internal = "INTERNAL"_fs; + +// Errors where the root cause of the problem is the result of a dependency or +// an environment failures. +inline constexpr auto External = "EXTERNAL"_fs; +} // namespace error_source + +namespace error_code { +using namespace folly::string_literals; + +// An error raised when an argument verification fails +inline constexpr auto InvalidArgument = "INVALID_ARGUMENT"_fs; + +// An error raised when the current state of a component is invalid. +inline constexpr auto InvalidState = "INVALID_STATE"_fs; + +// An error raised when unreachable code point was executed. +inline constexpr auto UnreachableCode = "UNREACHABLE_CODE"_fs; + +// An error raised when a requested operation is not yet implemented. +inline constexpr auto NotImplemented = "NOT_IMPLEMENTED"_fs; + +// An error raised when a requested operation is not supported. +inline constexpr auto NotSupported = "NOT_SUPPORTED"_fs; + +// As error raised during encoding optimization, when incompatible encoding is +// attempted. +inline constexpr auto IncompatibleEncoding = "INCOMPATIBLE_ENCODING"_fs; + +// An error raised if a corrupted file is detected +inline constexpr auto CorruptedFile = "CORRUPTED_FILE"_fs; + +// We do not know how to classify it yet. +inline constexpr auto Unknown = "UNKNOWN"_fs; +} // namespace error_code + +namespace external_source { +using namespace folly::string_literals; + +// Warm Storage +inline constexpr auto WarmStorage = "WARM_STORAGE"_fs; + +// Local File System +inline constexpr auto LocalFileSystem = "FILE_SYSTEM"_fs; + +} // namespace external_source + +// Base exception used by all other Nimble exceptions. Provides common +// functionality for all other exception types. +class NimbleException : public std::exception { + public: + explicit NimbleException( + std::string_view exceptionName, + const char* fileName, + size_t fileLine, + const char* functionName, + std::string_view failingExpression, + std::string_view errorMessage, + std::string_view errorCode, + bool retryable, + velox::VeloxException::Type veloxExceptionType); + + const char* what() const noexcept override; + + const std::string& exceptionName() const { + return exceptionName_; + } + + const char* fileName() const { + return fileName_; + } + + size_t fileLine() const { + return fileLine_; + } + + const char* functionName() const { + return functionName_; + } + + const std::string& failingExpression() const { + return failingExpression_; + } + + const std::string& errorMessage() const { + return errorMessage_; + } + + virtual const std::string_view errorSource() const = 0; + + const std::string& errorCode() const { + return errorCode_; + } + + bool retryable() const { + return retryable_; + } + + const std::string& context() const { + return context_; + } + + protected: + virtual void appendMessage(std::string& /* message */) const {} + + private: + void captureStackTraceFrames(); + void finalizeMessage() const; + + std::vector exceptionFrames_; + const std::string stackTrace_; + const std::string exceptionName_; + const char* fileName_; + const size_t fileLine_; + const char* functionName_; + const std::string failingExpression_; + const std::string errorMessage_; + const std::string errorCode_; + const bool retryable_; + const std::string context_; + + mutable folly::once_flag once_; + mutable std::string finalizedMessage_; +}; + +// Exception representing an error originating by a user misusing Nimble. +class NimbleUserError : public NimbleException { + public: + NimbleUserError( + const char* fileName, + size_t fileLine, + const char* functionName, + std::string_view failingExpression, + std::string_view errorMessage, + std::string_view errorCode, + bool retryable); + + const std::string_view errorSource() const override; +}; + +// Exception representing an internal error within the Nimble library. +class NimbleInternalError : public NimbleException { + public: + NimbleInternalError( + const char* fileName, + size_t fileLine, + const char* functionName, + std::string_view failingExpression, + std::string_view errorMessage, + std::string_view errorCode, + bool retryable); + + const std::string_view errorSource() const override; +}; +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/StatsUtil.h b/velox/dwio/nimble/common/StatsUtil.h new file mode 100644 index 00000000000..1ac466950d3 --- /dev/null +++ b/velox/dwio/nimble/common/StatsUtil.h @@ -0,0 +1,144 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "velox/common/base/SimdUtil.h" + +namespace facebook::nimble { + +template +struct MinMax { + T min; + T max; +}; + +template +constexpr bool kIntegralMinMaxType = std::is_integral_v> && + !std::is_same_v, bool>; + +template +constexpr bool kFloatingPointMinMaxType = + std::is_floating_point_v>; + +template +constexpr bool kStringMinMaxType = + std::is_same_v, std::string_view>; + +template +std::enable_if_t, MinMax>> +findMinMax(std::span values) { + using Value = std::remove_cv_t; + + Value minValue{values.front()}; + Value maxValue{values.front()}; + + for (const Value value : values) { + minValue = std::min(minValue, value); + maxValue = std::max(maxValue, value); + } + + return { + .min = minValue, + .max = maxValue, + }; +} + +template +std::enable_if_t< + kFloatingPointMinMaxType, + std::optional>>> +findMinMax(std::span values) { + using Value = std::remove_cv_t; + using Batch = xsimd::batch; + + const auto* rawValues = values.data(); + size_t index{0}; + Value minValue{values.front()}; + Value maxValue{values.front()}; + + if (std::isnan(minValue)) { + return std::nullopt; + } + + if (values.size() >= Batch::size) { + auto minBatch = Batch::load_unaligned(rawValues); + if (xsimd::any(minBatch != minBatch)) { + return std::nullopt; + } + auto maxBatch = minBatch; + index = Batch::size; + for (; index + Batch::size <= values.size(); index += Batch::size) { + const auto batch = Batch::load_unaligned(rawValues + index); + if (xsimd::any(batch != batch)) { + return std::nullopt; + } + minBatch = xsimd::min(minBatch, batch); + maxBatch = xsimd::max(maxBatch, batch); + } + minValue = xsimd::reduce_min(minBatch); + maxValue = xsimd::reduce_max(maxBatch); + } + + for (; index < values.size(); ++index) { + const auto value = values[index]; + if (std::isnan(value)) { + return std::nullopt; + } + minValue = std::min(minValue, value); + maxValue = std::max(maxValue, value); + } + + return MinMax{ + .min = minValue, + .max = maxValue, + }; +} + +/// Returns bounds as views into 'values'; copy them before the buffer is +/// recycled. Branches instead of calling std::min and std::max unconditionally +/// because string comparison is not free, and a value below the running minimum +/// cannot also be above the running maximum. +template + requires(kStringMinMaxType) +MinMax> findMinMax(std::span values) { + using Value = std::remove_cv_t; + + Value minValue{values.front()}; + Value maxValue{values.front()}; + + for (const Value value : values) { + if (value < minValue) { + minValue = value; + } else if (value > maxValue) { + maxValue = value; + } + } + + return { + .min = minValue, + .max = maxValue, + }; +} + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/Types.cpp b/velox/dwio/nimble/common/Types.cpp new file mode 100644 index 00000000000..5eb0c0bee1a --- /dev/null +++ b/velox/dwio/nimble/common/Types.cpp @@ -0,0 +1,214 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "velox/dwio/nimble/common/Types.h" + +#include +#include + +#include "velox/dwio/nimble/common/Exceptions.h" + +namespace facebook::nimble { +namespace { + +constexpr auto kEncodingTypes = + std::to_array>({ + {EncodingType::Trivial, "Trivial"}, + {EncodingType::RLE, "RLE"}, + {EncodingType::Dictionary, "Dictionary"}, + {EncodingType::FixedBitWidth, "FixedBitWidth"}, + {EncodingType::Sentinel, "Sentinel"}, + {EncodingType::Nullable, "Nullable"}, + {EncodingType::SparseBool, "SparseBool"}, + {EncodingType::Varint, "Varint"}, + {EncodingType::Delta, "Delta"}, + {EncodingType::Constant, "Constant"}, + {EncodingType::MainlyConstant, "MainlyConstant"}, + {EncodingType::Prefix, "Prefix"}, + {EncodingType::ALP, "ALP"}, + {EncodingType::PFOR, "PFOR"}, + {EncodingType::SimdForBitpack, "SimdForBitpack"}, + {EncodingType::BlockBitPacking, "BlockBitPacking"}, + {EncodingType::SubIntSplit, "SubIntSplit"}, + {EncodingType::FrequencyPartition, "FrequencyPartition"}, + {EncodingType::FOR, "FOR"}, + {EncodingType::Fsst, "Fsst"}, + {EncodingType::Huffman, "Huffman"}, + {EncodingType::DeltaBlock, "DeltaBlock"}, + {EncodingType::SharedDictionary, "SharedDictionary"}, + {EncodingType::Slice, "Slice"}, + }); + +constexpr auto kReadOnlyEncodingTypes = std::to_array({ + EncodingType::FOR, +}); + +constexpr auto kCompressionTypes = + std::to_array>({ + {CompressionType::Uncompressed, "Uncompressed"}, + {CompressionType::Zstd, "Zstd"}, + {CompressionType::MetaInternal, "MetaInternal"}, + {CompressionType::Lz4, "Lz4"}, + {CompressionType::OpenZL, "OpenZL"}, + }); + +} // namespace + +std::ostream& operator<<(std::ostream& out, EncodingType encodingType) { + return out << toString(encodingType); +} + +std::string toString(EncodingType encodingType) { + for (const auto& [type, name] : kEncodingTypes) { + if (encodingType == type) { + return std::string{name}; + } + } + return fmt::format( + "Unknown encoding type: {}", static_cast(encodingType)); +} + +EncodingType toEncodingType(std::string_view name) { + for (const auto& [type, candidate] : kEncodingTypes) { + if (name == candidate) { + return type; + } + } + NIMBLE_USER_FAIL("Unknown encoding type: {}", name); +} + +bool isReadOnlyEncoding(EncodingType encodingType) { + return std::find( + kReadOnlyEncodingTypes.begin(), + kReadOnlyEncodingTypes.end(), + encodingType) != kReadOnlyEncodingTypes.end(); +} + +bool isReadOnlyEncoding(std::string_view name) { + return std::any_of( + kReadOnlyEncodingTypes.begin(), + kReadOnlyEncodingTypes.end(), + [name](const auto encodingType) { + return name == toString(encodingType); + }); +} + +std::ostream& operator<<(std::ostream& out, DataType dataType) { + return out << toString(dataType); +} + +std::string toString(DataType dataType) { + switch (dataType) { + case DataType::Int8: + return "Int8"; + case DataType::Int16: + return "Int16"; + case DataType::Uint8: + return "Uint8"; + case DataType::Uint16: + return "Uint16"; + case DataType::Int32: + return "Int32"; + case DataType::Int64: + return "Int64"; + case DataType::Uint32: + return "Uint32"; + case DataType::Uint64: + return "Uint64"; + case DataType::Float: + return "Float"; + case DataType::Double: + return "Double"; + case DataType::Bool: + return "Bool"; + case DataType::String: + return "String"; + default: + return fmt::format( + "Unknown data type: {}", static_cast(dataType)); + } +} + +uint32_t decodedValueWidth(DataType dataType) { + switch (dataType) { + case DataType::Bool: + case DataType::Int8: + case DataType::Uint8: + return 1; + case DataType::Int16: + case DataType::Uint16: + return 2; + case DataType::Int32: + case DataType::Uint32: + case DataType::Float: + return 4; + case DataType::Int64: + case DataType::Uint64: + case DataType::Double: + return 8; + case DataType::String: + return sizeof(std::string_view); + case DataType::Undefined: + NIMBLE_UNSUPPORTED("Unsupported data type: {}", dataType); + } + NIMBLE_UNREACHABLE("Unknown data type: {}", dataType); +} + +std::string toString(CompressionType compressionType) { + for (const auto& [type, name] : kCompressionTypes) { + if (compressionType == type) { + return std::string{name}; + } + } + return fmt::format( + "Unknown compression type: {}", static_cast(compressionType)); +} + +CompressionType toCompressionType(std::string_view name) { + for (const auto& [type, candidate] : kCompressionTypes) { + if (name == candidate) { + return type; + } + } + NIMBLE_USER_FAIL("Unknown compression type: {}", name); +} + +std::ostream& operator<<(std::ostream& out, CompressionType compressionType) { + return out << toString(compressionType); +} + +template <> +void Variant::set( + VariantType& target, + std::string_view source) { + target = std::string(source); +} + +template <> +std::string_view Variant::get(VariantType& source) { + return std::get(source); +} + +std::string toString(ChecksumType type) { + switch (type) { + case ChecksumType::XXH3_64: + return "XXH3_64"; + default: + return fmt::format( + "Unknown checksum type: {}", static_cast(type)); + } +} + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/Types.h b/velox/dwio/nimble/common/Types.h new file mode 100644 index 00000000000..d6a413b422a --- /dev/null +++ b/velox/dwio/nimble/common/Types.h @@ -0,0 +1,460 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include + +// Single file containing all the types and enums in nimble, as well as some +// templates for mapping between those types and C++ types. +// +// A note on types: For all of our enum classes, we assume the number of +// types will stay small, so that each one is representable by a single byte. +// Don't violate this! + +namespace facebook::nimble { + +using VariantType = std::variant< + int8_t, + uint8_t, + int16_t, + uint16_t, + int32_t, + uint32_t, + int64_t, + uint64_t, + float, + double, + bool, + std::string>; + +template +class Variant { + public: + static void set(VariantType& target, T source) { + target = source; + } + + static T get(VariantType& source) { + return std::get(source); + } +}; + +template <> +void Variant::set( + VariantType& target, + std::string_view source); + +template <> +std::string_view Variant::get(VariantType& source); + +enum class EncodingType { + // Native encoding for numerics, simple packed chars with offsets for strings, + // bitpacked for bools. All data types supported. + Trivial = 0, + // Run length encoded data. The runs lengths are bit packed, and the run + // values are encoded like the trivial encoding. All data types supported. + RLE = 1, + // Data with the uniques encoded in a dictionary and the indices into that + // dictionary. All data types except bools supported. + Dictionary = 2, + // Stores integer types packed into a fixed number of bits (namely, the + // smallest required to represent the largest element). Currently only + // works with non-negative values, we may add ZigZag encoding later. + FixedBitWidth = 3, + // Stores nullable data using a 'sentinel' value to represent nulls in a + // single non-nullable encoding. + Sentinel = 4, + // Stores nullable data by wrapping one subencoding representing the non-nulls + // with another subencoding marking which rows are null. + Nullable = 5, + // Stores indices to set (or unset) bits. Useful for storing sparse data, such + // as when only a few rows in a encoding are non-null. + SparseBool = 6, + // Stores integer types via varint encoding. Currently only + // works with non-negative values, we may add ZigZag encoding later. + Varint = 7, + // Stores integer types with a delta encoding. Currently only supports + // positive deltas. + Delta = 8, + // Stores constant (i.e. only 1 unique value) data. + Constant = 9, + // Stores 'mainly constant' data, i.e. treats one particular value as special, + // using a bool child vector to store whether each row is that special value, + // and stores the non-special values as a separate encoding. + MainlyConstant = 10, + // Stores sorted string data with prefix compression. Common prefixes are + // shared across consecutive entries to reduce storage. Supports seek + // operations for efficient random access. + Prefix = 11, + // Adaptive Lossless floating-Point compression for numeric types. + ALP = 12, + // Patched Frame-of-Reference. Subtracts a min baseline, bitpacks ~90% of + // residuals at a narrow base bit width, and stores the remaining outliers + // ("exceptions") as a parallel position+value array. + // EXPERIMENTAL: Not production-ready. Do not enable for production tables + // without consulting the Nimble team (oncall: dwios). + PFOR = 13, + // SIMD Frame-of-Reference bitpacking. Subtracts baseline (global min), + // packs residuals in groups of 32 via Lemire FastPFor SIMD bitpacking. + // EXPERIMENTAL: Not production-ready. Do not enable for production tables + // without consulting the Nimble team (oncall: dwios). + SimdForBitpack = 14, + // Decomposes each value into bit-range sub-streams and encodes each + // independently. Optimal splits are chosen via a sample-driven DP algorithm. + // Only supported for 32- and 64-bit numeric types. + // EXPERIMENTAL: Not production-ready. Do not enable for production tables + // without consulting the Nimble team (oncall: dwios). + SubIntSplit = 16, + // Stores integer data in fixed-size chunks (default 1024 rows), each with + // its own baseline and bit width. Adapts per-chunk variable bit widths for + // better compression on data with locally narrow value ranges. + // EXPERIMENTAL: Not production-ready. Do not enable for production tables + // without consulting the Nimble team (oncall: dwios). + BlockBitPacking = 15, + // Partitions data by value frequency. Also, frequent values get shorter + // bit-width + // codes. Rows are reordered to group values with same code length. + // EXPERIMENTAL: Not production-ready. Do not enable for production tables + // without consulting the Nimble team (oncall: dwios). + FrequencyPartition = 17, + // Frame of Reference: stores offsets from per-frame minimum values. + // Supports O(1) random access. Preserves row order. + // EXPERIMENTAL: Not production-ready. Do not enable for production tables + // without consulting the Nimble team (oncall: dwios). + FOR = 18, + // Stores string data compressed with FSST (Fast Static Symbol Table). + // Trains a per-stripe symbol table and compresses each string independently, + // enabling random access decompression without touching neighboring strings. + // EXPERIMENTAL: Not production-ready. Do not enable for production tables + // without consulting the Nimble team (oncall: dwios). + Fsst = 19, + // Canonical Huffman coding for integral values. Preserves row order and + // stores periodic bit offsets for bounded random access. + // EXPERIMENTAL: Not production-ready. Do not enable for production tables + // without consulting the Nimble team (oncall: dwios). + Huffman = 20, + // Stores sorted integer streams as block checkpoints plus bit-packed deltas. + // EXPERIMENTAL: Not production-ready. Do not enable for production tables + // without consulting the Nimble team (oncall: dwios). + DeltaBlock = 21, + // Stores indices into an integer alphabet owned by a stripe, file, or + // external provider. The alphabet is resolved independently from the + // encoded index stream. + SharedDictionary = 22, + // A slice of another encoding: carries the source encoding verbatim plus + // the row offset at which the slice begins, deferring the slice work to + // decode time. Produced only by EncodingSliceFactory, never by encoding + // selection. + Slice = 23, +}; +std::string toString(EncodingType encodingType); +/// Returns the encoding type for 'name'. Throws if 'name' is unknown. +EncodingType toEncodingType(std::string_view name); +/// Returns true if the encoding is retained only for reading existing data. +bool isReadOnlyEncoding(EncodingType encodingType); +/// Returns true if the name identifies a read-only encoding. +bool isReadOnlyEncoding(std::string_view name); +std::ostream& operator<<(std::ostream& out, EncodingType encodingType); + +enum class DataType : uint8_t { + Undefined = 0, + Int8 = 1, + Uint8 = 2, + Int16 = 3, + Uint16 = 4, + Int32 = 5, + Uint32 = 6, + Int64 = 7, + Uint64 = 8, + Float = 9, + Double = 10, + Bool = 11, + String = 12, +}; + +std::string toString(DataType dataType); +std::ostream& operator<<(std::ostream& out, DataType dataType); + +/// Returns the byte width of the C++ value materialized for a data type. +/// String values materialize as std::string_view. +uint32_t decodedValueWidth(DataType dataType); + +// General string compression. Make sure values here match those in the footer +// specification +enum class CompressionType : uint8_t { + Uncompressed = 0, + // Zstd doesn't require us to externally store level or any other info. + Zstd = 1, + MetaInternal = 2, + Lz4 = 3, + OpenZL = 4, +}; + +std::string toString(CompressionType compressionType); +/// Returns the compression type for 'name'. Throws if 'name' is unknown. +CompressionType toCompressionType(std::string_view name); +std::ostream& operator<<(std::ostream& out, CompressionType compressionType); + +enum class ChecksumType : uint8_t { XXH3_64 = 0 }; + +std::string toString(ChecksumType type); + +/// A CompressionType and any type-specific configuration params. +struct CompressionParams { + CompressionType type; + + /// For zstd. + int zstdLevel = 1; + + /// Keep the compressed result only when compressedSize <= rawSize * + /// acceptRatio; otherwise fall back to Uncompressed. + float acceptRatio = 0.8f; +}; + +// Parameters controlling the search for the optimal encoding on a data set. +struct OptimalSearchParams { + // Whether recursive structures are allowed. E.g. a encoding may use another + // encoding as a subencoding, and may use the encoding factory to find the + // best subencoding. We must terminate the recursion at some depth. With the + // default of 1 allowed recursion the top level encoding may use recursive + // encodings, but its subencodings may not. + int allowedRecursions = 1; + + // Whether to log debug info during the search (such as estimated sizes of + // each encoding considered, etc.); + bool logSearch = false; + + // Helps align log messages when log_search=true to help distinguish + // subencoding log messages from higher-level ones. + int logDepth = 0; + + // Entropy encodings, such as HuffmanEncoding, can be much more compact than + // others, but are quite a bit slower. However, compared to applying a general + // string compression on top of another encoding they are relatively fast. + // In general the entropy encodings will also be smaller than GSC on top of + // a non-entropy encoding. + bool enableEntropyEncodings = true; + + // For some dimension encodings that will frequently be grouped by, we may + // want to force dictionary enabled encodings so grouping by that can will be + // fast. + bool requireDictionaryEnabled = false; +}; + +template +struct TypeTraits {}; + +template <> +struct TypeTraits { + using physicalType = uint8_t; + using sumType = int64_t; + static constexpr DataType dataType = DataType::Int8; +}; + +template <> +struct TypeTraits { + using physicalType = uint8_t; + using sumType = uint64_t; + static constexpr DataType dataType = DataType::Uint8; +}; + +template <> +struct TypeTraits { + using physicalType = uint16_t; + using sumType = int64_t; + static constexpr DataType dataType = DataType::Int16; +}; + +template <> +struct TypeTraits { + using physicalType = uint16_t; + using sumType = uint64_t; + static constexpr DataType dataType = DataType::Uint16; +}; + +template <> +struct TypeTraits { + using physicalType = uint32_t; + using sumType = int64_t; + static constexpr DataType dataType = DataType::Int32; +}; + +template <> +struct TypeTraits { + using physicalType = uint32_t; + using sumType = uint64_t; + static constexpr DataType dataType = DataType::Uint32; +}; + +template <> +struct TypeTraits { + using physicalType = uint64_t; + using sumType = int64_t; + static constexpr DataType dataType = DataType::Int64; +}; + +template <> +struct TypeTraits { + using physicalType = uint64_t; + using sumType = uint64_t; + static constexpr DataType dataType = DataType::Uint64; +}; + +template <> +struct TypeTraits { + using physicalType = uint32_t; + using sumType = double; + static constexpr DataType dataType = DataType::Float; +}; + +template <> +struct TypeTraits { + using physicalType = uint64_t; + using sumType = double; + static constexpr DataType dataType = DataType::Double; +}; + +template <> +struct TypeTraits { + using physicalType = bool; + static constexpr DataType dataType = DataType::Bool; +}; + +template <> +struct TypeTraits { + using physicalType = std::string; + static constexpr DataType dataType = DataType::String; +}; + +template <> +struct TypeTraits { + using physicalType = std::string_view; + static constexpr DataType dataType = DataType::String; +}; + +template +constexpr bool isOneByteIntegralType() { + return std::is_same_v || std::is_same_v; +} + +template +constexpr bool isTwoByteIntegralType() { + return std::is_same_v || std::is_same_v; +} + +template +constexpr bool isFourByteIntegralType() { + return std::is_same_v || std::is_same_v; +} + +template +constexpr bool isEightByteIntegralType() { + return std::is_same_v || std::is_same_v; +} + +template +constexpr bool isSignedIntegralType() { + return std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v; +} + +template +constexpr bool isUnsignedIntegralType() { + return std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v; +} + +template +constexpr bool isIntegralType() { + return isSignedIntegralType() || isUnsignedIntegralType(); +} + +template +constexpr bool isFloatingPointType() { + return std::is_same_v || std::is_same_v; +} + +template +constexpr bool isNumericType() { + return isIntegralType() || isFloatingPointType(); +} + +template +constexpr bool isStringType() { + return std::is_same_v || std::is_same_v; +} + +template +constexpr bool isSharedDictionaryType() { + return isIntegralType() || std::is_same_v; +} + +constexpr bool isSharedDictionaryType(DataType dataType) { + switch (dataType) { + case DataType::Int8: + case DataType::Uint8: + case DataType::Int16: + case DataType::Uint16: + case DataType::Int32: + case DataType::Uint32: + case DataType::Int64: + case DataType::Uint64: + case DataType::String: + return true; + case DataType::Undefined: + case DataType::Float: + case DataType::Double: + case DataType::Bool: + return false; + } + return false; +} + +template +constexpr bool isBoolType() { + return std::is_same_v; +} + +} // namespace facebook::nimble + +template <> +struct fmt::formatter + : formatter { + auto format(facebook::nimble::CompressionType s, format_context& ctx) const { + return formatter::format(facebook::nimble::toString(s), ctx); + } +}; + +template <> +struct fmt::formatter : formatter { + auto format(facebook::nimble::EncodingType s, format_context& ctx) const { + return formatter::format(facebook::nimble::toString(s), ctx); + } +}; + +template <> +struct fmt::formatter : formatter { + auto format(facebook::nimble::DataType s, format_context& ctx) const { + return formatter::format(facebook::nimble::toString(s), ctx); + } +}; diff --git a/velox/dwio/nimble/common/Varint.cpp b/velox/dwio/nimble/common/Varint.cpp new file mode 100644 index 00000000000..8bcc63fa6cd --- /dev/null +++ b/velox/dwio/nimble/common/Varint.cpp @@ -0,0 +1,563 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifdef __x86_64__ +#include +#define VARINT_SVE2 0 +#elif defined(__aarch64__) +#include +#if defined(__ARM_FEATURE_SVE) && defined(__ARM_FEATURE_SVE2_BITPERM) && \ + __has_include() +#include +#include +#define VARINT_SVE2 1 +#else +#define VARINT_SVE2 0 +#endif +#endif //__x86_64__ + +#include +#include + +#include + +#include "folly/Likely.h" +#include "velox/dwio/nimble/common/Exceptions.h" +#include "velox/dwio/nimble/common/Varint.h" + +namespace facebook::nimble::varint { + +const char* bulkVarintSkip(uint64_t n, const char* pos) { + const uint64_t* word = reinterpret_cast(pos); + while (n >= 8) { + // Zeros in the 8 * ith bits indicate termination of a varint. + n -= __builtin_popcountll(~(*word++) & 0x8080808080808080ULL); + } + pos = reinterpret_cast(word); + while (n--) { + skipVarint(&pos); + } + return pos; +} + +uint64_t bulkVarintSize32(std::span values) { + constexpr uint8_t kLookupSizeTable32[32] = {5, 5, 5, 5, 4, 4, 4, 4, 4, 4, 4, + 3, 3, 3, 3, 3, 3, 3, 2, 2, 2, 2, + 2, 2, 2, 1, 1, 1, 1, 1, 1, 1}; + uint64_t size = 0; + for (uint32_t value : values) { + size += kLookupSizeTable32[__builtin_clz(value | 1U)]; + } + return size; +} + +uint64_t bulkVarintSize64(std::span values) { + constexpr uint8_t kLookupSizeTable64[64] = { + 10, 9, 9, 9, 9, 9, 9, 9, 8, 8, 8, 8, 8, 8, 8, 7, 7, 7, 7, 7, 7, 7, + 6, 6, 6, 6, 6, 6, 6, 5, 5, 5, 5, 5, 5, 5, 4, 4, 4, 4, 4, 4, 4, 3, + 3, 3, 3, 3, 3, 3, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1}; + uint64_t size = 0; + for (uint64_t value : values) { + size += kLookupSizeTable64[__builtin_clzll(value | 1ULL)]; + } + return size; +} + +#if !VARINT_SVE2 + +// Declaration of the function we build via generated code below. +template +#ifdef __x86_64__ +__attribute__((__target__("bmi2"))) +#endif +// __attribute__ ((optimize("Os"))) +const char* bulkVarintDecodeBmi2(uint64_t n, const char* pos, T* output); + +#endif // !VARINT_SVE2 + +// Zero-extend 8 consecutive bytes into T-sized output elements using xsimd +// batch construction and store. +template +inline void expandByteWord(const uint8_t* bytes, T* output) { + using batch_type = xsimd::batch; + constexpr auto kBatchSize = batch_type::size; + + if constexpr (kBatchSize >= 8) { + batch_type( + static_cast(bytes[0]), + static_cast(bytes[1]), + static_cast(bytes[2]), + static_cast(bytes[3]), + static_cast(bytes[4]), + static_cast(bytes[5]), + static_cast(bytes[6]), + static_cast(bytes[7])) + .store_unaligned(output); + } else if constexpr (kBatchSize == 4) { + batch_type( + static_cast(bytes[0]), + static_cast(bytes[1]), + static_cast(bytes[2]), + static_cast(bytes[3])) + .store_unaligned(output); + batch_type( + static_cast(bytes[4]), + static_cast(bytes[5]), + static_cast(bytes[6]), + static_cast(bytes[7])) + .store_unaligned(output + 4); + } else if constexpr (kBatchSize == 2) { + batch_type(static_cast(bytes[0]), static_cast(bytes[1])) + .store_unaligned(output); + batch_type(static_cast(bytes[2]), static_cast(bytes[3])) + .store_unaligned(output + 2); + batch_type(static_cast(bytes[4]), static_cast(bytes[5])) + .store_unaligned(output + 4); + batch_type(static_cast(bytes[6]), static_cast(bytes[7])) + .store_unaligned(output + 6); + } +} + +// Process runs of single-byte varints using xsimd for both the high-bit +// check and byte-to-element widening. Works with uint8_t* throughout, +// avoiding reinterpret_cast to uint64_t* (alignment/strict-aliasing issues). +// Returns the number of elements remaining after processing. +template +inline uint64_t +bulkDecodeSingleByteRun(uint64_t n, const char*& pos, T*& output) { + using u8_batch = xsimd::batch; + constexpr auto kU8Size = u8_batch::size; + constexpr uint64_t wordSize = 8; + constexpr uint64_t kHighBits = 0x8080808080808080ULL; + + const auto* src = reinterpret_cast(pos); + uint64_t remaining = n; + + // Process kU8BatchSize bytes at a time. + // Single wide load + vptest + while (remaining >= kU8Size) { + auto bytes = u8_batch::load_unaligned(src); + if (xsimd::any((bytes & u8_batch(0x80)) != u8_batch(0))) { + break; + } + for (size_t i = 0; i < kU8Size; i += wordSize) { + expandByteWord(src + i, output + i); + } + src += kU8Size; + output += kU8Size; + remaining -= kU8Size; + } + + // Process 8 bytes at a time. Use memcpy for the high-bit check to avoid + // reinterpret_cast strict-aliasing/alignment issues. + while (remaining >= wordSize) { + uint64_t word; + std::memcpy(&word, src, sizeof(word)); + + if (word & kHighBits) { + break; + } + expandByteWord(src, output); + src += wordSize; + output += wordSize; + remaining -= wordSize; + } + + // Handle trailing single-byte varints one at a time. + while (remaining > 0 && !(src[0] & 0x80)) { + *output++ = static_cast(src[0]); + ++src; + --remaining; + } + + pos = reinterpret_cast(src); + return remaining; +} + +constexpr std::size_t kCacheLineBytes = 64; +constexpr std::size_t kMaxControlBitsValue = 64; +constexpr std::size_t kMaskLength = 6; + +// Process runs of two-byte varints. Each 2-byte varint has continuation bit +// set on byte 0 and clear on byte 1. We detect this pattern in 8-byte words +// and decode 4 varints at a time using simple scalar ops. +// Returns the number of elements remaining after processing. +template +inline uint64_t bulkDecodeTwoByteRun(uint64_t n, const char*& pos, T*& output) { + auto remaining = n; + const auto* src = reinterpret_cast(pos); + + // In a run of 2-byte varints, each 8-byte word has alternating high bits: + // bytes 0,2,4,6 have 0x80 set (continuation), bytes 1,3,5,7 have 0x80 + // clear (terminator). This gives the pattern 0x0080008000800080. + constexpr uint64_t kHighBits = 0x8080808080808080ULL; + constexpr uint64_t kTwoBytePattern = 0x0080008000800080ULL; + constexpr uint64_t kWordSize = 8; + + // Process 8 bytes at a time (4 two-byte varints). + while (remaining >= 4) { + uint64_t word; + std::memcpy(&word, src, sizeof(word)); + + if ((word & kHighBits) != kTwoBytePattern) { + break; + } + + output[0] = static_cast((src[0] & 0x7f) | (uint32_t(src[1]) << 7)); + output[1] = static_cast((src[2] & 0x7f) | (uint32_t(src[3]) << 7)); + output[2] = static_cast((src[4] & 0x7f) | (uint32_t(src[5]) << 7)); + output[3] = static_cast((src[6] & 0x7f) | (uint32_t(src[7]) << 7)); + + src += kWordSize; + output += 4; + remaining -= 4; + } + + // Handle trailing 2-byte varints one at a time. After bulkDecodeSingleByteRun + // we know the first byte has high bit set (otherwise it would have been + // consumed as a 1-byte varint), so we just check that the second byte is a + // terminator. + while (remaining > 0 && (src[0] & 0x80) && !(src[1] & 0x80)) { + *output++ = static_cast((src[0] & 0x7f) | (uint32_t(src[1]) << 7)); + src += 2; + --remaining; + } + + pos = reinterpret_cast(src); + return remaining; +} + +// Lookup table entry for table-driven BMI2 varint decode. +struct alignas(kCacheLineBytes) VarintLookupEntry { + // Extraction masks for up to 6 completed varints. Unused slots are 0 + uint64_t valueMasks[kMaskLength]; + // Extraction mask for carryover bytes (partial varint at end of chunk), and + // zero when the chunk ends on a clean varint boundary. + uint64_t carryOverMask; + uint8_t numCompleted; + uint8_t carryOverBits; + uint8_t padding[6]; +}; + +static_assert( + sizeof(VarintLookupEntry) == kCacheLineBytes, + "Must fit one cache line"); + +// We build the full 64-entry/kMaxControlBitsValue for control bit lookup table +// at compile time. +static constexpr auto kDecodeTable = [] { + std::array table{}; + for (int i = 0; i < kMaxControlBitsValue; ++i) { + VarintLookupEntry entry{}; + uint64_t currentMask = 0; + + int lastZero = -1, numCompleted = 0; + const uint8_t controlBits = static_cast(i); + for (int j = 0; j < kMaskLength; ++j) { + currentMask |= uint64_t(0x7f) << (j * 8); + if (!((controlBits >> j) & 1)) { + entry.valueMasks[numCompleted] = currentMask; + ++numCompleted; + + lastZero = j; + currentMask = 0; + } + } + + entry.numCompleted = static_cast(numCompleted); + + entry.carryOverMask = 0; + entry.carryOverBits = 0; + + if (lastZero < 5) { + // Partial varint at end of chunk. + entry.carryOverMask = currentMask; + entry.carryOverBits = static_cast(7 * (5 - lastZero)); + } else if (lastZero == -1) { + // All 6 bytes are continuation bytes (case 63). Accumulate carryover. + // This is a rare case + entry.carryOverMask = currentMask; + entry.carryOverBits = 42; + } + table[i] = entry; + } + return table; +}(); + +unsigned long long pext_u64(unsigned long long __X, unsigned long long __M) { +#ifdef __x86_64__ + return _pext_u64(__X, __M); +#else + unsigned long p = 0x4040404040404040UL; // initial bit permute control + const unsigned long mask = 0x8000000000000000UL; + unsigned long m = __M; + unsigned long c; + unsigned long result; + + p = 64 - __builtin_popcountl(__M); + result = 0; + /* We could a use a for loop here, but that combined with + -funroll-loops can expand to a lot of code. The while + loop avoids unrolling and the compiler commons the xor + from clearing the mask bit with the (m != 0) test. The + result is a more compact loop setup and body. */ + while (m != 0) { + unsigned long t; + c = __builtin_clzl(m); + t = (__X & (mask >> c)) >> (p - c); + m ^= (mask >> c); + result |= (t); + p++; + } + return (result); +#endif +} + +#if !VARINT_SVE2 + +// Table-driven BMI2 varint decode. Reads extraction masks from a lookup table. +// Takes n and output by reference so the caller can re-dispatch to fast paths +// after this function yields on a single-byte or two-byte run boundary. +template +#ifdef __x86_64__ +__attribute__((__target__("bmi2"))) +#endif +const char* +bulkVarintDecodeBmi2Table(uint64_t& n, const char* pos, T*& output) { + constexpr uint64_t kControlMask = 0x0000808080808080ULL; + constexpr int kChunkLen = 6; + // Control bit pattern for uniform single-byte chunks. + // cb=0: all 6 bytes are terminators (6 single-byte varints). + constexpr uint64_t kAllSingleByteCb = 0; + + uint64_t carryover = 0; + int carryoverBits = 0; + pos -= kChunkLen; + + while (n >= 8) { + pos += kChunkLen; + + uint64_t word; + std::memcpy(&word, pos, sizeof(word)); + const uint64_t cb = pext_u64(word, kControlMask); + + // If there is no carryover from a previous chunk and we see a run of + // single-byte varints, break out so the caller can use the dedicated + // fast path. + if (carryoverBits == 0 && cb == kAllSingleByteCb) { + return pos; + } + + // Case 63 (all continuation bytes) requires accumulating carryover + // rather than replacing it. This case is extremely rare + if (FOLLY_UNLIKELY(cb == 63)) { + carryover |= pext_u64(word, 0x00007f7f7f7f7f7fULL) << carryoverBits; + carryoverBits += 42; + continue; + } + + const auto& info = kDecodeTable[cb]; + + // Extract and store up to 6 values. Unused mask slots are 0, producing + // harmless zero writes that will be overwritten by subsequent iterations + output[0] = static_cast( + (pext_u64(word, info.valueMasks[0]) << carryoverBits) | carryover); + output[1] = static_cast(pext_u64(word, info.valueMasks[1])); + output[2] = static_cast(pext_u64(word, info.valueMasks[2])); + output[3] = static_cast(pext_u64(word, info.valueMasks[3])); + output[4] = static_cast(pext_u64(word, info.valueMasks[4])); + output[5] = static_cast(pext_u64(word, info.valueMasks[5])); + + output += info.numCompleted; + n -= info.numCompleted; + + // Update carryover. When carryoverMask is 0, _pext returns 0 and + // carryoverBits is 0, effectively clearing the carryover state. + carryover = pext_u64(word, info.carryOverMask); + carryoverBits = info.carryOverBits; + } + + pos += kChunkLen; + if (n > 0) { + if constexpr (std::is_same_v) { + *output++ = readVarint32(&pos) << carryoverBits | carryover; + for (uint64_t i = 1; i < n; ++i) { + *output++ = readVarint32(&pos); + } + } else { + *output++ = readVarint64(&pos) << carryoverBits | carryover; + for (uint64_t i = 1; i < n; ++i) { + *output++ = readVarint64(&pos); + } + } + n = 0; + } + return pos; +} + +#else + +// Table-driven BEXT varint decode. Reads extraction masks from a lookup table. +// Takes n and output by reference so the caller can re-dispatch to fast paths +// after this function yields on a single-byte or two-byte run boundary. +template +const char* +bulkVarintDecodeBmi2Table(uint64_t& n, const char* pos, T*& output) { + svuint64_t kControlMask = svdup_n_u64(0x0000808080808080ULL); + svuint64_t kCarryoverMask = svdup_n_u64(0x00007f7f7f7f7f7fULL); + constexpr int kChunkLen = 6; + // Control bit pattern for uniform single-byte chunks. + // cb=0: all 6 bytes are terminators (6 single-byte varints). + constexpr double kAllSingleByteCb = 0.0; + + svuint64_t carryover = svdup_n_u64(0); + svuint64_t carryoverBits = svdup_n_u64(0); + pos -= kChunkLen; + + while (n >= 8) { + pos += kChunkLen; + + svuint64_t svWord = svdup_n_u64(*reinterpret_cast(pos)); + uint64x2_t cb = svget_neonq(svbext_u64(svWord, kControlMask)); + float64x2_t fb = vreinterpretq_f64_u64(cb); + float64x2_t fob = vreinterpretq_f64_u64(svget_neonq(carryoverBits)); + + // If there is no carryover from a previous chunk and we see a run of + // single-byte varints, break out so the caller can use the dedicated + // fast path. + if ((fob[0] == 0.0) && (fb[0] == kAllSingleByteCb)) { + return pos; + } + + uint64_t gb = cb[0]; + + // Case 63 (all continuation bytes) requires accumulating carryover + // rather than replacing it. This case is extremely rare + if (FOLLY_UNLIKELY(gb == 63)) { + svuint64_t cbits = svbext_u64(svWord, kCarryoverMask); + carryover |= cbits << carryoverBits; + carryoverBits += 42; + continue; + } + + const auto& info = kDecodeTable[gb]; + + // Extract and store up to 6 values. Unused mask slots are 0, producing + // harmless zero writes that will be overwritten by subsequent iterations + + svuint64_t valueMasks0 = + svset_neonq(svundef_u64(), vld1q_u64(&(info.valueMasks[0]))); + svuint64_t valueMasks1 = + svset_neonq(svundef_u64(), vld1q_u64(&(info.valueMasks[2]))); + svuint64_t valueMasks2 = + svset_neonq(svundef_u64(), vld1q_u64(&(info.valueMasks[4]))); + + svuint64_t svo0 = svbext_u64(svWord, valueMasks0); + svuint64_t svo1 = svbext_u64(svWord, valueMasks1); + svuint64_t svo2 = svbext_u64(svWord, valueMasks2); + + uint64x2_t vo0 = svget_neonq(svo0); + uint64x2_t vo1 = svget_neonq(svo1); + uint64x2_t vo2 = svget_neonq(svo2); + + auto outputPtr = output; + + uint64x2_t cvo0 = svget_neonq((svo0 << carryoverBits) | carryover); + + if constexpr (sizeof(T) == 8) { + vst1q_lane_u64(outputPtr, cvo0, 0); + vst1q_lane_u64(outputPtr + 1, vo0, 1); + vst1q_u64(outputPtr + 2, vo1); + vst1q_u64(outputPtr + 4, vo2); + } else if constexpr (sizeof(T) == 4) { + vst1q_lane_u32(outputPtr, vreinterpretq_u32_u64(cvo0), 0); + vst1q_lane_u32(outputPtr + 1, vreinterpretq_u32_u64(vo0), 2); + svst1w_u64(svwhilelt_b32_u64(0, 4), outputPtr + 2, svo1); + svst1w_u64(svwhilelt_b32_u64(0, 4), outputPtr + 4, svo2); + } else if constexpr (sizeof(T) == 2) { + vst1q_lane_u16(outputPtr, vreinterpretq_u16_u64(cvo0), 0); + vst1q_lane_u16(outputPtr + 1, vreinterpretq_u16_u64(vo0), 4); + svst1h_u64(svwhilelt_b16_u64(0, 8), outputPtr + 2, svo1); + svst1h_u64(svwhilelt_b16_u64(0, 8), outputPtr + 4, svo2); + } else if constexpr (sizeof(T) == 1) { + vst1q_lane_u8(outputPtr, vreinterpretq_u8_u64(cvo0), 0); + vst1q_lane_u8(outputPtr + 1, vreinterpretq_u8_u64(vo0), 8); + svst1b_u64(svwhilelt_b8_u64(0, 16), outputPtr + 2, svo1); + svst1b_u64(svwhilelt_b8_u64(0, 16), outputPtr + 4, svo2); + } + + output += info.numCompleted; + n -= info.numCompleted; + + // Update carryover. When carryoverMask is 0, _pext returns 0 and + // carryoverBits is 0, effectively clearing the carryover state. + carryover = svbext_u64(svWord, svdup_n_u64(info.carryOverMask)); + carryoverBits = svdup_n_u64(info.carryOverBits); + } + + uint64_t carryoverVar = carryover[0]; + uint64_t carryoverBitsVar = carryoverBits[0]; + + pos += kChunkLen; + if (n > 0) { + if constexpr (std::is_same_v) { + *output++ = readVarint32(&pos) << carryoverBitsVar | carryoverVar; + for (uint64_t i = 1; i < n; ++i) { + *output++ = readVarint32(&pos); + } + } else { + *output++ = readVarint64(&pos) << carryoverBitsVar | carryoverVar; + for (uint64_t i = 1; i < n; ++i) { + *output++ = readVarint64(&pos); + } + } + n = 0; + } + return pos; +} + +#endif + +// Dispatch loop: cycles between fast paths and the general BMI2 decoder. +// When the BMI2 decoder detects a uniform single-byte or two-byte chunk +// boundary (with no carryover), it yields back here so the dedicated fast +// paths can handle the run efficiently. +template +inline const char* +bulkVarintDecodeDispatch(uint64_t n, const char* pos, T* output) { + auto remaining = n; + while (remaining > 0) { + remaining = bulkDecodeSingleByteRun(remaining, pos, output); + if (remaining == 0) { + break; + } + remaining = bulkDecodeTwoByteRun(remaining, pos, output); + if (remaining == 0) { + break; + } + pos = bulkVarintDecodeBmi2Table(remaining, pos, output); + } + return pos; +} + +const char* bulkVarintDecode32(uint64_t n, const char* pos, uint32_t* output) { + return bulkVarintDecodeDispatch(n, pos, output); +} + +const char* bulkVarintDecode64(uint64_t n, const char* pos, uint64_t* output) { + return bulkVarintDecodeDispatch(n, pos, output); +} + +} // namespace facebook::nimble::varint diff --git a/velox/dwio/nimble/common/Varint.h b/velox/dwio/nimble/common/Varint.h new file mode 100644 index 00000000000..f76ecefce50 --- /dev/null +++ b/velox/dwio/nimble/common/Varint.h @@ -0,0 +1,145 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include + +// Varint-related encoding methods. Same binary encoding as folly/Varint.h +// but with a different API, a faster decode method thanks to +// to not checking bounds, and bulk decoding methods. No functions in this +// library check bounds or deal with buffer overflow. +// +// The bulk decoding methods are particularly interesting. They range from +// 4x faster than the non-bulk in the best case (all 1 byte varints) to +// about 50% faster in the worst case (random byte lengths), assuming that +// bmi2 is available. See varintBenchmark.h for some details. + +namespace facebook::nimble::varint { + +// Decode n varints at once. Makes use of bmi2 instruction set if its +// available. Returns pos updated past the n varints. +const char* bulkVarintDecode32(uint64_t n, const char* pos, uint32_t* output); +const char* bulkVarintDecode64(uint64_t n, const char* pos, uint64_t* output); + +// Skips n varints, returning pos updated past the n varints. +const char* bulkVarintSkip(uint64_t n, const char* pos); + +// Returns the number of bytes the |values| will occupy after varint encoding. +uint64_t bulkVarintSize32(std::span values); +uint64_t bulkVarintSize64(std::span values); + +/// Inline non-bulk methods follow below. + +/// Returns the number of bytes needed to varint-encode a single value. +/// Uses CLZ (count leading zeros) for O(1) computation instead of a loop. +template +inline constexpr uint32_t varintSize(T val) noexcept { + if constexpr (sizeof(T) <= 4) { + // `| 1` avoids undefined behavior from __builtin_clz(0) and correctly + // returns 1 byte for val=0. + uint32_t bitsNeeded = 32 - __builtin_clz(static_cast(val) | 1); + return (bitsNeeded + 6) / 7; + } else { + uint32_t bitsNeeded = 64 - __builtin_clzll(static_cast(val) | 1); + return (bitsNeeded + 6) / 7; + } +} + +/// Returns the maximum number of bytes needed to varint-encode any value that +/// fits in `bitWidth` bits. +inline constexpr uint32_t maxVarintSizeForBitWidth(uint32_t bitWidth) noexcept { + return (std::min(bitWidth, 64u) + 6) / 7; +} + +template +inline void writeVarint(T val, char** pos) noexcept { + while (val >= 128) { + *((*pos)++) = 0x80 | (val & 0x7f); + val >>= 7; + } + *((*pos)++) = val; +} + +inline void skipVarint(const char** pos) noexcept { + while (*((*pos)++) & 128) { + } +} + +inline uint32_t readVarint32(const char** pos) noexcept { + uint32_t value = (**pos) & 127; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= (**pos & 127) << 7; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= (**pos & 127) << 14; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= (**pos & 127) << 21; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= (*((*pos)++) & 127) << 28; + return value; +} + +inline uint64_t readVarint64(const char** pos) noexcept { + uint64_t value = (**pos) & 127; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= (**pos & 127) << 7; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= (**pos & 127) << 14; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= (**pos & 127) << 21; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= static_cast(**pos & 127) << 28; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= static_cast(**pos & 127) << 35; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= static_cast(**pos & 127) << 42; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= static_cast(**pos & 127) << 49; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= static_cast(**pos & 127) << 56; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= static_cast(*((*pos)++) & 127) << 63; + return value; +} + +} // namespace facebook::nimble::varint diff --git a/velox/dwio/nimble/common/Vector.h b/velox/dwio/nimble/common/Vector.h new file mode 100644 index 00000000000..02d28be4ac5 --- /dev/null +++ b/velox/dwio/nimble/common/Vector.h @@ -0,0 +1,502 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "velox/buffer/Buffer.h" +#include "velox/buffer/BufferPool.h" +#include "velox/common/memory/Memory.h" +#include "velox/dwio/nimble/common/Exceptions.h" + +#include +#include +#include +#include +#include + +namespace facebook::nimble { + +/// A vector-like container similar to std::vector, but without the edge case +/// for booleans. Unlike std::vector, data() returns T* for all T, +/// allowing implicit conversion to std::span for all T. +template +class Vector { + using InnerType = + typename std::conditional, uint8_t, T>::type; + + public: + /// Constructs a vector with the given size, filled with the specified value. + Vector(velox::memory::MemoryPool* pool, size_t size, T value) : pool_{pool} { + init(size); + std::fill(dataRawPtr_, dataRawPtr_ + size_, value); + } + + /// Constructs a vector with the given size, with uninitialized elements. + Vector(velox::memory::MemoryPool* pool, size_t size) : pool_{pool} { + init(size); + } + + /// Constructs an empty vector. + explicit Vector(velox::memory::MemoryPool* pool) : pool_{pool} { + capacity_ = 0; + size_ = 0; + data_ = nullptr; + dataRawPtr_ = nullptr; +#ifndef NDEBUG + dataRawPtr_ = placeholder_.data(); +#endif + } + + /// Constructs a vector by adopting an existing buffer. The pool is extracted + /// from the buffer for future allocations. Size is set to 0 (no valid data). + explicit Vector(velox::BufferPtr buf) + : pool_{buf->pool()}, + data_{std::move(buf)}, + capacity_{data_->capacity() / sizeof(InnerType)}, + size_{0}, + dataRawPtr_{reinterpret_cast(data_->asMutable())} {} + + /// Constructs a vector from an iterator range. + template + Vector(velox::memory::MemoryPool* pool, It first, It last) : pool_{pool} { + auto size = last - first; + init(size); + std::copy(first, last, dataRawPtr_); + } + + /// Copy constructor. + Vector(const Vector& other) { + *this = other; + } + + /// Copy assignment operator. + Vector& operator=(const Vector& other) { + if (this != &other) { + size_ = other.size(); + capacity_ = other.capacity_; + pool_ = other.pool_; + allocateBuffer(); + std::copy(other.dataRawPtr_, other.dataRawPtr_ + size_, dataRawPtr_); + } + return *this; + } + + /// Move constructor. + Vector(Vector&& other) noexcept { + *this = std::move(other); + } + + /// Move assignment operator. + Vector& operator=(Vector&& other) noexcept { + if (this != &other) { + size_ = other.size(); + capacity_ = other.capacity_; + data_ = std::move(other.data_); +#ifndef NDEBUG + dataRawPtr_ = placeholder_.data(); +#endif + if (data_ != nullptr) { + dataRawPtr_ = reinterpret_cast(data_->asMutable()); + } + pool_ = other.pool_; + other.size_ = 0; + other.capacity_ = 0; + } + return *this; + } + + /// Constructs a vector from an initializer list. + Vector(velox::memory::MemoryPool* pool, std::initializer_list l) + : pool_{pool} { + init(l.size()); + std::copy(l.begin(), l.end(), dataRawPtr_); + } + + /// Returns the memory pool used by this vector. + inline velox::memory::MemoryPool* pool() { + return pool_; + } + + /// Returns the number of elements in the vector. + uint64_t size() const { + return size_; + } + + /// Returns true if the vector is empty. + bool empty() const { + return size_ == 0; + } + + /// Returns the current capacity of the vector. + uint64_t capacity() const { + return capacity_; + } + /// Returns a reference to the element at the given index. + T& operator[](uint64_t i) { + return dataRawPtr_[i]; + } + + /// Returns a const reference to the element at the given index. + const T& operator[](uint64_t i) const { + return dataRawPtr_[i]; + } + + /// Returns a pointer to the first element. + T* begin() { + return dataRawPtr_; + } + + /// Returns a pointer past the last element. + T* end() { + return dataRawPtr_ + size_; + } + + /// Returns a const pointer to the first element. + const T* begin() const { + return dataRawPtr_; + } + + /// Returns a const pointer past the last element. + const T* end() const { + return dataRawPtr_ + size_; + } + + /// Returns a reference to the last element. + T& back() { + return dataRawPtr_[size_ - 1]; + } + + /// Returns a const reference to the last element. + const T& back() const { + return dataRawPtr_[size_ - 1]; + } + + /// Directly updates the size to the given value. + /// Useful if you've filled in data directly using the underlying raw + /// pointers. + void update_size(uint64_t size) { + size_ = size; + } + + /// Fills all elements with the default value T(). + void zero_out() { + std::fill(dataRawPtr_, dataRawPtr_ + size_, T()); + } + + /// Fills all elements with the given value. + void fill(T value) { + std::fill(dataRawPtr_, dataRawPtr_ + size_, value); + } + + /// Resets the vector to an empty state, releasing allocated memory. + void clear() { + capacity_ = 0; + size_ = 0; + data_.reset(); + dataRawPtr_ = nullptr; + } + + /// Inserts elements from [inputStart, inputEnd) at the given output + /// position. + void insert(T* output, const T* inputStart, const T* inputEnd) { + const uint64_t inputSize = inputEnd - inputStart; + const uint64_t distanceToEnd = end() - output; + if (inputSize > distanceToEnd) { + const uint64_t sizeChange = inputSize - distanceToEnd; + const uint64_t distanceFromBegin = output - begin(); + resize(size_ + sizeChange); + std::move(inputStart, inputEnd, begin() + distanceFromBegin); + } else { + std::move(inputStart, inputEnd, output); + } + } + + /// Appends the given number of copies of value to the end of the vector. + void extend(uint64_t copies, T value) { + reserve(size_ + copies); + std::fill(end(), end() + copies, value); + size_ += copies; + } + + /// Returns a pointer to the underlying data. + T* data() noexcept { + return dataRawPtr_; + } + + /// Returns a const pointer to the underlying data. + const T* data() const noexcept { + return dataRawPtr_; + } + + /// Appends the given value to the end of the vector. + void push_back(T value) { + if (size_ == capacity_) { + reserve(calculateNewSize(capacity_)); + } + dataRawPtr_[size_] = value; + ++size_; + } + + /// Constructs an element in-place at the end of the vector. + template + void emplace_back(Args&&... args) { + if (size_ == capacity_) { + reserve(calculateNewSize(capacity_)); + } + // use placement new to construct the object. + new (dataRawPtr_ + size_) InnerType(std::forward(args)...); + ++size_; + } + + /// Ensures the vector can hold at least the given number of elements. + /// Does NOT shrink if size is less than current size, and does NOT + /// initialize any new elements. + void reserve(uint64_t size) { + if (size > capacity_) { + auto newData = + velox::AlignedBuffer::allocateExact(size, pool_); + // AlignedBuffer can allocate a bit more than requested for the alignment + // purpose, let's leverage that by using its true capacity. + capacity_ = newData->capacity() / sizeof(InnerType); + NIMBLE_DCHECK_GE( + capacity_, size, "Allocated capacity is smaller than requested"); + if (data_ != nullptr && size_ > 0) { + std::move( + dataRawPtr_, + dataRawPtr_ + size_, + reinterpret_cast(newData->template asMutable())); + } + data_ = std::move(newData); + dataRawPtr_ = reinterpret_cast(data_->asMutable()); + } + } + + /// Changes the size of the vector. + /// Does NOT shrink capacity, and does NOT initialize new elements. + void resize(uint64_t size) { + reserve(size); + size_ = size; + } + + /// Changes the size of the vector, initializing new elements to value. + /// Does NOT shrink capacity. + void resize(uint64_t newSize, const T& value) { + auto initialSize = size_; + resize(newSize); + + if (size_ > initialSize) { + std::fill(dataRawPtr_ + initialSize, end(), value); + } + } + + /// Releases ownership of the underlying buffer and returns it. + /// Sets the buffer's size to match the vector's logical size. + /// The vector is left in an empty state after this call. + velox::BufferPtr releaseOwnership() { + velox::BufferPtr tmp = std::move(data_); + tmp->setSize(size_); + capacity_ = 0; + size_ = 0; + data_ = nullptr; + dataRawPtr_ = nullptr; +#ifndef NDEBUG + dataRawPtr_ = placeholder_.data(); +#endif + return tmp; + } + + /// Releases the underlying buffer without setting its size. + /// Returns nullptr if the vector has no buffer. + /// The vector is left in an empty state after this call. + velox::BufferPtr releaseBuffer() { + auto tmp = std::move(data_); + capacity_ = 0; + size_ = 0; + dataRawPtr_ = nullptr; +#ifndef NDEBUG + dataRawPtr_ = placeholder_.data(); +#endif + return tmp; + } + + const velox::BufferPtr& testingBuffer() const { + return data_; + } + + private: + inline void init(size_t size) { + capacity_ = size; + size_ = size; + allocateBuffer(); + } + + inline size_t calculateNewSize(size_t size) { + auto newSize = size <<= 1; + if (newSize == 0) { + return velox::AlignedBuffer::kSizeofAlignedBuffer; + } + + return newSize; + } + + inline void allocateBuffer() { + data_ = velox::AlignedBuffer::allocateExact(capacity_, pool_); + dataRawPtr_ = reinterpret_cast(data_->asMutable()); + uint64_t newCapacity = data_->capacity() / sizeof(InnerType); + NIMBLE_DCHECK_GE( + newCapacity, capacity_, "Allocated capacity is smaller than requested"); + capacity_ = newCapacity; + } + + velox::memory::MemoryPool* pool_; + velox::BufferPtr data_; + uint64_t capacity_; + uint64_t size_; + T* dataRawPtr_; +#ifndef NDEBUG + inline static std::array placeholder_; +#endif +}; + +/// Owns a temporary Vector and returns its allocation to a BufferPool. +template +class ScopedVector { + public: + ScopedVector( + uint64_t size, + velox::memory::MemoryPool* pool, + velox::BufferPool* bufferPool) + : bufferPool_{bufferPool}, + vector_{acquireVectorBuffer(size, pool, bufferPool)} { + vector_.resize(size); + } + + ~ScopedVector() { + if (bufferPool_ != nullptr) { + auto buffer = vector_.releaseBuffer(); + if (buffer != nullptr) { + bufferPool_->release(std::move(buffer)); + } + } + } + + ScopedVector(const ScopedVector&) = delete; + ScopedVector& operator=(const ScopedVector&) = delete; + + operator Vector&() { + return vector_; + } + + operator const Vector&() const { + return vector_; + } + + Vector* operator->() { + return &vector_; + } + + const Vector* operator->() const { + return &vector_; + } + + Vector& operator*() { + return vector_; + } + + const Vector& operator*() const { + return vector_; + } + + V& operator[](uint64_t i) { + return vector_[i]; + } + + const V& operator[](uint64_t i) const { + return vector_[i]; + } + + V* begin() { + return vector_.begin(); + } + + V* end() { + return vector_.end(); + } + + const V* begin() const { + return vector_.begin(); + } + + const V* end() const { + return vector_.end(); + } + + V* data() { + return vector_.data(); + } + + const V* data() const { + return vector_.data(); + } + + uint64_t size() const { + return vector_.size(); + } + + bool empty() const { + return vector_.empty(); + } + + uint64_t capacity() const { + return vector_.capacity(); + } + + void reserve(uint64_t size) { + vector_.reserve(size); + } + + void resize(uint64_t size) { + vector_.resize(size); + } + + void push_back(V value) { + vector_.push_back(value); + } + + template + void emplace_back(Args&&... args) { + vector_.emplace_back(std::forward(args)...); + } + + private: + static Vector acquireVectorBuffer( + uint64_t size, + velox::memory::MemoryPool* pool, + velox::BufferPool* bufferPool) { + using InnerType = + typename std::conditional, uint8_t, V>::type; + if (bufferPool != nullptr && size > 0) { + if (auto buffer = bufferPool->get(size * sizeof(InnerType))) { + return Vector{std::move(buffer)}; + } + } + return Vector{pool}; + } + + velox::BufferPool* const bufferPool_; + Vector vector_; +}; + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/Zigzag.h b/velox/dwio/nimble/common/Zigzag.h new file mode 100644 index 00000000000..5205fbded73 --- /dev/null +++ b/velox/dwio/nimble/common/Zigzag.h @@ -0,0 +1,52 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +// ZigZag encoding maps signed integers to unsigned integers so that values +// with small absolute value have small encoded values, making them efficient +// for varint encoding. +// +// zigzagEncode32(0) == 0 +// zigzagEncode32(-1) == 1 +// zigzagEncode32(1) == 2 +// zigzagEncode32(-2) == 3 +// ... +// +// In general: +// if x >= 0, zigzagEncode32(x) == 2*x +// if x < 0, zigzagEncode32(x) == -2*x - 1 + +namespace facebook::nimble::zigzag { + +inline constexpr uint32_t zigzagEncode32(int32_t val) noexcept { + return static_cast((val << 1) ^ (val >> 31)); +} + +inline constexpr int32_t zigzagDecode32(uint32_t val) noexcept { + return static_cast((val >> 1) ^ -(val & 1)); +} + +inline constexpr uint64_t zigzagEncode64(int64_t val) noexcept { + return static_cast((val << 1) ^ (val >> 63)); +} + +inline constexpr int64_t zigzagDecode64(uint64_t val) noexcept { + return static_cast((val >> 1) ^ -(val & 1)); +} + +} // namespace facebook::nimble::zigzag diff --git a/velox/dwio/nimble/common/benchmarks/FixedBitArrayBenchmark.cpp b/velox/dwio/nimble/common/benchmarks/FixedBitArrayBenchmark.cpp new file mode 100644 index 00000000000..a2ecf822c06 --- /dev/null +++ b/velox/dwio/nimble/common/benchmarks/FixedBitArrayBenchmark.cpp @@ -0,0 +1,227 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "folly/Benchmark.h" +#include "folly/init/Init.h" +#include "velox/dwio/nimble/common/FixedBitArray.h" + +using namespace ::facebook; + +namespace { + +constexpr uint64_t kNumElements = 1000 * 1000; +constexpr uint64_t kBaseline = 12345; +constexpr uint64_t kStartOffset = 37; + +uint64_t maskForBitWidth(int bitWidth) { + return bitWidth == 64 ? ~0ULL : ((1ULL << bitWidth) - 1); +} + +std::vector makeValues(int bitWidth) { + const uint64_t mask = maskForBitWidth(bitWidth); + const uint64_t baseline = bitWidth == 64 ? 0 : kBaseline; + std::vector values(kNumElements); + for (uint64_t i = 0; i < kNumElements; ++i) { + const uint64_t residual = (i * 1000003ULL) & mask; + values[i] = residual + baseline; + } + return values; +} + +void clearBuffer(char* _Nonnull buffer, uint64_t bufferBytes) { + std::memset(buffer, 0, bufferBytes); +} + +void observeWrittenBuffer( + const char* buffer, + uint64_t elementCount, + int bitWidth) { + const uint64_t writtenBytes = + ((elementCount * static_cast(bitWidth)) + 7) >> 3; + folly::doNotOptimizeAway(buffer[writtenBytes - 1]); +} + +#define FIXED_BIT_ARRAY_BENCHMARKS(bitWidth) \ + BENCHMARK(BulkSet64WithBaseline_##bitWidth, iters) { \ + std::vector values; \ + std::unique_ptr buffer; \ + uint64_t bufferBytes; \ + BENCHMARK_SUSPEND { \ + values = makeValues(bitWidth); \ + bufferBytes = nimble::FixedBitArray::bufferSize(kNumElements, bitWidth); \ + buffer = std::make_unique(bufferBytes); \ + } \ + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); \ + const uint64_t baseline = bitWidth == 64 ? 0 : kBaseline; \ + while (iters--) { \ + BENCHMARK_SUSPEND { \ + clearBuffer(buffer.get(), bufferBytes); \ + } \ + fixedBitArray.bulkSetWithBaseline( \ + 0, kNumElements, values.data(), baseline); \ + observeWrittenBuffer(buffer.get(), kNumElements, bitWidth); \ + } \ + } \ + BENCHMARK_RELATIVE(ScalarSetWithBaseline_##bitWidth, iters) { \ + std::vector values; \ + std::unique_ptr buffer; \ + uint64_t bufferBytes; \ + BENCHMARK_SUSPEND { \ + values = makeValues(bitWidth); \ + bufferBytes = nimble::FixedBitArray::bufferSize(kNumElements, bitWidth); \ + buffer = std::make_unique(bufferBytes); \ + } \ + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); \ + const uint64_t baseline = bitWidth == 64 ? 0 : kBaseline; \ + while (iters--) { \ + BENCHMARK_SUSPEND { \ + clearBuffer(buffer.get(), bufferBytes); \ + } \ + const uint64_t* nextValue = values.data(); \ + for (uint64_t i = 0; i < kNumElements; ++i) { \ + fixedBitArray.set(i, *nextValue - baseline); \ + ++nextValue; \ + } \ + observeWrittenBuffer(buffer.get(), kNumElements, bitWidth); \ + } \ + } \ + BENCHMARK(BulkGet64WithBaseline_##bitWidth, iters) { \ + std::vector values; \ + std::unique_ptr buffer; \ + std::vector output; \ + BENCHMARK_SUSPEND { \ + values = makeValues(bitWidth); \ + const uint64_t bufferBytes = \ + nimble::FixedBitArray::bufferSize(kNumElements, bitWidth); \ + buffer = std::make_unique(bufferBytes); \ + clearBuffer(buffer.get(), bufferBytes); \ + output.resize(kNumElements); \ + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); \ + const uint64_t baseline = bitWidth == 64 ? 0 : kBaseline; \ + fixedBitArray.bulkSetWithBaseline( \ + 0, kNumElements, values.data(), baseline); \ + } \ + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); \ + const uint64_t baseline = bitWidth == 64 ? 0 : kBaseline; \ + while (iters--) { \ + fixedBitArray.bulkGetWithBaseline( \ + 0, kNumElements, output.data(), baseline); \ + folly::doNotOptimizeAway(*(output.data() + kNumElements - 1)); \ + } \ + } \ + BENCHMARK_RELATIVE(ScalarGetWithBaseline_##bitWidth, iters) { \ + std::vector values; \ + std::unique_ptr buffer; \ + std::vector output; \ + BENCHMARK_SUSPEND { \ + values = makeValues(bitWidth); \ + const uint64_t bufferBytes = \ + nimble::FixedBitArray::bufferSize(kNumElements, bitWidth); \ + buffer = std::make_unique(bufferBytes); \ + clearBuffer(buffer.get(), bufferBytes); \ + output.resize(kNumElements); \ + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); \ + const uint64_t baseline = bitWidth == 64 ? 0 : kBaseline; \ + fixedBitArray.bulkSetWithBaseline( \ + 0, kNumElements, values.data(), baseline); \ + } \ + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); \ + const uint64_t baseline = bitWidth == 64 ? 0 : kBaseline; \ + while (iters--) { \ + uint64_t* nextOutput = output.data(); \ + for (uint64_t i = 0; i < kNumElements; ++i) { \ + *nextOutput = fixedBitArray.get(i) + baseline; \ + ++nextOutput; \ + } \ + folly::doNotOptimizeAway(*(output.data() + kNumElements - 1)); \ + } \ + } + +#define FIXED_BIT_ARRAY_OFFSET_SET_BENCHMARKS(bitWidth) \ + BENCHMARK(BulkSet64WithBaselineOffset_##bitWidth, iters) { \ + std::vector values; \ + std::unique_ptr buffer; \ + uint64_t bufferBytes; \ + BENCHMARK_SUSPEND { \ + values = makeValues(bitWidth); \ + bufferBytes = nimble::FixedBitArray::bufferSize( \ + kStartOffset + kNumElements, bitWidth); \ + buffer = std::make_unique(bufferBytes); \ + } \ + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); \ + const uint64_t baseline = bitWidth == 64 ? 0 : kBaseline; \ + while (iters--) { \ + BENCHMARK_SUSPEND { \ + clearBuffer(buffer.get(), bufferBytes); \ + } \ + fixedBitArray.bulkSetWithBaseline( \ + kStartOffset, kNumElements, values.data(), baseline); \ + observeWrittenBuffer( \ + buffer.get(), kStartOffset + kNumElements, bitWidth); \ + } \ + } \ + BENCHMARK_RELATIVE(ScalarSetWithBaselineOffset_##bitWidth, iters) { \ + std::vector values; \ + std::unique_ptr buffer; \ + uint64_t bufferBytes; \ + BENCHMARK_SUSPEND { \ + values = makeValues(bitWidth); \ + bufferBytes = nimble::FixedBitArray::bufferSize( \ + kStartOffset + kNumElements, bitWidth); \ + buffer = std::make_unique(bufferBytes); \ + } \ + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); \ + const uint64_t baseline = bitWidth == 64 ? 0 : kBaseline; \ + while (iters--) { \ + BENCHMARK_SUSPEND { \ + clearBuffer(buffer.get(), bufferBytes); \ + } \ + const uint64_t* nextValue = values.data(); \ + for (uint64_t i = 0; i < kNumElements; ++i) { \ + fixedBitArray.set(kStartOffset + i, *nextValue - baseline); \ + ++nextValue; \ + } \ + observeWrittenBuffer( \ + buffer.get(), kStartOffset + kNumElements, bitWidth); \ + } \ + } + +FIXED_BIT_ARRAY_BENCHMARKS(16) +FIXED_BIT_ARRAY_BENCHMARKS(32) +FIXED_BIT_ARRAY_BENCHMARKS(33) +FIXED_BIT_ARRAY_BENCHMARKS(40) +FIXED_BIT_ARRAY_BENCHMARKS(48) +FIXED_BIT_ARRAY_BENCHMARKS(56) +FIXED_BIT_ARRAY_BENCHMARKS(57) +FIXED_BIT_ARRAY_BENCHMARKS(58) +FIXED_BIT_ARRAY_BENCHMARKS(60) +FIXED_BIT_ARRAY_BENCHMARKS(64) + +FIXED_BIT_ARRAY_OFFSET_SET_BENCHMARKS(32) +FIXED_BIT_ARRAY_OFFSET_SET_BENCHMARKS(40) +FIXED_BIT_ARRAY_OFFSET_SET_BENCHMARKS(48) +FIXED_BIT_ARRAY_OFFSET_SET_BENCHMARKS(56) +FIXED_BIT_ARRAY_OFFSET_SET_BENCHMARKS(64) + +} // namespace + +int main(int argc, char** argv) { + folly::Init init(&argc, &argv); + folly::runBenchmarks(); +} diff --git a/velox/dwio/nimble/common/benchmarks/VarintBenchmark.cpp b/velox/dwio/nimble/common/benchmarks/VarintBenchmark.cpp new file mode 100644 index 00000000000..ac3fa6fd27e --- /dev/null +++ b/velox/dwio/nimble/common/benchmarks/VarintBenchmark.cpp @@ -0,0 +1,281 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "folly/Benchmark.h" +#include "folly/Random.h" +#include "folly/Varint.h" +#include "velox/dwio/nimble/common/Varint.h" + +using namespace ::facebook; + +const int kNumElements = 1000 * 1000; + +// Basically same code as dwrf::IntDecoder::readVuLong. +uint64_t DwrfRead(const char** bufferStart, const char* bufferEnd) { + if (LIKELY(bufferEnd - *bufferStart >= folly::kMaxVarintLength64)) { + const char* p = *bufferStart; + uint64_t val; + do { + int64_t b; + b = *p++; + val = (b & 0x7f); + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 7; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 14; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 21; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 28; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 35; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 42; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 49; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 56; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x01) << 63; + if (LIKELY(b >= 0)) { + break; + } else { + throw std::runtime_error{"invalid encoding: likely corrupt data"}; + } + } while (false); + *bufferStart = p; + return val; + } else { + // this part isn't the same, but doesn't measurably effect time. + return nimble::varint::readVarint64(bufferStart); + } +} + +// Makes random data uniform over in bit width over 32 bits. +std::vector MakeUniformData(int num_elements = kNumElements) { + std::vector data(num_elements); + for (int i = 0; i < num_elements; ++i) { + const int bit_shift = 1 + folly::Random::secureRand32() % 32; + data[i] = folly::Random::secureRand32() % (1 << bit_shift); + } + return data; +} + +// Makes 95% 1 byte, 5% 2 byte data. +std::vector MakeSkewedData(int num_elements = kNumElements) { + std::vector data(num_elements); + for (int i = 0; i < num_elements; ++i) { + if (folly::Random::secureRand32() % 20) { + data[i] = folly::Random::secureRand32() % (1 << 7); + } else { + data[i] = folly::Random::secureRand32() % (1 << 14); + } + } + return data; +} + +// Makes data where all values fit in exactly `numBytes` varint bytes. +std::vector MakeFixedWidthData32( + int numBytes, + int num_elements = kNumElements) { + std::vector data(num_elements); + uint32_t lo = (numBytes == 1) ? 0 : (1u << (7 * (numBytes - 1))); + uint32_t hi = (1u << (7 * numBytes)) - 1; + if (numBytes == 5) { + hi = UINT32_MAX; + } + for (int i = 0; i < num_elements; ++i) { + data[i] = lo + folly::Random::secureRand32() % (hi - lo + 1); + } + return data; +} + +// Makes 64-bit data where all values fit in exactly `numBytes` varint bytes. +std::vector MakeFixedWidthData64( + int numBytes, + int num_elements = kNumElements) { + std::vector data(num_elements); + uint64_t lo = (numBytes == 1) ? 0 : (1ull << (7 * (numBytes - 1))); + uint64_t hi = (numBytes >= 10) ? UINT64_MAX : ((1ull << (7 * numBytes)) - 1); + for (int i = 0; i < num_elements; ++i) { + data[i] = lo + folly::Random::secureRand64() % (hi - lo + 1); + } + return data; +} + +// Encode data into a varint buffer, returns total encoded size. +template +std::unique_ptr EncodeData( + const std::vector& data, + uint64_t& encodedSize) { + auto buf = std::make_unique(data.size() * folly::kMaxVarintLength64); + char* pos = buf.get(); + for (auto val : data) { + nimble::varint::writeVarint(val, &pos); + } + encodedSize = pos - buf.get(); + return buf; +} + +// ============================================================================ +// Original benchmarks (uniform + skewed, 32-bit) +// ============================================================================ + +BENCHMARK(Encode, iters) { + std::vector data; + std::unique_ptr buf; + BENCHMARK_SUSPEND { + data = MakeUniformData(); + buf = std::make_unique(kNumElements * folly::kMaxVarintLength32); + } + while (iters--) { + char* pos = buf.get(); + for (int i = 0; i < kNumElements; ++i) { + nimble::varint::writeVarint(data[i], &pos); + } + CHECK_GE(pos - buf.get(), kNumElements); + } +} + +BENCHMARK(NimbleBulkDecodeUniform, iters) { + std::vector data; + std::unique_ptr buf; + std::vector recovered; + BENCHMARK_SUSPEND { + recovered.resize(kNumElements); + data = MakeUniformData(); + buf = std::make_unique(kNumElements * folly::kMaxVarintLength32); + char* pos = buf.get(); + for (int i = 0; i < kNumElements; ++i) { + nimble::varint::writeVarint(data[i], &pos); + } + } + while (iters--) { + const char* cpos = buf.get(); + nimble::varint::bulkVarintDecode32(kNumElements, cpos, recovered.data()); + CHECK_EQ(recovered.back(), data.back()); + } +} + +// ============================================================================ +// Fixed byte-width benchmarks (32-bit): isolate per-width performance +// ============================================================================ + +BENCHMARK_DRAW_LINE(); + +BENCHMARK(BulkDecode_1byte, iters) { + std::vector data; + std::unique_ptr buf; + std::vector recovered; + BENCHMARK_SUSPEND { + recovered.resize(kNumElements); + data = MakeFixedWidthData32(1); + uint64_t sz; + buf = EncodeData(data, sz); + } + while (iters--) { + const char* cpos = buf.get(); + nimble::varint::bulkVarintDecode32(kNumElements, cpos, recovered.data()); + CHECK_EQ(recovered.back(), data.back()); + } +} + +BENCHMARK_DRAW_LINE(); + +BENCHMARK(BulkDecode_2byte, iters) { + std::vector data; + std::unique_ptr buf; + std::vector recovered; + BENCHMARK_SUSPEND { + recovered.resize(kNumElements); + data = MakeFixedWidthData32(2); + uint64_t sz; + buf = EncodeData(data, sz); + } + while (iters--) { + const char* cpos = buf.get(); + nimble::varint::bulkVarintDecode32(kNumElements, cpos, recovered.data()); + CHECK_EQ(recovered.back(), data.back()); + } +} + +BENCHMARK_DRAW_LINE(); + +// ============================================================================ +// Batch size benchmarks: how does bulk decode scale with n? +// ============================================================================ + +BENCHMARK_DRAW_LINE(); + +#define BATCH_SIZE_BENCH(N) \ + BENCHMARK(BulkDecode_batch##N, iters) { \ + std::vector data; \ + std::unique_ptr buf; \ + std::vector recovered; \ + BENCHMARK_SUSPEND { \ + recovered.resize(N); \ + data = MakeUniformData(N); \ + uint64_t sz; \ + buf = EncodeData(data, sz); \ + } \ + while (iters--) { \ + const char* cpos = buf.get(); \ + nimble::varint::bulkVarintDecode32(N, cpos, recovered.data()); \ + folly::doNotOptimizeAway(recovered.back()); \ + } \ + } + +BATCH_SIZE_BENCH(4) +BATCH_SIZE_BENCH(8) +BATCH_SIZE_BENCH(16) +BATCH_SIZE_BENCH(64) +BATCH_SIZE_BENCH(256) +BATCH_SIZE_BENCH(1024) +BATCH_SIZE_BENCH(4096) + +int main() { + folly::runBenchmarks(); +} diff --git a/velox/dwio/nimble/common/tests/BitEncoderTests.cpp b/velox/dwio/nimble/common/tests/BitEncoderTests.cpp new file mode 100644 index 00000000000..90effadf752 --- /dev/null +++ b/velox/dwio/nimble/common/tests/BitEncoderTests.cpp @@ -0,0 +1,78 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include "folly/Random.h" +#include "velox/dwio/nimble/common/BitEncoder.h" + +using namespace ::facebook; + +TEST(BitEncoderTests, writeThenReadDifferentBitLengths) { + constexpr int elementCount = 1000; + std::vector buffer(4 * elementCount); + nimble::BitEncoder bitEncoder(buffer.data()); + for (int i = 0; i < elementCount; ++i) { + bitEncoder.putBits(i % 31, (i % 31) + 1); + } + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(i % 31, bitEncoder.getBits((i % 31) + 1)); + } +} + +TEST(BitEncoderTests, writeAndReadIntermixed) { + constexpr int elementCount = 1000; + std::vector data; + for (int i = 0; i < elementCount; ++i) { + data.push_back(2 * i); + } + std::vector bitLengths; + for (int i = 0; i < elementCount; ++i) { + bitLengths.push_back(11 + (i % 10)); + } + std::vector buffer(4 * elementCount); + nimble::BitEncoder bitEncoder(buffer.data()); + for (int i = 0; i < elementCount; i += 2) { + bitEncoder.putBits(data[i], bitLengths[i]); + bitEncoder.putBits(data[i + 1], bitLengths[i + 1]); + ASSERT_EQ(data[i >> 1], bitEncoder.getBits(bitLengths[i >> 1])); + } + for (int i = elementCount / 2; i < elementCount; ++i) { + ASSERT_EQ(data[i], bitEncoder.getBits(bitLengths[i])); + } +} + +TEST(BitEncoderTests, writeThenReadFullBitRange) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + constexpr int elementCount = 1000; + std::vector data; + for (int i = 0; i < elementCount; ++i) { + data.push_back(folly::Random::rand64(rng)); + } + for (int bits = 1; bits < 65; ++bits) { + const uint64_t mask = bits == 64 ? (~0ULL) : ((1ULL << bits) - 1); + std::vector buffer(8 * elementCount); + nimble::BitEncoder bitEncoder(buffer.data()); + for (int i = 0; i < elementCount; ++i) { + bitEncoder.putBits(data[i] & mask, bits); + } + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(bitEncoder.getBits(bits), data[i] & mask); + } + } +} diff --git a/velox/dwio/nimble/common/tests/BitsTests.cpp b/velox/dwio/nimble/common/tests/BitsTests.cpp new file mode 100644 index 00000000000..8be58d98db5 --- /dev/null +++ b/velox/dwio/nimble/common/tests/BitsTests.cpp @@ -0,0 +1,112 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include "folly/Random.h" +#include "velox/common/base/BitUtil.h" + +using namespace facebook::velox::bits; + +template +void repeat(int32_t times, const T& t) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + while (times-- > 0) { + t(rng); + } +} + +TEST(BitsTests, setBits) { + repeat(10, [](auto& rng) { + auto size = folly::Random::rand32(64 * 1024, rng) + 1; + auto begin = folly::Random::rand32(size, rng); + auto end = folly::Random::rand32(begin, size, rng); + std::vector bitmap(divRoundUp(size, 64) * 8, 0); + fillBits(reinterpret_cast(bitmap.data()), begin, end, true); + for (auto i = 0; i < size; ++i) { + bool expected = (i >= begin && i < end); + EXPECT_EQ( + expected, + isBitSet(reinterpret_cast(bitmap.data()), i)) + << i; + } + }); +} + +TEST(BitsTests, clearBits) { + repeat(10, [](auto& rng) { + auto size = folly::Random::rand32(64 * 1024, rng) + 1; + auto begin = folly::Random::rand32(size, rng); + auto end = folly::Random::rand32(begin, size, rng); + std::vector bitmap(divRoundUp(size, 64) * 8, 0xff); + fillBits(reinterpret_cast(bitmap.data()), begin, end, false); + for (auto i = 0; i < size; ++i) { + bool expected = (i < begin || i >= end); + EXPECT_EQ( + expected, + isBitSet(reinterpret_cast(bitmap.data()), i)) + << i; + } + }); +} + +TEST(BitsTests, findSetBit) { + repeat(10, [](auto& rng) { + auto size = folly::Random::rand32(64 * 1024, rng) + 1; + std::vector bitmap(divRoundUp(size, 64) * 8, 0); + auto begin = folly::Random::rand32(size, rng); + auto n = (size - begin) / 3; + auto numSetBits = 0; + auto pos = size; + for (auto i = begin; i < size; ++i) { + if (folly::Random::oneIn(3, rng)) { + setBit(reinterpret_cast(bitmap.data()), i); + if (++numSetBits == n) { + pos = i; + } + } + } + EXPECT_EQ(pos, findSetBit(bitmap.data(), begin, size, n)); + }); +} + +TEST(BitsTests, copy) { + repeat(1, [](auto& rng) { + auto size = folly::Random::rand32(64 * 1024, rng) + 1; + auto begin = folly::Random::rand32(size, rng); + auto end = folly::Random::oneIn(2) ? folly::Random::rand32(begin, size, rng) + : size; + std::vector src(divRoundUp(size, 64) * 8, 0); + std::vector dst(src.size(), 0); + for (auto i = begin; i < end; ++i) { + maybeSetBit(src.data(), i, folly::Random::oneIn(2, rng)); + } + Bitmap srcBitmap{src.data(), static_cast(size)}; + BitmapBuilder dstBitmap{dst.data(), static_cast(size)}; + dstBitmap.copy(srcBitmap, begin, end); + for (auto i = 0; i < end; ++i) { + if (i < begin) { + EXPECT_FALSE(isBitSet(reinterpret_cast(dst.data()), i)); + } else { + EXPECT_EQ( + isBitSet(reinterpret_cast(src.data()), i), + isBitSet(reinterpret_cast(dst.data()), i)) + << i; + } + } + }); +} diff --git a/velox/dwio/nimble/common/tests/BufferTest.cpp b/velox/dwio/nimble/common/tests/BufferTest.cpp new file mode 100644 index 00000000000..b3046c3fe25 --- /dev/null +++ b/velox/dwio/nimble/common/tests/BufferTest.cpp @@ -0,0 +1,336 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include +#include +#include + +#include "velox/common/memory/Memory.h" +#include "velox/dwio/nimble/common/Buffer.h" +#include "velox/dwio/nimble/common/NimbleException.h" +#include "velox/dwio/nimble/common/tests/GTestUtils.h" + +namespace facebook::nimble::test { + +class BufferTest : public ::testing::Test { + protected: + void SetUp() override { + pool_ = velox::memory::deprecatedAddDefaultLeafMemoryPool(); + } + + std::shared_ptr pool_; +}; + +TEST_F(BufferTest, reserveSmall) { + Buffer buffer(*pool_); + char* ptr = buffer.reserve(10); + ASSERT_NE(ptr, nullptr); + // Write and read back + std::memset(ptr, 'A', 10); + EXPECT_EQ(ptr[0], 'A'); + EXPECT_EQ(ptr[9], 'A'); +} + +TEST_F(BufferTest, reserveMultipleWithinChunk) { + Buffer buffer(*pool_); + char* ptr1 = buffer.reserve(100); + char* ptr2 = buffer.reserve(200); + ASSERT_NE(ptr1, nullptr); + ASSERT_NE(ptr2, nullptr); + + // Both pointers should be valid and non-overlapping + std::memset(ptr1, 'A', 100); + std::memset(ptr2, 'B', 200); + EXPECT_EQ(ptr1[0], 'A'); + EXPECT_EQ(ptr2[0], 'B'); +} + +TEST_F(BufferTest, reserveOversized) { + // Request more than kMinChunkSize (1 MB) to trigger a new chunk allocation + Buffer buffer(*pool_); + constexpr uint64_t bigSize = 2 * 1024 * 1024; // 2 MB + char* ptr = buffer.reserve(bigSize); + ASSERT_NE(ptr, nullptr); + // Write at boundaries + ptr[0] = 'X'; + ptr[bigSize - 1] = 'Y'; + EXPECT_EQ(ptr[0], 'X'); + EXPECT_EQ(ptr[bigSize - 1], 'Y'); +} +TEST_F(BufferTest, writeStringBasic) { + // Single write + { + Buffer buffer(*pool_); + std::string_view input = "hello world"; + auto result = buffer.writeString(input); + EXPECT_EQ(result, "hello world"); + EXPECT_EQ(result.size(), input.size()); + } + + // Multiple writes and view validity + { + Buffer buffer(*pool_); + auto r1 = buffer.writeString("aaa"); + auto r2 = buffer.writeString("bbb"); + auto r3 = buffer.writeString("ccc"); + + EXPECT_EQ(r1, "aaa"); + EXPECT_EQ(r2, "bbb"); + EXPECT_EQ(r3, "ccc"); + + // All views should remain valid + EXPECT_EQ(r1, "aaa"); + } + + // Empty write + { + Buffer buffer(*pool_); + auto result = buffer.writeString(""); + EXPECT_EQ(result, ""); + EXPECT_EQ(result.size(), 0); + } +} +TEST_F(BufferTest, takeOwnership) { + Buffer buffer(*pool_); + + auto simple_view = [&]() { + auto bufferPtr = velox::AlignedBuffer::allocate(128, pool_.get()); + char* raw = bufferPtr->asMutable(); + std::memcpy(raw, "simple extended test data", 25); + return buffer.takeOwnership(std::move(bufferPtr)); + }(); + EXPECT_GE(simple_view.size(), 25); + EXPECT_EQ( + std::string_view(simple_view.data(), 25), "simple extended test data"); +} + +TEST_F(BufferTest, transferBuffers) { + Buffer buffer(*pool_, 4096); + + auto* firstChunk = buffer.reserve(10); + std::memcpy(firstChunk, "abcdefghij", 10); + + constexpr uint64_t bigSize{2 * 1024 * 1024}; + auto* secondChunk = buffer.reserve(bigSize); + std::memset(secondChunk, 'B', bigSize); + + auto chunks = buffer.transferBuffers(); + ASSERT_EQ(chunks.size(), 2); + + EXPECT_GE(chunks[0]->capacity(), 10); + EXPECT_EQ(std::string_view(chunks[0]->as(), 10), "abcdefghij"); + + EXPECT_GE(chunks[1]->capacity(), bigSize); + EXPECT_EQ(chunks[1]->as()[0], 'B'); + EXPECT_EQ(chunks[1]->as()[bigSize - 1], 'B'); + + EXPECT_EQ(buffer.testingChunkCount(), 0); + auto emptyChunks = buffer.transferBuffers(); + EXPECT_TRUE(emptyChunks.empty()); + EXPECT_EQ(buffer.testingChunkCount(), 0); +} + +TEST_F(BufferTest, getMemoryPool) { + Buffer buffer(*pool_); + auto& poolRef = buffer.getMemoryPool(); + // Just verify we get a valid reference back + EXPECT_GT(poolRef.capacity(), 0); +} + +TEST_F(BufferTest, resetReusesFirstChunk) { + Buffer buffer(*pool_, 4096); + EXPECT_EQ(buffer.testingChunkCount(), 1); + EXPECT_EQ(buffer.testingCurrentChunkIndex(), 0); + + // Reserve within the first chunk. + char* ptr1 = buffer.reserve(100); + std::memset(ptr1, 'A', 100); + + // Reset and reserve again — should reuse the same chunk. + buffer.reset(); + EXPECT_EQ(buffer.testingChunkCount(), 1); + EXPECT_EQ(buffer.testingCurrentChunkIndex(), 0); + + char* ptr2 = buffer.reserve(100); + ASSERT_NE(ptr2, nullptr); + // Should get the same pointer since we reset to the beginning. + EXPECT_EQ(ptr1, ptr2); +} + +TEST_F(BufferTest, resetReusesMultipleChunks) { + // Use a small initial chunk so we can force multiple chunks. + Buffer buffer(*pool_, 4096); + EXPECT_EQ(buffer.testingChunkCount(), 1); + + // Force a second chunk by reserving more than kMinChunkSize. + constexpr uint64_t bigSize = 2 * 1024 * 1024; + char* bigPtr1 = buffer.reserve(bigSize); + std::memset(bigPtr1, 'B', bigSize); + EXPECT_EQ(buffer.testingChunkCount(), 2); + EXPECT_EQ(buffer.testingCurrentChunkIndex(), 1); + + // Reset — should go back to chunk 0, but keep both chunks. + buffer.reset(); + EXPECT_EQ(buffer.testingChunkCount(), 2); + EXPECT_EQ(buffer.testingCurrentChunkIndex(), 0); + + // Small reserve fits in first chunk. + char* smallPtr = buffer.reserve(100); + ASSERT_NE(smallPtr, nullptr); + EXPECT_EQ(buffer.testingCurrentChunkIndex(), 0); + + // Large reserve should reuse the second chunk via tryAdvanceToNextChunk, + // not allocate a third. + char* bigPtr2 = buffer.reserve(bigSize); + ASSERT_NE(bigPtr2, nullptr); + EXPECT_EQ(buffer.testingChunkCount(), 2); + EXPECT_EQ(buffer.testingCurrentChunkIndex(), 1); + // Should be the same underlying memory. + EXPECT_EQ(bigPtr1, bigPtr2); +} + +TEST_F(BufferTest, tryAdvanceSkipsTooSmallChunks) { + // kMinChunkSize is 1MB. Initial chunk is 1MB. + Buffer buffer(*pool_, 4096); + EXPECT_EQ(buffer.testingChunkCount(), 1); + + // Force a second 2MB chunk. + constexpr uint64_t bigSize = 2 * 1024 * 1024; + buffer.reserve(bigSize); + EXPECT_EQ(buffer.testingChunkCount(), 2); + + // Reset, then request 3MB — neither chunk (1MB or 2MB) can satisfy it, + // so a third chunk must be allocated. + constexpr uint64_t hugeSize = 3 * 1024 * 1024; + buffer.reset(); + buffer.reserve(hugeSize); + EXPECT_EQ(buffer.testingChunkCount(), 3); + EXPECT_EQ(buffer.testingCurrentChunkIndex(), 2); +} + +TEST_F(BufferTest, resetMultipleTimes) { + Buffer buffer(*pool_, 4096); + + // Cycle reset+reserve multiple times — chunk count should stay stable. + for (int i = 0; i < 5; ++i) { + buffer.reset(); + char* ptr = buffer.reserve(2048); + ASSERT_NE(ptr, nullptr); + std::memset(ptr, static_cast('A' + i), 2048); + EXPECT_EQ(ptr[0], static_cast('A' + i)); + } + // Only the initial chunk should exist — no growth. + EXPECT_EQ(buffer.testingChunkCount(), 1); +} + +TEST_F(BufferTest, encodingBufferPool) { + struct TestCase { + const char* name; + uint32_t maxCachedBuffers; + uint32_t numReleaseBuffers; + std::optional expectedReusedBufferIndex; + }; + + for (const auto testCase : { + TestCase{"no cached buffers", 0, 1, std::nullopt}, + TestCase{"single cached buffer", 1, 1, 0}, + TestCase{"cache limit drops extra released buffer", 1, 2, 0}, + TestCase{"multi-buffer cache reuses last cached buffer", 2, 2, 1}, + }) { + SCOPED_TRACE(testCase.name); + + EncodingBufferPool bufferPool{ + pool_.get(), /*maxCachedBuffers=*/testCase.maxCachedBuffers}; + + std::vector> buffers; + buffers.reserve(testCase.numReleaseBuffers); + std::vector bufferAddresses; + bufferAddresses.reserve(testCase.numReleaseBuffers); + + constexpr uint64_t bigSize = 2 * 1024 * 1024; + for (uint32_t i = 0; i < testCase.numReleaseBuffers; ++i) { + buffers.push_back(bufferPool.acquire()); + bufferAddresses.push_back(buffers.back().get()); + + buffers.back()->reserve(bigSize); + EXPECT_EQ(buffers.back()->testingChunkCount(), 2); + EXPECT_EQ(buffers.back()->testingCurrentChunkIndex(), 1); + } + + for (auto& buffer : buffers) { + bufferPool.release(std::move(buffer)); + } + + auto reusedBuffer = bufferPool.acquire(); + if (testCase.expectedReusedBufferIndex.has_value()) { + EXPECT_EQ( + reusedBuffer.get(), + bufferAddresses[*testCase.expectedReusedBufferIndex]); + EXPECT_EQ(reusedBuffer->testingChunkCount(), 2); + } else { + EXPECT_EQ(reusedBuffer->testingChunkCount(), 1); + } + EXPECT_EQ(reusedBuffer->testingCurrentChunkIndex(), 0); + + reusedBuffer->reserve(128); + EXPECT_EQ(reusedBuffer->testingCurrentChunkIndex(), 0); + } +} + +TEST_F(BufferTest, encodingBufferPoolRejectsNullBufferRelease) { + EncodingBufferPool bufferPool{pool_.get()}; + + NIMBLE_ASSERT_THROW(bufferPool.release(nullptr), "Buffer cannot be null"); +} + +TEST_F(BufferTest, scopedEncodingBuffer) { + { + SCOPED_TRACE("returns buffer to pool"); + EncodingBufferPool bufferPool{pool_.get(), /*maxCachedBuffers=*/1}; + Buffer* scopedBufferAddress; + + { + ScopedEncodingBuffer scopedBuffer{pool_.get(), &bufferPool}; + scopedBufferAddress = &scopedBuffer.get(); + scopedBuffer.get().reserve(128); + } + + auto reusedBuffer = bufferPool.acquire(); + EXPECT_EQ(reusedBuffer.get(), scopedBufferAddress); + } + + { + SCOPED_TRACE("rejects null memory pool"); + EncodingBufferPool bufferPool{pool_.get()}; + + NIMBLE_ASSERT_THROW( + ScopedEncodingBuffer(nullptr, &bufferPool), + "Memory pool cannot be null"); + } + + { + SCOPED_TRACE("uses memory pool without buffer pool"); + ScopedEncodingBuffer scopedBuffer{pool_.get(), nullptr}; + + auto* ptr = scopedBuffer.get().reserve(128); + ASSERT_NE(ptr, nullptr); + EXPECT_EQ(&scopedBuffer.get().getMemoryPool(), pool_.get()); + } +} + +} // namespace facebook::nimble::test diff --git a/velox/dwio/nimble/common/tests/CMakeLists.txt b/velox/dwio/nimble/common/tests/CMakeLists.txt new file mode 100644 index 00000000000..7143dfd78eb --- /dev/null +++ b/velox/dwio/nimble/common/tests/CMakeLists.txt @@ -0,0 +1,46 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +add_library(nimble_common_file_writer NimbleFileWriter.cpp NimbleFileWriter.h) + +target_link_libraries(nimble_common_file_writer nimble_common nimble_writer velox_vector) + +add_executable( + nimble_common_tests + BitEncoderTests.cpp + BitsTests.cpp + ChunkHeaderTest.cpp + ConstantsTest.cpp + DataTypeDispatchTest.cpp + ExceptionHelperTest.cpp + ExceptionTests.cpp + FeatureGateTest.cpp + FixedBitArrayTests.cpp + VarintTests.cpp + VectorTest.cpp + TestUtils.h + NimbleCompare.h + GTestUtils.h +) + +add_test(nimble_common_tests nimble_common_tests) + +target_link_libraries( + nimble_common_tests + nimble_common + velox_common_base + gtest + gtest_main + glog::glog + Folly::folly +) diff --git a/velox/dwio/nimble/common/tests/ChecksumTest.cpp b/velox/dwio/nimble/common/tests/ChecksumTest.cpp new file mode 100644 index 00000000000..448289edaff --- /dev/null +++ b/velox/dwio/nimble/common/tests/ChecksumTest.cpp @@ -0,0 +1,128 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "velox/dwio/nimble/common/Checksum.h" +#include + +using namespace facebook::nimble; + +TEST(ChecksumTests, createXxh3_64) { + auto checksum = ChecksumFactory::create(ChecksumType::XXH3_64); + ASSERT_NE(checksum, nullptr); + EXPECT_EQ(checksum->getType(), ChecksumType::XXH3_64); +} + +TEST(ChecksumTests, emptyDataChecksum) { + auto checksum = ChecksumFactory::create(ChecksumType::XXH3_64); + auto result = checksum->getChecksum(); + EXPECT_NE(result, 0); +} + +TEST(ChecksumTests, singleUpdateChecksum) { + auto checksum = ChecksumFactory::create(ChecksumType::XXH3_64); + checksum->update("hello world"); + auto result = checksum->getChecksum(); + EXPECT_NE(result, 0); +} + +TEST(ChecksumTests, multipleUpdatesChecksum) { + auto checksum = ChecksumFactory::create(ChecksumType::XXH3_64); + checksum->update("hello"); + checksum->update(" "); + checksum->update("world"); + auto result = checksum->getChecksum(); + EXPECT_NE(result, 0); +} + +TEST(ChecksumTests, sameDataProducesSameChecksum) { + auto checksum1 = ChecksumFactory::create(ChecksumType::XXH3_64); + auto checksum2 = ChecksumFactory::create(ChecksumType::XXH3_64); + + checksum1->update("test data"); + checksum2->update("test data"); + + EXPECT_EQ(checksum1->getChecksum(), checksum2->getChecksum()); +} + +TEST(ChecksumTests, differentDataProducesDifferentChecksum) { + auto checksum1 = ChecksumFactory::create(ChecksumType::XXH3_64); + auto checksum2 = ChecksumFactory::create(ChecksumType::XXH3_64); + + checksum1->update("data1"); + checksum2->update("data2"); + + EXPECT_NE(checksum1->getChecksum(), checksum2->getChecksum()); +} + +TEST(ChecksumTests, getChecksumWithoutReset) { + auto checksum = ChecksumFactory::create(ChecksumType::XXH3_64); + checksum->update("test"); + + auto result1 = checksum->getChecksum(false); + auto result2 = checksum->getChecksum(false); + + EXPECT_EQ(result1, result2); +} + +TEST(ChecksumTests, getChecksumWithReset) { + auto checksum = ChecksumFactory::create(ChecksumType::XXH3_64); + checksum->update("test"); + + auto resultBeforeReset = checksum->getChecksum(true); + auto resultAfterReset = checksum->getChecksum(); + + EXPECT_NE(resultBeforeReset, resultAfterReset); +} + +TEST(ChecksumTests, resetAllowsReuse) { + auto checksum = ChecksumFactory::create(ChecksumType::XXH3_64); + + checksum->update("first data"); + auto firstChecksum = checksum->getChecksum(true); + + checksum->update("first data"); + auto secondChecksum = checksum->getChecksum(); + + EXPECT_EQ(firstChecksum, secondChecksum); +} + +TEST(ChecksumTests, incrementalUpdateMatchesSingleUpdate) { + auto incrementalChecksum = ChecksumFactory::create(ChecksumType::XXH3_64); + auto singleChecksum = ChecksumFactory::create(ChecksumType::XXH3_64); + + incrementalChecksum->update("abc"); + incrementalChecksum->update("def"); + incrementalChecksum->update("ghi"); + + singleChecksum->update("abcdefghi"); + + EXPECT_EQ(incrementalChecksum->getChecksum(), singleChecksum->getChecksum()); +} + +TEST(ChecksumTests, binaryDataChecksum) { + auto checksum = ChecksumFactory::create(ChecksumType::XXH3_64); + std::string binaryData = {'\x00', '\x01', '\x02', '\xff', '\xfe'}; + checksum->update(binaryData); + auto result = checksum->getChecksum(); + EXPECT_NE(result, 0); +} + +TEST(ChecksumTests, largeDataChecksum) { + auto checksum = ChecksumFactory::create(ChecksumType::XXH3_64); + std::string largeData(1024 * 1024, 'x'); + checksum->update(largeData); + auto result = checksum->getChecksum(); + EXPECT_NE(result, 0); +} diff --git a/velox/dwio/nimble/common/tests/ChunkHeaderTest.cpp b/velox/dwio/nimble/common/tests/ChunkHeaderTest.cpp new file mode 100644 index 00000000000..122d617f5e9 --- /dev/null +++ b/velox/dwio/nimble/common/tests/ChunkHeaderTest.cpp @@ -0,0 +1,63 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include "velox/dwio/nimble/common/ChunkHeader.h" + +namespace facebook::nimble::test { + +TEST(ChunkHeaderTest, size) { + EXPECT_EQ(kChunkHeaderSize, 5); + EXPECT_EQ(kChunkHeaderSize, sizeof(uint32_t) + sizeof(CompressionType)); +} + +TEST(ChunkHeaderTest, writeAndRead) { + char buffer[kChunkHeaderSize]; + auto* writePos = buffer; + writeChunkHeader(1'234, CompressionType::Zstd, writePos); + EXPECT_EQ(writePos - buffer, kChunkHeaderSize); + + const char* readPos = buffer; + const auto header = readChunkHeader(readPos); + EXPECT_EQ(readPos - buffer, kChunkHeaderSize); + EXPECT_EQ(header.length, 1'234); + EXPECT_EQ(header.compressionType, CompressionType::Zstd); +} + +TEST(ChunkHeaderTest, uncompressed) { + char buffer[kChunkHeaderSize]; + auto* writePos = buffer; + writeChunkHeader(42, CompressionType::Uncompressed, writePos); + + const char* readPos = buffer; + const auto header = readChunkHeader(readPos); + EXPECT_EQ(header.length, 42); + EXPECT_EQ(header.compressionType, CompressionType::Uncompressed); +} + +TEST(ChunkHeaderTest, structuredBinding) { + char buffer[kChunkHeaderSize]; + auto* writePos = buffer; + writeChunkHeader(100'000, CompressionType::Zstd, writePos); + + const char* readPos = buffer; + const auto [length, compressionType] = readChunkHeader(readPos); + EXPECT_EQ(length, 100'000); + EXPECT_EQ(compressionType, CompressionType::Zstd); +} + +} // namespace facebook::nimble::test diff --git a/velox/dwio/nimble/common/tests/ConstantsTest.cpp b/velox/dwio/nimble/common/tests/ConstantsTest.cpp new file mode 100644 index 00000000000..28b2efd4c4f --- /dev/null +++ b/velox/dwio/nimble/common/tests/ConstantsTest.cpp @@ -0,0 +1,21 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +namespace facebook::nimble::test { + +// kChunkHeaderSize tests moved to ChunkHeaderTest.cpp. + +} // namespace facebook::nimble::test diff --git a/velox/dwio/nimble/common/tests/DataTypeDispatchTest.cpp b/velox/dwio/nimble/common/tests/DataTypeDispatchTest.cpp new file mode 100644 index 00000000000..2a1a3aa6cc0 --- /dev/null +++ b/velox/dwio/nimble/common/tests/DataTypeDispatchTest.cpp @@ -0,0 +1,178 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "velox/dwio/nimble/common/DataTypeDispatch.h" + +#include +#include +#include + +#include + +#include "velox/dwio/nimble/common/tests/GTestUtils.h" + +using namespace facebook; + +namespace { + +struct TypeName { + template + std::string operator()() const { + if constexpr (std::is_same_v) { + return "int8"; + } else if constexpr (std::is_same_v) { + return "uint8"; + } else if constexpr (std::is_same_v) { + return "int16"; + } else if constexpr (std::is_same_v) { + return "uint16"; + } else if constexpr (std::is_same_v) { + return "int32"; + } else if constexpr (std::is_same_v) { + return "uint32"; + } else if constexpr (std::is_same_v) { + return "int64"; + } else if constexpr (std::is_same_v) { + return "uint64"; + } else if constexpr (std::is_same_v) { + return "float"; + } else if constexpr (std::is_same_v) { + return "double"; + } else if constexpr (std::is_same_v) { + return "bool"; + } else if constexpr (std::is_same_v) { + return "string_view"; + } + } +}; + +std::string dispatchDataType(nimble::DataType dataType) { + NIMBLE_RETURN_BY_DATA_TYPE(dataType, T, TypeName{}.operator()()); +} + +std::string tryDispatchDataType(nimble::DataType dataType) { + NIMBLE_TRY_RETURN_BY_DATA_TYPE(dataType, T, TypeName{}.operator()()); +} + +std::string dispatchVarintDataType(nimble::DataType dataType) { + NIMBLE_RETURN_BY_VARINT_DATA_TYPE(dataType, T, TypeName{}.operator()()); +} + +std::string dispatchNonBoolDataType(nimble::DataType dataType) { + NIMBLE_RETURN_BY_NON_BOOL_DATA_TYPE(dataType, T, TypeName{}.operator()()); +} + +std::string dispatchNumericDataType(nimble::DataType dataType) { + NIMBLE_RETURN_BY_NUMERIC_DATA_TYPE(dataType, T, TypeName{}.operator()()); +} + +std::string dispatchFloatingPointDataType(nimble::DataType dataType) { + NIMBLE_RETURN_BY_FLOATING_POINT_DATA_TYPE( + dataType, T, TypeName{}.operator()()); +} + +std::string dispatchIntegerDataType(nimble::DataType dataType) { + NIMBLE_RETURN_BY_INTEGER_DATA_TYPE(dataType, T, TypeName{}.operator()()); +} + +} // namespace + +TEST(DataTypeDispatchTest, dispatchesAllDataTypes) { + EXPECT_EQ(dispatchDataType(nimble::DataType::Int8), "int8"); + EXPECT_EQ(dispatchDataType(nimble::DataType::Uint8), "uint8"); + EXPECT_EQ(dispatchDataType(nimble::DataType::Int16), "int16"); + EXPECT_EQ(dispatchDataType(nimble::DataType::Uint16), "uint16"); + EXPECT_EQ(dispatchDataType(nimble::DataType::Int32), "int32"); + EXPECT_EQ(dispatchDataType(nimble::DataType::Uint32), "uint32"); + EXPECT_EQ(dispatchDataType(nimble::DataType::Int64), "int64"); + EXPECT_EQ(dispatchDataType(nimble::DataType::Uint64), "uint64"); + EXPECT_EQ(dispatchDataType(nimble::DataType::Float), "float"); + EXPECT_EQ(dispatchDataType(nimble::DataType::Double), "double"); + EXPECT_EQ(dispatchDataType(nimble::DataType::Bool), "bool"); + EXPECT_EQ(dispatchDataType(nimble::DataType::String), "string_view"); +} + +TEST(DataTypeDispatchTest, tryDispatchReturnsDefaultForUndefined) { + EXPECT_EQ(tryDispatchDataType(nimble::DataType::Undefined), ""); +} + +TEST(DataTypeDispatchTest, dispatchRejectsUndefined) { + NIMBLE_ASSERT_THROW( + dispatchDataType(nimble::DataType::Undefined), "Unsupported data type"); +} + +TEST(DataTypeDispatchTest, dispatchesVarintDataTypes) { + EXPECT_EQ(dispatchVarintDataType(nimble::DataType::Int32), "int32"); + EXPECT_EQ(dispatchVarintDataType(nimble::DataType::Uint32), "uint32"); + EXPECT_EQ(dispatchVarintDataType(nimble::DataType::Int64), "int64"); + EXPECT_EQ(dispatchVarintDataType(nimble::DataType::Uint64), "uint64"); + EXPECT_EQ(dispatchVarintDataType(nimble::DataType::Float), "float"); + EXPECT_EQ(dispatchVarintDataType(nimble::DataType::Double), "double"); + + NIMBLE_ASSERT_THROW( + dispatchVarintDataType(nimble::DataType::Int16), + "Unsupported varint data type"); + NIMBLE_ASSERT_THROW( + dispatchVarintDataType(nimble::DataType::String), + "Unsupported varint data type"); +} + +TEST(DataTypeDispatchTest, dispatchesNonBoolDataTypes) { + EXPECT_EQ(dispatchNonBoolDataType(nimble::DataType::Int8), "int8"); + EXPECT_EQ(dispatchNonBoolDataType(nimble::DataType::String), "string_view"); + + NIMBLE_ASSERT_THROW( + dispatchNonBoolDataType(nimble::DataType::Bool), + "Unsupported non-bool data type"); + NIMBLE_ASSERT_THROW( + dispatchNonBoolDataType(nimble::DataType::Undefined), + "Unsupported non-bool data type"); +} + +TEST(DataTypeDispatchTest, dispatchesNumericDataTypes) { + EXPECT_EQ(dispatchNumericDataType(nimble::DataType::Int8), "int8"); + EXPECT_EQ(dispatchNumericDataType(nimble::DataType::Double), "double"); + + NIMBLE_ASSERT_THROW( + dispatchNumericDataType(nimble::DataType::Bool), + "Unsupported numeric data type"); + NIMBLE_ASSERT_THROW( + dispatchNumericDataType(nimble::DataType::String), + "Unsupported numeric data type"); +} + +TEST(DataTypeDispatchTest, dispatchesFloatingPointDataTypes) { + EXPECT_EQ(dispatchFloatingPointDataType(nimble::DataType::Float), "float"); + EXPECT_EQ(dispatchFloatingPointDataType(nimble::DataType::Double), "double"); + + NIMBLE_ASSERT_THROW( + dispatchFloatingPointDataType(nimble::DataType::Int32), + "Unsupported floating point data type"); + NIMBLE_ASSERT_THROW( + dispatchFloatingPointDataType(nimble::DataType::String), + "Unsupported floating point data type"); +} + +TEST(DataTypeDispatchTest, dispatchesIntegerDataTypes) { + EXPECT_EQ(dispatchIntegerDataType(nimble::DataType::Int8), "int8"); + EXPECT_EQ(dispatchIntegerDataType(nimble::DataType::Uint64), "uint64"); + + NIMBLE_ASSERT_THROW( + dispatchIntegerDataType(nimble::DataType::Float), + "Unsupported integer data type"); + NIMBLE_ASSERT_THROW( + dispatchIntegerDataType(nimble::DataType::Bool), + "Unsupported integer data type"); +} diff --git a/velox/dwio/nimble/common/tests/EncodingPrimitivesTest.cpp b/velox/dwio/nimble/common/tests/EncodingPrimitivesTest.cpp new file mode 100644 index 00000000000..d71c99bbc60 --- /dev/null +++ b/velox/dwio/nimble/common/tests/EncodingPrimitivesTest.cpp @@ -0,0 +1,219 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include +#include + +#include "velox/dwio/nimble/encodings/common/EncodingPrimitives.h" + +namespace facebook::nimble::encoding::test { + +// --- write/read round-trips for typed helpers --- + +TEST(EncodingPrimitivesTest, writeReadUint32) { + char buf[16]; + char* wpos = buf; + writeUint32(12345, wpos); + EXPECT_EQ(wpos, buf + sizeof(uint32_t)); + + const char* rpos = buf; + EXPECT_EQ(readUint32(rpos), 12345u); + EXPECT_EQ(rpos, buf + sizeof(uint32_t)); +} + +TEST(EncodingPrimitivesTest, writeReadUint64) { + char buf[16]; + char* wpos = buf; + writeUint64(0xDEADBEEFCAFEull, wpos); + EXPECT_EQ(wpos, buf + sizeof(uint64_t)); + + const char* rpos = buf; + EXPECT_EQ(readUint64(rpos), 0xDEADBEEFCAFEull); + EXPECT_EQ(rpos, buf + sizeof(uint64_t)); +} + +TEST(EncodingPrimitivesTest, writeReadChar) { + char buf[4]; + char* wpos = buf; + writeChar('Z', wpos); + EXPECT_EQ(wpos, buf + sizeof(char)); + + const char* rpos = buf; + EXPECT_EQ(readChar(rpos), 'Z'); + EXPECT_EQ(rpos, buf + sizeof(char)); +} + +// --- write/read template round-trips --- + +TEST(EncodingPrimitivesTest, writeReadTemplateInt32) { + char buf[16]; + char* wpos = buf; + write(-42, wpos); + + const char* rpos = buf; + EXPECT_EQ((read(rpos)), -42); +} + +TEST(EncodingPrimitivesTest, writeReadTemplateFloat) { + char buf[16]; + char* wpos = buf; + write(3.14f, wpos); + + const char* rpos = buf; + EXPECT_FLOAT_EQ((read(rpos)), 3.14f); +} + +TEST(EncodingPrimitivesTest, writeReadTemplateDouble) { + char buf[16]; + char* wpos = buf; + write(2.718281828, wpos); + + const char* rpos = buf; + EXPECT_DOUBLE_EQ((read(rpos)), 2.718281828); +} + +TEST(EncodingPrimitivesTest, writeReadTemplateUint64) { + char buf[16]; + char* wpos = buf; + write(999999999999ull, wpos); + + const char* rpos = buf; + EXPECT_EQ((read(rpos)), 999999999999ull); +} + +// --- writeString/readString with 4-byte length prefix --- + +TEST(EncodingPrimitivesTest, writeReadString) { + char buf[128]; + char* wpos = buf; + std::string_view input = "hello"; + writeString(input, wpos); + // Should advance by 4 (length prefix) + 5 (chars) + EXPECT_EQ(wpos, buf + sizeof(uint32_t) + 5); + + const char* rpos = buf; + auto result = readString(rpos); + EXPECT_EQ(result, "hello"); + EXPECT_EQ(rpos, buf + sizeof(uint32_t) + 5); +} + +TEST(EncodingPrimitivesTest, writeReadStringEmpty) { + char buf[16]; + char* wpos = buf; + writeString("", wpos); + EXPECT_EQ(wpos, buf + sizeof(uint32_t)); + + const char* rpos = buf; + auto result = readString(rpos); + EXPECT_EQ(result, ""); + EXPECT_EQ(result.size(), 0); +} + +// --- readOwnedString --- + +TEST(EncodingPrimitivesTest, readOwnedString) { + char buf[128]; + char* wpos = buf; + writeString("owned", wpos); + + const char* rpos = buf; + std::string result = readOwnedString(rpos); + EXPECT_EQ(result, "owned"); + // Result should be an independent std::string + EXPECT_EQ(rpos, buf + sizeof(uint32_t) + 5); +} + +// --- writeBytes --- + +TEST(EncodingPrimitivesTest, writeBytes) { + char buf[32]; + char* wpos = buf; + std::string_view data = "raw bytes"; + writeBytes(data, wpos); + EXPECT_EQ(wpos, buf + data.size()); + EXPECT_EQ(std::string_view(buf, data.size()), "raw bytes"); +} + +// --- peek --- + +TEST(EncodingPrimitivesTest, peekDoesNotAdvance) { + char buf[16]; + char* wpos = buf; + write(42, wpos); + + const char* rpos = buf; + auto val = peek(rpos); + EXPECT_EQ(val, 42u); + // peek should not advance the pointer + EXPECT_EQ(rpos, buf); +} + +TEST(EncodingPrimitivesTest, peekFloat) { + char buf[16]; + char* wpos = buf; + write(1.5f, wpos); + + const char* rpos = buf; + EXPECT_FLOAT_EQ((peek(rpos)), 1.5f); + EXPECT_EQ(rpos, buf); +} + +// --- Sequential writes and pointer advancement --- + +TEST(EncodingPrimitivesTest, sequentialWritesThenReads) { + char buf[128]; + char* wpos = buf; + write(100, wpos); + write(-200, wpos); + write(1.5f, wpos); + writeString("test", wpos); + + const char* rpos = buf; + EXPECT_EQ((read(rpos)), 100u); + EXPECT_EQ((read(rpos)), -200); + EXPECT_FLOAT_EQ((read(rpos)), 1.5f); + EXPECT_EQ(readString(rpos), "test"); + + // Both pointers should have advanced the same amount + EXPECT_EQ(wpos - buf, rpos - buf); +} + +// --- write/read via string_view template specialization --- + +TEST(EncodingPrimitivesTest, writeReadTemplateStringView) { + char buf[128]; + char* wpos = buf; + write(std::string_view("via template"), wpos); + + const char* rpos = buf; + auto result = read(rpos); + EXPECT_EQ(result, "via template"); +} + +TEST(EncodingPrimitivesTest, writeReadTemplateStdString) { + char buf[128]; + char* wpos = buf; + std::string input = "std string"; + write(input, wpos); + + const char* rpos = buf; + auto result = read(rpos); + EXPECT_EQ(result, "std string"); +} + +} // namespace facebook::nimble::encoding::test diff --git a/velox/dwio/nimble/common/tests/EncodingTypeTest.cpp b/velox/dwio/nimble/common/tests/EncodingTypeTest.cpp new file mode 100644 index 00000000000..0bf69210b4e --- /dev/null +++ b/velox/dwio/nimble/common/tests/EncodingTypeTest.cpp @@ -0,0 +1,190 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include +#include +#include +#include + +#include "velox/dwio/nimble/encodings/common/EncodingType.h" + +namespace facebook::nimble::test { + +// --- EncodingPhysicalType: identity mapping --- + +TEST(EncodingTypeTest, int32IdentityMapping) { + int32_t value = 42; + auto physical = EncodingPhysicalType::asEncodingPhysicalType(value); + static_assert(std::is_same_v); + auto logical = EncodingPhysicalType::asEncodingLogicalType(physical); + EXPECT_EQ(logical, 42); +} + +TEST(EncodingTypeTest, int32NegativeRoundTrip) { + int32_t value = -12345; + auto physical = EncodingPhysicalType::asEncodingPhysicalType(value); + auto logical = EncodingPhysicalType::asEncodingLogicalType(physical); + EXPECT_EQ(logical, -12345); +} + +TEST(EncodingTypeTest, int32Limits) { + { + int32_t value = std::numeric_limits::min(); + auto physical = + EncodingPhysicalType::asEncodingPhysicalType(value); + auto logical = + EncodingPhysicalType::asEncodingLogicalType(physical); + EXPECT_EQ(logical, std::numeric_limits::min()); + } + { + int32_t value = std::numeric_limits::max(); + auto physical = + EncodingPhysicalType::asEncodingPhysicalType(value); + auto logical = + EncodingPhysicalType::asEncodingLogicalType(physical); + EXPECT_EQ(logical, std::numeric_limits::max()); + } +} + +// --- EncodingPhysicalType: float <-> uint32_t bitwise conversion --- + +TEST(EncodingTypeTest, floatRoundTrip) { + float value = 3.14f; + auto physical = EncodingPhysicalType::asEncodingPhysicalType(value); + static_assert(std::is_same_v); + auto logical = EncodingPhysicalType::asEncodingLogicalType(physical); + EXPECT_FLOAT_EQ(logical, 3.14f); +} + +TEST(EncodingTypeTest, floatNegativeZero) { + float value = -0.0f; + auto physical = EncodingPhysicalType::asEncodingPhysicalType(value); + auto logical = EncodingPhysicalType::asEncodingLogicalType(physical); + // -0.0 and +0.0 compare equal but have different bit patterns + EXPECT_EQ(logical, 0.0f); + uint32_t bits; + std::memcpy(&bits, &logical, sizeof(float)); + EXPECT_NE(bits, 0u); // -0.0f has sign bit set +} + +TEST(EncodingTypeTest, floatNaN) { + float value = std::numeric_limits::quiet_NaN(); + auto physical = EncodingPhysicalType::asEncodingPhysicalType(value); + auto logical = EncodingPhysicalType::asEncodingLogicalType(physical); + EXPECT_TRUE(std::isnan(logical)); +} + +TEST(EncodingTypeTest, floatInfinity) { + float value = std::numeric_limits::infinity(); + auto physical = EncodingPhysicalType::asEncodingPhysicalType(value); + auto logical = EncodingPhysicalType::asEncodingLogicalType(physical); + EXPECT_EQ(logical, std::numeric_limits::infinity()); +} + +// --- EncodingPhysicalType: double <-> uint64_t bitwise conversion --- + +TEST(EncodingTypeTest, doubleRoundTrip) { + double value = 2.718281828; + auto physical = EncodingPhysicalType::asEncodingPhysicalType(value); + static_assert(std::is_same_v); + auto logical = EncodingPhysicalType::asEncodingLogicalType(physical); + EXPECT_DOUBLE_EQ(logical, 2.718281828); +} + +TEST(EncodingTypeTest, doubleNaN) { + double value = std::numeric_limits::quiet_NaN(); + auto physical = EncodingPhysicalType::asEncodingPhysicalType(value); + auto logical = EncodingPhysicalType::asEncodingLogicalType(physical); + EXPECT_TRUE(std::isnan(logical)); +} + +TEST(EncodingTypeTest, doubleNegativeZero) { + double value = -0.0; + auto physical = EncodingPhysicalType::asEncodingPhysicalType(value); + auto logical = EncodingPhysicalType::asEncodingLogicalType(physical); + EXPECT_EQ(logical, 0.0); + uint64_t bits; + std::memcpy(&bits, &logical, sizeof(double)); + EXPECT_NE(bits, 0u); // -0.0 has sign bit set +} + +// --- asEncodingPhysicalTypeSpan --- + +TEST(EncodingTypeTest, asEncodingPhysicalTypeSpanFloat) { + std::vector values = {1.0f, 2.0f, 3.0f, -0.0f}; + auto span = std::span(values); + auto physicalSpan = + EncodingPhysicalType::asEncodingPhysicalTypeSpan(span); + + EXPECT_EQ(physicalSpan.size(), values.size()); + // Verify the bitwise reinterpretation is consistent + for (size_t i = 0; i < values.size(); ++i) { + auto expected = + EncodingPhysicalType::asEncodingPhysicalType(values[i]); + EXPECT_EQ(physicalSpan[i], expected); + } +} + +TEST(EncodingTypeTest, asEncodingPhysicalTypeSpanInt32) { + std::vector values = {-1, 0, 1, 100}; + auto span = std::span(values); + auto physicalSpan = + EncodingPhysicalType::asEncodingPhysicalTypeSpan(span); + + EXPECT_EQ(physicalSpan.size(), values.size()); + for (size_t i = 0; i < values.size(); ++i) { + auto expected = + EncodingPhysicalType::asEncodingPhysicalType(values[i]); + EXPECT_EQ(physicalSpan[i], expected); + } +} + +TEST(EncodingTypeTest, asEncodingPhysicalType) { + using PhysicalType = EncodingPhysicalType::type; + + std::vector values = {1.0f, 2.0f, 3.0f}; + auto physicalSpan = EncodingPhysicalType::asEncodingPhysicalTypeSpan( + std::span{values}); + auto physicalValues = [&] { + return std::vector{physicalSpan.begin(), physicalSpan.end()}; + }; + + const std::vector expectedBefore{ + EncodingPhysicalType::asEncodingPhysicalType(1.0f), + EncodingPhysicalType::asEncodingPhysicalType(2.0f), + EncodingPhysicalType::asEncodingPhysicalType(3.0f)}; + EXPECT_EQ(expectedBefore, physicalValues()); + + physicalSpan[1] = EncodingPhysicalType::asEncodingPhysicalType(-3.5f); + const std::vector expectedAfter{ + EncodingPhysicalType::asEncodingPhysicalType(1.0f), + EncodingPhysicalType::asEncodingPhysicalType(-3.5f), + EncodingPhysicalType::asEncodingPhysicalType(3.0f)}; + EXPECT_EQ(expectedAfter, physicalValues()); + EXPECT_FLOAT_EQ(values[1], -3.5f); +} + +TEST(EncodingTypeTest, asEncodingPhysicalTypeSpanEmpty) { + std::vector values; + auto span = std::span(values); + auto physicalSpan = + EncodingPhysicalType::asEncodingPhysicalTypeSpan(span); + EXPECT_EQ(physicalSpan.size(), 0); +} + +} // namespace facebook::nimble::test diff --git a/velox/dwio/nimble/common/tests/ExceptionHelperTest.cpp b/velox/dwio/nimble/common/tests/ExceptionHelperTest.cpp new file mode 100644 index 00000000000..473cd3a40d9 --- /dev/null +++ b/velox/dwio/nimble/common/tests/ExceptionHelperTest.cpp @@ -0,0 +1,56 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "velox/dwio/nimble/common/ExceptionHelper.h" +#include + +using namespace facebook::nimble; + +TEST(ExceptionHelperTest, noArgs) { + auto result = errorMessage(); + // CompileTimeEmptyString is convertible to const char*, string_view, string. + const char* asCharPtr = result; + EXPECT_STREQ(asCharPtr, ""); + + std::string_view asSv = result; + EXPECT_EQ(asSv, ""); + EXPECT_EQ(asSv.size(), 0); + + std::string asStr = result; + EXPECT_EQ(asStr, ""); +} + +TEST(ExceptionHelperTest, constCharPtr) { + const char* input = "hello"; + const char* result = errorMessage(input); + EXPECT_EQ(result, input); + EXPECT_STREQ(result, "hello"); +} + +TEST(ExceptionHelperTest, stdString) { + std::string input = "test message"; + std::string result = errorMessage(input); + EXPECT_EQ(result, "test message"); +} + +TEST(ExceptionHelperTest, formatArgs) { + std::string result = errorMessage("val={}", 42); + EXPECT_EQ(result, "val=42"); +} + +TEST(ExceptionHelperTest, formatArgsMultiple) { + std::string result = errorMessage("{} + {} = {}", 1, 2, 3); + EXPECT_EQ(result, "1 + 2 = 3"); +} diff --git a/velox/dwio/nimble/common/tests/ExceptionTests.cpp b/velox/dwio/nimble/common/tests/ExceptionTests.cpp new file mode 100644 index 00000000000..f5476ae170d --- /dev/null +++ b/velox/dwio/nimble/common/tests/ExceptionTests.cpp @@ -0,0 +1,477 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include "velox/common/base/VeloxException.h" +#include "velox/dwio/nimble/common/Exceptions.h" + +namespace facebook { +namespace { + +template +void verifyException( + const T& e, + const std::string& exceptionName, + const std::string& fileName, + const std::string& fileLine, + const std::string& functionName, + const std::string& failingExpression, + const std::string& errorMessage, + const std::string& errorSource, + const std::string& errorCode, + const std::string& retryable, + const std::string& additionalMessage = "") { + EXPECT_EQ(fileName, e.fileName()); + if (!fileLine.empty()) { + EXPECT_EQ(fileLine, folly::to(e.fileLine())); + } + EXPECT_EQ(functionName, e.functionName()); + EXPECT_EQ(failingExpression, e.failingExpression()); + EXPECT_EQ(errorMessage, e.errorMessage()); + EXPECT_EQ(errorSource, e.errorSource()); + EXPECT_EQ(errorCode, e.errorCode()); + EXPECT_EQ(retryable, e.retryable() ? "True" : "False"); + + EXPECT_NE( + std::string(e.what()).find(exceptionName + "\n"), std::string::npos); + EXPECT_NE( + std::string(e.what()).find("Error Source: " + errorSource + "\n"), + std::string::npos); + EXPECT_NE( + std::string(e.what()).find("Error Code: " + errorCode + "\n"), + std::string::npos); + if (!errorMessage.empty()) { + EXPECT_NE( + std::string(e.what()).find("Error Message: " + errorMessage + "\n"), + std::string::npos); + } + EXPECT_NE( + std::string(e.what()).find("Retryable: " + retryable + "\n"), + std::string::npos); + EXPECT_NE( + std::string(e.what()).find( + "Location: " + + folly::to(functionName, '@', fileName, ':', fileLine)), + std::string::npos); + if (!failingExpression.empty()) { + EXPECT_NE( + std::string(e.what()).find("Expression: " + failingExpression + "\n"), + std::string::npos); + } + EXPECT_NE(std::string(e.what()).find("Stack Trace:\n"), std::string::npos); + + if (!additionalMessage.empty()) { + EXPECT_NE(std::string(e.what()).find(additionalMessage), std::string::npos); + } +} + +template +void verifyFileComparisonFailure( + Check&& check, + std::string_view comparedValues, + std::string_view message) { + try { + check(); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleUserError& exception) { + EXPECT_EQ(exception.errorCode(), "CORRUPTED_FILE"); + EXPECT_EQ(exception.errorSource(), "USER"); + EXPECT_NE(exception.errorMessage().find(comparedValues), std::string::npos); + EXPECT_NE(exception.errorMessage().find(message), std::string::npos); + } +} + +TEST(ExceptionTests, format) { + verifyException( + nimble::NimbleUserError( + "file1", 23, "func1", "expr1", "err1", "code1", true), + "NimbleUserError", + "file1", + "23", + "func1", + "expr1", + "err1", + "USER", + "code1", + "True"); + + verifyException( + nimble::NimbleInternalError( + "file2", 24, "func2", "expr2", "err2", "code2", false), + "NimbleInternalError", + "file2", + "24", + "func2", + "expr2", + "err2", + "INTERNAL", + "code2", + "False"); +} + +TEST(ExceptionTests, check) { + int a = 5; + try { + NIMBLE_CHECK(a < 3, "error message1"); + } catch (const nimble::NimbleInternalError& e) { + verifyException( + e, + "NimbleInternalError", + __FILE__, + "", + "TestBody", + "a < 3", + "error message1", + "INTERNAL", + "INVALID_ARGUMENT", + "False"); + } +} + +TEST(ExceptionTests, unreachable) { + try { + NIMBLE_UNREACHABLE("error message"); + } catch (const nimble::NimbleInternalError& e) { + verifyException( + e, + "NimbleInternalError", + __FILE__, + "", + "TestBody", + "", + "error message", + "INTERNAL", + "UNREACHABLE_CODE", + "False"); + } +} + +TEST(ExceptionTests, notImplemented) { + try { + NIMBLE_NOT_IMPLEMENTED("error message7"); + } catch (const nimble::NimbleInternalError& e) { + verifyException( + e, + "NimbleInternalError", + __FILE__, + "", + "TestBody", + "", + "error message7", + "INTERNAL", + "NOT_IMPLEMENTED", + "False"); + } +} + +TEST(ExceptionTests, notSupported) { + try { + NIMBLE_UNSUPPORTED("error message6"); + } catch (const nimble::NimbleUserError& e) { + verifyException( + e, + "NimbleUserError", + __FILE__, + "", + "TestBody", + "", + "error message6", + "USER", + "NOT_SUPPORTED", + "False"); + } +} + +TEST(ExceptionTests, stackTraceThreads) { + // Make sure captured stack trace doesn't need anything from thread local + // storage + std::exception_ptr e; + auto throwFunc = []() { NIMBLE_CHECK(false, "Test."); }; + std::thread t([&]() { + try { + throwFunc(); + } catch (...) { + e = std::current_exception(); + } + }); + + t.join(); + + ASSERT_NE(nullptr, e); + const auto trace = folly::exceptionStr(e).toStdString(); + if (trace.find("_ZN8facebook6nimble") != std::string::npos) { + // folly symbolized the frames but left them mangled, which happens when + // folly is built without demangling support. Velox neither declares nor + // installs libiberty, so a stock OSS build lands here. File and line + // resolution still works; only the spelling of the symbol differs. + GTEST_SKIP() << "folly emitted mangled frames; demangling is unavailable " + "in this build."; + } + EXPECT_NE( + std::string::npos, + trace.find("facebook::nimble::NimbleException::NimbleException")); +} + +TEST(ExceptionTests, context) { + auto messageFunc = [](velox::VeloxException::Type exceptionType, void* arg) { + auto msg = *static_cast(arg); + switch (exceptionType) { + case velox::VeloxException::Type::kUser: + return fmt::format("USER {}", msg); + case velox::VeloxException::Type::kSystem: + return fmt::format("SYSTEM {}", msg); + } + NIMBLE_UNREACHABLE( + "Unknown exception type: {}", static_cast(exceptionType)); + }; + std::string context1Message = "1"; + velox::ExceptionContextSetter context1({messageFunc, &context1Message, true}); + std::string context2Message = "2"; + velox::ExceptionContextSetter context2( + {messageFunc, &context2Message, false}); + std::string context3Message = "3"; + velox::ExceptionContextSetter context3( + {messageFunc, &context3Message, false}); + + try { + NIMBLE_UNSUPPORTED(""); + FAIL(); + } catch (const nimble::NimbleException& e) { + ASSERT_EQ(e.context(), "USER 3 USER 1"); + ASSERT_NE( + std::string(e.what()).find("Context: " + e.context()), + std::string::npos); + } + + try { + NIMBLE_UNKNOWN(""); + FAIL(); + } catch (const nimble::NimbleException& e) { + ASSERT_EQ(e.context(), "SYSTEM 3 SYSTEM 1"); + ASSERT_NE( + std::string(e.what()).find("Context: " + e.context()), + std::string::npos); + } +} + +TEST(ExceptionTests, checkComparisons) { + // Test NIMBLE_CHECK_GT + try { + NIMBLE_CHECK_GT(5, 10); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleInternalError& e) { + EXPECT_NE(std::string(e.what()).find("(5 vs. 10)"), std::string::npos); + EXPECT_EQ(e.errorCode(), "INVALID_ARGUMENT"); + } + + // Test NIMBLE_CHECK_GE + try { + NIMBLE_CHECK_GE(3, 5); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleInternalError& e) { + EXPECT_NE(std::string(e.what()).find("(3 vs. 5)"), std::string::npos); + } + + // Test NIMBLE_CHECK_LT + try { + NIMBLE_CHECK_LT(10, 5); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleInternalError& e) { + EXPECT_NE(std::string(e.what()).find("(10 vs. 5)"), std::string::npos); + } + + // Test NIMBLE_CHECK_LE + try { + NIMBLE_CHECK_LE(8, 3); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleInternalError& e) { + EXPECT_NE(std::string(e.what()).find("(8 vs. 3)"), std::string::npos); + } + + // Test NIMBLE_CHECK_EQ + try { + NIMBLE_CHECK_EQ(5, 10); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleInternalError& e) { + EXPECT_NE(std::string(e.what()).find("(5 vs. 10)"), std::string::npos); + } + + // Test NIMBLE_CHECK_NE + try { + NIMBLE_CHECK_NE(7, 7); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleInternalError& e) { + EXPECT_NE(std::string(e.what()).find("(7 vs. 7)"), std::string::npos); + } + + // Test successful comparison (should not throw) + NIMBLE_CHECK_GT(10, 5); + NIMBLE_CHECK_GE(5, 5); + NIMBLE_CHECK_LT(3, 8); + NIMBLE_CHECK_LE(4, 4); + NIMBLE_CHECK_EQ(7, 7); + NIMBLE_CHECK_NE(3, 5); +} + +TEST(ExceptionTests, checkComparisonsWithCustomMessage) { + // Test with custom format message + try { + NIMBLE_CHECK_GT(5, 10, "custom message: {} items", 42); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleInternalError& e) { + std::string what = e.what(); + EXPECT_NE(what.find("(5 vs. 10)"), std::string::npos); + EXPECT_NE(what.find("custom message: 42 items"), std::string::npos); + } + + // Test NIMBLE_CHECK_EQ with custom message + try { + NIMBLE_CHECK_EQ(100, 200, "Size mismatch"); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleInternalError& e) { + std::string what = e.what(); + EXPECT_NE(what.find("(100 vs. 200)"), std::string::npos); + EXPECT_NE(what.find("Size mismatch"), std::string::npos); + } +} + +TEST(ExceptionTests, checkNull) { + int* nullPtr = nullptr; + int value = 42; + int* validPtr = &value; + + // Test NIMBLE_CHECK_NULL - should throw when pointer is not null + try { + NIMBLE_CHECK_NULL(validPtr); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleInternalError& e) { + EXPECT_EQ(e.errorCode(), "INVALID_ARGUMENT"); + } + + // Test NIMBLE_CHECK_NULL - should not throw when pointer is null + NIMBLE_CHECK_NULL(nullPtr); + + // Test NIMBLE_CHECK_NOT_NULL - should throw when pointer is null + try { + NIMBLE_CHECK_NOT_NULL(nullPtr); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleInternalError& e) { + EXPECT_EQ(e.errorCode(), "INVALID_ARGUMENT"); + } + + // Test NIMBLE_CHECK_NOT_NULL - should not throw when pointer is valid + NIMBLE_CHECK_NOT_NULL(validPtr); +} + +TEST(ExceptionTests, checkFileNotNull) { + int* nullPointer = nullptr; + int value = 42; + int* validPointer = &value; + + try { + NIMBLE_CHECK_FILE_NOT_NULL(nullPointer, "Invalid file pointer"); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleUserError& exception) { + EXPECT_EQ(exception.errorCode(), "CORRUPTED_FILE"); + EXPECT_EQ(exception.errorSource(), "USER"); + EXPECT_EQ(exception.errorMessage(), "Invalid file pointer"); + } + + NIMBLE_CHECK_FILE_NOT_NULL(validPointer); +} + +TEST(ExceptionTests, checkFileComparisons) { + verifyFileComparisonFailure( + [] { NIMBLE_CHECK_FILE_GT(5, 10, "Invalid greater-than value"); }, + "(5 vs. 10)", + "Invalid greater-than value"); + verifyFileComparisonFailure( + [] { NIMBLE_CHECK_FILE_GE(3, 5, "Invalid greater-or-equal value"); }, + "(3 vs. 5)", + "Invalid greater-or-equal value"); + verifyFileComparisonFailure( + [] { NIMBLE_CHECK_FILE_LT(10, 5, "Invalid less-than value"); }, + "(10 vs. 5)", + "Invalid less-than value"); + verifyFileComparisonFailure( + [] { NIMBLE_CHECK_FILE_LE(8, 3, "Invalid less-or-equal value"); }, + "(8 vs. 3)", + "Invalid less-or-equal value"); + verifyFileComparisonFailure( + [] { NIMBLE_CHECK_FILE_EQ(5, 10, "Mismatched file value"); }, + "(5 vs. 10)", + "Mismatched file value"); + verifyFileComparisonFailure( + [] { NIMBLE_CHECK_FILE_NE(7, 7, "Duplicate file value"); }, + "(7 vs. 7)", + "Duplicate file value"); + + NIMBLE_CHECK_FILE_GT(10, 5); + NIMBLE_CHECK_FILE_GE(5, 5); + NIMBLE_CHECK_FILE_LT(3, 8); + NIMBLE_CHECK_FILE_LE(4, 4); + NIMBLE_CHECK_FILE_EQ(7, 7); + NIMBLE_CHECK_FILE_NE(7, 8); +} + +TEST(ExceptionTests, checkNullWithMessage) { + int* nullPtr = nullptr; + int value = 42; + int* validPtr = &value; + + // Test with custom message + try { + NIMBLE_CHECK_NULL(validPtr, "Expected null pointer"); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleInternalError& e) { + EXPECT_NE( + std::string(e.what()).find("Expected null pointer"), std::string::npos); + } + + try { + NIMBLE_CHECK_NOT_NULL(nullPtr, "Pointer should not be null"); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleInternalError& e) { + EXPECT_NE( + std::string(e.what()).find("Pointer should not be null"), + std::string::npos); + } +} + +TEST(ExceptionTests, failMacros) { + // Test NIMBLE_FAIL + try { + NIMBLE_FAIL("Internal error: {}", 42); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleInternalError& e) { + EXPECT_NE( + std::string(e.what()).find("Internal error: 42"), std::string::npos); + EXPECT_EQ(e.errorCode(), "INVALID_STATE"); + } + + // Test NIMBLE_USER_FAIL + try { + NIMBLE_USER_FAIL("User error: {} is invalid", "input"); + FAIL() << "Should have thrown"; + } catch (const nimble::NimbleUserError& e) { + EXPECT_NE( + std::string(e.what()).find("User error: input is invalid"), + std::string::npos); + EXPECT_EQ(e.errorCode(), "INVALID_ARGUMENT"); + } +} + +} // namespace +} // namespace facebook diff --git a/velox/dwio/nimble/common/tests/FeatureGateTest.cpp b/velox/dwio/nimble/common/tests/FeatureGateTest.cpp new file mode 100644 index 00000000000..fd1b4392a05 --- /dev/null +++ b/velox/dwio/nimble/common/tests/FeatureGateTest.cpp @@ -0,0 +1,82 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "velox/dwio/nimble/common/FeatureGate.h" + +#include +#include + +#include + +namespace facebook::nimble { +namespace { + +// A gate that returns a fixed answer regardless of the requested default, +// modeling an internal gate that force-enables or gates-off a feature. +class FixedFeatureGate : public FeatureGate { + public: + explicit FixedFeatureGate(bool value) : value_{value} {} + + bool enabled(std::string_view /*feature*/, bool /*defaultValue*/) + const override { + return value_; + } + + private: + const bool value_; +}; + +class FeatureGateTest : public ::testing::Test { + protected: + // Restore the default no-op gate so tests don't leak process-wide state. + void TearDown() override { + registerFeatureGate(nullptr); + } +}; + +TEST_F(FeatureGateTest, defaultGatePassesRequestedValueThrough) { + const auto gate = featureGate(); + EXPECT_TRUE(gate->enabled("any_feature", true)); + EXPECT_FALSE(gate->enabled("any_feature", false)); +} + +TEST_F(FeatureGateTest, registeredGateOverridesDefault) { + registerFeatureGate(std::make_shared(false)); + EXPECT_FALSE(featureGate()->enabled("feature", true)); + + registerFeatureGate(std::make_shared(true)); + EXPECT_TRUE(featureGate()->enabled("feature", false)); +} + +TEST_F(FeatureGateTest, registerNullptrRestoresDefault) { + registerFeatureGate(std::make_shared(false)); + ASSERT_FALSE(featureGate()->enabled("feature", true)); + + registerFeatureGate(nullptr); + EXPECT_TRUE(featureGate()->enabled("feature", true)); + EXPECT_FALSE(featureGate()->enabled("feature", false)); +} + +TEST_F(FeatureGateTest, returnedGateOutlivesReRegistration) { + // featureGate() hands out an owning pointer, so an already-fetched gate stays + // valid (and unchanged) even after a different gate is registered. + const auto original = featureGate(); + registerFeatureGate(std::make_shared(false)); + EXPECT_TRUE(original->enabled("feature", true)); + EXPECT_FALSE(featureGate()->enabled("feature", true)); +} + +} // namespace +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/tests/FixedBitArrayTests.cpp b/velox/dwio/nimble/common/tests/FixedBitArrayTests.cpp new file mode 100644 index 00000000000..ac3045caa22 --- /dev/null +++ b/velox/dwio/nimble/common/tests/FixedBitArrayTests.cpp @@ -0,0 +1,702 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "fmt/format.h" +#include "folly/Benchmark.h" +#include "folly/Random.h" +#include "velox/dwio/nimble/common/FixedBitArray.h" + +using namespace ::facebook; + +TEST(FixedBitArrayTests, setThenGet) { + constexpr int bitWidth = 10; + constexpr int elementCount = 10000; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.set(i, (i + i * i) % (1 << bitWidth)); + } + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(fixedBitArray.get(i), (i + i * i) % (1 << bitWidth)); + } +} + +TEST(FixedBitArrayTests, zeroAndSet) { + constexpr int bitWidth = 4; + constexpr int elementCount = 1234; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.set(i, (i + i * i) % (1 << bitWidth)); + } + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(fixedBitArray.get(i), (i + i * i) % (1 << bitWidth)); + } + for (int i = 0; i < elementCount; ++i) { + if (i % 3 == 0) { + fixedBitArray.zeroAndSet(i, i % (1 << bitWidth)); + } + } + for (int i = 0; i < elementCount; ++i) { + if (i % 3 == 0) { + ASSERT_EQ(fixedBitArray.get(i), i % (1 << bitWidth)); + } else { + ASSERT_EQ(fixedBitArray.get(i), (i + i * i) % (1 << bitWidth)); + } + } +} + +constexpr int kNumTestsPerBitWidth = 10; +constexpr int kMaxElements = 1000; + +TEST(FixedBitArrayTests, setThenGetRandom) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int bitWidth = 1; bitWidth <= 64; ++bitWidth) { + const uint64_t elementMask = + bitWidth == 64 ? (~0ULL) : ((1ULL << bitWidth) - 1); + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = folly::Random::rand32(rng) % kMaxElements; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = folly::Random::rand64(rng) & elementMask; + } + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.set(i, randomValues[i]); + } + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(fixedBitArray.get(i), randomValues[i]); + } + } + } +} + +TEST(FixedBitArrayTests, zeroAndSetThenGetRandom) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int bitWidth = 1; bitWidth <= 64; ++bitWidth) { + const uint64_t elementMask = + bitWidth == 64 ? (~0ULL) : ((1ULL << bitWidth) - 1); + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = folly::Random::rand32(rng) % kMaxElements; + std::vector buffer( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth), 0xFF); + nimble::FixedBitArray fixedBitArray(buffer.data(), bitWidth); + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = folly::Random::rand64(rng) & elementMask; + } + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.zeroAndSet(i, randomValues[i]); + } + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(fixedBitArray.get(i), randomValues[i]) << bitWidth; + } + } + } +} + +TEST(FixedBitArrayTests, set32ThenGet32Random) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int bitWidth = 1; bitWidth <= 32; ++bitWidth) { + const uint64_t elementMask = (1ULL << bitWidth) - 1; + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = folly::Random::rand32(rng) % kMaxElements; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = folly::Random::rand32(rng) & elementMask; + } + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.set32(i, randomValues[i]); + } + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(fixedBitArray.get32(i), randomValues[i]) + << i << " " << bitWidth; + } + } + } +} + +TEST(FixedBitArrayTests, bulkGet32) { + constexpr int kBitWidth = 7; + constexpr int kNumElements = 27; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(kNumElements, kBitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), kBitWidth); + for (int i = 0; i < kNumElements; ++i) { + fixedBitArray.set32(i, 100 - i); + } + std::vector values(kNumElements); + fixedBitArray.bulkGet32(0, kNumElements, values.data()); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(values[i], 100 - i); + } + std::vector values64(kNumElements); + fixedBitArray.bulkGet32Into64(0, kNumElements, values64.data()); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(values64[i], 100 - i); + } +} + +TEST(FixedBitArrayTests, bulkGet32Random) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int bitWidth = 1; bitWidth <= 32; ++bitWidth) { + const uint64_t elementMask = (1ULL << bitWidth) - 1; + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = folly::Random::rand32(rng) % kMaxElements; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = folly::Random::rand32(rng) & elementMask; + } + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.set32(i, randomValues[i]); + } + std::vector values(elementCount); + fixedBitArray.bulkGet32(0, elementCount, values.data()); + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(values[i], randomValues[i]); + } + for (int i = 0; i < elementCount; ++i) { + uint32_t element; + fixedBitArray.bulkGet32(i, 1, &element); + ASSERT_EQ(element, values[i]); + } + std::vector values64(elementCount); + fixedBitArray.bulkGet32Into64(0, elementCount, values64.data()); + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(values64[i], randomValues[i]); + } + for (int i = 0; i < elementCount; ++i) { + uint64_t element; + fixedBitArray.bulkGet32Into64(i, 1, &element); + ASSERT_EQ(element, values64[i]); + } + } + } +} + +TEST(FixedBitArrayTests, bulkGetWithBaseline32Random) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int bitWidth = 4; bitWidth <= 32; ++bitWidth) { + const uint64_t maxElement = (1ULL << bitWidth); + const uint64_t baseline = folly::Random::rand32(rng) % maxElement; + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = folly::Random::rand32(rng) % kMaxElements; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = folly::Random::rand32(rng) % (maxElement - baseline); + } + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.set32(i, randomValues[i]); + } + std::vector values(elementCount); + fixedBitArray.bulkGetWithBaseline( + 0, elementCount, values.data(), baseline); + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(values[i], randomValues[i] + baseline) + << "Bit Width: " << bitWidth << ", i: " << i; + } + for (int i = 0; i < elementCount; ++i) { + uint32_t element; + fixedBitArray.bulkGetWithBaseline(i, 1, &element, baseline); + ASSERT_EQ(element, randomValues[i] + baseline); + } + std::vector values64(elementCount); + fixedBitArray.bulkGetWithBaseline32Into64( + 0, elementCount, values64.data(), baseline); + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(values64[i], randomValues[i] + baseline); + } + for (int i = 0; i < elementCount; ++i) { + uint64_t element; + fixedBitArray.bulkGetWithBaseline32Into64(i, 1, &element, baseline); + ASSERT_EQ(element, randomValues[i] + baseline); + } + } + } +} + +TEST(FixedBitArrayTests, bulkGet64WithBaseline) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + // Test all three code paths: + // - bitWidth <= 32: delegates to bulkGetWithBaseline32Into64 + // - bitWidth 33-58: branchless byte-aligned loads + // - bitWidth > 58: inline cross-word boundary handling + for (int bitWidth = 1; bitWidth <= 64; ++bitWidth) { + SCOPED_TRACE(fmt::format("bitWidth={}", bitWidth)); + const uint64_t maxElement = bitWidth == 64 + ? std::numeric_limits::max() + : (1ULL << bitWidth) - 1; + const uint64_t baseline = folly::Random::rand64(rng) % (maxElement / 2 + 1); + const uint64_t valueRange = maxElement - baseline; + + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = folly::Random::rand32(rng) % kMaxElements; + const auto bufferBytes = + nimble::FixedBitArray::bufferSize(elementCount, bitWidth); + auto buffer = std::make_unique(bufferBytes); + std::memset(buffer.get(), 0, bufferBytes); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = valueRange == std::numeric_limits::max() + ? folly::Random::rand64(rng) + : folly::Random::rand64(rng) % (valueRange + 1); + } + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.set(i, randomValues[i]); + } + + // Bulk read all elements. + std::vector values(elementCount); + fixedBitArray.bulkGetWithBaseline( + 0, elementCount, values.data(), baseline); + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(values[i], randomValues[i] + baseline) + << "bitWidth: " << bitWidth << ", i: " << i; + } + + // Single-element reads at each position. + for (int i = 0; i < elementCount; ++i) { + uint64_t element; + fixedBitArray.bulkGetWithBaseline(i, 1, &element, baseline); + ASSERT_EQ(element, randomValues[i] + baseline); + } + + // Read from a random offset. + if (elementCount > 1) { + const int offset = folly::Random::rand32(rng) % (elementCount - 1); + const int count = elementCount - offset; + std::vector partial(count); + fixedBitArray.bulkGetWithBaseline( + offset, count, partial.data(), baseline); + for (int i = 0; i < count; ++i) { + ASSERT_EQ(partial[i], randomValues[offset + i] + baseline); + } + } + } + } +} + +TEST(FixedBitArrayTests, bulkGet64WithBaselineZeroBaseline) { + // Verify bulkGet64WithBaseline with baseline=0 matches per-element get(). + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int bitWidth : {1, 8, 16, 32, 33, 40, 48, 56, 57, 58, 60, 64}) { + SCOPED_TRACE(fmt::format("bitWidth={}", bitWidth)); + const int elementCount = 100 + folly::Random::rand32(rng) % 200; + const uint64_t maxElement = bitWidth == 64 + ? std::numeric_limits::max() + : (1ULL << bitWidth) - 1; + const auto bufferBytes = + nimble::FixedBitArray::bufferSize(elementCount, bitWidth); + auto buffer = std::make_unique(bufferBytes); + std::memset(buffer.get(), 0, bufferBytes); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + + for (int i = 0; i < elementCount; ++i) { + const uint64_t value = bitWidth == 64 + ? folly::Random::rand64(rng) + : folly::Random::rand64(rng) % (maxElement + 1); + fixedBitArray.set(i, value); + } + + std::vector bulkValues(elementCount); + fixedBitArray.bulkGetWithBaseline(0, elementCount, bulkValues.data(), 0); + + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(bulkValues[i], fixedBitArray.get(i)) + << "bitWidth: " << bitWidth << ", i: " << i; + } + } +} + +TEST(FixedBitArrayTests, bulkSet32Random) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + for (int bitWidth = 1; bitWidth <= 32; ++bitWidth) { + const uint64_t maxElement = (1ULL << bitWidth); + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = 1 + folly::Random::rand32(rng) % kMaxElements; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = folly::Random::rand32(rng) % maxElement; + } + const int offset = folly::Random::rand32(rng) % elementCount; + const int size = elementCount - offset; + fixedBitArray.bulkSet32(offset, size, randomValues.data() + offset); + for (int i = 0; i < size; ++i) { + ASSERT_EQ(fixedBitArray.get32(offset + i), randomValues[offset + i]); + } + } + } +} + +TEST(FixedBitArrayTests, bulkSet32WithBaselineRandom) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + for (int bitWidth = 4; bitWidth <= 32; ++bitWidth) { + const uint64_t maxElement = (1ULL << bitWidth); + const uint64_t baseline = folly::Random::rand32(rng) % maxElement; + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = 1 + folly::Random::rand32(rng) % kMaxElements; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + std::vector randomValues(elementCount); + std::vector randomValuesWithBaseline(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = folly::Random::rand32(rng) % (maxElement - baseline); + randomValuesWithBaseline[i] = randomValues[i] + baseline; + } + const int offset = folly::Random::rand32(rng) % elementCount; + const int size = elementCount - offset; + fixedBitArray.bulkSetWithBaseline( + offset, size, randomValuesWithBaseline.data() + offset, baseline); + for (int i = 0; i < size; ++i) { + ASSERT_EQ(fixedBitArray.get32(offset + i), randomValues[offset + i]) + << "Bit Width: " << bitWidth << ", i: " << i; + } + } + } +} + +TEST(FixedBitArrayTests, equals32Random) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int bitWidth = 1; bitWidth <= 32; ++bitWidth) { + const uint64_t elementMask = (1ULL << bitWidth) - 1; + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = 1 + folly::Random::rand32(rng) % kMaxElements; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + // Test both the full range and also a case where equals are actually + // likely. + randomValues[i] = folly::Random::rand32(rng) & elementMask; + if (test % 2 == 1) { + randomValues[i] = randomValues[i] % 10; + } + fixedBitArray.set32(i, randomValues[i]); + } + // Remember start needs to be a multiple of 64. + const int start = + (folly::Random::rand32(rng) % elementCount) & (0xFFFFFF00); + const int length = elementCount - start; + auto equalsBuffer = std::make_unique( + nimble::FixedBitArray::bufferSize(length, 1)); + const uint32_t equalsValue = + randomValues[folly::Random::rand32(rng) % elementCount]; + fixedBitArray.equals32(start, length, equalsValue, equalsBuffer.get()); + for (int i = 0; i < length; ++i) { + ASSERT_EQ( + velox::bits::isBitSet( + reinterpret_cast(equalsBuffer.get()), i), + randomValues[start + i] == equalsValue); + } + } + } +} + +TEST(FixedBitArrayTests, bulkGet64WithBaselineBitWidth58Boundary) { + constexpr int bitWidth = 58; + constexpr int elementCount = 5; + constexpr int start = 3; + constexpr uint64_t baseline = 17; + constexpr uint64_t maxElement = (1ULL << bitWidth) - 1; + + const auto bufferBytes = + nimble::FixedBitArray::bufferSize(elementCount, bitWidth); + auto buffer = std::make_unique(bufferBytes); + std::memset(buffer.get(), 0, bufferBytes); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + + // start * bitWidth leaves a bit remainder of 6. For bitWidth 58 this is the + // highest possible remainder, and 58 + 6 exactly fills one 64-bit load. + fixedBitArray.set(start, maxElement); + + uint64_t value = 0; + fixedBitArray.bulkGetWithBaseline(start, 1, &value, baseline); + ASSERT_EQ(value, maxElement + baseline); +} + +TEST(FixedBitArrayTests, bulkSet64WithBaseline) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + // Test all three code paths: + // - bitWidth <= 32: delegates to template-unrolled bulkSetInternal32 + // - bitWidth 33-58: branchless byte-aligned stores + // - bitWidth > 58: inline cross-word boundary handling + for (int bitWidth = 1; bitWidth <= 64; ++bitWidth) { + SCOPED_TRACE(fmt::format("bitWidth={}", bitWidth)); + const uint64_t maxElement = bitWidth == 64 + ? std::numeric_limits::max() + : (1ULL << bitWidth) - 1; + const uint64_t baseline = folly::Random::rand64(rng) % (maxElement / 2 + 1); + const uint64_t valueRange = maxElement - baseline; + + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = 1 + folly::Random::rand32(rng) % kMaxElements; + const auto bufferBytes = + nimble::FixedBitArray::bufferSize(elementCount, bitWidth); + auto buffer = std::make_unique(bufferBytes); + std::memset(buffer.get(), 0, bufferBytes); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + + std::vector inputValues(elementCount); + std::vector expectedResiduals(elementCount); + for (int i = 0; i < elementCount; ++i) { + expectedResiduals[i] = + valueRange == std::numeric_limits::max() + ? folly::Random::rand64(rng) + : folly::Random::rand64(rng) % (valueRange + 1); + inputValues[i] = expectedResiduals[i] + baseline; + } + + fixedBitArray.bulkSetWithBaseline( + 0, elementCount, inputValues.data(), baseline); + + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(fixedBitArray.get(i), expectedResiduals[i]) << "i=" << i; + } + + std::vector recovered(elementCount); + fixedBitArray.bulkGetWithBaseline( + 0, elementCount, recovered.data(), baseline); + ASSERT_EQ(recovered, inputValues); + } + } +} + +TEST(FixedBitArrayTests, bulkSet64WithBaselinePartialRange) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int bitWidth : {4, 16, 32, 33, 40, 48, 56, 57, 58, 64}) { + SCOPED_TRACE(fmt::format("bitWidth={}", bitWidth)); + const uint64_t maxElement = bitWidth == 64 + ? std::numeric_limits::max() + : (1ULL << bitWidth) - 1; + const uint64_t baseline = folly::Random::rand64(rng) % (maxElement / 2 + 1); + const uint64_t valueRange = maxElement - baseline; + + const int elementCount = 200; + const auto bufferBytes = + nimble::FixedBitArray::bufferSize(elementCount, bitWidth); + auto buffer = std::make_unique(bufferBytes); + std::memset(buffer.get(), 0, bufferBytes); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + + const int offset = 50; + const int count = 100; + std::vector inputValues(count); + for (int i = 0; i < count; ++i) { + const uint64_t residual = + valueRange == std::numeric_limits::max() + ? folly::Random::rand64(rng) + : folly::Random::rand64(rng) % (valueRange + 1); + inputValues[i] = residual + baseline; + } + + fixedBitArray.bulkSetWithBaseline( + offset, count, inputValues.data(), baseline); + + for (int i = 0; i < count; ++i) { + ASSERT_EQ(fixedBitArray.get(offset + i), inputValues[i] - baseline) + << "i=" << i; + } + + // Round-trip the partial range through the bulk read path. A non-zero start + // offset exercises the byte-aligned masked load at a non-zero base pointer + // (its last element also reads into bufferSize's trailing slop). + std::vector recovered(count); + fixedBitArray.bulkGetWithBaseline( + offset, count, recovered.data(), baseline); + ASSERT_EQ(recovered, inputValues); + + for (int i = 0; i < offset; ++i) { + ASSERT_EQ(fixedBitArray.get(i), 0) + << "pre-range slot " << i << " should be zero"; + } + for (int i = offset + count; i < elementCount; ++i) { + ASSERT_EQ(fixedBitArray.get(i), 0) + << "post-range slot " << i << " should be zero"; + } + } +} + +// Packs random values of element type T through the unified +// bulkSetWithBaseline and confirms the stored residuals round-trip. bitWidth is +// chosen so values fit in T. +template +void verifyBulkSetWithBaseline( + int bitWidth, + int elementCount, + std::mt19937& rng) { + SCOPED_TRACE( + fmt::format( + "typeBytes={} bitWidth={} elementCount={}", + sizeof(T), + bitWidth, + elementCount)); + const auto bufferBytes = + nimble::FixedBitArray::bufferSize(elementCount, bitWidth); + auto buffer = std::make_unique(bufferBytes); + std::memset(buffer.get(), 0, bufferBytes); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + + const T baseline = static_cast(folly::Random::rand32(rng) % 7); + const uint64_t residualMask = + bitWidth == 64 ? ~0ULL : ((1ULL << bitWidth) - 1); + std::vector inputValues(elementCount); + std::vector expectedResiduals(elementCount); + for (int i = 0; i < elementCount; ++i) { + expectedResiduals[i] = folly::Random::rand64(rng) & residualMask; + inputValues[i] = static_cast(expectedResiduals[i] + baseline); + } + + fixedBitArray.bulkSetWithBaseline( + 0, elementCount, inputValues.data(), baseline); + + std::vector actualResiduals(elementCount); + for (int i = 0; i < elementCount; ++i) { + actualResiduals[i] = fixedBitArray.get(i); + } + EXPECT_EQ(actualResiduals, expectedResiduals); +} + +TEST(FixedBitArrayTests, bulkSetWithBaselineDispatchesByElementWidth) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + const int smallCount = 1 + folly::Random::rand32(rng) % kMaxElements; + // One element type per compile-time branch: 1- and 2-byte exercise the widen + // path, 4-byte delegates to bulkSet32, 8-byte to bulkSet64. + verifyBulkSetWithBaseline(/*bitWidth=*/5, smallCount, rng); + verifyBulkSetWithBaseline(/*bitWidth=*/12, smallCount, rng); + verifyBulkSetWithBaseline(/*bitWidth=*/20, smallCount, rng); + verifyBulkSetWithBaseline(/*bitWidth=*/40, smallCount, rng); + + // Exercise the narrow widen path across multiple stack chunks + // (length > kWidenChunk = 1024), including a partial final chunk. + verifyBulkSetWithBaseline( + /*bitWidth=*/5, /*elementCount=*/2500, rng); + verifyBulkSetWithBaseline( + /*bitWidth=*/12, /*elementCount=*/2500, rng); +} + +// Round-trips values of element type T through bulkSetWithBaseline + +// bulkGetWithBaseline and confirms each compile-time branch recovers them. +template +void verifyBulkGetWithBaseline( + int bitWidth, + int elementCount, + std::mt19937& rng) { + SCOPED_TRACE( + fmt::format( + "typeBytes={} bitWidth={} elementCount={}", + sizeof(T), + bitWidth, + elementCount)); + const auto bufferBytes = + nimble::FixedBitArray::bufferSize(elementCount, bitWidth); + auto buffer = std::make_unique(bufferBytes); + std::memset(buffer.get(), 0, bufferBytes); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + + const T baseline = static_cast(folly::Random::rand32(rng) % 7); + const uint64_t residualMask = + bitWidth == 64 ? ~0ULL : ((1ULL << bitWidth) - 1); + std::vector inputValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + inputValues[i] = + static_cast((folly::Random::rand64(rng) & residualMask) + baseline); + } + + fixedBitArray.bulkSetWithBaseline( + 0, elementCount, inputValues.data(), baseline); + + std::vector actual(elementCount); + fixedBitArray.bulkGetWithBaseline(0, elementCount, actual.data(), baseline); + EXPECT_EQ(actual, inputValues); +} + +TEST(FixedBitArrayTests, bulkGetWithBaselineDispatchesByElementWidth) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + const int smallCount = 1 + folly::Random::rand32(rng) % kMaxElements; + verifyBulkGetWithBaseline(/*bitWidth=*/5, smallCount, rng); + verifyBulkGetWithBaseline(/*bitWidth=*/12, smallCount, rng); + verifyBulkGetWithBaseline(/*bitWidth=*/20, smallCount, rng); + verifyBulkGetWithBaseline(/*bitWidth=*/40, smallCount, rng); + + // Narrow read path across multiple stack chunks (length > kWidenChunk). + verifyBulkGetWithBaseline( + /*bitWidth=*/5, /*elementCount=*/2500, rng); + verifyBulkGetWithBaseline( + /*bitWidth=*/12, /*elementCount=*/2500, rng); +} diff --git a/velox/dwio/nimble/common/tests/GTestUtils.h b/velox/dwio/nimble/common/tests/GTestUtils.h new file mode 100644 index 00000000000..d5ae3881ef6 --- /dev/null +++ b/velox/dwio/nimble/common/tests/GTestUtils.h @@ -0,0 +1,87 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +// gtest v1.10 deprecated *_TEST_CASE in favor of *_TEST_SUITE. These +// macros are provided for portability between different gtest versions. +#ifdef TYPED_TEST_SUITE +#define NIMBLE_TYPED_TEST_SUITE TYPED_TEST_SUITE +#else +#define NIMBLE_TYPED_TEST_SUITE TYPED_TEST_CASE +#endif + +#ifdef INSTANTIATE_TEST_SUITE_P +#define NIMBLE_INSTANTIATE_TEST_SUITE_P INSTANTIATE_TEST_SUITE_P +#else +#define NIMBLE_INSTANTIATE_TEST_SUITE_P INSTANTIATE_TEST_CASE_P +#endif + +// The void static cast supresses the "unused expression result" warning in +// clang. +#define NIMBLE_ASSERT_THROW_IMPL(_type, _expression, _errorMessage) \ + try { \ + static_cast(_expression); \ + FAIL() << "Expected an exception"; \ + } catch (const _type& e) { \ + ASSERT_TRUE(e.errorMessage().find(_errorMessage) != std::string::npos) \ + << "Expected error message to contain '" << (_errorMessage) \ + << "', but received '" << e.errorMessage() << "'."; \ + } + +#define NIMBLE_ASSERT_THROW(_expression, _errorMessage) \ + NIMBLE_ASSERT_THROW_IMPL( \ + facebook::nimble::NimbleException, _expression, _errorMessage) + +#define NIMBLE_ASSERT_USER_THROW(_expression, _errorMessage) \ + NIMBLE_ASSERT_THROW_IMPL( \ + facebook::nimble::NimbleUserError, _expression, _errorMessage) + +#define NIMBLE_ASSERT_FILE_THROW(_expression, _errorMessage) \ + try { \ + static_cast(_expression); \ + FAIL() << "Expected a corrupted file exception"; \ + } catch (const facebook::nimble::NimbleUserError& exception) { \ + ASSERT_EQ(exception.errorCode(), "CORRUPTED_FILE"); \ + ASSERT_NE(exception.errorMessage().find(_errorMessage), std::string::npos) \ + << "Expected error message to contain '" << (_errorMessage) \ + << "', but received '" << exception.errorMessage() << "'."; \ + } + +#define NIMBLE_ASSERT_RUNTIME_THROW(_expression, _errorMessage) \ + NIMBLE_ASSERT_THROW_IMPL( \ + facebook::nimble::NimbleRuntimeError, _expression, _errorMessage) + +#ifndef NDEBUG +#define DEBUG_ONLY_TEST(test_fixture, test_name) TEST(test_fixture, test_name) +#define DEBUG_ONLY_TEST_F(test_fixture, test_name) \ + TEST_F(test_fixture, test_name) +#define DEBUG_ONLY_TEST_P(test_fixture, test_name) \ + TEST_P(test_fixture, test_name) +#define DEBUG_ONLY_CO_TEST_F(test_fixture, test_name) \ + CO_TEST_F(test_fixture, test_name) +#else +#define DEBUG_ONLY_TEST(test_fixture, test_name) \ + TEST(test_fixture, DISABLED_##test_name) +#define DEBUG_ONLY_TEST_F(test_fixture, test_name) \ + TEST_F(test_fixture, DISABLED_##test_name) +#define DEBUG_ONLY_TEST_P(test_fixture, test_name) \ + TEST_P(test_fixture, DISABLED_##test_name) +#define DEBUG_ONLY_CO_TEST_F(test_fixture, test_name) \ + CO_TEST_F(test_fixture, DISABLED_test_name) +#endif diff --git a/velox/dwio/nimble/common/tests/MetricsLoggerTest.cpp b/velox/dwio/nimble/common/tests/MetricsLoggerTest.cpp new file mode 100644 index 00000000000..dce62d806ec --- /dev/null +++ b/velox/dwio/nimble/common/tests/MetricsLoggerTest.cpp @@ -0,0 +1,139 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include "velox/dwio/nimble/common/MetricsLogger.h" + +namespace facebook::nimble::test { + +// --- StripeLoadMetrics::serialize --- + +TEST(MetricsLoggerTest, stripeLoadMetricsSerialize) { + StripeLoadMetrics metrics{ + .stripeIndex = 3, + .rowsInStripe = 1000, + .streamCount = 5, + .totalStreamSize = 4096, + .cpuUsec = 100, + .wallTimeUsec = 200, + }; + + auto obj = metrics.serialize(); + EXPECT_EQ(obj["stripeIndex"].asInt(), 3); + EXPECT_EQ(obj["rowsInStripe"].asInt(), 1000); + EXPECT_EQ(obj["streamCount"].asInt(), 5); + EXPECT_EQ(obj["totalStreamSize"].asInt(), 4096); +} + +TEST(MetricsLoggerTest, stripeLoadMetricsSerializeDefaults) { + StripeLoadMetrics metrics{ + .stripeIndex = 0, + .rowsInStripe = 0, + .cpuUsec = 0, + .wallTimeUsec = 0, + }; + + auto obj = metrics.serialize(); + EXPECT_EQ(obj["streamCount"].asInt(), 0); + EXPECT_EQ(obj["totalStreamSize"].asInt(), 0); +} + +// --- StripeFlushMetrics::serialize --- + +TEST(MetricsLoggerTest, stripeFlushMetricsSerialize) { + StripeFlushMetrics metrics{ + .inputSize = 50000, + .rowCount = 10000, + .stripeSize = 30000, + .trackedMemory = 65536, + .flushCpuUsec = 500, + .flushWallTimeUsec = 600, + }; + + auto obj = metrics.serialize(); + EXPECT_EQ(obj["inputSize"].asInt(), 50000); + EXPECT_EQ(obj["rowCount"].asInt(), 10000); + EXPECT_EQ(obj["stripeSize"].asInt(), 30000); + EXPECT_EQ(obj["trackedMemory"].asInt(), 65536); +} + +// --- FileCloseMetrics::serialize --- + +TEST(MetricsLoggerTest, fileCloseMetricsSerialize) { + FileCloseMetrics metrics{ + .rowCount = 100000, + .inputSize = 500000, + .stripeCount = 5, + .fileSize = 300000, + .encodingCpuNs = 2000000, + .encodingWallNs = 3000000, + }; + + auto obj = metrics.serialize(); + EXPECT_EQ(obj["rowCount"].asInt(), 100000); + EXPECT_EQ(obj["inputSize"].asInt(), 500000); + EXPECT_EQ(obj["stripeCount"].asInt(), 5); + EXPECT_EQ(obj["fileSize"].asInt(), 300000); + EXPECT_EQ(obj["encodingCpuNs"].asInt(), 2000000); + EXPECT_EQ(obj["encodingWallNs"].asInt(), 3000000); +} + +// --- LoggingScope --- + +TEST(MetricsLoggerTest, loggingScopeNullInitially) { + // Without any LoggingScope, getLogger should return nullptr + EXPECT_EQ(LoggingScope::getLogger(), nullptr); +} + +TEST(MetricsLoggerTest, loggingScopeSetAndClear) { + MetricsLogger logger; + { + LoggingScope scope(logger); + EXPECT_EQ(LoggingScope::getLogger(), &logger); + } + // After scope destruction, logger should be null again + EXPECT_EQ(LoggingScope::getLogger(), nullptr); +} + +TEST(MetricsLoggerTest, loggingScopeNestedScopes) { + MetricsLogger logger1; + MetricsLogger logger2; + { + LoggingScope scope1(logger1); + EXPECT_EQ(LoggingScope::getLogger(), &logger1); + { + LoggingScope scope2(logger2); + EXPECT_EQ(LoggingScope::getLogger(), &logger2); + } + // After inner scope destruction, logger is null (not restored to logger1) + EXPECT_EQ(LoggingScope::getLogger(), nullptr); + } +} + +// --- Default MetricsLogger virtual methods are safe no-ops --- + +TEST(MetricsLoggerTest, defaultMethodsAreNoOps) { + MetricsLogger logger; + // These should all execute without error + logger.logException(LogOperation::Write, "test error"); + logger.logStripeLoad(StripeLoadMetrics{}); + logger.logStripeFlush(StripeFlushMetrics{}); + logger.logFileClose(FileCloseMetrics{}); + logger.logCompressionContext("test context"); +} + +} // namespace facebook::nimble::test diff --git a/velox/dwio/nimble/common/tests/NimbleCompare.h b/velox/dwio/nimble/common/tests/NimbleCompare.h new file mode 100644 index 00000000000..4c1a9225d6a --- /dev/null +++ b/velox/dwio/nimble/common/tests/NimbleCompare.h @@ -0,0 +1,85 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +namespace facebook::nimble { + +template +class NimbleCompare { + public: + static bool equals(const T& a, const T& b); +}; + +template +inline bool NimbleCompare::equals(const T& a, const T& b) { + return a == b; +} + +template +class NimbleCompare< + T, + std::enable_if_t || std::is_same_v>> { + public: + // double or float + using FloatingType = + typename std::conditional, double, float>::type; + // 64bit integer or 32 bit integer + using IntegralType = typename std:: + conditional, int64_t, int32_t>::type; + + static_assert(sizeof(FloatingType) == sizeof(IntegralType)); + + static bool equals(const T& a, const T& b); + + // This will be convenient when for debug, logging. + static IntegralType asInteger(const T& a); +}; + +template +inline bool NimbleCompare< + T, + std::enable_if_t || std::is_same_v>>:: + equals(const T& a, const T& b) { + // For floating point types, we do bit-wise comparison, for other types, + // just use the original ==. + // TODO: handle NaN. + return *(reinterpret_cast(&(a))) == + *(reinterpret_cast(&(b))); +} + +template +inline typename NimbleCompare< + T, + std::enable_if_t< + std::is_same_v || std::is_same_v>>::IntegralType +NimbleCompare< + T, + std::enable_if_t || std::is_same_v>>:: + asInteger(const T& a) { + return *(reinterpret_cast(&(a))); +} + +template +struct NimbleComparator { + constexpr bool operator()(const T& a, const T& b) const { + return NimbleCompare::equals(a, b); + } +}; + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/common/tests/NimbleFileWriter.cpp b/velox/dwio/nimble/common/tests/NimbleFileWriter.cpp new file mode 100644 index 00000000000..d6abae43ef2 --- /dev/null +++ b/velox/dwio/nimble/common/tests/NimbleFileWriter.cpp @@ -0,0 +1,65 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "velox/dwio/nimble/common/tests/NimbleFileWriter.h" +#include "velox/dwio/nimble/common/Exceptions.h" +#include "velox/dwio/nimble/writer/Writer.h" +#include "velox/dwio/nimble/writer/WriterOptions.h" + +namespace facebook::nimble::test { + +std::string createNimbleFile( + velox::memory::MemoryPool& memoryPool, + const velox::VectorPtr& vector, + nimble::WriterOptions writerOptions, + bool flushAfterWrite) { + return createNimbleFile( + memoryPool, + std::vector{vector}, + std::move(writerOptions), + flushAfterWrite); +} + +std::string createNimbleFile( + velox::memory::MemoryPool& memoryPool, + const std::vector& vectors, + nimble::WriterOptions writerOptions, + bool flushAfterWrite) { + std::string file; + auto writeFile = std::make_unique(&file); + + NIMBLE_CHECK_GT(vectors.size(), 0, "Expecting at least one input vector."); + auto& type = vectors[0]->type(); + + for (int i = 1; i < vectors.size(); ++i) { + NIMBLE_CHECK( + vectors[i]->type()->kindEquals(type), + "All vectors should have the same schema."); + } + + nimble::Writer writer( + type, std::move(writeFile), memoryPool, std::move(writerOptions)); + for (const auto& vector : vectors) { + writer.write(vector); + if (flushAfterWrite) { + writer.flush(); + } + } + writer.close(); + + return file; +} + +} // namespace facebook::nimble::test diff --git a/velox/dwio/nimble/common/tests/NimbleFileWriter.h b/velox/dwio/nimble/common/tests/NimbleFileWriter.h new file mode 100644 index 00000000000..e33ab2a5c97 --- /dev/null +++ b/velox/dwio/nimble/common/tests/NimbleFileWriter.h @@ -0,0 +1,35 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "velox/dwio/nimble/writer/WriterOptions.h" +#include "velox/vector/BaseVector.h" + +namespace facebook::nimble::test { + +std::string createNimbleFile( + velox::memory::MemoryPool& memoryPool, + const std::vector& vectors, + nimble::WriterOptions writerOptions = {}, + bool flushAfterWrite = true); + +std::string createNimbleFile( + velox::memory::MemoryPool& memoryPool, + const velox::VectorPtr& vector, + nimble::WriterOptions writerOptions = {}, + bool flushAfterWrite = true); + +} // namespace facebook::nimble::test diff --git a/velox/dwio/nimble/common/tests/StatsUtilTest.cpp b/velox/dwio/nimble/common/tests/StatsUtilTest.cpp new file mode 100644 index 00000000000..489f651fa56 --- /dev/null +++ b/velox/dwio/nimble/common/tests/StatsUtilTest.cpp @@ -0,0 +1,128 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "velox/dwio/nimble/common/StatsUtil.h" + +using namespace facebook::nimble; + +namespace { + +template +void expectIntegralMinMax(T expectedMin, T expectedMax, std::vector values) { + const auto minMax = findMinMax(std::span{values}); + EXPECT_EQ(minMax.min, expectedMin); + EXPECT_EQ(minMax.max, expectedMax); +} + +template +void expectFloatingPointMinMax( + T expectedMin, + T expectedMax, + std::vector values) { + const auto minMax = findMinMax(std::span{values}); + ASSERT_TRUE(minMax.has_value()); + if constexpr (std::is_same_v) { + EXPECT_FLOAT_EQ(minMax->min, expectedMin); + EXPECT_FLOAT_EQ(minMax->max, expectedMax); + } else { + EXPECT_DOUBLE_EQ(minMax->min, expectedMin); + EXPECT_DOUBLE_EQ(minMax->max, expectedMax); + } +} + +} // namespace + +TEST(StatsUtilTest, findMinMaxHandlesIntegralTypes) { + std::vector signedValues(257); + for (size_t i = 0; i < signedValues.size(); ++i) { + signedValues[i] = static_cast(i) - 128; + } + signedValues[13] = std::numeric_limits::min(); + signedValues[199] = std::numeric_limits::max(); + expectIntegralMinMax( + std::numeric_limits::min(), + std::numeric_limits::max(), + signedValues); + + std::vector unsignedValues(257); + for (size_t i = 0; i < unsignedValues.size(); ++i) { + unsignedValues[i] = 1000 + i; + } + unsignedValues[17] = 0; + unsignedValues[211] = std::numeric_limits::max(); + expectIntegralMinMax( + 0, std::numeric_limits::max(), unsignedValues); +} + +TEST(StatsUtilTest, findMinMaxHandlesStringViews) { + using namespace std::string_view_literals; + + // Bounds trail the first value, so seeding from values.front() and skipping + // the max check when a value lowers the min must still find both. + std::vector values = {"mmm"sv, "aaa"sv, "zzz"sv, "qqq"sv}; + auto minMax = findMinMax(std::span(values)); + EXPECT_EQ(minMax.min, "aaa"sv); + EXPECT_EQ(minMax.max, "zzz"sv); + + // Ordering is by bytes and length, not NUL termination. + std::vector withNuls = {"a\0a"sv, "a\0"sv, "a\0z"sv}; + minMax = findMinMax(std::span(withNuls)); + EXPECT_EQ(minMax.min, "a\0"sv); + EXPECT_EQ(minMax.max, "a\0z"sv); + + std::vector single = {"only"sv}; + minMax = findMinMax(std::span(single)); + EXPECT_EQ(minMax.min, "only"sv); + EXPECT_EQ(minMax.max, "only"sv); + + // The returned bounds are views into the input, not copies. + std::vector aliased = {"beta"sv, "alpha"sv}; + minMax = findMinMax(std::span(aliased)); + EXPECT_EQ(minMax.min.data(), aliased[1].data()); + EXPECT_EQ(minMax.max.data(), aliased[0].data()); +} + +TEST(StatsUtilTest, findMinMaxHandlesFloatingPointTypes) { + std::vector floatValues(257, 1.25F); + floatValues[11] = -123.5F; + floatValues[211] = 456.75F; + expectFloatingPointMinMax(-123.5F, 456.75F, floatValues); + + std::vector doubleValues(257, 2.5); + doubleValues[31] = -9876.5; + doubleValues[199] = 54321.25; + expectFloatingPointMinMax(-9876.5, 54321.25, doubleValues); +} + +TEST(StatsUtilTest, findMinMaxReturnsEmptyForFloatingPointNaN) { + std::vector floatValues(257, 1.25F); + floatValues[211] = std::numeric_limits::quiet_NaN(); + EXPECT_EQ(findMinMax(std::span{floatValues}), std::nullopt); + + std::vector doubleValues(257, 2.5); + doubleValues[31] = std::numeric_limits::quiet_NaN(); + EXPECT_EQ(findMinMax(std::span{doubleValues}), std::nullopt); +} diff --git a/velox/dwio/nimble/common/tests/TestUtils.h b/velox/dwio/nimble/common/tests/TestUtils.h new file mode 100644 index 00000000000..edfa99d7bc3 --- /dev/null +++ b/velox/dwio/nimble/common/tests/TestUtils.h @@ -0,0 +1,548 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "folly/Random.h" +#include "folly/Synchronized.h" +#include "velox/common/file/File.h" +#include "velox/common/file/FileSystems.h" +#include "velox/common/memory/Memory.h" +#include "velox/dwio/nimble/common/Buffer.h" +#include "velox/dwio/nimble/common/Types.h" +#include "velox/dwio/nimble/common/Vector.h" + +// Utilities to support testing in nimble. + +namespace facebook::nimble::testing { + +// Adds random data covering the whole range of the data type. +template +void addRandomData(RNG&& rng, int rowCount, Vector* data, Buffer* buffer); + +// Draw an int uniformly from [min, max). +class Util { + public: + Util(velox::memory::MemoryPool& memoryPool) : memoryPool_(memoryPool) {} + + // Makes between 1 and maxRows random values over T's data range. + // For strings we improvise a bit. + template + Vector makeRandomData(RNG&& rng, uint32_t maxRows, Buffer* buffer); + + // A random length vector of all the same data. + template + Vector makeConstantData(RNG&& rng, uint32_t maxRows, Buffer* buffer); + + // A random length vector of runs of various lengths. + template + Vector makeRLEData(RNG&& rng, uint32_t maxRows, Buffer* buffer); + + // Makes data compatible with our HuffmanColumn (namely in the + // range [0, 4096)) for integer data only. + template + Vector makeHuffmanData(RNG&& rng, uint32_t maxRows, Buffer* buffer); + + // A Vector of all the Make* Vectors. + template + std::vector> + makeDataPatterns(RNG&& rng, uint32_t maxRows, Buffer* buffer); + + // Trims data down so that it will be friendly to sums, i.e. so it won't + // overflow the sum type. This mainly applies to 64-bit integers. + template + Vector sumFriendlyData(const Vector& data); + + inline uint64_t uniformRandom(uint64_t min, uint64_t max) { + CHECK_LT(min, max); + return min + (folly::Random::rand32() % (max - min)); + } + + // Draw an int uniformly from [0, max); + inline uint64_t uniformRandom(uint64_t max) { + return uniformRandom(0, max); + } + + private: + velox::memory::MemoryPool& memoryPool_; +}; + +// +// End of public API. Implementation follows. +// + +template +inline void +addRandomData(RNG&& rng, int rowCount, Vector* data, Buffer* /* buffer */) { + // Half the time only add positive data. + if (!std::is_signed_v || folly::Random::rand32() % 2) { + for (int i = 0; i < rowCount; ++i) { + if constexpr (sizeof(T) > 4) { + const uint64_t rand = folly::Random::rand64(std::forward(rng)); + data->push_back(*reinterpret_cast(&rand)); + + } else { + const uint32_t rand = folly::Random::rand32(std::forward(rng)); + data->push_back(*reinterpret_cast(&rand)); + } + } + } else { + for (int i = 0; i < rowCount; ++i) { + if constexpr (sizeof(T) > 4) { + const uint64_t rand = + folly::Random::rand64(std::forward(rng)) & ((1ULL << 63) - 1); + data->push_back(*reinterpret_cast(&rand)); + } else { + const uint32_t rand = + folly::Random::rand32(std::forward(rng)) & ((1U << 31) - 1); + data->push_back(*reinterpret_cast(&rand)); + } + } + } +} + +template +inline void addRandomData( + RNG&& rng, + int rowCount, + Vector* data, + Buffer* /* buffer */) { + for (int i = 0; i < rowCount; ++i) { + const uint32_t rand = folly::Random::rand32(std::forward(rng)); + data->push_back(static_cast(rand)); + } +} + +template +inline void addRandomData( + RNG&& rng, + int rowCount, + Vector* data, + Buffer* /* buffer */) { + for (int i = 0; i < rowCount; ++i) { + const uint64_t rand = folly::Random::rand64(std::forward(rng)); + data->push_back(static_cast(rand)); + } +} + +template +inline void addRandomData( + RNG&& rng, + int rowCount, + Vector* data, + Buffer* buffer) { + for (int i = 0; i < rowCount; ++i) { + // This is a bit arbitrary, but lets stick to 50 char max string length. + // We can have some separate tests for large strings if we want. + const int len = folly::Random::rand32(std::forward(rng)) % 50; + char* pos = buffer->reserve(len); + for (int j = 0; j < len; ++j) { + pos[j] = folly::Random::rand32(std::forward(rng)) % 256; + } + data->emplace_back(pos, len); + } +} + +template +inline void addRandomData( + RNG&& rng, + int rowCount, + Vector* data, + Buffer* /* buffer */) { + for (int i = 0; i < rowCount; ++i) { + data->push_back(folly::Random::rand32(std::forward(rng)) & 1); + } +} + +template +Vector Util::makeRandomData(RNG&& rng, uint32_t maxRows, Buffer* buffer) { + Vector randomData(&memoryPool_); + const int rowCount = + 1 + folly::Random::rand32(std::forward(rng)) % maxRows; + randomData.reserve(rowCount); + addRandomData(std::forward(rng), rowCount, &randomData, buffer); + return randomData; +} + +template +Vector Util::makeConstantData(RNG&& rng, uint32_t maxRows, Buffer* buffer) { + Vector data = makeRandomData(std::forward(rng), maxRows, buffer); + for (int i = 1; i < data.size(); ++i) { + data[i] = data[0]; + } + return data; +} + +template +Vector Util::makeRLEData(RNG&& rng, uint32_t maxRows, Buffer* buffer) { + Vector data = makeRandomData(std::forward(rng), maxRows, buffer); + int index = 0; + while (index < data.size()) { + const uint32_t runLength = data.size() - index == 1 ? 1 + : 1 + + folly::Random::rand32(std::forward(rng)) % + (data.size() - index - 1); + for (int i = 0; i < runLength; ++i) { + data[index + i] = data[index]; + } + index += runLength; + } + return data; +} + +template +Vector +Util::makeHuffmanData(RNG&& rng, uint32_t maxRows, Buffer* /* buffer */) { + const int rowCount = + 1 + folly::Random::rand32(std::forward(rng)) % maxRows; + Vector huffmanData(&memoryPool_); + huffmanData.reserve(rowCount); + // Half the time draw all the symbols from within [0, symbolCount), + // emulating encoding a dictionary index, and the other half of the time + // emulate encoding a normal stream by drawing from the full [0, 4096) range. + const int symbolCount = + 1 + folly::Random::rand32(std::forward(rng)) % rowCount; + const int maxValue = + (folly::Random::rand32(std::forward(rng)) & 1) ? symbolCount : 4096; + for (uint32_t i = 0; i < rowCount; ++i) { + huffmanData.push_back( + folly::Random::rand32(std::forward(rng)) % maxValue); + } + return huffmanData; +} + +template +std::vector> +Util::makeDataPatterns(RNG&& rng, uint32_t maxRows, Buffer* buffer) { + std::vector> patterns; + patterns.push_back( + makeRandomData(std::forward(rng), maxRows, buffer)); + patterns.push_back( + makeConstantData(std::forward(rng), maxRows, buffer)); + patterns.push_back(makeRLEData(std::forward(rng), maxRows, buffer)); + if constexpr (isIntegralType()) { + patterns.push_back( + makeHuffmanData(std::forward(rng), maxRows, buffer)); + } + return patterns; +} + +template +Vector sumFriendlyDataHelper( + velox::memory::MemoryPool&, + const Vector& data) { + return data; +} + +template +Vector Util::sumFriendlyData(const Vector& data) { + return data; +} + +template <> +inline Vector Util::sumFriendlyData(const Vector& data) { + Vector sumData(&memoryPool_); + // 15 is somewhat arbitrary shift, but this is testing code, so meh. + for (int64_t datum : data) { + sumData.push_back(datum >> 15); + } + return sumData; +} + +template <> +inline Vector Util::sumFriendlyData(const Vector& data) { + Vector sumData(&memoryPool_); + // 15 is somewhat arbitrary shift, but this is testing code, so meh. + for (uint64_t datum : data) { + sumData.push_back(datum >> 15); + } + return sumData; +} + +struct Chunk { + uint64_t offset; + uint64_t size; +}; + +// Wrapper around InMemoryReadFile (can't inherit, as InMemoryReadFile is final) +// which tracks all offsets and sizes being read. This is used to verify Nimble +// reader coalese behavior. +class InMemoryTrackableReadFile final : public velox::ReadFile { + public: + explicit InMemoryTrackableReadFile( + std::string_view file, + bool shouldProduceChainedBuffers) + : file_{file}, + shouldProduceChainedBuffers_{shouldProduceChainedBuffers} {} + + std::string_view pread( + uint64_t offset, + uint64_t length, + void* buf, + const velox::FileIoContext& context = {}) const final { + chunks_.wlock()->push_back({offset, length}); + return file_.pread(offset, length, buf, context); + } + + std::string pread( + uint64_t offset, + uint64_t length, + const velox::FileIoContext& context = {}) const final { + chunks_.wlock()->push_back({offset, length}); + return file_.pread(offset, length, context); + } + + uint64_t preadv( + uint64_t offset, + const std::vector>& buffers, + const velox::FileIoContext& /*context*/ = {}) const final { + uint64_t totalRead = 0; + for (const auto& range : buffers) { + if (range.data() != nullptr) { + file_.pread(offset, range.size(), range.data()); + } + chunks_.wlock()->push_back({offset, range.size()}); + offset += range.size(); + totalRead += range.size(); + } + return totalRead; + } + + uint64_t preadv( + folly::Range regions, + folly::Range iobufs, + const velox::FileIoContext& context = {}) const override { + VELOX_CHECK_EQ(regions.size(), iobufs.size()); + uint64_t length = 0; + for (size_t i = 0; i < regions.size(); ++i) { + const auto& region = regions[i]; + length += region.length; + auto& output = iobufs[i]; + if (shouldProduceChainedBuffers_) { + chunks_.wlock()->push_back({region.offset, region.length}); + uint64_t splitPoint = region.length / 2; + output = folly::IOBuf(folly::IOBuf::CREATE, splitPoint); + file_.pread(region.offset, splitPoint, output.writableData()); + output.append(splitPoint); + const uint64_t nextLength = region.length - splitPoint; + auto next = folly::IOBuf::create(nextLength); + file_.pread( + region.offset + splitPoint, nextLength, next->writableData()); + next->append(nextLength); + output.appendChain(std::move(next)); + } else { + output = folly::IOBuf(folly::IOBuf::CREATE, region.length); + pread(region.offset, region.length, output.writableData(), context); + output.append(region.length); + } + } + + return length; + } + + uint64_t size() const final { + return file_.size(); + } + + uint64_t memoryUsage() const final { + return file_.memoryUsage(); + } + + // Mainly for testing. Coalescing isn't helpful for in memory data. + void setShouldCoalesce(bool shouldCoalesce) { + file_.setShouldCoalesce(shouldCoalesce); + } + + bool shouldCoalesce() const final { + return file_.shouldCoalesce(); + } + + std::vector chunks() { + return *chunks_.rlock(); + } + + void resetChunks() { + chunks_.wlock()->clear(); + } + + std::string getName() const override { + return ""; + } + + uint64_t getNaturalReadSize() const override { + return 1024; + } + + private: + velox::InMemoryReadFile file_; + bool shouldProduceChainedBuffers_; + mutable folly::Synchronized> chunks_; +}; + +// Wraps any ReadFile and tracks the maximum read offset across all pread calls. +// Used to verify that certain reads don't access regions beyond a boundary +// (e.g. metadata regions). +class TrackingReadFile : public velox::ReadFile { + public: + explicit TrackingReadFile(std::shared_ptr delegate) + : delegate_(std::move(delegate)) {} + + void setReadDelayUs(uint64_t delayUs) { + readDelayUs_ = delayUs; + } + + void setReadError(std::exception_ptr error) { + std::lock_guard l(mu_); + readError_ = std::move(error); + } + + void clearReadError() { + std::lock_guard l(mu_); + readError_ = nullptr; + } + + std::string_view pread( + uint64_t offset, + uint64_t length, + void* buf, + const velox::FileIoContext& context = {}) const override { + updateMaxReadOffset(offset + length); + maybeThrow(); + maybeDelay(); + return delegate_->pread(offset, length, buf, context); + } + + std::string pread( + uint64_t offset, + uint64_t length, + const velox::FileIoContext& context = {}) const override { + updateMaxReadOffset(offset + length); + maybeThrow(); + maybeDelay(); + return delegate_->pread(offset, length, context); + } + + uint64_t preadv( + uint64_t offset, + const std::vector>& buffers, + const velox::FileIoContext& context = {}) const override { + uint64_t totalLength = 0; + for (const auto& buffer : buffers) { + totalLength += buffer.size(); + } + updateMaxReadOffset(offset + totalLength); + ioGroups_.wlock()->push_back({offset, totalLength}); + maybeThrow(); + maybeDelay(); + return delegate_->preadv(offset, buffers, context); + } + + uint64_t preadv( + folly::Range regions, + folly::Range iobufs, + const velox::FileIoContext& context = {}) const override { + for (const auto& region : regions) { + updateMaxReadOffset(region.offset + region.length); + } + maybeThrow(); + maybeDelay(); + return delegate_->preadv(regions, iobufs, context); + } + + uint64_t size() const override { + return delegate_->size(); + } + + uint64_t memoryUsage() const override { + return delegate_->memoryUsage(); + } + + bool shouldCoalesce() const override { + return delegate_->shouldCoalesce(); + } + + std::string getName() const override { + return delegate_->getName(); + } + + uint64_t getNaturalReadSize() const override { + return delegate_->getNaturalReadSize(); + } + + // Returns the maximum end offset (offset + length) across all reads. + uint64_t maxReadOffset() const { + return maxReadOffset_.load(std::memory_order_relaxed); + } + + void resetMaxReadOffset() { + maxReadOffset_.store(0, std::memory_order_relaxed); + } + + struct IoGroup { + uint64_t offset; + uint64_t size; + + std::string debugString() const { + return fmt::format("[offset={}, size={}]", offset, size); + } + }; + + std::vector ioGroups() const { + return *ioGroups_.rlock(); + } + + void resetIoGroups() { + ioGroups_.wlock()->clear(); + } + + private: + void updateMaxReadOffset(uint64_t endOffset) const { + auto current = maxReadOffset_.load(std::memory_order_relaxed); + while (endOffset > current && + !maxReadOffset_.compare_exchange_weak( + current, endOffset, std::memory_order_relaxed)) { + } + } + + void maybeThrow() const { + std::lock_guard l(mu_); + if (readError_ != nullptr) { + std::rethrow_exception(readError_); + } + } + + void maybeDelay() const { + const auto delayUs = readDelayUs_.load(std::memory_order_relaxed); + if (delayUs > 0) { + std::this_thread::sleep_for(std::chrono::microseconds(delayUs)); + } + } + + const std::shared_ptr delegate_; + mutable std::atomic_uint64_t maxReadOffset_{0}; + mutable folly::Synchronized> ioGroups_; + std::atomic_uint64_t readDelayUs_{0}; + mutable std::mutex mu_; + std::exception_ptr readError_; +}; + +} // namespace facebook::nimble::testing diff --git a/velox/dwio/nimble/common/tests/TypesTest.cpp b/velox/dwio/nimble/common/tests/TypesTest.cpp new file mode 100644 index 00000000000..333e0a9ef7f --- /dev/null +++ b/velox/dwio/nimble/common/tests/TypesTest.cpp @@ -0,0 +1,385 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include +#include + +#include "velox/dwio/nimble/common/NimbleException.h" +#include "velox/dwio/nimble/common/Types.h" +#include "velox/dwio/nimble/common/tests/GTestUtils.h" + +namespace facebook::nimble::test { + +// --- toString for EncodingType --- + +TEST(TypesTest, encodingTypeStringConversion) { + const std::vector> testCases{ + {EncodingType::Trivial, "Trivial"}, + {EncodingType::RLE, "RLE"}, + {EncodingType::Dictionary, "Dictionary"}, + {EncodingType::FixedBitWidth, "FixedBitWidth"}, + {EncodingType::Sentinel, "Sentinel"}, + {EncodingType::Nullable, "Nullable"}, + {EncodingType::SparseBool, "SparseBool"}, + {EncodingType::Varint, "Varint"}, + {EncodingType::Delta, "Delta"}, + {EncodingType::Constant, "Constant"}, + {EncodingType::MainlyConstant, "MainlyConstant"}, + {EncodingType::Prefix, "Prefix"}, + {EncodingType::ALP, "ALP"}, + {EncodingType::PFOR, "PFOR"}, + {EncodingType::SimdForBitpack, "SimdForBitpack"}, + {EncodingType::BlockBitPacking, "BlockBitPacking"}, + {EncodingType::SubIntSplit, "SubIntSplit"}, + {EncodingType::FrequencyPartition, "FrequencyPartition"}, + {EncodingType::FOR, "FOR"}, + {EncodingType::Fsst, "Fsst"}, + {EncodingType::Huffman, "Huffman"}, + {EncodingType::DeltaBlock, "DeltaBlock"}, + }; + for (const auto& [type, name] : testCases) { + SCOPED_TRACE(name); + EXPECT_EQ(toString(type), name); + EXPECT_EQ(toEncodingType(name), type); + } +} + +TEST(TypesTest, encodingTypeToStringUnknown) { + auto result = toString(static_cast(255)); + EXPECT_NE(result.find("Unknown"), std::string::npos); +} + +TEST(TypesTest, toEncodingTypeUnknown) { + EXPECT_ANY_THROW(toEncodingType("unknown")); +} + +TEST(TypesTest, readOnlyEncoding) { + EXPECT_FALSE(isReadOnlyEncoding(EncodingType::PFOR)); + EXPECT_FALSE(isReadOnlyEncoding("PFOR")); + EXPECT_TRUE(isReadOnlyEncoding(EncodingType::FOR)); + EXPECT_TRUE(isReadOnlyEncoding("FOR")); + EXPECT_FALSE(isReadOnlyEncoding(EncodingType::Trivial)); + EXPECT_FALSE(isReadOnlyEncoding("Trivial")); + EXPECT_FALSE(isReadOnlyEncoding("unknown")); +} + +TEST(TypesTest, encodingTypeStreamOperator) { + std::ostringstream ss; + ss << EncodingType::Trivial; + EXPECT_EQ(ss.str(), "Trivial"); +} + +// --- toString for DataType --- + +TEST(TypesTest, dataTypeToString) { + EXPECT_EQ(toString(DataType::Int8), "Int8"); + EXPECT_EQ(toString(DataType::Uint8), "Uint8"); + EXPECT_EQ(toString(DataType::Int16), "Int16"); + EXPECT_EQ(toString(DataType::Uint16), "Uint16"); + EXPECT_EQ(toString(DataType::Int32), "Int32"); + EXPECT_EQ(toString(DataType::Uint32), "Uint32"); + EXPECT_EQ(toString(DataType::Int64), "Int64"); + EXPECT_EQ(toString(DataType::Uint64), "Uint64"); + EXPECT_EQ(toString(DataType::Float), "Float"); + EXPECT_EQ(toString(DataType::Double), "Double"); + EXPECT_EQ(toString(DataType::Bool), "Bool"); + EXPECT_EQ(toString(DataType::String), "String"); +} + +TEST(TypesTest, dataTypeToStringUnknown) { + auto result = toString(static_cast(200)); + EXPECT_NE(result.find("Unknown"), std::string::npos); +} + +TEST(TypesTest, dataTypeStreamOperator) { + std::ostringstream ss; + ss << DataType::Int32; + EXPECT_EQ(ss.str(), "Int32"); +} + +TEST(TypesTest, dataTypeFmtFormat) { + EXPECT_EQ(fmt::format("{}", DataType::Int32), "Int32"); + EXPECT_EQ(fmt::format("{}", DataType::String), "String"); +} + +// --- toString for CompressionType --- + +TEST(TypesTest, compressionTypeStringConversion) { + const std::vector> testCases{ + {CompressionType::Uncompressed, "Uncompressed"}, + {CompressionType::Zstd, "Zstd"}, + {CompressionType::MetaInternal, "MetaInternal"}, + {CompressionType::Lz4, "Lz4"}, + {CompressionType::OpenZL, "OpenZL"}, + }; + for (const auto& [type, name] : testCases) { + SCOPED_TRACE(name); + EXPECT_EQ(toString(type), name); + EXPECT_EQ(toCompressionType(name), type); + } +} + +TEST(TypesTest, compressionTypeToStringUnknown) { + auto result = toString(static_cast(200)); + EXPECT_NE(result.find("Unknown"), std::string::npos); +} + +TEST(TypesTest, toCompressionTypeUnknown) { + EXPECT_ANY_THROW(toCompressionType("unknown")); +} + +TEST(TypesTest, compressionTypeStreamOperator) { + std::ostringstream ss; + ss << CompressionType::Zstd; + EXPECT_EQ(ss.str(), "Zstd"); +} + +TEST(TypesTest, compressionTypeFmtFormatter) { + auto str = fmt::format("{}", CompressionType::Zstd); + EXPECT_EQ(str, "Zstd"); +} + +TEST(TypesTest, encodingTypeFmtFormatter) { + auto str = fmt::format("{}", EncodingType::Dictionary); + EXPECT_EQ(str, "Dictionary"); +} + +// --- toString for ChecksumType --- + +TEST(TypesTest, checksumTypeToString) { + EXPECT_EQ(toString(ChecksumType::XXH3_64), "XXH3_64"); +} + +TEST(TypesTest, checksumTypeToStringUnknown) { + auto result = toString(static_cast(200)); + EXPECT_NE(result.find("Unknown"), std::string::npos); +} + +// --- Variant --- + +TEST(TypesTest, variantSetGetInt) { + VariantType v; + Variant::set(v, 42); + EXPECT_EQ(Variant::get(v), 42); +} + +TEST(TypesTest, variantSetGetNegativeInt) { + VariantType v; + Variant::set(v, -12345); + EXPECT_EQ(Variant::get(v), -12345); +} + +TEST(TypesTest, variantSetGetDouble) { + VariantType v; + Variant::set(v, 3.14); + EXPECT_DOUBLE_EQ(Variant::get(v), 3.14); +} + +TEST(TypesTest, variantSetGetFloat) { + VariantType v; + Variant::set(v, 2.5f); + EXPECT_FLOAT_EQ(Variant::get(v), 2.5f); +} + +TEST(TypesTest, variantSetGetBool) { + VariantType v; + Variant::set(v, true); + EXPECT_TRUE(Variant::get(v)); + + Variant::set(v, false); + EXPECT_FALSE(Variant::get(v)); +} + +TEST(TypesTest, variantSetGetString) { + VariantType v; + Variant::set(v, std::string("hello")); + EXPECT_EQ(Variant::get(v), "hello"); +} + +TEST(TypesTest, variantSetGetStringView) { + VariantType v; + // string_view specialization stores as std::string internally + Variant::set(v, std::string_view("world")); + auto result = Variant::get(v); + EXPECT_EQ(result, "world"); +} + +TEST(TypesTest, variantSetGetUint8) { + VariantType v; + Variant::set(v, 200); + EXPECT_EQ(Variant::get(v), 200); +} + +TEST(TypesTest, variantSetGetInt16) { + VariantType v; + Variant::set(v, -300); + EXPECT_EQ(Variant::get(v), -300); +} + +// --- TypeTraits --- + +TEST(TypesTest, typeTraitsDataType) { + EXPECT_EQ(TypeTraits::dataType, DataType::Int8); + EXPECT_EQ(TypeTraits::dataType, DataType::Uint8); + EXPECT_EQ(TypeTraits::dataType, DataType::Int16); + EXPECT_EQ(TypeTraits::dataType, DataType::Uint16); + EXPECT_EQ(TypeTraits::dataType, DataType::Int32); + EXPECT_EQ(TypeTraits::dataType, DataType::Uint32); + EXPECT_EQ(TypeTraits::dataType, DataType::Int64); + EXPECT_EQ(TypeTraits::dataType, DataType::Uint64); + EXPECT_EQ(TypeTraits::dataType, DataType::Float); + EXPECT_EQ(TypeTraits::dataType, DataType::Double); + EXPECT_EQ(TypeTraits::dataType, DataType::Bool); + EXPECT_EQ(TypeTraits::dataType, DataType::String); + EXPECT_EQ(TypeTraits::dataType, DataType::String); +} + +TEST(TypesTest, typeTraitsPhysicalType) { + static_assert(std::is_same_v::physicalType, uint8_t>); + static_assert(std::is_same_v::physicalType, uint8_t>); + static_assert(std::is_same_v::physicalType, uint16_t>); + static_assert(std::is_same_v::physicalType, uint16_t>); + static_assert(std::is_same_v::physicalType, uint32_t>); + static_assert(std::is_same_v::physicalType, uint32_t>); + static_assert(std::is_same_v::physicalType, uint64_t>); + static_assert(std::is_same_v::physicalType, uint64_t>); + static_assert(std::is_same_v::physicalType, uint32_t>); + static_assert(std::is_same_v::physicalType, uint64_t>); + static_assert(std::is_same_v::physicalType, bool>); + static_assert( + std::is_same_v::physicalType, std::string>); + static_assert(std::is_same_v< + TypeTraits::physicalType, + std::string_view>); +} + +TEST(TypesTest, decodedValueWidth) { + EXPECT_EQ(decodedValueWidth(DataType::Bool), sizeof(bool)); + EXPECT_EQ(decodedValueWidth(DataType::Int8), sizeof(int8_t)); + EXPECT_EQ(decodedValueWidth(DataType::Uint8), sizeof(uint8_t)); + EXPECT_EQ(decodedValueWidth(DataType::Int16), sizeof(int16_t)); + EXPECT_EQ(decodedValueWidth(DataType::Uint16), sizeof(uint16_t)); + EXPECT_EQ(decodedValueWidth(DataType::Int32), sizeof(int32_t)); + EXPECT_EQ(decodedValueWidth(DataType::Uint32), sizeof(uint32_t)); + EXPECT_EQ(decodedValueWidth(DataType::Int64), sizeof(int64_t)); + EXPECT_EQ(decodedValueWidth(DataType::Uint64), sizeof(uint64_t)); + EXPECT_EQ(decodedValueWidth(DataType::Float), sizeof(float)); + EXPECT_EQ(decodedValueWidth(DataType::Double), sizeof(double)); + EXPECT_EQ(decodedValueWidth(DataType::String), sizeof(std::string_view)); + NIMBLE_ASSERT_THROW( + decodedValueWidth(DataType::Undefined), "Unsupported data type"); +} + +// --- Type predicates --- + +TEST(TypesTest, isIntegralType) { + EXPECT_TRUE(isIntegralType()); + EXPECT_TRUE(isIntegralType()); + EXPECT_TRUE(isIntegralType()); + EXPECT_TRUE(isIntegralType()); + EXPECT_TRUE(isIntegralType()); + EXPECT_TRUE(isIntegralType()); + EXPECT_TRUE(isIntegralType()); + EXPECT_TRUE(isIntegralType()); + EXPECT_FALSE(isIntegralType()); + EXPECT_FALSE(isIntegralType()); + EXPECT_FALSE(isIntegralType()); + EXPECT_FALSE(isIntegralType()); +} + +TEST(TypesTest, isSignedIntegralType) { + EXPECT_TRUE(isSignedIntegralType()); + EXPECT_TRUE(isSignedIntegralType()); + EXPECT_TRUE(isSignedIntegralType()); + EXPECT_TRUE(isSignedIntegralType()); + EXPECT_FALSE(isSignedIntegralType()); + EXPECT_FALSE(isSignedIntegralType()); +} + +TEST(TypesTest, isUnsignedIntegralType) { + EXPECT_TRUE(isUnsignedIntegralType()); + EXPECT_TRUE(isUnsignedIntegralType()); + EXPECT_TRUE(isUnsignedIntegralType()); + EXPECT_TRUE(isUnsignedIntegralType()); + EXPECT_FALSE(isUnsignedIntegralType()); + EXPECT_FALSE(isUnsignedIntegralType()); +} + +TEST(TypesTest, isOneByteIntegralType) { + EXPECT_TRUE(isOneByteIntegralType()); + EXPECT_TRUE(isOneByteIntegralType()); + EXPECT_FALSE(isOneByteIntegralType()); + EXPECT_FALSE(isOneByteIntegralType()); + EXPECT_FALSE(isOneByteIntegralType()); +} + +TEST(TypesTest, isTwoByteIntegralType) { + EXPECT_TRUE(isTwoByteIntegralType()); + EXPECT_TRUE(isTwoByteIntegralType()); + EXPECT_FALSE(isTwoByteIntegralType()); + EXPECT_FALSE(isTwoByteIntegralType()); + EXPECT_FALSE(isTwoByteIntegralType()); +} + +TEST(TypesTest, isFourByteIntegralType) { + EXPECT_TRUE(isFourByteIntegralType()); + EXPECT_TRUE(isFourByteIntegralType()); + EXPECT_FALSE(isFourByteIntegralType()); + EXPECT_FALSE(isFourByteIntegralType()); + EXPECT_FALSE(isFourByteIntegralType()); +} + +TEST(TypesTest, isEightByteIntegralType) { + EXPECT_TRUE(isEightByteIntegralType()); + EXPECT_TRUE(isEightByteIntegralType()); + EXPECT_FALSE(isEightByteIntegralType()); + EXPECT_FALSE(isEightByteIntegralType()); + EXPECT_FALSE(isEightByteIntegralType()); + EXPECT_FALSE(isEightByteIntegralType()); +} + +TEST(TypesTest, isFloatingPointType) { + EXPECT_TRUE(isFloatingPointType()); + EXPECT_TRUE(isFloatingPointType()); + EXPECT_FALSE(isFloatingPointType()); + EXPECT_FALSE(isFloatingPointType()); +} + +TEST(TypesTest, isNumericType) { + EXPECT_TRUE(isNumericType()); + EXPECT_TRUE(isNumericType()); + EXPECT_TRUE(isNumericType()); + EXPECT_TRUE(isNumericType()); + EXPECT_FALSE(isNumericType()); + EXPECT_FALSE(isNumericType()); +} + +TEST(TypesTest, isStringType) { + EXPECT_TRUE(isStringType()); + EXPECT_TRUE(isStringType()); + EXPECT_FALSE(isStringType()); + EXPECT_FALSE(isStringType()); +} + +TEST(TypesTest, isBoolType) { + EXPECT_TRUE(isBoolType()); + EXPECT_FALSE(isBoolType()); + EXPECT_FALSE(isBoolType()); +} + +} // namespace facebook::nimble::test diff --git a/velox/dwio/nimble/common/tests/VarintTests.cpp b/velox/dwio/nimble/common/tests/VarintTests.cpp new file mode 100644 index 00000000000..a7d2f1fe378 --- /dev/null +++ b/velox/dwio/nimble/common/tests/VarintTests.cpp @@ -0,0 +1,798 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include + +#include "folly/Random.h" +#include "folly/Range.h" +#include "folly/Varint.h" +#include "velox/dwio/nimble/common/Varint.h" + +using namespace ::facebook; + +namespace { +const int kNumElements = 10000; + +// Encode a vector of values into a varint buffer, returning the buffer and its +// size. +template +std::pair, size_t> encodeValues( + const std::vector& values) { + auto buf = + std::make_unique(values.size() * folly::kMaxVarintLength64); + char* pos = buf.get(); + for (auto val : values) { + nimble::varint::writeVarint(val, &pos); + } + return {std::move(buf), static_cast(pos - buf.get())}; +} + +// Bulk-decode and verify the result matches the expected values. +template +void verifyBulkDecode(const std::vector& expected, const char* encoded) { + std::vector decoded(expected.size()); + if constexpr (sizeof(T) == 4) { + nimble::varint::bulkVarintDecode32( + expected.size(), encoded, decoded.data()); + } else { + nimble::varint::bulkVarintDecode64( + expected.size(), encoded, decoded.data()); + } + for (size_t i = 0; i < expected.size(); ++i) { + ASSERT_EQ(expected[i], decoded[i]) + << "mismatch at index " << i << " of " << expected.size(); + } +} + +} // namespace + +TEST(VarintTests, varintSize32) { + // Boundary values for varint encoding. + EXPECT_EQ(nimble::varint::varintSize(uint32_t{0}), 1); + EXPECT_EQ(nimble::varint::varintSize(uint32_t{1}), 1); + EXPECT_EQ(nimble::varint::varintSize(uint32_t{127}), 1); + EXPECT_EQ(nimble::varint::varintSize(uint32_t{128}), 2); + EXPECT_EQ(nimble::varint::varintSize(uint32_t{16383}), 2); + EXPECT_EQ(nimble::varint::varintSize(uint32_t{16384}), 3); + EXPECT_EQ(nimble::varint::varintSize(uint32_t{2097151}), 3); + EXPECT_EQ(nimble::varint::varintSize(uint32_t{2097152}), 4); + EXPECT_EQ(nimble::varint::varintSize(uint32_t{268435455}), 4); + EXPECT_EQ(nimble::varint::varintSize(uint32_t{268435456}), 5); + EXPECT_EQ( + nimble::varint::varintSize(std::numeric_limits::max()), 5); + + // Verify consistency with writeVarint for random values. + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + char buf[folly::kMaxVarintLength32]; + for (int i = 0; i < kNumElements; ++i) { + const int bitShift = folly::Random::rand32(rng) % 32; + uint32_t val = folly::Random::rand32(rng) >> bitShift; + char* pos = buf; + nimble::varint::writeVarint(val, &pos); + ASSERT_EQ(nimble::varint::varintSize(val), static_cast(pos - buf)) + << "mismatch for val=" << val; + } +} + +TEST(VarintTests, varintSize64) { + // Boundary values for varint encoding. + EXPECT_EQ(nimble::varint::varintSize(uint64_t{0}), 1); + EXPECT_EQ(nimble::varint::varintSize(uint64_t{127}), 1); + EXPECT_EQ(nimble::varint::varintSize(uint64_t{128}), 2); + EXPECT_EQ(nimble::varint::varintSize(uint64_t{16383}), 2); + EXPECT_EQ(nimble::varint::varintSize(uint64_t{16384}), 3); + EXPECT_EQ( + nimble::varint::varintSize(std::numeric_limits::max()), 10); + + // Verify consistency with writeVarint for random values. + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + char buf[folly::kMaxVarintLength64]; + for (int i = 0; i < kNumElements; ++i) { + const int bitShift = folly::Random::rand64(rng) % 64; + uint64_t val = folly::Random::rand64(rng) >> bitShift; + char* pos = buf; + nimble::varint::writeVarint(val, &pos); + ASSERT_EQ(nimble::varint::varintSize(val), static_cast(pos - buf)) + << "mismatch for val=" << val; + } +} + +TEST(VarintTests, maxVarintSizeForBitWidth) { + EXPECT_EQ(nimble::varint::maxVarintSizeForBitWidth(0), 0); + EXPECT_EQ(nimble::varint::maxVarintSizeForBitWidth(1), 1); + EXPECT_EQ(nimble::varint::maxVarintSizeForBitWidth(7), 1); + EXPECT_EQ(nimble::varint::maxVarintSizeForBitWidth(8), 2); + EXPECT_EQ(nimble::varint::maxVarintSizeForBitWidth(63), 9); + EXPECT_EQ(nimble::varint::maxVarintSizeForBitWidth(64), 10); + EXPECT_EQ(nimble::varint::maxVarintSizeForBitWidth(65), 10); +} + +TEST(VarintTests, writeRead32) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + std::vector data; + // Generate data with uniform bit length. + for (int i = 0; i < kNumElements; ++i) { + const int bitShift = folly::Random::rand32(rng) % 32; + data.push_back(folly::Random::rand32(rng) >> bitShift); + } + + auto buffer = + std::make_unique(kNumElements * folly::kMaxVarintLength32); + char* pos = buffer.get(); + for (int i = 0; i < kNumElements; ++i) { + nimble::varint::writeVarint(data[i], &pos); + } + + auto follyBuffer = + std::make_unique(kNumElements * folly::kMaxVarintLength32); + uint8_t* fpos = follyBuffer.get(); + for (int i = 0; i < kNumElements; ++i) { + fpos += folly::encodeVarint(data[i], fpos); + } + + ASSERT_EQ(pos - buffer.get(), fpos - follyBuffer.get()); + + ASSERT_EQ(nimble::varint::bulkVarintSize32(data), pos - buffer.get()); + + const char* cpos = buffer.get(); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(data[i], nimble::varint::readVarint32(&cpos)); + } + + const uint8_t* fstart = follyBuffer.get(); + const uint8_t* fend = follyBuffer.get() + (fpos - follyBuffer.get()); + folly::Range frange(fstart, fend); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(data[i], folly::decodeVarint(frange)); + } + + std::vector bulk(kNumElements); + cpos = buffer.get(); + nimble::varint::bulkVarintDecode32(kNumElements, cpos, bulk.data()); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(data[i], bulk[i]); + } +} + +TEST(VarintTests, writeRead64) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + std::vector data; + // Generate data with uniform bit length. + for (int i = 0; i < kNumElements; ++i) { + const int bitShift = folly::Random::rand64(rng) % 64; + data.push_back(folly::Random::rand64(rng) >> bitShift); + } + + auto buffer = + std::make_unique(kNumElements * folly::kMaxVarintLength64); + char* pos = buffer.get(); + for (int i = 0; i < kNumElements; ++i) { + nimble::varint::writeVarint(data[i], &pos); + } + + auto follyBuffer = + std::make_unique(kNumElements * folly::kMaxVarintLength64); + uint8_t* fpos = follyBuffer.get(); + for (int i = 0; i < kNumElements; ++i) { + fpos += folly::encodeVarint(data[i], fpos); + } + + ASSERT_EQ(pos - buffer.get(), fpos - follyBuffer.get()); + + ASSERT_EQ(nimble::varint::bulkVarintSize64(data), pos - buffer.get()); + + const char* cpos = buffer.get(); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(data[i], nimble::varint::readVarint64(&cpos)); + } + + const uint8_t* fstart = follyBuffer.get(); + const uint8_t* fend = follyBuffer.get() + (fpos - follyBuffer.get()); + folly::Range frange(fstart, fend); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(data[i], folly::decodeVarint(frange)); + } + + std::vector bulk(kNumElements); + cpos = buffer.get(); + nimble::varint::bulkVarintDecode64(kNumElements, cpos, bulk.data()); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(data[i], bulk[i]); + } +} + +// ============================================================================ +// Single-byte varint tests: exercise the SIMD decodeSingleByteRun path. +// The function has three loops: +// 1. Wide loop: processes kU8BatchSize bytes (32 on AVX2, 16 on SSE/NEON) +// 2. 8-byte loop: processes 8 bytes at a time +// 3. Tail loop: processes 1 byte at a time +// These tests cover boundary conditions for all three loops. +// ============================================================================ + +// All 128 single-byte values (0-127) decode correctly for uint32_t. +TEST(VarintTests, singleByte32AllValues) { + std::vector data(128); + std::iota(data.begin(), data.end(), 0); + auto [buf, size] = encodeValues(data); + ASSERT_EQ(size, 128u); + verifyBulkDecode(data, buf.get()); +} + +// All 128 single-byte values (0-127) decode correctly for uint64_t. +TEST(VarintTests, singleByte64AllValues) { + std::vector data(128); + std::iota(data.begin(), data.end(), 0); + auto [buf, size] = encodeValues(data); + ASSERT_EQ(size, 128u); + verifyBulkDecode(data, buf.get()); +} + +// Test every count from 0 to 100 with all-zero values. +// Exercises exact boundary transitions between wide/8-byte/tail loops. +TEST(VarintTests, singleByte32AllCountsZero) { + for (int count = 0; count <= 100; ++count) { + std::vector data(count, 0); + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } +} + +// Test every count from 0 to 100 with value 127 (max single-byte varint). +TEST(VarintTests, singleByte32AllCountsMax) { + for (int count = 0; count <= 100; ++count) { + std::vector data(count, 127); + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } +} + +// Test counts at specific SIMD boundaries with uint64_t. +TEST(VarintTests, singleByte64SimdBoundaries) { + for (int count : {0, 1, 2, 7, 8, 9, 15, 16, 17, 31, + 32, 33, 63, 64, 65, 96, 127, 128, 256, 1000}) { + std::vector data(count); + for (int i = 0; i < count; ++i) { + data[i] = i % 128; + } + auto [buf, size] = encodeValues(data); + ASSERT_EQ(size, static_cast(count)); + verifyBulkDecode(data, buf.get()); + } +} + +// A multi-byte varint (>=128) interrupts the single-byte run at each position +// within a 64-element window. Verifies the SIMD path correctly bails out and +// the remaining elements are decoded by the fallback path. +TEST(VarintTests, singleByte32MultiByteInterrupt) { + for (int interruptPos = 0; interruptPos < 64; ++interruptPos) { + const int count = 64; + std::vector data(count); + for (int i = 0; i < count; ++i) { + data[i] = (i == interruptPos) ? 200 : (i % 128); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } +} + +TEST(VarintTests, singleByte64MultiByteInterrupt) { + for (int interruptPos = 0; interruptPos < 64; ++interruptPos) { + const int count = 64; + std::vector data(count); + for (int i = 0; i < count; ++i) { + data[i] = (i == interruptPos) ? 200 : (i % 128); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } +} + +// Single-byte values followed by progressively longer multi-byte varints. +// Tests the transition from decodeSingleByteRun into the BMI2/scalar path. +TEST(VarintTests, singleByte32TransitionToMultiByte) { + for (int singleCount : {0, 1, 7, 8, 15, 16, 31, 32, 33, 64}) { + for (int multiCount : {0, 1, 5, 10}) { + std::vector data; + data.reserve(singleCount + multiCount); + for (int i = 0; i < singleCount; ++i) { + data.push_back(i % 128); + } + for (int i = 0; i < multiCount; ++i) { + data.push_back(128 + i * 1000); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } + } +} + +TEST(VarintTests, singleByte64TransitionToMultiByte) { + for (int singleCount : {0, 1, 7, 8, 15, 16, 31, 32, 33, 64}) { + for (int multiCount : {0, 1, 5, 10}) { + std::vector data; + data.reserve(singleCount + multiCount); + for (int i = 0; i < singleCount; ++i) { + data.push_back(i % 128); + } + for (int i = 0; i < multiCount; ++i) { + data.push_back(128 + static_cast(i) * 1000); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } + } +} + +// Alternating single-byte and multi-byte varints. The SIMD path must +// correctly handle frequent bail-outs and re-entries. +TEST(VarintTests, singleByte32AlternatingSingleMulti) { + std::vector data; + data.reserve(200); + for (int i = 0; i < 200; ++i) { + data.push_back(i % 2 == 0 ? (i % 128) : (128 + i)); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); +} + +// Large run of single-byte varints to stress the wide SIMD loop. +TEST(VarintTests, singleByte32LargeRun) { + const int count = 100000; + std::vector data(count); + for (int i = 0; i < count; ++i) { + data[i] = i % 128; + } + auto [buf, size] = encodeValues(data); + ASSERT_EQ(size, static_cast(count)); + verifyBulkDecode(data, buf.get()); +} + +TEST(VarintTests, singleByte64LargeRun) { + const int count = 100000; + std::vector data(count); + for (int i = 0; i < count; ++i) { + data[i] = i % 128; + } + auto [buf, size] = encodeValues(data); + ASSERT_EQ(size, static_cast(count)); + verifyBulkDecode(data, buf.get()); +} + +// Random mix: ~80% single-byte, ~20% multi-byte, with a random seed. +TEST(VarintTests, singleByte32RandomMix) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + std::vector data(kNumElements); + for (int i = 0; i < kNumElements; ++i) { + if (folly::Random::rand32(rng) % 5 != 0) { + data[i] = folly::Random::rand32(rng) % 128; + } else { + data[i] = 128 + folly::Random::rand32(rng) % 10000; + } + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); +} + +TEST(VarintTests, singleByte64RandomMix) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + std::vector data(kNumElements); + for (int i = 0; i < kNumElements; ++i) { + if (folly::Random::rand32(rng) % 5 != 0) { + data[i] = folly::Random::rand32(rng) % 128; + } else { + data[i] = 128 + folly::Random::rand64(rng) % 1000000; + } + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); +} + +// Constant value runs for each single-byte value. +TEST(VarintTests, singleByte32ConstantRuns) { + for (uint32_t val = 0; val < 128; ++val) { + std::vector data(37, val); + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } +} + +// Verify that a single multi-byte varint at the very start works. +TEST(VarintTests, singleByte32MultiByteFirst) { + std::vector data = {300}; + for (int i = 0; i < 50; ++i) { + data.push_back(i % 128); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); +} + +// Verify that a single multi-byte varint at the very end works. +TEST(VarintTests, singleByte32MultiByteLast) { + std::vector data; + data.reserve(51); + for (int i = 0; i < 50; ++i) { + data.push_back(i % 128); + } + data.push_back(300); + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); +} + +// ============================================================================ +// Two-byte varint tests: exercise the bulkDecodeTwoByteRun fast path. +// Two-byte varints have values in [128, 16383]. The fast path processes +// 4 varints (8 bytes) at a time, then handles trailing varints one at a time. +// ============================================================================ + +// All two-byte boundary values decode correctly for uint32_t. +TEST(VarintTests, twoByte32BoundaryValues) { + std::vector data = {128, 255, 256, 1000, 8191, 16383}; + auto [buf, size] = encodeValues(data); + ASSERT_EQ(size, 12u); + verifyBulkDecode(data, buf.get()); +} + +// All two-byte boundary values decode correctly for uint64_t. +TEST(VarintTests, twoByte64BoundaryValues) { + std::vector data = {128, 255, 256, 1000, 8191, 16383}; + auto [buf, size] = encodeValues(data); + ASSERT_EQ(size, 12u); + verifyBulkDecode(data, buf.get()); +} + +// Test every count from 0 to 100 with uniform two-byte varints. +// Exercises the 4-at-a-time wide loop and trailing scalar loop boundaries. +TEST(VarintTests, twoByte32AllCounts) { + for (int count = 0; count <= 100; ++count) { + std::vector data(count, 200); + auto [buf, size] = encodeValues(data); + ASSERT_EQ(size, static_cast(count * 2)); + verifyBulkDecode(data, buf.get()); + } +} + +// Test counts at multiples of 4 (wide loop boundary) with uint64_t. +TEST(VarintTests, twoByte64WideBoundaries) { + for (int count : {0, 1, 2, 3, 4, 5, 7, 8, 9, 12, 16, 32, 64, 100, 1000}) { + std::vector data(count); + for (int i = 0; i < count; ++i) { + data[i] = 128 + (i % 16256); + } + auto [buf, size] = encodeValues(data); + ASSERT_EQ(size, static_cast(count * 2)); + verifyBulkDecode(data, buf.get()); + } +} + +// Constant value runs for representative two-byte values. +TEST(VarintTests, twoByte32ConstantRuns) { + for (uint32_t val : {128u, 200u, 1000u, 8192u, 16383u}) { + std::vector data(37, val); + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } +} + +// Large run of two-byte varints to stress the wide loop. +TEST(VarintTests, twoByte32LargeRun) { + const int count = 100000; + std::vector data(count); + for (int i = 0; i < count; ++i) { + data[i] = 128 + (i % 16256); + } + auto [buf, size] = encodeValues(data); + ASSERT_EQ(size, static_cast(count * 2)); + verifyBulkDecode(data, buf.get()); +} + +TEST(VarintTests, twoByte64LargeRun) { + const int count = 100000; + std::vector data(count); + for (int i = 0; i < count; ++i) { + data[i] = 128 + (i % 16256); + } + auto [buf, size] = encodeValues(data); + ASSERT_EQ(size, static_cast(count * 2)); + verifyBulkDecode(data, buf.get()); +} + +// A 3+ byte varint interrupts the two-byte run at each position within a +// 32-element window. Verifies the fast path correctly bails out. +TEST(VarintTests, twoByte32MultiByteInterrupt) { + for (int interruptPos = 0; interruptPos < 32; ++interruptPos) { + const int count = 32; + std::vector data(count); + for (int i = 0; i < count; ++i) { + data[i] = (i == interruptPos) ? 20000 : 200; + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } +} + +// Two-byte values followed by multi-byte values. +TEST(VarintTests, twoByte32TransitionToMultiByte) { + for (int twoByteCount : {0, 1, 3, 4, 5, 8, 16, 32}) { + for (int multiCount : {0, 1, 5, 10}) { + std::vector data; + data.reserve(twoByteCount + multiCount); + for (int i = 0; i < twoByteCount; ++i) { + data.push_back(128 + (i % 16256)); + } + for (int i = 0; i < multiCount; ++i) { + data.push_back(20000 + i * 1000); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } + } +} + +// ============================================================================ +// Dispatch loop tests: exercise the bulkVarintDecodeDispatch re-entry logic. +// The dispatch loop cycles: single-byte fast path -> two-byte fast path -> +// BMI2 general decoder, and the BMI2 decoder yields back when it detects +// a single-byte boundary (no carryover). These tests verify correct decoding +// across multiple dispatch loop iterations. +// ============================================================================ + +// Single-byte run, then two-byte run, then single-byte run again. +// The BMI2 decoder should not be entered at all. +TEST(VarintTests, dispatch32SingleTwoSingle) { + std::vector data; + data.reserve(150); + for (int i = 0; i < 50; ++i) { + data.push_back(i % 128); + } + for (int i = 0; i < 50; ++i) { + data.push_back(128 + (i % 16256)); + } + for (int i = 0; i < 50; ++i) { + data.push_back(i % 128); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); +} + +// A large 5-byte varint followed by single-byte varints. +// Tests the dispatch loop re-entering single-byte fast path after BMI2. +TEST(VarintTests, dispatch32LargeHeadThenSingleByte) { + std::vector data; + data.reserve(201); + data.push_back(UINT32_MAX); + for (int i = 0; i < 200; ++i) { + data.push_back(i % 128); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); +} + +// A large 5-byte varint followed by two-byte varints. +// Tests the dispatch loop re-entering two-byte fast path after BMI2. +TEST(VarintTests, dispatch32LargeHeadThenTwoByte) { + std::vector data; + data.reserve(201); + data.push_back(UINT32_MAX); + for (int i = 0; i < 200; ++i) { + data.push_back(128 + (i % 16256)); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); +} + +TEST(VarintTests, dispatch64LargeHeadThenSingleByte) { + std::vector data; + data.reserve(201); + data.push_back(UINT64_MAX); + for (int i = 0; i < 200; ++i) { + data.push_back(i % 128); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); +} + +TEST(VarintTests, dispatch64LargeHeadThenTwoByte) { + std::vector data; + data.reserve(201); + data.push_back(UINT64_MAX); + for (int i = 0; i < 200; ++i) { + data.push_back(128 + (i % 16256)); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); +} + +// Sporadic large values every N elements among single-byte values. +// Each large value forces the BMI2 decoder, then the dispatch loop must +// re-enter the single-byte fast path. +TEST(VarintTests, dispatch32SporadicLargeAmongSingleByte) { + for (int interval : {4, 8, 16, 64, 256}) { + const int count = 1024; + std::vector data(count); + for (int i = 0; i < count; ++i) { + data[i] = (i % interval == 0) ? UINT32_MAX : (i % 128); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } +} + +// Sporadic large values every N elements among two-byte values. +TEST(VarintTests, dispatch32SporadicLargeAmongTwoByte) { + for (int interval : {4, 8, 16, 64, 256}) { + const int count = 1024; + std::vector data(count); + for (int i = 0; i < count; ++i) { + data[i] = (i % interval == 0) ? UINT32_MAX : (128 + i % 16256); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } +} + +TEST(VarintTests, dispatch64SporadicLargeAmongSingleByte) { + for (int interval : {4, 8, 16, 64, 256}) { + const int count = 1024; + std::vector data(count); + for (int i = 0; i < count; ++i) { + data[i] = (i % interval == 0) ? UINT64_MAX : uint64_t(i % 128); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } +} + +// Repeating pattern: single-byte block, two-byte block, large value. +// Exercises all three dispatch loop phases in sequence, multiple times. +TEST(VarintTests, dispatch32RepeatingThreePhase) { + std::vector data; + data.reserve(20 * 21); + for (int round = 0; round < 20; ++round) { + for (int i = 0; i < 10; ++i) { + data.push_back(i % 128); + } + for (int i = 0; i < 10; ++i) { + data.push_back(128 + (i % 16256)); + } + data.push_back(UINT32_MAX); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); +} + +TEST(VarintTests, dispatch64RepeatingThreePhase) { + std::vector data; + data.reserve(20 * 21); + for (int round = 0; round < 20; ++round) { + for (int i = 0; i < 10; ++i) { + data.push_back(i % 128); + } + for (int i = 0; i < 10; ++i) { + data.push_back(128 + (i % 16256)); + } + data.push_back(UINT64_MAX); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); +} + +// All 3-byte, 4-byte, and 5-byte varints (no fast path applies). +// Everything goes through BMI2 general decoder. +TEST(VarintTests, dispatch32AllMultiByte) { + for (auto [lo, hi] : std::vector>{ + {16384, 2097151}, {2097152, 268435455}, {268435456, UINT32_MAX}}) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + const int count = 500; + std::vector data(count); + for (int i = 0; i < count; ++i) { + data[i] = lo + folly::Random::rand32(rng) % (hi - lo + 1); + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } +} + +// Random mix of all varint widths with dispatch loop stress. +TEST(VarintTests, dispatch32RandomAllWidths) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + std::vector data(kNumElements); + for (int i = 0; i < kNumElements; ++i) { + int width = folly::Random::rand32(rng) % 5; + switch (width) { + case 0: + data[i] = folly::Random::rand32(rng) % 128; + break; + case 1: + data[i] = 128 + folly::Random::rand32(rng) % 16256; + break; + case 2: + data[i] = 16384 + folly::Random::rand32(rng) % 2080768; + break; + case 3: + data[i] = 2097152 + folly::Random::rand32(rng) % 266338304; + break; + case 4: + data[i] = + 268435456 + folly::Random::rand32(rng) % (UINT32_MAX - 268435456); + break; + } + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); +} + +TEST(VarintTests, dispatch64RandomAllWidths) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + std::vector data(kNumElements); + for (int i = 0; i < kNumElements; ++i) { + int width = folly::Random::rand32(rng) % 5; + switch (width) { + case 0: + data[i] = folly::Random::rand32(rng) % 128; + break; + case 1: + data[i] = 128 + folly::Random::rand32(rng) % 16256; + break; + case 2: + data[i] = 16384 + folly::Random::rand32(rng) % 2080768; + break; + case 3: + data[i] = uint64_t(2097152) + folly::Random::rand64(rng) % 266338304; + break; + case 4: + data[i] = folly::Random::rand64(rng); + break; + } + } + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); +} + +// Edge case: exactly n=1 for each varint width. +TEST(VarintTests, dispatch32SingleElement) { + for (uint32_t val : {0u, 127u, 128u, 16383u, 16384u, 2097151u, UINT32_MAX}) { + std::vector data = {val}; + auto [buf, size] = encodeValues(data); + verifyBulkDecode(data, buf.get()); + } +} diff --git a/velox/dwio/nimble/common/tests/VectorTest.cpp b/velox/dwio/nimble/common/tests/VectorTest.cpp new file mode 100644 index 00000000000..563dc8e7b32 --- /dev/null +++ b/velox/dwio/nimble/common/tests/VectorTest.cpp @@ -0,0 +1,426 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "velox/dwio/nimble/common/Vector.h" +#include +#include +#include "velox/buffer/BufferPool.h" +#include "velox/common/memory/Memory.h" + +DECLARE_bool(velox_enable_memory_usage_track_in_default_memory_pool); + +using namespace ::facebook; + +namespace { + +// Verifies that a Vector is in a clean empty state with no buffer. +template +void checkVectorEmpty(const nimble::Vector& v) { + EXPECT_EQ(v.size(), 0); + EXPECT_EQ(v.capacity(), 0); + EXPECT_EQ(v.testingBuffer(), nullptr); +} + +// Verifies that a Vector holds the expected buffer with the given capacity +// and zero logical size. +template +void checkVectorBuffer( + const nimble::Vector& v, + const velox::Buffer* expectedBuffer, + uint64_t expectedCapacity) { + EXPECT_EQ(v.size(), 0); + EXPECT_EQ(v.capacity(), expectedCapacity); + EXPECT_EQ(v.testingBuffer().get(), expectedBuffer); +} + +} // namespace + +class VectorTest : public ::testing::Test { + protected: + static void SetUpTestCase() { + FLAGS_velox_enable_memory_usage_track_in_default_memory_pool = true; + velox::memory::MemoryManager::testingSetInstance({}); + } + + void SetUp() override { + rootPool_ = velox::memory::memoryManager()->addRootPool("VectorTest"); + pool_ = rootPool_->addLeafChild("leaf"); + } + + std::shared_ptr rootPool_; + std::shared_ptr pool_; +}; + +TEST_F(VectorTest, fromRange) { + std::vector source{4, 5, 6}; + nimble::Vector v1(pool_.get(), source.begin(), source.end()); + EXPECT_EQ(3, v1.size()); + EXPECT_EQ(4, v1[0]); + EXPECT_EQ(5, v1[1]); + EXPECT_EQ(6, v1[2]); +} + +TEST_F(VectorTest, equalOp1) { + nimble::Vector v1(pool_.get()); + v1.push_back(1); + v1.emplace_back(2); + v1.push_back(3); + + nimble::Vector v2(pool_.get()); + v2.push_back(4); + v2.emplace_back(5); + + EXPECT_EQ(3, v1.size()); + EXPECT_EQ(1, v1[0]); + EXPECT_EQ(2, v1[1]); + EXPECT_EQ(3, v1[2]); + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(4, v2[0]); + EXPECT_EQ(5, v2[1]); + + v1 = v2; + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(4, v1[0]); + EXPECT_EQ(5, v1[1]); + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(4, v2[0]); + EXPECT_EQ(5, v2[1]); +} + +TEST_F(VectorTest, explicitMoveEqualOp) { + nimble::Vector v1(pool_.get()); + v1.push_back(1); + v1.emplace_back(2); + v1.push_back(3); + + nimble::Vector v2(pool_.get()); + v2.push_back(4); + v2.emplace_back(5); + + EXPECT_EQ(3, v1.size()); + ASSERT_FALSE(v1.empty()); + EXPECT_EQ(1, v1[0]); + EXPECT_EQ(2, v1[1]); + EXPECT_EQ(3, v1[2]); + EXPECT_EQ(2, v2.size()); + ASSERT_FALSE(v2.empty()); + EXPECT_EQ(4, v2[0]); + EXPECT_EQ(5, v2[1]); + + v1 = std::move(v2); + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(4, v1[0]); + EXPECT_EQ(5, v1[1]); + // @lint-ignore CLANGTIDY bugprone-use-after-move + EXPECT_EQ(0, v2.size()); + ASSERT_TRUE(v2.empty()); +} + +TEST_F(VectorTest, moveEqualOp1) { + nimble::Vector v1(pool_.get()); + v1.push_back(1); + v1.emplace_back(2); + v1.push_back(3); + EXPECT_EQ(3, v1.size()); + EXPECT_EQ(1, v1[0]); + EXPECT_EQ(2, v1[1]); + EXPECT_EQ(3, v1[2]); + v1 = nimble::Vector(pool_.get(), {4, 5}); + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(4, v1[0]); + EXPECT_EQ(5, v1[1]); +} + +TEST_F(VectorTest, copyCtr) { + nimble::Vector v2(pool_.get()); + v2.push_back(3); + v2.emplace_back(4); + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(3, v2[0]); + EXPECT_EQ(4, v2[1]); + nimble::Vector v1(v2); + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(3, v1[0]); + EXPECT_EQ(4, v1[1]); + + // make sure they do not share buffer + v1[0] = 1; + v1[1] = 2; + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(1, v1[0]); + EXPECT_EQ(2, v1[1]); + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(3, v2[0]); + EXPECT_EQ(4, v2[1]); +} + +TEST_F(VectorTest, boolInitializerList) { + nimble::Vector v1(pool_.get(), {true, false, true}); + EXPECT_EQ(3, v1.size()); + EXPECT_EQ(true, v1[0]); + EXPECT_EQ(false, v1[1]); + EXPECT_EQ(true, v1[2]); +} + +TEST_F(VectorTest, boolEqualOp1) { + nimble::Vector v1(pool_.get()); + v1.push_back(false); + v1.emplace_back(true); + v1.push_back(true); + + EXPECT_EQ(3, v1.size()); + EXPECT_EQ(false, v1[0]); + EXPECT_EQ(true, v1[1]); + EXPECT_EQ(true, v1[2]); + + nimble::Vector v2(pool_.get()); + v2.push_back(true); + v2.emplace_back(false); + + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(true, v2[0]); + EXPECT_EQ(false, v2[1]); + + v1 = v2; + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(true, v1[0]); + EXPECT_EQ(false, v1[1]); + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(true, v2[0]); + EXPECT_EQ(false, v2[1]); +} + +TEST_F(VectorTest, boolMoveEqualOp1) { + nimble::Vector v1(pool_.get()); + v1.push_back(true); + v1.emplace_back(false); + v1.push_back(false); + + EXPECT_EQ(3, v1.size()); + EXPECT_EQ(true, v1[0]); + EXPECT_EQ(false, v1[1]); + EXPECT_EQ(false, v1[2]); + + v1 = nimble::Vector(pool_.get(), {false, true}); + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(false, v1[0]); + EXPECT_EQ(true, v1[1]); +} + +TEST_F(VectorTest, boolCopyCtr) { + nimble::Vector v2(pool_.get()); + v2.push_back(true); + v2.emplace_back(false); + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(true, v2[0]); + EXPECT_EQ(false, v2[1]); + nimble::Vector v1(v2); + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(true, v1[0]); + EXPECT_EQ(false, v1[1]); + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(true, v2[0]); + EXPECT_EQ(false, v2[1]); +} + +TEST_F(VectorTest, scopedVectorReleasesBufferToPool) { + velox::BufferPool bufferPool{velox::BufferPool::kDefaultCapacity}; + + { + nimble::ScopedVector scratch{ + /*size=*/4, pool_.get(), &bufferPool}; + ASSERT_EQ(scratch.size(), 4); + scratch[0] = 123; + EXPECT_EQ(scratch[0], 123); + EXPECT_EQ(bufferPool.size(), 0); + } + EXPECT_EQ(bufferPool.size(), 1); + + { + nimble::ScopedVector scratch{ + /*size=*/2, pool_.get(), &bufferPool}; + EXPECT_EQ(scratch.size(), 2); + EXPECT_EQ(bufferPool.size(), 0); + } + EXPECT_EQ(bufferPool.size(), 1); +} + +TEST_F(VectorTest, scopedVectorWorksWithoutBufferPool) { + nimble::ScopedVector scratch{ + /*size=*/3, pool_.get(), /*bufferPool=*/nullptr}; + ASSERT_EQ(scratch.size(), 3); + scratch[0] = true; + scratch[1] = false; + EXPECT_TRUE(scratch[0]); + EXPECT_FALSE(scratch[1]); +} + +TEST_F(VectorTest, scopedVectorOperations) { + nimble::ScopedVector scratch{ + /*size=*/0, pool_.get(), /*bufferPool=*/nullptr}; + EXPECT_TRUE(scratch.empty()); + + scratch.reserve(3); + EXPECT_GE(scratch.capacity(), 3); + scratch.push_back(1); + scratch.emplace_back(2); + scratch.resize(3); + scratch[2] = 3; + + EXPECT_EQ(scratch.size(), 3); + + nimble::Vector& values = scratch; + EXPECT_EQ(values.size(), scratch.size()); + EXPECT_EQ(values.data(), scratch.data()); + EXPECT_EQ(values[0], 1); + + EXPECT_EQ(scratch->capacity(), scratch.capacity()); + EXPECT_EQ((*scratch)[1], 2); + EXPECT_EQ(scratch[2], 3); + EXPECT_EQ(std::accumulate(scratch.begin(), scratch.end(), uint32_t{0}), 6); + + const auto& constScratch = scratch; + const nimble::Vector& constValues = constScratch; + EXPECT_EQ(constValues.data(), constScratch.data()); + EXPECT_EQ(constValues[0], 1); + EXPECT_EQ(constScratch->size(), 3); + EXPECT_EQ((*constScratch)[1], 2); + EXPECT_EQ(constScratch[2], 3); + EXPECT_EQ( + std::accumulate(constScratch.begin(), constScratch.end(), uint32_t{0}), + 6); +} + +TEST_F(VectorTest, memoryCleanup) { + EXPECT_EQ(0, pool_->usedBytes()); + { + nimble::Vector v(pool_.get()); + EXPECT_EQ(0, pool_->usedBytes()); + v.resize(1000, 10); + EXPECT_NE(0, pool_->usedBytes()); + } + EXPECT_EQ(0, pool_->usedBytes()); + { + nimble::Vector v(pool_.get()); + EXPECT_EQ(0, pool_->usedBytes()); + v.resize(1000, 10); + EXPECT_NE(0, pool_->usedBytes()); + + auto vCopy(v); + } + EXPECT_EQ(0, pool_->usedBytes()); + { + nimble::Vector v(pool_.get()); + EXPECT_EQ(0, pool_->usedBytes()); + v.resize(1000, 10); + EXPECT_NE(0, pool_->usedBytes()); + + auto vCopy(std::move(v)); + } + EXPECT_EQ(0, pool_->usedBytes()); +} + +TEST_F(VectorTest, reserveActualSize) { + EXPECT_EQ(0, pool_->usedBytes()); + + // There is no good way to assert the exact expected size because of padding + // logic in the AlignedBuffer::allocate. + // We know that without the exactSize flag being passed into the + // AlignedBuffer, it would allocate 12,582,912 bytes for any requested size in + // the range [8,388,609 - 12,582,912]. What we can do is to pick some value + // in the middle of this range and assert that it's roughly what we expect it + // to be, e.g allocatedSize is in [X, X+1MB]. + { + // 1 byte type + nimble::Vector v(pool_.get()); + EXPECT_EQ(0, pool_->usedBytes()); + + const size_t lowerBound = 9 * 1024 * 1024; + const size_t upperBound = lowerBound + 1024 * 1024; + v.reserve(lowerBound); + EXPECT_GE(pool_->usedBytes(), lowerBound); + EXPECT_LE(pool_->usedBytes(), upperBound); + } + + { + // 4 byte type + nimble::Vector v(pool_.get()); + EXPECT_EQ(0, pool_->usedBytes()); + const size_t lowerBound = 9 * 1024 * 1024; + const size_t upperBound = lowerBound + 1024 * 1024; + const uint64_t valueCount = lowerBound / sizeof(int32_t); + v.reserve(valueCount); + EXPECT_GE(pool_->usedBytes(), lowerBound); + EXPECT_LE(pool_->usedBytes(), upperBound); + } +} + +TEST_F(VectorTest, releaseBuffer) { + nimble::Vector v(pool_.get()); + v.push_back(10); + v.push_back(20); + v.push_back(30); + EXPECT_EQ(v.size(), 3); + + auto buf = v.releaseBuffer(); + EXPECT_NE(buf, nullptr); + EXPECT_GE(buf->capacity(), 3 * sizeof(int32_t)); + + // Vector is fully reset after release. + checkVectorEmpty(v); + + // Releasing again returns nullptr. + auto buf2 = v.releaseBuffer(); + EXPECT_EQ(buf2, nullptr); +} + +TEST_F(VectorTest, constructFromBuffer) { + auto buf = velox::AlignedBuffer::allocate(100, pool_.get()); + const auto expectedCapacity = buf->capacity() / sizeof(int32_t); + auto* rawBuf = buf.get(); + + nimble::Vector v(std::move(buf)); + + checkVectorBuffer(v, rawBuf, expectedCapacity); + EXPECT_EQ(v.pool(), pool_.get()); + + // Can use the vector normally. + v.push_back(42); + EXPECT_EQ(v.size(), 1); + EXPECT_EQ(v[0], 42); +} + +TEST_F(VectorTest, releaseAndConstructRoundTrip) { + nimble::Vector v1(pool_.get()); + for (int i = 0; i < 100; ++i) { + v1.push_back(i); + } + const auto originalCapacity = v1.capacity(); + + // Release from v1, construct v2 from the buffer. + auto buf = v1.releaseBuffer(); + checkVectorEmpty(v1); + + auto* rawBuf = buf.get(); + nimble::Vector v2(std::move(buf)); + + checkVectorBuffer(v2, rawBuf, originalCapacity); + EXPECT_EQ(v2.pool(), pool_.get()); + + // v2 can be used normally. + v2.resize(50, 7); + EXPECT_EQ(v2.size(), 50); + EXPECT_EQ(v2[0], 7); +} diff --git a/velox/dwio/nimble/common/tests/ZigzagTests.cpp b/velox/dwio/nimble/common/tests/ZigzagTests.cpp new file mode 100644 index 00000000000..22cbd35e2f5 --- /dev/null +++ b/velox/dwio/nimble/common/tests/ZigzagTests.cpp @@ -0,0 +1,87 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include + +#include + +#include "velox/dwio/nimble/common/Zigzag.h" + +using namespace ::facebook::nimble::zigzag; + +TEST(ZigzagTest, encode32) { + EXPECT_EQ(zigzagEncode32(0), 0); + EXPECT_EQ(zigzagEncode32(-1), 1); + EXPECT_EQ(zigzagEncode32(1), 2); + EXPECT_EQ(zigzagEncode32(-2), 3); + EXPECT_EQ(zigzagEncode32(2), 4); + EXPECT_EQ(zigzagEncode32(-3), 5); + EXPECT_EQ(zigzagEncode32(std::numeric_limits::max()), 4294967294U); + EXPECT_EQ(zigzagEncode32(std::numeric_limits::min()), 4294967295U); +} + +TEST(ZigzagTest, decode32) { + EXPECT_EQ(zigzagDecode32(0), 0); + EXPECT_EQ(zigzagDecode32(1), -1); + EXPECT_EQ(zigzagDecode32(2), 1); + EXPECT_EQ(zigzagDecode32(3), -2); + EXPECT_EQ(zigzagDecode32(4), 2); + EXPECT_EQ(zigzagDecode32(5), -3); + EXPECT_EQ(zigzagDecode32(4294967294U), std::numeric_limits::max()); + EXPECT_EQ(zigzagDecode32(4294967295U), std::numeric_limits::min()); +} + +TEST(ZigzagTest, roundTrip32) { + for (int32_t val : + {0, + 1, + -1, + 100, + -100, + 1'000'000, + -1'000'000, + std::numeric_limits::max(), + std::numeric_limits::min()}) { + EXPECT_EQ(zigzagDecode32(zigzagEncode32(val)), val); + } +} + +TEST(ZigzagTest, encode64) { + EXPECT_EQ(zigzagEncode64(0), 0); + EXPECT_EQ(zigzagEncode64(-1), 1); + EXPECT_EQ(zigzagEncode64(1), 2); + EXPECT_EQ(zigzagEncode64(-2), 3); + EXPECT_EQ( + zigzagEncode64(std::numeric_limits::max()), + std::numeric_limits::max() - 1); + EXPECT_EQ( + zigzagEncode64(std::numeric_limits::min()), + std::numeric_limits::max()); +} + +TEST(ZigzagTest, roundTrip64) { + for (int64_t val : + {int64_t{0}, + int64_t{1}, + int64_t{-1}, + int64_t{100}, + int64_t{-100}, + int64_t{1'000'000'000'000}, + int64_t{-1'000'000'000'000}, + std::numeric_limits::max(), + std::numeric_limits::min()}) { + EXPECT_EQ(zigzagDecode64(zigzagEncode64(val)), val); + } +} diff --git a/velox/dwio/nimble/compression/CMakeLists.txt b/velox/dwio/nimble/compression/CMakeLists.txt new file mode 100644 index 00000000000..fd2f7710f00 --- /dev/null +++ b/velox/dwio/nimble/compression/CMakeLists.txt @@ -0,0 +1,27 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +add_library( + nimble_compression + Compression.cpp + CompressionPolicy.cpp + Lz4Compressor.cpp + OpenZLCompressor.cpp + ZstdCompressor.cpp + ZstdCompressor.h + OpenZLCompressor.h + Lz4Compressor.h + Compression.h +) + +target_link_libraries(nimble_compression nimble_common Folly::folly lz4 ${VELOX_OPENZL_LIBRARIES}) diff --git a/velox/dwio/nimble/compression/COMPRESSION.md b/velox/dwio/nimble/compression/COMPRESSION.md new file mode 100644 index 00000000000..47b6c24f833 --- /dev/null +++ b/velox/dwio/nimble/compression/COMPRESSION.md @@ -0,0 +1,349 @@ +# Nimble Compression Guide + +This document explains how compression works in Nimble, how to adjust compression settings, and how to add a new compression codec. + +## Architecture Overview + +Nimble's compression system has three layers: + +1. **Codec layer** — individual compressor implementations (`ICompressor`) that perform the actual compress/decompress work. +2. **Policy layer** — `CompressionPolicy` decides *which* codec to use, with what parameters, and whether to accept the result. +3. **Configuration layer** — `CompressionOptions` (exposed via `WriterOptions`) lets users tune compression behavior. + +Compression is applied at **leaf encoding nodes** in the encoding tree (e.g. `Trivial`, `FixedBitWidth`). The `CompressionEncoder` template class in `Compression.h` handles the common pattern of trying compression and falling back to uncompressed if the policy rejects the result. + +``` +WriterOptions::compressionOptions + → ManualEncodingSelectionPolicy (encoding selection) + → ConfiguredCompressionPolicy (created per-stream at leaf level) + → Compression::compress() (static dispatch) + → CompressorRegistry → ICompressor::compress() +``` + +## Supported Codecs + +| Codec | Enum Value | Class | Files | Default Settings | Min Size | OSS | +|-------|-----------|-------|-------|-----------------|----------|-----| +| Uncompressed | 0 | — | — | — | — | Yes | +| Zstd | 1 | `ZstdCompressor` | `ZstdCompressor.{h,cpp}` | level=3 | 25 bytes | Yes | +| MetaInternal | 2 | `MetaInternalCompressor` | `fb/MetaInternal{,MC}Compressor.{h,cpp}` | comp=4, decomp=2 | 40 bytes | No | +| LZ4 | 3 | `Lz4Compressor` | `Lz4Compressor.{h,cpp}` | accel=1 | 12 bytes | Yes | + +The default codec is **MetaInternal** internally and **Zstd** in OSS builds (controlled by `DISABLE_META_INTERNAL_COMPRESSOR`). + +## Adjusting Compression Settings + +### Option 1: Direct `WriterOptions` (programmatic) + +Set fields on `WriterOptions::compressionOptions` before constructing a `Writer`: + +```cpp +#include "velox/dwio/nimble/writer/WriterOptions.h" + +WriterOptions options; + +// Switch to Zstd with higher compression level +options.compressionOptions.compressionType = CompressionType::Zstd; +options.compressionOptions.zstdCompressionLevel = 7; + +// Reject compressed output unless it saves at least 5% +options.compressionOptions.compressionAcceptRatio = 0.95f; + +// Increase minimum stream size before attempting compression +options.compressionOptions.zstdMinCompressionSize = 64; + +Writer writer(type, std::move(file), pool, std::move(options)); +``` + +### Option 2: Serde Parameters (Hive table properties) + +When writing through the DWIO API layer (`NimbleWriterOptionBuilder`), compression settings are read from Hive serde parameters: + +| Serde Parameter | Field | +|---|---| +| `alpha.encodingselection.compression.accept.ratio` | `compressionAcceptRatio` | +| `alpha.zstd.compression.size.min` | `zstdMinCompressionSize` | +| `alpha.zstrong.compression.size.min` | `internalMinCompressionSize` | +| `alpha.zstrong.compression.level` | `internalCompressionLevel` | +| `alpha.zstrong.decompression.level` | `internalDecompressionLevel` | +| `alpha.zstrong.enable.variable.bit.width.compressor` | `useVariableBitWidthCompressor` | +| `alpha.openzl.managed.compression.key` | `metaInternalCompressionKey` | + +These are defined in `dwio/api/NimbleConfig.{h,cpp}`. + +### Option 3: `WriterOptionOverrides` (high-level DWIO FileWriter) + +The `WriterOptionOverrides` lambda lets you override options that were already populated from serde parameters: + +```cpp +#include "dwio/api/FileWriter.h" + +WriterOptionOverrides overrides; +overrides.nimbleOverrides = [](nimble::WriterOptions& opts) { + opts.compressionOptions.compressionType = CompressionType::Lz4; + opts.compressionOptions.lz4AccelerationLevel = 2; +}; +``` + +### `CompressionOptions` Reference + +Defined in `velox/dwio/nimble/encodings/selection/EncodingSelectionPolicy.h`: + +```cpp +struct CompressionOptions { + // Reject compression if compressedSize > uncompressedSize * ratio + float compressionAcceptRatio = 0.98f; + + // Which codec to use + CompressionType compressionType = CompressionType::MetaInternal; // Zstd in OSS + + // Zstd settings + uint64_t zstdMinCompressionSize = 25; // skip streams smaller than this + uint32_t zstdCompressionLevel = 3; // zstd level (1=fast, 22=max) + + // LZ4 settings + uint64_t lz4MinCompressionSize = 12; + uint32_t lz4AccelerationLevel = 1; // higher = faster but less compression + + // MetaInternal (Zstrong) settings + uint64_t internalMinCompressionSize = 40; + uint32_t internalCompressionLevel = 4; + uint32_t internalDecompressionLevel = 2; + bool useVariableBitWidthCompressor = false; + MetaInternalCompressionKey metaInternalCompressionKey; +}; +``` + +### Metadata Compression + +Separate from per-stream data compression, metadata sections (stripe groups, optional sections) in the file footer are compressed with Zstd when they exceed `metadataCompressionThreshold` (default: 64 KB). Configure via `WriterOptions::metadataCompressionThreshold`. + +## How Compression Selection Works + +Compression is **not** per-column — a single `CompressionOptions` applies uniformly to all streams. The per-stream decision is accept/reject: + +1. If the stream's raw size is below `minCompressionSize`, skip compression. +2. Compress the data with the configured codec. +3. Call `CompressionPolicy::shouldAccept()` — if `compressedSize > uncompressedSize * compressionAcceptRatio`, reject and keep uncompressed. +4. Individual codecs also reject internally (e.g. Zstd returns `Uncompressed` if `ZSTD_compress` produces output larger than input). + +## Adding a New Compression Codec + +### Step 1: Add the enum value + +Add a new entry to `CompressionType` in **both** locations. The values must match. + +**`velox/dwio/nimble/common/Types.h`:** +```cpp +enum class CompressionType : uint8_t { + Uncompressed = 0, + Zstd = 1, + MetaInternal = 2, + Lz4 = 3, + MyCodec = 4, // <-- new +}; +``` + +**`velox/dwio/nimble/tablet/Footer.fbs`:** +```fbs +enum CompressionType:uint8 { + Uncompressed = 0, + Zstd = 1, + MetaInternal = 2, + Lz4 = 3, + MyCodec = 4, +} +``` + +Also update the `toString(CompressionType)` function in `velox/dwio/nimble/common/Types.cpp`. + +### Step 2: Implement the `ICompressor` interface + +Create `MyCodecCompressor.h` and `MyCodecCompressor.cpp` in `velox/dwio/nimble/compression/`. The interface is defined in `Compression.h`: + +```cpp +#pragma once + +#include "velox/dwio/nimble/compression/Compression.h" + +namespace facebook::nimble { + +class MyCodecCompressor : public ICompressor { + public: + /// Compress data. Return CompressionType::Uncompressed with std::nullopt + /// buffer if compression is not beneficial. + CompressionResult compress( + velox::memory::MemoryPool& pool, + std::string_view data, + DataType dataType, + int bitWidth, + const CompressionPolicy& compressionPolicy) override; + + /// Decompress data. The compressionType parameter identifies which codec + /// produced the data. + velox::BufferPtr uncompress( + velox::memory::MemoryPool& pool, + CompressionType compressionType, + DataType dataType, + std::string_view data, + velox::BufferPool* bufferPool = nullptr) override; + + /// Return the uncompressed size from the compressed data header, if + /// available. + std::optional uncompressedSize(std::string_view data) const override; + + CompressionType compressionType() override; +}; + +} // namespace facebook::nimble +``` + +Implementation notes (follow the pattern in `ZstdCompressor.cpp` or `Lz4Compressor.cpp`): + +- **Prepend the uncompressed size** as a `uint32_t` header before the compressed payload. Both Zstd and LZ4 do this so `uncompressedSize()` can read it without decompressing. +- **Return `Uncompressed`** if the compressed output is not smaller than the input — return `{CompressionType::Uncompressed, std::nullopt}`. +- **Call `compressionPolicy.shouldAccept()`** to let the policy decide whether the compression ratio is good enough. If it returns `false`, return `Uncompressed`. +- **Define a minimum compression size constant** in `velox/dwio/nimble/common/Constants.h` (e.g. `constexpr uint64_t kMyCodecMinCompressionSize = 20;`). +- Use `allocateBuffer()` from `Compression.h` to allocate output buffers. + +### Step 3: Register the compressor + +Add the new compressor to the `CompressorRegistry` in `Compression.cpp`: + +```cpp +#include "velox/dwio/nimble/compression/MyCodecCompressor.h" + +struct CompressorRegistry { + CompressorRegistry() { + compressors.reserve(4); // <-- update count + compressors.emplace(CompressionType::Zstd, std::make_unique()); + compressors.emplace(CompressionType::Lz4, std::make_unique()); + compressors.emplace(CompressionType::MyCodec, std::make_unique()); // <-- new +#ifndef DISABLE_META_INTERNAL_COMPRESSOR + compressors.emplace(CompressionType::MetaInternal, std::make_unique()); +#endif + } + // ... +}; +``` + +### Step 4: Add compression parameters + +**`velox/dwio/nimble/compression/CompressionPolicy.h`** — add a parameter struct and include it in `CompressionParameters`: + +```cpp +struct MyCodecCompressionParameters { + int16_t myOption = 5; +}; + +struct CompressionParameters { + ZstdCompressionParameters zstd{}; + Lz4CompressionParameters lz4{}; + MetaInternalCompressionParameters metaInternal{}; + MyCodecCompressionParameters myCodec{}; // <-- new +}; +``` + +### Step 5: Add configuration fields + +**`velox/dwio/nimble/compression/CompressionPolicy.h`** — add fields to `CompressionOptions`: + +```cpp +struct CompressionOptions { + // ... existing fields ... + uint64_t myCodecMinCompressionSize = kMyCodecMinCompressionSize; + int32_t myCodecOption = 5; +}; +``` + +**`velox/dwio/nimble/compression/CompressionPolicy.cpp`** — update `ConfiguredCompressionPolicy::config()` to handle the new codec: + +```cpp +CompressionConfig config() const override { + // ... existing Zstd and Lz4 branches ... + if (compressionOptions_.compressionType == CompressionType::MyCodec) { + CompressionConfig information{ + .compressionType = CompressionType::MyCodec, + .minCompressionSize = compressionOptions_.myCodecMinCompressionSize}; + information.parameters.myCodec.myOption = + compressionOptions_.myCodecOption; + return information; + } + // ... MetaInternal fallthrough ... +} +``` + +### Step 6: Update the BUCK file + +Add the new source and header files to `velox/dwio/nimble/compression/BUCK`, plus any third-party dependency: + +```python +cpp_library( + name = "compression", + srcs = [ + "Compression.cpp", + "Lz4Compressor.cpp", + "MyCodecCompressor.cpp", # <-- new + "ZstdCompressor.cpp", + "fb/MetaInternalCompressor.cpp", + "fb/MetaInternalMCCompressor.cpp", + ], + headers = [ + # ... existing headers ... + "MyCodecCompressor.h", # <-- new + ], + deps = [ + # ... existing deps ... + "fbsource//third-party/mycodec:mycodec", # <-- if needed + ], + # ... +) +``` + +### Step 7 (optional): Wire up serde parameters + +If the codec should be configurable via Hive table properties, add config entries in `dwio/api/NimbleConfig.{h,cpp}` and populate them in `dwio/api/NimbleWriterOptionBuilder.cpp`. + +### Step 8: Add tests + +Add compression round-trip tests in `velox/dwio/nimble/encodings/tests/` or `velox/dwio/nimble/compression/tests/`. At minimum, verify: + +- Compress → decompress round-trip produces identical data. +- Streams below `minCompressionSize` are not compressed. +- Streams that don't compress well are kept uncompressed. +- `uncompressedSize()` returns the correct value from the header. +- End-to-end: write a Nimble file with the new codec and read it back. + +## Checklist for Adding a New Codec + +- [ ] Add enum value to `CompressionType` in `Types.h` and `Footer.fbs` (must match) +- [ ] Update `toString(CompressionType)` in `Types.cpp` +- [ ] Implement `ICompressor` subclass (`MyCodecCompressor.{h,cpp}`) +- [ ] Add min compression size constant to `Constants.h` +- [ ] Register in `CompressorRegistry` in `Compression.cpp` +- [ ] Add parameter struct to `CompressionPolicy.h` +- [ ] Add fields to `CompressionOptions` in `CompressionPolicy.h` +- [ ] Handle new codec in `ConfiguredCompressionPolicy::config()` +- [ ] Add source/header to `compression/BUCK` +- [ ] (Optional) Add serde parameters in `NimbleConfig.{h,cpp}` and `NimbleWriterOptionBuilder.cpp` +- [ ] Add unit tests + +## Key Files Reference + +| File | Purpose | +|------|---------| +| `common/Types.h` | `CompressionType` enum | +| `common/Constants.h` | Min compression size constants | +| `compression/Compression.h` | `ICompressor` interface, `CompressionEncoder`, `Compression` static API | +| `compression/Compression.cpp` | `CompressorRegistry`, dispatch logic | +| `compression/CompressionPolicy.h` | `CompressionPolicy` interface, parameter structs, compression options | +| `compression/CompressionPolicy.cpp` | `ConfiguredCompressionPolicy` implementation | +| `compression/ZstdCompressor.{h,cpp}` | Zstd codec (good reference implementation) | +| `compression/Lz4Compressor.{h,cpp}` | LZ4 codec | +| `compression/fb/MetaInternalCompressor.{h,cpp}` | Zstrong/Managed Compression (internal only) | +| `compression/BUCK` | Build target for all compression code | +| `encodings/selection/EncodingSelectionPolicy.h` | Encoding selection policy | +| `velox/WriterOptions.h` | `WriterOptions::compressionOptions` | +| `tablet/Footer.fbs` | FlatBuffers schema (stores `CompressionType` per metadata section) | +| `dwio/api/NimbleConfig.{h,cpp}` | Serde parameter definitions | +| `dwio/api/NimbleWriterOptionBuilder.cpp` | Populates `CompressionOptions` from serde params | diff --git a/velox/dwio/nimble/compression/Compression.cpp b/velox/dwio/nimble/compression/Compression.cpp new file mode 100644 index 00000000000..2213001d3eb --- /dev/null +++ b/velox/dwio/nimble/compression/Compression.cpp @@ -0,0 +1,94 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "velox/dwio/nimble/compression/Compression.h" +#include "velox/dwio/common/Statistics.h" +#include "velox/dwio/nimble/common/Exceptions.h" +#include "velox/dwio/nimble/compression/Lz4Compressor.h" +#include "velox/dwio/nimble/compression/OpenZLCompressor.h" +#include "velox/dwio/nimble/compression/ZstdCompressor.h" +#include "velox/dwio/nimble/encodings/common/EncodingPrimitives.h" + +#ifndef DISABLE_META_INTERNAL_COMPRESSOR +#include "velox/dwio/nimble/compression/fb/MetaInternalCompressor.h" +#endif + +namespace facebook::nimble { + +namespace { + +struct CompressorRegistry { + CompressorRegistry() { + compressors.reserve(4); + compressors.emplace( + CompressionType::Zstd, std::make_unique()); + compressors.emplace( + CompressionType::Lz4, std::make_unique()); + compressors.emplace( + CompressionType::OpenZL, std::make_unique()); +#ifndef DISABLE_META_INTERNAL_COMPRESSOR + compressors.emplace( + CompressionType::MetaInternal, + std::make_unique()); +#endif + } + + std::unordered_map> compressors; +}; + +ICompressor& getCompressor(CompressionType compressionType) { + static CompressorRegistry registry; + auto it = registry.compressors.find(compressionType); + NIMBLE_CHECK( + it != registry.compressors.end(), + "Compressor for type {} is not registered.", + toString(compressionType)); + return *it->second; +} +} // namespace + +/* static */ CompressionResult Compression::compress( + velox::memory::MemoryPool& pool, + std::string_view data, + DataType dataType, + int bitWidth, + const CompressionPolicy& compressionPolicy) { + auto compression = compressionPolicy.config(); + + return getCompressor(compression.compressionType) + .compress(pool, data, dataType, bitWidth, compressionPolicy); +} + +/* static */ velox::BufferPtr Compression::uncompress( + velox::memory::MemoryPool& pool, + CompressionType compressionType, + DataType dataType, + std::string_view data, + velox::io::IoCounter* decompressCounter, + velox::BufferPool* bufferPool) { + auto& compressor = getCompressor(compressionType); + return velox::dwio::common::withDecompressStats(decompressCounter, [&] { + return compressor.uncompress( + pool, compressionType, dataType, data, bufferPool); + }); +} + +/* static */ std::optional Compression::uncompressedSize( + CompressionType compressionType, + std::string_view data) { + return getCompressor(compressionType).uncompressedSize(data); +} + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/compression/Compression.h b/velox/dwio/nimble/compression/Compression.h new file mode 100644 index 00000000000..3726e4ed8e7 --- /dev/null +++ b/velox/dwio/nimble/compression/Compression.h @@ -0,0 +1,224 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "folly/io/IOBuf.h" +#include "velox/buffer/Buffer.h" +#include "velox/buffer/BufferPool.h" +#include "velox/common/base/IoCounter.h" +#include "velox/dwio/nimble/common/Types.h" +#include "velox/dwio/nimble/common/Vector.h" +#include "velox/dwio/nimble/compression/CompressionPolicy.h" +#include "velox/dwio/nimble/encodings/common/EncodingPrimitives.h" + +namespace facebook::nimble { + +struct CompressionResult { + CompressionType compressionType; + std::optional> buffer; +}; + +struct ICompressor { + virtual CompressionResult compress( + velox::memory::MemoryPool& pool, + std::string_view data, + DataType dataType, + int bitWidth, + const CompressionPolicy& compressionPolicy) = 0; + + virtual velox::BufferPtr uncompress( + velox::memory::MemoryPool& pool, + CompressionType compressionType, + DataType dataType, + std::string_view data, + velox::BufferPool* bufferPool = nullptr) = 0; + + virtual std::optional uncompressedSize( + std::string_view data) const = 0; + + virtual CompressionType compressionType() = 0; + + virtual ~ICompressor() = default; +}; + +/// Allocates a buffer of at least 'bytes' capacity, trying the BufferPool +/// first if available, falling back to MemoryPool allocation. +inline velox::BufferPtr allocateBuffer( + velox::memory::MemoryPool& memoryPool, + velox::BufferPool* bufferPool, + uint64_t bytes) { + if (bufferPool != nullptr) { + if (auto buf = bufferPool->get(bytes)) { + buf->setSize(bytes); + return buf; + } + } + return velox::AlignedBuffer::allocate(bytes, &memoryPool); +} + +class Compression { + public: + static CompressionResult compress( + velox::memory::MemoryPool& pool, + std::string_view data, + DataType dataType, + int bitWidth, + const CompressionPolicy& compressionPolicy); + + static velox::BufferPtr uncompress( + velox::memory::MemoryPool& pool, + CompressionType compressionType, + DataType dataType, + std::string_view data, + velox::io::IoCounter* decompressCounter, + velox::BufferPool* bufferPool = nullptr); + + static std::optional uncompressedSize( + CompressionType compressionType, + std::string_view data); + + static void registerCompressor(std::unique_ptr&& compressor); +}; + +// Encodings using compression repeat the same pattern, which involves trying to +// compress the data, and throwing it away if compression policy decides to +// throw it away. This class tries to extract this common logic into one place. +// Note: There are actually two sub-patterns, therefore, there are two CTors in +// this class. More on this below. +template +class CompressionEncoder { + public: + // This CTor handles the sub-pattern where the source data is already encoded + // correctly, so no extra encoding is needed. + CompressionEncoder( + velox::memory::MemoryPool& pool, + const CompressionPolicy& compressionPolicy, + DataType dataType, + std::string_view uncompressedBuffer, + int bitWidth = 0) + : dataSize_{uncompressedBuffer.size()}, + compressionType_{CompressionType::Uncompressed} { + if (dataSize_ == 0 || + dataSize_ < compressionPolicy.config().minCompressionSize || + compressionPolicy.config().compressionType == + CompressionType::Uncompressed) { + // No compression, just use the original buffer. + data_ = uncompressedBuffer; + return; + } + + auto compressionResult = Compression::compress( + pool, uncompressedBuffer, dataType, bitWidth, compressionPolicy); + + if (compressionResult.compressionType == CompressionType::Uncompressed) { + // Compression declined. Use the original buffer. + data_ = uncompressedBuffer; + return; + } + + // Compression accepted. Use the compressed buffer. + compressed_ = std::move(compressionResult.buffer); + data_ = {compressed_->data(), compressed_->size()}; + dataSize_ = compressed_->size(); + compressionType_ = compressionResult.compressionType; + } + + // This CTor handles the sub-pattern where the source data requires special + // encoding before it is compressed/written. + // Note that in this case, the target buffer for the newly encoded data is + // different if compression is applied (it is written to a temp buffer), or if + // compression is skipped (written directly to the stream buffer). + CompressionEncoder( + velox::memory::MemoryPool& pool, + const CompressionPolicy& compressionPolicy, + DataType dataType, + int bitWidth, + size_t uncompressedSize, + const std::function()>& allocateUncompressedBuffer, + std::function encoder) + : encoder_{std::move(encoder)}, + dataSize_{uncompressedSize}, + compressionType_{CompressionType::Uncompressed} { + if (uncompressedSize == 0 || + uncompressedSize < compressionPolicy.config().minCompressionSize || + compressionPolicy.config().compressionType == + CompressionType::Uncompressed) { + // No compression. Do not encode the data yet. It will be encoded later + // (in write()) directly into the output buffer. + return; + } + + // Compression is attempted. Encode the data before compressing it, into a + // temp buffer. + auto uncompressed = allocateUncompressedBuffer(); + char* pos = uncompressed.data(); + encoder_(pos); + auto compressionResult = Compression::compress( + pool, + {uncompressed.data(), uncompressed.size()}, + dataType, + bitWidth, + compressionPolicy); + + if (compressionResult.compressionType == CompressionType::Uncompressed) { + // Compression declined. Since we already encoded the data, remember the + // temp buffer for later. + // Note: data size is still the same uncompressed size. + data_ = uncompressed; + return; + } + + // Compression accepted. Use the compressed buffer. + compressed_ = std::move(compressionResult.buffer); + data_ = {compressed_->data(), compressed_->size()}; + dataSize_ = compressed_->size(); + compressionType_ = compressionResult.compressionType; + } + + size_t getSize() { + return dataSize_; + } + + void write(char*& pos) { + if (!data_.has_value()) { + // If we are here, it means we handle uncompressed data that needs to be + // encoded directly into the target buffer. + encoder_(pos); + return; + } + + if (data_->data() == nullptr) { + return; + } + + std::copy(data_->begin(), data_->end(), pos); + pos += data_->size(); + } + + CompressionType compressionType() { + return compressionType_; + } + + private: + const std::function encoder_; + + size_t dataSize_; + std::optional> data_; + std::optional> compressed_; + CompressionType compressionType_; +}; + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/compression/CompressionPolicy.cpp b/velox/dwio/nimble/compression/CompressionPolicy.cpp new file mode 100644 index 00000000000..56e30240063 --- /dev/null +++ b/velox/dwio/nimble/compression/CompressionPolicy.cpp @@ -0,0 +1,175 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "velox/dwio/nimble/compression/CompressionPolicy.h" + +namespace facebook::nimble { + +ConfiguredCompressionPolicy::ConfiguredCompressionPolicy( + CompressionOptions compressionOptions, + EncodingType encodingType) + : compressionOptions_{std::move(compressionOptions)}, + effectiveAcceptRatio_{getAcceptRatio(encodingType)} {} + +CompressionConfig ConfiguredCompressionPolicy::config() const { + if (compressionOptions_.compressionType == CompressionType::Uncompressed) { + return {.compressionType = CompressionType::Uncompressed}; + } + + if (compressionOptions_.compressionType == CompressionType::Zstd) { + CompressionConfig config{ + .compressionType = CompressionType::Zstd, + .minCompressionSize = compressionOptions_.zstdMinCompressionSize}; + config.parameters.zstd.compressionLevel = + compressionOptions_.zstdCompressionLevel; + return config; + } + + if (compressionOptions_.compressionType == CompressionType::Lz4) { + CompressionConfig config{ + .compressionType = CompressionType::Lz4, + .minCompressionSize = compressionOptions_.lz4MinCompressionSize}; + config.parameters.lz4.accelerationLevel = + compressionOptions_.lz4AccelerationLevel; + return config; + } + + if (compressionOptions_.compressionType == CompressionType::OpenZL) { + CompressionConfig config{ + .compressionType = CompressionType::OpenZL, + .minCompressionSize = compressionOptions_.openzlMinCompressionSize}; + config.parameters.openzl.compressionLevel = + compressionOptions_.openzlCompressionLevel; + config.parameters.openzl.decompressionLevel = + compressionOptions_.openzlDecompressionLevel; + config.parameters.openzl.formatVersion = + compressionOptions_.openzlFormatVersion; + return config; + } + + CompressionConfig config{ + .compressionType = CompressionType::MetaInternal, + .minCompressionSize = compressionOptions_.internalMinCompressionSize}; + config.parameters.metaInternal.compressionLevel = + compressionOptions_.internalCompressionLevel; + config.parameters.metaInternal.decompressionLevel = + compressionOptions_.internalDecompressionLevel; + config.parameters.metaInternal.useVariableBitWidthCompressor = + compressionOptions_.useVariableBitWidthCompressor; + config.parameters.metaInternal.compressionKey = + compressionOptions_.metaInternalCompressionKey; + return config; +} + +bool ConfiguredCompressionPolicy::shouldAccept( + CompressionType /* compressionType */, + uint64_t uncompressedSize, + uint64_t compressedSize) const { + if (uncompressedSize * effectiveAcceptRatio_ < compressedSize) { + return false; + } + + return true; +} + +float ConfiguredCompressionPolicy::getAcceptRatio( + EncodingType encodingType) const { + for (const auto& [enc, ratio] : + compressionOptions_.compressionAcceptRatioOverrides) { + if (enc == encodingType) { + return ratio; + } + } + return compressionOptions_.compressionAcceptRatio; +} + +ReplayedCompressionPolicy::ReplayedCompressionPolicy( + CompressionType compressionType, + CompressionOptions compressionOptions) + : compressionType_{compressionType}, + compressionOptions_{std::move(compressionOptions)} {} + +CompressionConfig ReplayedCompressionPolicy::config() const { + if (compressionType_ == CompressionType::Uncompressed) { + return {.compressionType = CompressionType::Uncompressed}; + } + + if (compressionType_ == CompressionType::Zstd) { + CompressionConfig config{ + .compressionType = CompressionType::Zstd, + .minCompressionSize = compressionOptions_.zstdMinCompressionSize}; + config.parameters.zstd.compressionLevel = + compressionOptions_.zstdCompressionLevel; + return config; + } + + if (compressionType_ == CompressionType::Lz4) { + CompressionConfig config{ + .compressionType = CompressionType::Lz4, + .minCompressionSize = compressionOptions_.lz4MinCompressionSize}; + config.parameters.lz4.accelerationLevel = + compressionOptions_.lz4AccelerationLevel; + return config; + } + + if (compressionType_ == CompressionType::OpenZL) { + CompressionConfig config{ + .compressionType = CompressionType::OpenZL, + .minCompressionSize = compressionOptions_.openzlMinCompressionSize}; + config.parameters.openzl.compressionLevel = + compressionOptions_.openzlCompressionLevel; + config.parameters.openzl.decompressionLevel = + compressionOptions_.openzlDecompressionLevel; + config.parameters.openzl.formatVersion = + compressionOptions_.openzlFormatVersion; + return config; + } + +#ifdef DISABLE_META_INTERNAL_COMPRESSOR + // Replaying a layout recorded by an internal writer can ask for the Meta + // internal compressor, which has no OSS implementation and is not put in the + // registry by Compression.cpp. Falling through to it would throw from + // getCompressor(), so redirect to Zstd. + CompressionConfig config{ + .compressionType = CompressionType::Zstd, + .minCompressionSize = compressionOptions_.zstdMinCompressionSize}; + config.parameters.zstd.compressionLevel = + compressionOptions_.zstdCompressionLevel; + return config; +#else + CompressionConfig config{ + .compressionType = CompressionType::MetaInternal, + .minCompressionSize = compressionOptions_.internalMinCompressionSize}; + config.parameters.metaInternal.compressionLevel = + compressionOptions_.internalCompressionLevel; + config.parameters.metaInternal.decompressionLevel = + compressionOptions_.internalDecompressionLevel; + config.parameters.metaInternal.useVariableBitWidthCompressor = + compressionOptions_.useVariableBitWidthCompressor; + config.parameters.metaInternal.compressionKey = + compressionOptions_.metaInternalCompressionKey; + return config; +#endif +} + +bool ReplayedCompressionPolicy::shouldAccept( + CompressionType /* compressionType */, + uint64_t uncompressedSize, + uint64_t compressedSize) const { + return compressedSize <= + uncompressedSize * compressionOptions_.compressionAcceptRatio; +} + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/compression/CompressionPolicy.h b/velox/dwio/nimble/compression/CompressionPolicy.h new file mode 100644 index 00000000000..10e22d7dd2d --- /dev/null +++ b/velox/dwio/nimble/compression/CompressionPolicy.h @@ -0,0 +1,216 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include +#include + +#include "folly/json/json.h" +#include "velox/dwio/nimble/common/Constants.h" +#include "velox/dwio/nimble/common/Types.h" + +namespace facebook::nimble { + +/// +/// Compression policy type definitions: +/// A compression policy defines which compression algorithm to apply on the +/// data (if any) and what parameters to use for this compression algorithm. In +/// addition, once compression is applied to the data, the compression policy +/// can decide if the compressed result is statisfactory or if it should be +/// discarded. +struct ZstdCompressionParameters { + int16_t compressionLevel = 3; +}; + +struct Lz4CompressionParameters { + int16_t accelerationLevel = 1; +}; + +/// An identifier for the meta internal compression policy. +class MetaInternalCompressionKey { + public: + MetaInternalCompressionKey() = default; + + MetaInternalCompressionKey( + std::string ns, + std::string tableName, + std::string columnName) + : ns_{std::move(ns)}, + tableName_{std::move(tableName)}, + columnName_{std::move(columnName)} {} + + const std::string& ns() const { + return ns_; + } + + const std::string& tableName() const { + return tableName_; + } + + const std::string& columnName() const { + return columnName_; + } + + std::string toString() const { + folly::dynamic json = folly::dynamic::object("ns", ns_)( + "tableName", tableName_)("columnName", columnName_); + return folly::toJson(json); + } + + static MetaInternalCompressionKey fromString(const std::string& str) { + auto json = folly::parseJson(str); + return MetaInternalCompressionKey{ + json["ns"].asString(), + json["tableName"].asString(), + json["columnName"].asString()}; + } + + private: + std::string ns_; + std::string tableName_; + std::string columnName_; +}; + +struct MetaInternalCompressionParameters { + int16_t compressionLevel = 0; + int16_t decompressionLevel = 0; + bool useVariableBitWidthCompressor = true; + MetaInternalCompressionKey compressionKey; +}; + +struct OpenZLCompressionParameters { + int compressionLevel = 6; + int decompressionLevel = 3; + int formatVersion = 25; // current prod max, as of 2026-06-10 +}; + +struct CompressionParameters { + ZstdCompressionParameters zstd{}; + Lz4CompressionParameters lz4{}; + MetaInternalCompressionParameters metaInternal{}; + OpenZLCompressionParameters openzl{}; +}; + +struct CompressionConfig { + CompressionType compressionType{}; + CompressionParameters parameters{}; + uint64_t minCompressionSize = 0; +}; + +struct CompressionOptions { + /// Rejects compression when compressedSize exceeds uncompressedSize + /// multiplied by this ratio. Enforced by every compressor through + /// CompressionPolicy::shouldAccept(). + float compressionAcceptRatio = 0.98f; +#ifndef DISABLE_META_INTERNAL_COMPRESSOR + CompressionType compressionType = CompressionType::MetaInternal; +#else + CompressionType compressionType = CompressionType::Zstd; +#endif + uint64_t zstdMinCompressionSize = kZstdMinCompressionSize; + uint32_t zstdCompressionLevel = 3; + uint64_t lz4MinCompressionSize = kLz4MinCompressionSize; + uint32_t lz4AccelerationLevel = 1; + uint64_t internalMinCompressionSize = kMetaInternalMinCompressionSize; + uint32_t internalCompressionLevel = 4; + uint32_t internalDecompressionLevel = 2; + bool useVariableBitWidthCompressor = false; + MetaInternalCompressionKey metaInternalCompressionKey{}; + uint64_t openzlMinCompressionSize = kOpenZLMinCompressionSize; + int32_t openzlCompressionLevel = 6; + int32_t openzlDecompressionLevel = 3; + int32_t openzlFormatVersion = 25; + /// Per-encoding overrides for compressionAcceptRatio. + /// BlockBitPacking default 0.7: data is already well-packed via per-block + /// baselines, so compression must save at least 30% to justify CPU cost. + std::vector> compressionAcceptRatioOverrides = + {{EncodingType::BlockBitPacking, 0.7f}}; +}; + +class CompressionPolicy { + public: + virtual CompressionConfig config() const = 0; + virtual bool shouldAccept( + CompressionType /* compressionType */, + uint64_t /* uncompressedSize */, + uint64_t /* compressedSize */) const = 0; + + virtual ~CompressionPolicy() = default; +}; + +/// Default compression policy. Default behavior (if not compression policy is +/// provided) is to not compress. +class NoCompressionPolicy : public CompressionPolicy { + public: + CompressionConfig config() const override { + return {.compressionType = CompressionType::Uncompressed}; + } + + virtual bool shouldAccept( + CompressionType /* compressionType */, + uint64_t /* uncompressedSize */, + uint64_t /* compressedSize */) const override { + return false; + } +}; + +/// Compression policy backed by CompressionOptions. +/// Requests the configured compression type and uses +/// CompressionOptions::compressionAcceptRatio to decide whether to keep the +/// compressed result. +class ConfiguredCompressionPolicy : public CompressionPolicy { + public: + ConfiguredCompressionPolicy( + CompressionOptions compressionOptions, + EncodingType encodingType); + + CompressionConfig config() const override; + + virtual bool shouldAccept( + CompressionType compressionType, + uint64_t uncompressedSize, + uint64_t compressedSize) const override; + + private: + float getAcceptRatio(EncodingType encodingType) const; + + const CompressionOptions compressionOptions_; + const float effectiveAcceptRatio_; +}; + +/// Compression policy for replaying an encoding layout's captured compression +/// type while using current CompressionOptions parameters. +class ReplayedCompressionPolicy : public CompressionPolicy { + public: + ReplayedCompressionPolicy( + CompressionType compressionType, + CompressionOptions compressionOptions); + + CompressionConfig config() const override; + + virtual bool shouldAccept( + CompressionType compressionType, + uint64_t uncompressedSize, + uint64_t compressedSize) const override; + + private: + const CompressionType compressionType_; + const CompressionOptions compressionOptions_; +}; + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/compression/Lz4Compressor.cpp b/velox/dwio/nimble/compression/Lz4Compressor.cpp new file mode 100644 index 00000000000..dbae78a3499 --- /dev/null +++ b/velox/dwio/nimble/compression/Lz4Compressor.cpp @@ -0,0 +1,81 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/dwio/nimble/compression/Lz4Compressor.h" + +#include +#include + +namespace facebook::nimble { + +CompressionResult Lz4Compressor::compress( + velox::memory::MemoryPool& memoryPool, + std::string_view data, + DataType /* dataType */, + int /* bitWidth */, + const CompressionPolicy& compressionPolicy) { + auto parameters = compressionPolicy.config().parameters.lz4; + Vector buffer{&memoryPool, data.size() + sizeof(uint32_t)}; + auto pos = buffer.data(); + encoding::writeUint32(data.size(), pos); + auto ret = LZ4_compress_fast( + data.data(), pos, data.size(), data.size(), parameters.accelerationLevel); + const auto compressedSize = static_cast(ret) + sizeof(uint32_t); + if (ret == 0 || + !compressionPolicy.shouldAccept( + CompressionType::Lz4, data.size(), compressedSize)) { + return { + .compressionType = CompressionType::Uncompressed, + .buffer = std::nullopt, + }; + } + + buffer.resize(ret + sizeof(uint32_t)); + return { + .compressionType = CompressionType::Lz4, + .buffer = std::move(buffer), + }; +} + +velox::BufferPtr Lz4Compressor::uncompress( + velox::memory::MemoryPool& memoryPool, + const CompressionType /* compressionType */, + const DataType /* dataType */, + std::string_view data, + velox::BufferPool* bufferPool) { + auto pos = data.data(); + const uint32_t uncompressedSize = encoding::readUint32(pos); + auto buffer = allocateBuffer(memoryPool, bufferPool, uncompressedSize); + auto ret = LZ4_decompress_safe( + pos, + buffer->asMutable(), + data.size() - sizeof(uint32_t), + uncompressedSize); + NIMBLE_CHECK(ret >= 0, "Error decompressing LZ4 data."); + return buffer; +} + +std::optional Lz4Compressor::uncompressedSize( + std::string_view data) const { + auto* pos = data.data(); + return encoding::readUint32(pos); +} + +CompressionType Lz4Compressor::compressionType() { + return CompressionType::Lz4; +} + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/compression/Lz4Compressor.h b/velox/dwio/nimble/compression/Lz4Compressor.h new file mode 100644 index 00000000000..950a4bf3ddb --- /dev/null +++ b/velox/dwio/nimble/compression/Lz4Compressor.h @@ -0,0 +1,44 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "velox/dwio/nimble/compression/Compression.h" + +namespace facebook::nimble { + +class Lz4Compressor : public ICompressor { + public: + CompressionResult compress( + velox::memory::MemoryPool& memoryPool, + std::string_view data, + DataType dataType, + int bitWidth, + const CompressionPolicy& compressionPolicy) override; + + velox::BufferPtr uncompress( + velox::memory::MemoryPool& memoryPool, + CompressionType compressionType, + DataType dataType, + std::string_view data, + velox::BufferPool* bufferPool = nullptr) override; + + std::optional uncompressedSize(std::string_view data) const override; + + CompressionType compressionType() override; +}; + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/compression/OpenZLCompressor.cpp b/velox/dwio/nimble/compression/OpenZLCompressor.cpp new file mode 100644 index 00000000000..711ffdbd81d --- /dev/null +++ b/velox/dwio/nimble/compression/OpenZLCompressor.cpp @@ -0,0 +1,430 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/dwio/nimble/compression/OpenZLCompressor.h" + +#include +#include +#include + +#include "velox/dwio/nimble/common/Exceptions.h" + +#include "openzl/cpp/CCtx.hpp" +#include "openzl/cpp/CParam.hpp" +#include "openzl/cpp/Codecs.hpp" +#include "openzl/cpp/Compressor.hpp" +#include "openzl/cpp/DCtx.hpp" +#include "openzl/cpp/Input.hpp" +#include "openzl/cpp/Selector.hpp" +#include "openzl/cpp/Type.hpp" + +#include "openzl/zl_compressor.h" +#include "openzl/zl_data.h" +#include "openzl/zl_decompress.h" +#include "openzl/zl_errors.h" +#include "openzl/zl_localParams.h" +#include "openzl/zl_segmenter.h" +#include "openzl/zl_version.h" + +#include "openzl/common/assertion.h" +#include "openzl/shared/bits.h" +#include "openzl/shared/estimate.h" +#include "openzl/shared/numeric_operations.h" + +namespace facebook::nimble { +namespace { + +// Local parameter ID for element byte width, consumed by the chunk segmenter. +constexpr int kNimbleElementByteWidthParamId = 1; +// Target chunk size: ~16MB. +constexpr size_t kDefaultChunkByteSizeTarget = 16 << 20; + +// Multi-chunk segmenter. It splits the serial input into +// ~16MB chunks (aligned to element width) and routes each through the inner +// numeric graph. There is no C++ wrapper for segmenters, so this is registered +// via the C API. +ZL_Report defaultSegmenter(ZL_Segmenter* sctx) { + ZL_RESULT_DECLARE_SCOPE_REPORT(sctx); + + size_t const numInputs = ZL_Segmenter_numInputs(sctx); + ZL_ERR_IF_NE(numInputs, 1, node_invalid_input); + + const ZL_Input* const input = ZL_Segmenter_getInput(sctx, 0); + ZL_ASSERT_NN(input); + size_t const totalBytes = ZL_Input_contentSize(input); + + if (totalBytes == 0) { + ZL_ERR_IF_ERR(ZL_Segmenter_processChunk( + sctx, &totalBytes, 1, ZL_GRAPH_STORE, nullptr)); + return ZL_returnSuccess(); + } + + ZL_GraphIDList const customGraphs = ZL_Segmenter_getCustomGraphs(sctx); + ZL_ASSERT_EQ(customGraphs.nbGraphIDs, 1); + ZL_GraphID const innerGraph = customGraphs.graphids[0]; + + ZL_IntParam const eltWidthParam = + ZL_Segmenter_getLocalIntParam(sctx, kNimbleElementByteWidthParamId); + ZL_ERR_IF_EQ( + eltWidthParam.paramId, ZL_LP_INVALID_PARAMID, node_invalid_input); + size_t const eltWidth = (size_t)eltWidthParam.paramValue; + + size_t const chunkTarget = + (kDefaultChunkByteSizeTarget / eltWidth) * eltWidth; + + size_t remaining = totalBytes; + while (remaining > 0) { + size_t chunkSize = (remaining > chunkTarget) ? chunkTarget : remaining; + ZL_ERR_IF_ERR( + ZL_Segmenter_processChunk(sctx, &chunkSize, 1, innerGraph, nullptr)); + remaining -= chunkSize; + } + return ZL_returnSuccess(); +} + +// Decides if a tokenize transform can and should be used on @p input. It +// calculates an upper bound on the size of tokenization itself (size of +// alphabet + size of indices) and compares it to the original size times a +// threshold multiplier. The multiplier differs for floats and ints. +bool shouldTokenize( + const openzl::Input& input, + int compressionLevel, + bool isFloat) { + auto const eltWidth = input.eltWidth(); + auto const nbElts = input.numElts(); + auto const src = input.ptr(); + auto const srcSize = eltWidth * nbElts; + + if (compressionLevel == 1 && eltWidth == 8) { + // Disable tokenization for level 1 and 64-bit integers because it is + // quite slow. + return false; + } + + static constexpr size_t maxAlphabetSizeInBytes = + 64 * 1024; // see T246457068 for details + const auto maxAlphabetSizeInElts = maxAlphabetSizeInBytes / eltWidth; + + uint64_t const maxEltValue = + (eltWidth >= 8) ? UINT64_MAX : (1ull << (8 * eltWidth)) - 1; + uint64_t const maxCardValue = + std::min(std::min(maxEltValue, (uint64_t)nbElts), maxAlphabetSizeInElts); + auto const cardinality = + ZL_estimateCardinality_fixed(src, nbElts, eltWidth, maxCardValue); + if (cardinality.estimateUpperBound >= maxAlphabetSizeInElts) { + return false; + } + size_t const multiplier = (isFloat || compressionLevel == 1) ? 7 : 9; + auto const tokenizeEstimatedAlphabetSize = + cardinality.estimateUpperBound * eltWidth; + auto const tokenizeEstimatedIndicesSize = + nbElts * (size_t)ZL_nextPow2(cardinality.estimateUpperBound) / 8; + auto const tokenizeEstimatedUpperBounds = + tokenizeEstimatedAlphabetSize + tokenizeEstimatedIndicesSize; + return tokenizeEstimatedUpperBounds < srcSize * multiplier / 10; +} + +// Selects field-lz vs zstd as the backend LZ stage. Single-byte elements go to +// zstd; wider elements use the struct-aware field-lz graph. +class SelectLz : public openzl::Selector { + public: + explicit SelectLz(openzl::GraphID fieldLz) : fieldLz_{fieldLz} {} + + openzl::SelectorDescription selectorDescription() const override { + return { + .name = std::string{"nimble.select_lz"}, + .inputTypeMask = openzl::TypeMask::Numeric, + .customGraphs = {fieldLz_}, + .localParams = {}, + }; + } + + openzl::GraphID select( + openzl::SelectorState& /* state */, + const openzl::Input& input) const override { + if (input.eltWidth() == 1) { + return ZL_GRAPH_ZSTD; + } + return fieldLz_; + } + + private: + openzl::GraphID fieldLz_; +}; + +// Selects whether to apply the tokenize stage based on a cardinality estimate. +class SelectTokenize : public openzl::Selector { + public: + SelectTokenize( + openzl::GraphID tokenize, + openzl::GraphID noTokenize, + bool isFloat) + : tokenize_{tokenize}, noTokenize_{noTokenize}, isFloat_{isFloat} {} + + openzl::SelectorDescription selectorDescription() const override { + return { + .name = isFloat_ ? "nimble.select_tokenize_f32" + : "nimble.select_tokenize_int", + .inputTypeMask = openzl::TypeMask::Numeric, + .customGraphs = {tokenize_, noTokenize_}, + .localParams = {}, + }; + } + + openzl::GraphID select( + openzl::SelectorState& state, + const openzl::Input& input) const override { + const int compressionLevel = + state.getCParam(openzl::CParam::CompressionLevel); + return shouldTokenize(input, compressionLevel, isFloat_) ? tokenize_ + : noTokenize_; + } + + private: + openzl::GraphID tokenize_; + openzl::GraphID noTokenize_; + bool isFloat_; +}; + +// Selects whether to range-pack (subtract min and pack into the smallest +// integer width) based on the actual value range of the input. +class SelectRangePack : public openzl::Selector { + public: + SelectRangePack(openzl::GraphID rangePack, openzl::GraphID noRangePack) + : rangePack_{rangePack}, noRangePack_{noRangePack} {} + + openzl::SelectorDescription selectorDescription() const override { + return { + .name = std::string{"nimble.select_range_pack"}, + .inputTypeMask = openzl::TypeMask::Numeric, + .customGraphs = {rangePack_, noRangePack_}, + .localParams = {}, + }; + } + + openzl::GraphID select( + openzl::SelectorState& /* state */, + const openzl::Input& input) const override { + const auto eltWidth = input.eltWidth(); + const auto nbElts = input.numElts(); + const auto src = input.ptr(); + const ZL_ElementRange range = + ZL_computeUnsignedRange(src, nbElts, eltWidth); + const auto rangeSize = range.max - range.min; + if (NUMOP_numericWidthForValue(rangeSize) < eltWidth) { + return rangePack_; + } + return noRangePack_; + } + + private: + openzl::GraphID rangePack_; + openzl::GraphID noRangePack_; +}; + +// Registers the chunk segmenter on the raw compressor and wraps @p innerGraph +// with it, threading the element byte width through as a local parameter. This +// is the only place the C API is used directly, since segmenters have no C++ +// wrapper. +openzl::GraphID wrapWithDefaultSegmenter( + ZL_Compressor* cgraph, + openzl::GraphID innerGraph, + size_t elementByteWidth) { + ZL_Type inputType = ZL_Type_serial; + ZL_SegmenterDesc desc = { + .name = "!nimble.default_segmenter", + .segmenterFn = defaultSegmenter, + .inputTypeMasks = &inputType, + .numInputs = 1, + .lastInputIsVariable = false, + .customGraphs = nullptr, + .numCustomGraphs = 0, + .localParams = {}, + .opaque = {}, + .mparamMat = {}, + .mparam = {}, + }; + ZL_GraphID const segmenterBase = + ZL_Compressor_registerSegmenter(cgraph, &desc); + + ZL_IntParam intParams[] = {{ + .paramId = kNimbleElementByteWidthParamId, + .paramValue = static_cast(elementByteWidth), + }}; + ZL_LocalParams segParams = { + .intParams = {.intParams = intParams, .nbIntParams = 1}, + .copyParams = {}, + .refParams = {}, + }; + ZL_ParameterizedGraphDesc const segGraphDesc = { + .name = nullptr, + .graph = segmenterBase, + .customGraphs = &innerGraph, + .nbCustomGraphs = 1, + .customNodes = nullptr, + .nbCustomNodes = 0, + .localParams = &segParams, + .mparam = {}, + }; + return ZL_Compressor_registerParameterizedGraph(cgraph, &segGraphDesc); +} + +// Builds the default numeric graph: +// serial -> interpret-as-LE(width) -> range-pack? -> tokenize?+delta -> lz +// Wrapped by the multi-chunk segmenter. +// This decision-making is a default "should-be-good-enough" pipeline that can +// be modified or replaced to fit to your specific use case. +openzl::GraphID makeNumericGraph( + openzl::Compressor& compressor, + bool isFloat, + size_t elementBitWidth, + int formatVersion) { + const size_t elementByteWidth = elementBitWidth / 8; + + const openzl::GraphID fieldLz = openzl::graphs::FieldLz{}(compressor); + const openzl::GraphID lz = openzl::Selector::registerSelector( + compressor, std::make_shared(fieldLz)); + + const openzl::GraphID deltaLz = openzl::nodes::DeltaInt{}(compressor, lz); + const openzl::GraphID tokenize = + openzl::nodes::TokenizeNumeric{/* sort */ true}(compressor, deltaLz, lz); + + const openzl::GraphID tokenizeOrNot = openzl::Selector::registerSelector( + compressor, std::make_shared(tokenize, lz, isFloat)); + + const openzl::GraphID range = + openzl::nodes::RangePack{}(compressor, tokenizeOrNot); + const openzl::GraphID rangeOrNot = openzl::Selector::registerSelector( + compressor, std::make_shared(range, tokenizeOrNot)); + + openzl::nodes::ConvertSerialToNumLE toNumeric{ + static_cast(elementByteWidth)}; + const openzl::GraphID innerGraph = toNumeric(compressor, rangeOrNot); + return wrapWithDefaultSegmenter( + compressor.get(), innerGraph, elementByteWidth); +} + +// Maps the nimble DataType to the appropriate numeric graph. Types without +// a dedicated numeric pipeline fall back to a plain zstd graph (still a valid +// OpenZL frame). +openzl::GraphID createGraph( + openzl::Compressor& compressor, + DataType dataType, + int formatVersion) { + switch (dataType) { + case DataType::Int8: + case DataType::Uint8: + return makeNumericGraph( + compressor, /* isFloat */ false, 8, formatVersion); + case DataType::Int16: + case DataType::Uint16: + return makeNumericGraph( + compressor, /* isFloat */ false, 16, formatVersion); + case DataType::Int32: + case DataType::Uint32: + return makeNumericGraph( + compressor, /* isFloat */ false, 32, formatVersion); + case DataType::Int64: + case DataType::Uint64: + return makeNumericGraph( + compressor, /* isFloat */ false, 64, formatVersion); + case DataType::Float: + return makeNumericGraph( + compressor, /* isFloat */ true, 32, formatVersion); + case DataType::Double: + case DataType::Bool: + case DataType::String: + case DataType::Undefined: + default: + return ZL_GRAPH_ZSTD; + } +} + +} // namespace + +CompressionResult OpenZLCompressor::compress( + velox::memory::MemoryPool& pool, + std::string_view data, + DataType dataType, + int /* bitWidth */, + const CompressionPolicy& compressionPolicy) { + const auto parameters = compressionPolicy.config().parameters.openzl; + + openzl::Compressor compressor; + compressor.selectStartingGraph( + createGraph(compressor, dataType, parameters.formatVersion)); + + openzl::CCtx cctx; + cctx.setParameter( + openzl::CParam::CompressionLevel, parameters.compressionLevel); + cctx.setParameter( + openzl::CParam::DecompressionLevel, parameters.decompressionLevel); + cctx.setParameter(openzl::CParam::FormatVersion, parameters.formatVersion); + cctx.refCompressor(compressor); + + Vector buffer{&pool, openzl::compressBound(data.size())}; + const size_t compressedSize = + cctx.compressSerial({buffer.data(), buffer.size()}, data); + + if (!compressionPolicy.shouldAccept( + CompressionType::OpenZL, data.size(), compressedSize)) { + return { + .compressionType = CompressionType::Uncompressed, + .buffer = std::nullopt, + }; + } + + buffer.resize(compressedSize); + return { + .compressionType = CompressionType::OpenZL, + .buffer = std::move(buffer), + }; +} + +velox::BufferPtr OpenZLCompressor::uncompress( + velox::memory::MemoryPool& pool, + const CompressionType /* compressionType */, + const DataType /* dataType */, + std::string_view data, + velox::BufferPool* bufferPool) { + openzl::DCtx dctx; + const size_t uncompressedSize = + dctx.unwrap(ZL_getDecompressedSize(data.data(), data.size())); + auto buffer = allocateBuffer(pool, bufferPool, uncompressedSize); + const size_t written = + dctx.decompressSerial({buffer->asMutable(), buffer->size()}, data); + NIMBLE_CHECK( + written == uncompressedSize, + "Decompressed size mismatch: expected {}, got {}", + uncompressedSize, + written); + return buffer; +} + +std::optional OpenZLCompressor::uncompressedSize( + std::string_view data) const { + const ZL_Report report = ZL_getDecompressedSize(data.data(), data.size()); + if (ZL_isError(report)) { + return std::nullopt; + } + return ZL_validResult(report); +} + +CompressionType OpenZLCompressor::compressionType() { + return CompressionType::OpenZL; +} + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/compression/OpenZLCompressor.h b/velox/dwio/nimble/compression/OpenZLCompressor.h new file mode 100644 index 00000000000..20d4ea7e79a --- /dev/null +++ b/velox/dwio/nimble/compression/OpenZLCompressor.h @@ -0,0 +1,45 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "velox/dwio/nimble/compression/Compression.h" + +namespace facebook::nimble { + +/// ICompressor backed by custom graphs implemented using OpenZL +class OpenZLCompressor : public ICompressor { + public: + CompressionResult compress( + velox::memory::MemoryPool& pool, + std::string_view data, + DataType dataType, + int bitWidth, + const CompressionPolicy& compressionPolicy) override; + + velox::BufferPtr uncompress( + velox::memory::MemoryPool& pool, + CompressionType compressionType, + DataType dataType, + std::string_view data, + velox::BufferPool* bufferPool = nullptr) override; + + std::optional uncompressedSize(std::string_view data) const override; + + CompressionType compressionType() override; +}; + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/compression/ZstdCompressor.cpp b/velox/dwio/nimble/compression/ZstdCompressor.cpp new file mode 100644 index 00000000000..44dfb334103 --- /dev/null +++ b/velox/dwio/nimble/compression/ZstdCompressor.cpp @@ -0,0 +1,111 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/dwio/nimble/compression/ZstdCompressor.h" + +#include +#include +#include +#include + +namespace facebook::nimble { + +CompressionResult ZstdCompressor::compress( + velox::memory::MemoryPool& pool, + std::string_view data, + DataType dataType, + int /* bitWidth */, + const CompressionPolicy& compressionPolicy) { + auto parameters = compressionPolicy.config().parameters.zstd; + Vector buffer{&pool, data.size() + sizeof(uint32_t)}; + auto pos = buffer.data(); + encoding::writeUint32(data.size(), pos); + auto ret = ZSTD_compress( + pos, data.size(), data.data(), data.size(), parameters.compressionLevel); + if (ZSTD_isError(ret)) { + NIMBLE_CHECK( + ZSTD_getErrorCode(ret) == ZSTD_ErrorCode::ZSTD_error_dstSize_tooSmall, + "Error while compressing data: {}", + ZSTD_getErrorName(ret)); + return { + .compressionType = CompressionType::Uncompressed, + .buffer = std::nullopt, + }; + } + const auto compressedSize = ret + sizeof(uint32_t); + const bool shouldAccept = compressionPolicy.shouldAccept( + CompressionType::Zstd, data.size(), compressedSize); + if (!shouldAccept) { + return { + .compressionType = CompressionType::Uncompressed, + .buffer = std::nullopt, + }; + } + + buffer.resize(compressedSize); + return { + .compressionType = CompressionType::Zstd, + .buffer = std::move(buffer), + }; +} + +namespace { +ZSTD_DCtx* getThreadLocalDCtx() { + struct DCtxDeleter { + void operator()(ZSTD_DCtx* ctx) const { + ZSTD_freeDCtx(ctx); + } + }; + static thread_local std::unique_ptr ctx{ + ZSTD_createDCtx()}; + NIMBLE_CHECK(ctx != nullptr, "Failed to create ZSTD decompression context"); + return ctx.get(); +} +} // namespace + +velox::BufferPtr ZstdCompressor::uncompress( + velox::memory::MemoryPool& pool, + const CompressionType compressionType, + const DataType /* dataType */, + std::string_view data, + velox::BufferPool* bufferPool) { + auto pos = data.data(); + const uint32_t uncompressedSize = encoding::readUint32(pos); + auto buffer = allocateBuffer(pool, bufferPool, uncompressedSize); + auto ret = ZSTD_decompressDCtx( + getThreadLocalDCtx(), + buffer->asMutable(), + buffer->size(), + pos, + data.size() - sizeof(uint32_t)); + NIMBLE_CHECK( + !ZSTD_isError(ret), + "Error uncompressing data: {}", + ZSTD_getErrorName(ret)); + return buffer; +} + +std::optional ZstdCompressor::uncompressedSize( + std::string_view data) const { + auto* pos = data.data(); + return encoding::readUint32(pos); +} + +CompressionType ZstdCompressor::compressionType() { + return CompressionType::Zstd; +} + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/compression/ZstdCompressor.h b/velox/dwio/nimble/compression/ZstdCompressor.h new file mode 100644 index 00000000000..98d3ef6249e --- /dev/null +++ b/velox/dwio/nimble/compression/ZstdCompressor.h @@ -0,0 +1,44 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "velox/dwio/nimble/compression/Compression.h" + +namespace facebook::nimble { + +class ZstdCompressor : public ICompressor { + public: + CompressionResult compress( + velox::memory::MemoryPool& pool, + std::string_view data, + DataType dataType, + int bitWidth, + const CompressionPolicy& compressionPolicy) override; + + velox::BufferPtr uncompress( + velox::memory::MemoryPool& pool, + CompressionType compressionType, + DataType dataType, + std::string_view data, + velox::BufferPool* bufferPool = nullptr) override; + + std::optional uncompressedSize(std::string_view data) const override; + + CompressionType compressionType() override; +}; + +} // namespace facebook::nimble diff --git a/velox/dwio/nimble/docs/Encodings.md b/velox/dwio/nimble/docs/Encodings.md new file mode 100644 index 00000000000..a278501fc55 --- /dev/null +++ b/velox/dwio/nimble/docs/Encodings.md @@ -0,0 +1,121 @@ +# Nimble Encodings Documentation + + +This document describes Nimble's encoding system, including the available encoding +types, the cascading (recursive) and encoding selection policies. + +## Encoding Types + +Nimble supports various encoding types which are extensible. +The following table shows the list of already supported encodings. + + +| Encoding | Data Sequence | Description | +|----------|-------------|-------------| +| **Trivial** | Any | It is baseline encoding. Raw values with no transformation. | +| **Constant** | All identical | Stores a single value for the entire block. | +| **MainlyConstant** | Mostly identical | Stores a common value plus exceptions. | +| **RLE** | Repeated runs | Run-length encoding: (value, count) pairs. | +| **Dictionary** | Low cardinality | Alphabet of unique values + index array. | +| **FixedBitWidth** | Bounded integers | Bit-packs integers using the minimum bits needed. | +| **Varint** | Variable integers | Variable-length integer encoding. | +| **Delta** | Sequential/trending | Stores deltas between consecutive values. | +| **Sentinel** | Nullable data | It uses a sentinel value to represent nulls. | +| **Nullable** | Nullable data | Explicit null bitmap + non-null values stream. | +| **SparseBool** | Skewed on false/true | Stores positions of the minority boolean value. | +| **Prefix** | Shared prefixes | Prefix-based string encoding. | + +## Cascading (Recursive) Encoding + +Nimble's key differentiator is **cascading encoding**: encodings can contain other +encodings, and helps forming an encoding tree. +* Each nested encoding is independently selected based on the characteristics of its sub-stream data. +* When selecting an encoding for a nested sub-stream, the parent encoding type is +excluded from the candidate list to prevent infinite recursion. +* For example, a `DictionaryEncoding`'s alphabet sub-stream will not be encoded with another `DictionaryEncoding`. + +### Example: Dictionary Encoding Tree + +For an example, consider an input array `["apple", "banana", "apple", "cherry", "banana", "strawberry"]`: + +``` +DictionaryEncoding +├── Alphabet: TrivialEncoding +│ └── ["apple", "banana", "cherry", "strawberry"] +│ └── Compression: Zstd (if beneficial) +└── Indices: FixedBitWidthEncoding + └── [0, 1, 0, 2, 1, 3] (2 bits each) + └── Compression: Zstd (if beneficial) +``` + +The dictionary encoding produces two sub-streams: +1. **Alphabet**: The unique values, encoded with whatever encoding best fits them. +2. **Indices**: Integer references into the alphabet, typically bit-packed. + +Each sub-stream is recursively encoded using the same selection process, and +compression is independently applied at each level. + +## Encoding Selection + +Encoding selection is controlled by pluggable **EncodingSelectionPolicy** implementations. +The policy inspects the data and its statistics, then selects the best encoding. + +### Selection Flow + +1. Compute `Statistics` for the data (unique counts, min/max, value distribution) +2. Call `policy->select(values, statistics)` to get an `EncodingSelectionResult` +3. Create an `EncodingSelection` context with the result, statistics, and policy +4. Dispatch to the selected encoding's `encode()` method +5. If the encoding is nested, `encodeNested()` triggers recursive selection + +### ManualEncodingSelectionPolicy (Default) + +Uses **cost-based selection**: evaluates each candidate encoding, estimates the encoded +size using `EncodingSizeEstimation`, and applies a read performance factor: + +``` +cost = estimatedSize × readFactor +``` + +Read factors weight encodings by decode performance: + +| Encoding | Read Factor | Rationale | +|----------|------------|-----------| +| Trivial | 0.7 | Fast decode + compression benefits | +| FixedBitWidth | 0.9 | Faster than variable-width | +| Dictionary | 1.0 | Standard | +| RLE | 1.0 | Standard | +| MainlyConstant | 1.0 | Standard | + +The encoding with the lowest cost wins. + +### LearnedEncodingSelectionPolicy + +Uses an **ML model** to predict the best encoding choice based on data statistics. +Intended for workloads where the cost model doesn't capture the full picture + +### ReplayedEncodingSelectionPolicy + +Replays a previously captured **EncodingLayout** to reproduce an exact encoding tree. +Useful for: +- Testing: ensuring deterministic encoding across runs +- Debugging: reproducing a specific encoding configuration +- Migration: applying a known-good layout to new data + +## Compression + +Compression is applied **per encoding level**, after the encoding produces its byte +output. The compression policy decides whether to keep the compressed version based +on the compression ratio. + + +## Statistics + +The `Statistics` class computes per-block statistics used by encoding selection: +- Unique value count (NDV) +- Min/max repeat lengths (for RLE) +- Value distribution characteristics +- Min/max values +- BucketCounts +Column-level statistics (across stripes) are computed by the `velox/stats/` module +for use in query planning and data skipping. diff --git a/velox/dwio/nimble/docs/NimbleDSL.md b/velox/dwio/nimble/docs/NimbleDSL.md new file mode 100644 index 00000000000..120a5f92631 --- /dev/null +++ b/velox/dwio/nimble/docs/NimbleDSL.md @@ -0,0 +1,152 @@ +# NimbleDSL + +NimbleDSL is an interactive SQL-like interface for inspecting Nimble files. It provides a user-friendly way to query row data, view schema information, and inspect file metadata such as layout, encodings, and indexes. + +This document gives a brief overview of the commands available in NimbleDSL. For more detailed information. + + +## Usage + +```bash +./nimble_dsl -- +``` + +`` is the path to a Nimble file. + +Type `HELP` at the prompt for a summary of all commands. Type `QUIT`, `EXIT`, or press Ctrl-D to exit. + +## Commands + +All commands are **case-insensitive**. Trailing semicolons and commas are optional. + +### SELECT + +Reads and displays row data from the file. (The default limit is 20 rows) + +``` +SELECT * [LIMIT n] [OFFSET n] [STRIPE s] +SELECT col1, col2 [LIMIT n] [OFFSET n] [STRIPE s] +``` + + +- `*` selects all columns. Specific columns can be listed by name (space or comma separated). +- `LIMIT n` — return at most `n` rows. +- `OFFSET n` — skip the first `n` rows before returning results. +- `STRIPE s` — restrict reading to stripe `s` (0-indexed). The offset is relative to the start of that stripe. +- A `FROM` clause is accepted and silently ignored for SQL familiarity. + +Examples: + +``` +nimble> SELECT * +nimble> SELECT name, age LIMIT 5 +nimble> SELECT * LIMIT 10 OFFSET 100 +nimble> SELECT * LIMIT 50 STRIPE 0 +``` + +If an invalid column name is provided, an error message is printed suggesting the `DESCRIBE` command. +If an invalid stripe ID is provided, an error is printed showing the valid stripe count. + +### DESCRIBE + +``` +DESCRIBE +``` + +Shows a table of top-level column names, their types, and the corresponding Nimble stream offsets. + +Example output: + +``` +Column Type Stream +------------------------------------------------------------ +user_id BIGINT 0 +name VARCHAR 2 +scores ARRAY 4 +``` + +### SHOW SCHEMA + +``` +SHOW SCHEMA +``` + +Shows the full Nimble schema tree including nested types (arrays, maps, rows, flat maps) with stream offsets and type kinds. This is useful for understanding the complete structure of complex nested schemas. + +### SHOW INFO + +``` +SHOW INFO +``` + +Shows file-level metadata: + +- Nimble format version (major.minor) +- File size +- Checksum value and type (e.g. `XXH3_64`) +- Stripe count +- Row count +- User-defined metadata key-value pairs (e.g. `build.revision`, `hostname`) + +### SHOW STATS + +``` +SHOW STATS +``` + +Shows per-column statistics: value count, null count, min, max, logical size, and physical size. + +* Requires the file to have been written with `enableVectorizedStats = true`. +* If no statistics are available, a message is printed instead. + +### SHOW STRIPES + +``` +SHOW STRIPES +``` + +Shows stripe-level information: stripe ID, byte offset, byte size, and row count for each stripe in the file. + +### SHOW STREAMS + +``` +SHOW STREAMS [STRIPE s] +``` + +Shows stream-level information: stream ID, byte offset, byte length, item count, and stream label. Without `STRIPE`, shows streams across all stripes. +With `STRIPE s`, filters to a single stripe. + +### SHOW ENCODING + +``` +SHOW ENCODING [STRIPE s] +``` + +Shows the encoding tree for each stream: encoding types, data types, and compression used. Without `STRIPE`, shows encodings across all stripes. +With `STRIPE s`, filters to a single stripe. + +### SHOW INDEX + +``` +SHOW INDEX +``` + +Shows index information if the file has an index configured: index columns with sort orders (ASC/DESC NULLS LAST), number of stripes and index groups, per-group compression and offset details, and per-stripe key stream regions. +If no index is configured, a message is printed instead. + +### HELP + +``` +HELP +``` + +Prints a summary of all available commands. + +### QUIT / EXIT + +``` +QUIT +EXIT +``` + +Exits the REPL. Ctrl-D (EOF) also exits. diff --git a/velox/dwio/nimble/docs/architecture/nimble_write_path.html b/velox/dwio/nimble/docs/architecture/nimble_write_path.html new file mode 100644 index 00000000000..f073116414c --- /dev/null +++ b/velox/dwio/nimble/docs/architecture/nimble_write_path.html @@ -0,0 +1,1456 @@ + + + + + +Nimble Write Path: Velox → Nimble → File + + + + +

Nimble Write Path: Velox → Nimble → File

+

End-to-end write path from Velox HiveDataSink to Nimble physical file output.

+ +
+
Column — schema node
+
Stream — one StreamData object (data+nulls bundled)
+
Chunk — ~20MB partition of one stream
+
Encoding — recursive tree inside one chunk
+
+ + + + +
+

0. Write Path — Velox → Nimble File

+ + +
+ + +
+
HIVE CONNECTOR
+
+
+
HiveDataSink::appendData(RowVectorPtr)
+
+ +
│
+ +
+
FileDataSink::write(index, input)
+
+ +
│
+ +
+ writers_[index]→write(dataInput) +
+ +
│
+ + +
+
NIMBLE ADAPTER
+
+
+
NimbleWriterAdapter::write(VectorPtr)
+
fb_velox/dwio/nimble/writer/NimbleWriter.cpp
+
+ +
│
+ +
+ writer_.write(data) +
+ +
│
+ + +
+
VELOX WRITER
+
+
+
Nimble::Writer::write(VectorPtr)
+
+ +
│
+ + +
+
rootWriter_→write(input, OrderedRanges)
+
FieldWriter tree: Velox vectors → StreamData buffers
+
+ +
│
+ +
+
addIndexKey(input)
+
cluster / hash / sorted index writers
+
+ +
│
+ + +
+
evaluateFlushPolicy()
+
default: StripeRawSizeFlushPolicy @ 256MB
+
+
+ + +
+ + +
+ + +
+
┌──────────────────
+
┴
+
──────────────────┐
+
+ +
+ +
+
shouldChunk() = true
+
+
flushChunks()
+
soft → hard chunking
+
+
│
+
+
ENCODING
+
+
+
writeStreams() / writeChunks(lastChunk)
+
+
│
+
+
encodeStreamData()
+
EncodingSelectionPolicy → EncodingFactory::encode()
+
+
no disk write
+
+ + +
+
shouldFlush() = true
+
+
writeStripe()
+
+
│
+
←────── same encoding
+
│
+
+
TABLET WRITER
+
+
+
TabletWriter::writeStripe(rowCount, streams)
+
+
│
+
+
TabletWriter::close()
+
Footer → Postscript → Magic
+
+
+
+ +
+ + +
+
FILE LAYOUT
+ + +
+
+ writeStripe
──→ +
+
+
StripeGroup 0
+
Stripe 0
+
Stripe 1
+
StripeGroup 1
+
Stripe 2
+
Stripe 3
+
+
+ + +
+
+ close()
──→ +
+
+
StripeGroup Meta
+
Optional Sections
+
Footer
+
Postscript + Magic
+
+
+
+ +
+
+ + + + +
+

1. Column → Stream Mapping

+ +

+ Each schema node creates one StreamData object in FieldWriterContext::streams_[]. + Data and nulls are bundled inside the same object — NOT separate streams. +

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
Column type# streamsStreamData classWhat's inside
BOOL1NullableContentStreamData<bool>data_ + nonNulls_ bundled
INT321NullableContentStreamData<int32>data_ + nonNulls_ bundled
DOUBLE1NullableContentStreamData<double>data_ + nonNulls_ bundled
STRING1NullableContentStringStreamDatabuffer_ + lengths_ + nonNulls_ all bundled
ROW / STRUCT1 + childrenNullsStreamDatajust nonNulls_ (children create their own streams)
ARRAY<T>1 + childNullableContentStreamData<uint32>lengths (1 per row) + child has SUM(lengths) rows
+ MAP<K, V> + 1 + key + valNullableContentStreamData<uint32>lengths (1 per row) + key/val children have SUM(lengths) rows
+ + +
+
Data Pipeline (1000-row batch)
+ + +
+ + +
row
nulls
+
a
INT32
+
b
DOUBLE
+
c
STRING
+
d: MAP<STRING, INT>
+
e: ARRAY<INT>
+
f: STRUCT{x:INT, y:DOUBLE, z:STRING}
+ + +
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+ + +
+ +
RowFieldWriter (root)
+ +
+ +
SimpleFieldWriter
<int32>
+
SimpleFieldWriter
<double>
+
StringFieldWriter
+ +
+
MapFieldWriter
+
lengthsStream_
+
+
+
key_ = StringFieldWriter
+
value_ = SimpleFieldWriter<int32>
+
+
+ +
+
ArrayFieldWriter
+
lengthsStream_
+
+
+
elements_ = SimpleFieldWriter<int32>
+
+
+ +
+
RowFieldWriter (nested struct f)
+
nullsStream_
+
+
+
SimpleFieldWriter<int32> x
+
SimpleFieldWriter<double> y
+
StringFieldWriter z
+
+
+
+ + +
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+
↓
+ + +
[0] row nulls
NullsStreamData
1000
+
[1] a INT32
NullableContent<i32>
1000
+
[2] b DOUBLE
NullableContent<dbl>
1000
+
[3] c STRING
NullableContentString
1000
+
[4] d.lengths
NullableContent<u32>
1000
1 per row
+
[5] d.key
NullableContentString
~4000
SUM(lengths)
+
[6] d.val
NullableContent<i32>
~4000
SUM(lengths)
+
[7] e.lengths
NullableContent<u32>
1000
1 per row
+
[8] e.elem
NullableContent<i32>
~3000
SUM(lengths)
+
[9] f nulls
NullsStreamData
1000
+
[10] f.x
NullableContent<i32>
1000
+
[11] f.y
NullableContent<dbl>
1000
+
[12] f.z
NullableContentString
1000
+
+
MAP/ARRAY children: SUM(parent lengths) rows | struct/row: nulls only | chunk boundaries independent per stream
+
+
+ + + + +
+

1.1 StreamData Internals — What Each Stream Actually Stores

+ +
+ + +
+
+
Velox RowVector [10,000 rows]
+
+ +
+ +
+
│
+
+
column 0 (int64)
+
NullableContentStreamData<int64_t>
+
+ data_: Vector<int64_t>
+ [non-null values appended]
+ nonNulls_: Vector<bool>
+ [null bitmap appended] +
+
+
+ + +
+
│
+
+
column 1 (double)
+
NullableContentStreamData<double>
+
+ data_: Vector<double>
+ [non-null values appended]
+ nonNulls_: Vector<bool>
+ [null bitmap appended] +
+
+
+ + +
+
│
+
+
column 2 (string)
+
NullableContentStringStreamData
+
+ buffer_: Vector<char>
+ [string bytes concatenated]
+ lengths_: Vector<size_t>
+ [per-string lengths]
+ nonNulls_: Vector<bool>
+ [null bitmap appended] +
+
+ materialize() builds string_view[] from buffer_ + lengths_ before encoding +
+
+
+
+
+ + +
+
+ Append-only buffers
+ Each write() call appends to the existing vectors. No per-write allocation — just vector growth. +
+
+ Dictionary vectors are flattened
+ If the input is a Velox DictionaryVector, the FieldWriter uses DecodedVector to flatten it before appending raw values. No dictionary indices are preserved. +
+
+ String storage
+ Strings are stored as buffer_ (all chars concatenated) + lengths_ (per-string). Before encoding, materialize() builds string_view[] pointing into buffer_ — O(n), no copies. +
+
+
+
+ + + + +
+

2. Full Pipeline — All Types

+

+ Schema: struct row { a: BOOL, b: INT32, c: DOUBLE, d: STRING, e: MAP<STRING, INT>, f: ARRAY<INT>, g: STRUCT { x: INT, y: DOUBLE, z: STRING } } +

+ + +
+ Layer 1 — Schema nodes (15 total) +
+
+
row
STRUCT
+
a
BOOL
+
b
INT32
+
c
DOUBLE
+
d
STRING
+
e
MAP
+
e.key
STRING
+
e.val
INT
+
f
ARRAY
+
f.elem
INT
+
g
STRUCT
+
g.x
INT
+
g.y
DOUBLE
+
g.z
STRING
+
+ +
↓
+
each schema node → one StreamData entry (1:1). 15 nodes but root "row" counts too → 14 streams (row has no separate stream? no, it does: NullsStreamData)
+ + +
+ Layer 2 — Streams + FieldWriterContext::streams_[] — 14 entries +
+ +
+
+
[0] nulls
+
+ NullsStreamData · 0.6MB +
+
+
[1] a BOOL
+
+ NullableContent<bool> · 0.6MB +
+
+
[2] b INT32
+
+ NullableContent<i32> · 25MB +
+
+
[3] c DOUBLE
+
+ NullableContent<dbl> · 50MB +
+
+
[4] d STRING
+
+ NullableContentString · 80MB +
+
+
[5] e lengths
+
+ NullableContent<u32> · 2.5MB +
+
+
[6] e.key
+
+ NullableContentString · SUM(lengths) strings +
+
+
[7] e.val
+
+ NullableContent<i32> · SUM(lengths) ints +
+
+
[8] f lengths
+
+ NullableContent<u32> · 2.5MB +
+
+
[9] f.elem
+
+ NullableContent<i32> · SUM(lengths) ints +
+
+
[10] g
+
+ NullsStreamData · 0.6MB +
+
+
[11] g.x INT
+
+ NullableContent<i32> · 25MB +
+
+
[12] g.y DOUBLE
+
+ NullableContent<dbl> · 50MB +
+
+
[13] g.z STRING
+
+ NullableContentString · 80MB +
+
+ +
↓
+
each stream chunked independently by raw byte size (~20MB max)
+ + +
+ Layer 3 — Chunks + Chunker slices data+nulls together at the same row boundary within each stream. Boundaries don't align across streams. +
+ +
+
+
[0] nulls
+
C0
+ 0.6MB +
+
+
[1] a BOOL
+
C0
+ 0.6MB +
+
+
[2] b INT32
+
C0 20MB
C1 5MB
+ 25MB +
+
+
[3] c DOUBLE
+
C0 20MB
C1 20MB
C2 10MB
+ 50MB +
+
+
[4] d STRING
+
C0 20MB
C1 20MB
C2 20MB
C3 20MB
+ 80MB +
+
+
[5] e lengths
+
C0
+ 2.5MB +
+
+
[6] e.key
+
C0 20MB
C1 20MB
C2 8MB
+ SUM(lengths) strings +
+
+
[7] e.val
+
C0 20MB
C1 5MB
+ SUM(lengths) ints +
+
+
[8] f lengths
+
C0
+ 2.5MB +
+
+
[9] f.elem
+
C0 20MB
C1 5MB
+ SUM(lengths) ints +
+
+
[10] g
+
C0
+ 0.6MB +
+
+
[11] g.x
+
C0 20MB
C1 5MB
+ 25MB +
+
+
[12] g.y
+
C0 20MB
C1 20MB
C2 10MB
+ 50MB +
+
+
[13] g.z
+
C0 20MB
C1 20MB
C2 20MB
C3 20MB
+ 80MB +
+
+
+ + + + +
+

3. Encode Path (Chunking)

+ +
+ +
+
[2] INT32 layout
+
+
NullableEncoding prefix: [1B type][1B dataType][rowCount]
+
nulls → SparseBool → MetaInternal(Zstrong)
+
data → FixedBitWidth → MetaInternal(Zstrong)
+
+ +
+
+
Chunk Header
+
4B size
+
1B Uncompressed
+
+
+
+
Prefix
+
1B = 5 (Nullable)
+
1B Int32
+
varint rows
+
+
+
+
nulls
+
1B=6(SparseBool) | 1B Bool | rows | sparseVal | nested indices
+
+
+
+
+
data
+
1B=3(FBW) | 1B Int32 | rows | compr | 4B baseline | 1B bitWidth | packed bits
+
+
+
+
+
Self-describing: byte 0 = EncodingType enum, byte 1 = DataType. No external metadata.
+
+ +
+
▸ [2] b INT32, Chunk 0
+
+
encodeStreamData(chunk view)
+
+
│
+
+
NullableEncoding::encodeNullable()
+
+
┌────────────
┴
────────────┐
+
+
+
encodeNested<bool>(nulls)
+
SparseBool::encode()
+
│
+
CompressionEncoder
MetaInternal(Zstrong)(nulls)
+
+
+
encodeNested<int32>(data)
+
FixedBitWidth::encode()
+
│
+
CompressionEncoder
MetaInternal(Zstrong)(data)
+
+
+
└────────────
┬
────────────┘
+
│
+
ChunkedStreamWriter::encode(blob)
prepend 5B header: [4B size][1B Uncompressed]
+
+
+ + +
+ +
+
[4] STRING Dictionary layout
+
+
NullableEncoding prefix: [1B type][1B dataType][rowCount]
+
nulls → SparseBool → MetaInternal(Zstrong)
+
+
data → Dictionary
+
+
alphabet → Trivial
+
lengths → FBW → MetaInternal(Zstrong)
+
chars → raw bytes → MetaInternal(Zstrong)
+
indices → FBW (9b) → MetaInternal(Zstrong)
+
+
+
+ +
+
+
Chunk Header
+
4B size
+
1B Uncompressed
+
+
+
+
Prefix
+
1B = 5 (Nullable)
+
1B String
+
varint rows
+
+
+
+
nulls
+
1B=6(SparseBool) | 1B Bool | rows | sparseVal | nested indices
+
+
+
+
+
data
+
1B=2(Dictionary) | 1B String | rows
+
+
+
+
alphabet
+
1B=0(Trivial) | lengths(FBW) | chars(raw+Zstrong)
+
+
+
indices
+
1B=3(FBW) | 1B UInt32 | rows | compr | baseline | bitWidth | packed
+
+
+
+
+
+
Self-describing: byte 0 = EncodingType enum. Recursive nesting — Dictionary contains Trivial + FBW inside.
+
+ +
+
▸ [4] d STRING, Chunk 0 (500 unique) — Dictionary
+
+
encodeStreamData(chunk view)
+
+
│
+
+
NullableEncoding::encodeNullable()
+
+
┌──────────────
┴
──────────────┐
+
+
+
nulls
+
SparseBool
+
│
+
MetaInternal(Zstrong)
+
+
+
data
+
Dictionary::encode()
+
┌──────────
┴
──────────┐
+
+
+
alphabet
+
Trivial
+
┌────
┴
────┐
+
+
+
lengths
+
FBW
+
│
+
MetaInternal(Zstrong)
+
+
+
chars
+
raw bytes
+
│
+
MetaInternal(Zstrong)
+
+
+
+
+
indices
+
FBW (9b)
+
│
+
MetaInternal(Zstrong)
+
+
+
+
+
└──────────────
┬
──────────────┘
+
│
+
ChunkedStreamWriter::encode(blob)
prepend 5B header: [4B size][1B Uncompressed]
+
+
+ + +
+ +
+
[4] STRING Trivial layout — different encoding!
+
+
NullableEncoding prefix: [1B type][1B dataType][rowCount]
+
nulls → SparseBool → MetaInternal(Zstrong)
+
+
data → Trivial
+
+
lengths → FBW → MetaInternal(Zstrong)
+
chars → raw bytes → MetaInternal(Zstrong)
+
+
+
+ +
+
+
Chunk Header
+
4B size
+
1B Uncompressed
+
+
+
+
Prefix
+
1B = 5 (Nullable)
+
1B String
+
varint rows
+
+
+
+
nulls
+
1B=6(SparseBool) | 1B Bool | rows | sparseVal | nested indices
+
+
+
+
+
data
+
1B=0(Trivial) | 1B String | rows
+
+
+
lengths: FBW(bitWidth) + packed
+
chars: raw bytes (possibly compressed with Zstrong)
+
+
+
+
+
Same Nullable wrapper, but data uses Trivial (type=0) instead of Dictionary — chosen per chunk based on data characteristics.
+
+ +
+
▸ [4] d STRING, Chunk 1 (all unique) — Trivial different encoding!
+
+
encodeStreamData(chunk view)
+
+
│
+
+
NullableEncoding::encodeNullable()
+
+
┌────────────
┴
────────────┐
+
+
+
nulls
+
SparseBool
+
│
+
MetaInternal(Zstrong)
+
+
+
data
+
Trivial::encode()
+
┌────
┴
────┐
+
+
+
lengths
+
FBW
+
│
+
MetaInternal(Zstrong)
+
+
+
chars
+
raw bytes
+
│
+
MetaInternal(Zstrong)
+
+
+
+
+
└────────────
┬
────────────┘
+
│
+
ChunkedStreamWriter::encode(blob)
prepend 5B header: [4B size][1B Uncompressed]
+
+
+ + +
+ compressionAcceptRatio (default 0.98): each leaf encoding compresses independently via CompressionEncoder. + If compressed > 98% of original → falls back to uncompressed for that payload. + Nested sub-encodings (alphabet, indices) are inside the chunk blob — NOT separate streams. +
+ +
+
+ Schema node + → Stream (1 StreamData, data+nulls bundled) + → Chunks (~20MB slices) + → Encoding tree (recursive, per chunk) +
+
+ Leaf columns: 1 node = 1 stream = N chunks. Chunking is per column.
+ MAP/ARRAY: lengths is its own stream; children are separate streams with different row counts (SUM of lengths).
+ STRUCT: nulls-only stream; children are separate streams with same row count as parent. +
+
+
+ + + + +
+

4. Write Data Flow (3 columns: A INT64, B INT64, C STRING)

+ + +
+
+
Phase 1: write() calls — RowVector → StreamData buffers
+
+ +
+ 
TableWriter feeds RowVectors
+─────────────────────────────
+
+RowVector 1 ── Writer::write()
+RowVector 2 ──     │
+RowVector 3 ──     ├─ RowFieldWriter::write()
+   ...      ──     │       │
+RowVector N ──     │       ├─ FieldWriter A ──append── StreamData A (Vector<int64>)
+                   │       ├─ FieldWriter B ──append── StreamData B (Vector<int64>)
+                   │       └─ FieldWriter C ──append── StreamData C (Vector<string_view>)
+                   │
+                   │   (original RowVector can be freed now)
+                   │
+                   └─ evaluateFlushPolicy()
+
+ +
+
+
Phase 1.0 — Initial Writes
+ StreamData A (raw):   10MB
+ StreamData B (raw):    3MB
+ StreamData C (raw):    1MB
+ encodedStreams_[A-C]:  0MB
+ ────────────────────────────
+ memoryUsed = 14MB     stripeEncodedSize = 0MB +
+
+
Phase 1.1 — More write() Calls
+ stream[1] col A:  120MB (data_: 120MB, nonNulls_: ~0.1MB)
+ stream[2] col B:   90MB (data_: 90MB, no nulls)
+ stream[3] col C:   10MB (buffer_: 8MB, lengths_: 1.5MB)
+ ────────────────────────────
+ memoryUsed = 220MB    stripeEncodedSize = 0MB
+ (each stream is one StreamData object — data+nulls bundled inside) +
+
+
+ + +
+
+
Phase 2a: Soft Chunking — evaluateFlushPolicy() called after every write()
+
+ +
+ 
evaluateFlushPolicy()
+│
+├─ shouldChunk?  220MB + 0MB = 220MB ≥ 200MB?  → YES
+│
+├─ SOFT CHUNKING: pick streams ≥ 20MB
+│     A: 120MB ≥ 20MB → YES
+│     B:  90MB ≥ 20MB → YES
+│     C:  10MB < 20MB → SKIP  ← too small, untouched!
+│
+│  Chunk A (120MB raw, ensureFullChunks=true):
+ +
+
Stream A: 120MB raw → 6 encoded chunks (~30MB total)
+
+
raw:
+
+
20MB
+
20MB
+
20MB
+
20MB
+
20MB
+
20MB
+
+
+
↓ encode each
+
+
enc:
+
+
5MB
+
5MB
+
5MB
+
5MB
+
5MB
+
5MB
+
+
= 30MB
+
+
+ 
+│
+│  After chunking A, re-check shouldChunk?
+│     memoryUsed = 0 (A) + 90 (B) + 10 (C) = 100MB
+│     stripeEncodedSize = 30MB
+│     100 + 30 = 130MB > 100MB (low watermark) → CONTINUE chunking
+│
+│  Chunk B (90MB raw, ensureFullChunks=true):
+ +
+
Stream B: 90MB raw → 4 encoded chunks (~20MB) + 10MB leftover stays raw
+
+
raw:
+
+
20MB
+
20MB
+
20MB
+
20MB
+
+
10MB
+
+
+
+
+
enc:
+
+
5MB
+
5MB
+
5MB
+
5MB
+
+
10MB raw
+
+
+
= 20MB + 10MB raw
+
+
+ 
+│
+│  After chunking B, re-check shouldChunk?
+│     memoryUsed = 0 (A) + 10 (B leftover) + 10 (C) = 20MB
+│     stripeEncodedSize = 30 (A) + 20 (B) = 50MB
+│     20 + 50 = 70MB < 100MB (low watermark) → STOP chunking
+│
+│  C is NEVER chunked — global pressure relieved before we got to it
+│
+├─ shouldFlush?
+│     Estimated: 50 + 20/3.7 ≈ 55MB < 100MB target → NO
+│     → KEEP WRITING
+ + +
IN MEMORY: 20MB raw + 50MB encoded = 70MB total
+
+
+
A:
+
+
c1
+
c2
+
c3
+
c4
+
c5
+
c6
+
+
30MB enc
+
+
+
B:
+
+
10MB
+
c1
+
c2
+
c3
+
c4
+
+
10MB raw + 20MB enc
+
+
+
C:
+
+
10MB
+
+
10MB raw
+
+
+
raw
+
encoded
+
+
+
+
+ + +
+
+
Phase 2b: Hard Chunking — soft chunking not enough
+
+ +
+ 
... more write() calls → 207MB total → shouldChunk = YES
+
+Soft chunking runs first:
+  A: 80MB → 4 × 20MB → encode → freed
+  B: 60MB → 3 × 20MB → encode → freed
+  C: 17MB < 20MB → SKIP (too small for soft)
+
+Re-check: 17MB raw + 85MB encoded = 102MB > 100MB → still above!
+ + +
After soft chunking: 17MB raw + 85MB encoded = 102MB
+
+
+
A:
+
+
c1
+
c2
+
c3
+
c4
+
c5
+
c6
+
c7
+
c8
+
c9
+
c10
+
+
0MB raw + 50MB enc (10 chunks)
+
+
+
B:
+
+
c1
+
c2
+
c3
+
c4
+
c5
+
c6
+
c7
+
+
0MB raw + 35MB enc (7 chunks)
+
+
+
C:
+
+
17MB raw
+
+
← only stream with raw data!
+
+
+ + 
╔══════════════════════════════════════════════════════╗
+║  HARD CHUNKING: chunk ALL streams, ensureFullChunks=false  ║
+║  C is 17MB — only stream left with raw data                ║
+║  17MB < 20MB but allowed because ensureFullChunks=false    ║
+╚══════════════════════════════════════════════════════╝
+ + +
+
Stream C: 17MB raw → hard chunk (undersized allowed)
+
+
raw:
+
+
17MB (<20MB!)
+
+
+
+
↓ encode (undersized allowed)
+
+
enc:
+
+
~4MB
+
+
+
+
+ + 
+Re-check: 0 + 89 = 89MB < 100MB → STOP
+shouldFlush? 89MB < 100MB → NO → KEEP WRITING
+ + +
IN MEMORY: 0MB raw + 89MB encoded
+
+
+
A:
+
+
c1
+
c2
+
c3
+
c4
+
c5
+
c6
+
c7
+
c8
+
c9
+
c10
+
+
50MB (10 chunks)
+
+
+
B:
+
+
c1
+
c2
+
c3
+
c4
+
c5
+
c6
+
c7
+
+
+
35MB (7 chunks)
+
+
+
C:
+
+
c1
+
+
+
4MB (1 chunk)
+
+
+
+
+ + +
+
+
Phase 2c: Final writes trigger stripe flush
+
+ +
+ ... more write() calls ...
+ StreamData A (raw):   25MB  (new since last chunk)
+ StreamData B (raw):   15MB
+ StreamData C (raw):    3MB
+ encodedStreams_[A]:  50MB  (10 chunks)
+ encodedStreams_[B]:  35MB  (7 chunks)
+ encodedStreams_[C]:   4MB  (1 chunk)
+ ────────────────────────────
+ memoryUsed = 43MB      stripeEncodedSize = 89MB +
+ +
+ 
evaluateFlushPolicy()
+│
+├─ shouldChunk?  43 + 89 = 132MB < 200MB → NO
+│
+├─ shouldFlush?
+│     Estimated: 89 + 43/3.7 ≈ 100.6MB ≥ 100MB → YES!
+│
+└─→ writeStripe()!
+
+
+ + +
+
+
Phase 3: writeStripe() — flush everything to disk
+
+ +
+ 
writeStripe()
+│
+├─ 1. Encode ALL remaining raw data (lastChunk=true, minChunkSize=0)
+│
+│     A: 25MB raw → [view1: 20MB] [view2: 5MB]
+│        → encodeChunk(20MB) → ~5MB encoded
+│        → encodeChunk(5MB)  → ~1.3MB encoded  ← allowed because minChunkSize=0
+│     B: 15MB raw → [view1: 15MB]
+│        → encodeChunk(15MB) → ~3.8MB encoded
+│     C: 3MB raw → [view1: 3MB]
+│        → encodeChunk(3MB)  → ~0.8MB encoded
+│
+├─ 2. Final chunk inventory:
+│
+│     encodedStreams_[A].chunks: [c1..c10 from chunking rounds] + [c11, c12]
+│                                 = 12 chunks total
+│     encodedStreams_[B].chunks: [c1..c7 from chunking rounds] + [c8]
+│                                 = 8 chunks total
+│     encodedStreams_[C].chunks: [c1 from hard chunking] + [c2]
+│                                 = 2 chunks total
+│
+├─ 3. Compact encodedStreams_ (remove empty slots), pass to TabletWriter
+│
+├─ 4. TabletWriter::writeStripe() writes to disk:
+│
+│     for (stream : streams) {
+│         writeStreamWithChecksum(stream);  // all chunks contiguously
+│     }
+│
+└─ 5. Reset everything (clear all StreamData + encodedStreams_), start next stripe
+ + +
+
ON DISK (one stripe)
+
+
A-c1 A-c2 .. A-c12
◄── Stream A ──►
+
B-c1 B-c2 .. B-c8
◄── Stream B ──►
+
C-c1 C-c2
◄─ C ─►
+
+
◄──────────────────────────── ONE STRIPE ────────────────────────────►
+
+ + +
+ CHUNK INDEX (metadata):
+ Stream A: 12 chunks │ chunkRows, chunkOffsets, chunkMins, chunkMaxs
+ Stream B:  8 chunks │ chunkRows, chunkOffsets, chunkMins, chunkMaxs
+ Stream C:  2 chunks │ chunkRows, chunkOffsets, chunkMins, chunkMaxs
+ → Reader can skip individual chunks via min/max stats (filter pushdown) +
+
+
+
+ + + diff --git a/velox/dwio/nimble/docs/architecture/string_dictionary_optimization.html b/velox/dwio/nimble/docs/architecture/string_dictionary_optimization.html new file mode 100644 index 00000000000..bd0cec65622 --- /dev/null +++ b/velox/dwio/nimble/docs/architecture/string_dictionary_optimization.html @@ -0,0 +1,1313 @@ + + + + + + +StringColumnReader Dictionary Optimization + + + + + + +
+ +

StringColumnReader Dictionary Optimization

+

How Nimble reads string columns as dictionary indices + post-hoc SIMD filtering

+ + +
+
Architecture Overview
+ +

Batch Reader vs Selective Reader

+ +
+

+ Nimble has two completely separate read paths with different class hierarchies, + different encoding APIs, and different capabilities. + The dictionary optimization only exists in the selective reader path. +

+ +
+
+

Batch Reader (nimble/velox/)

+
+ BatchReader
+   → FieldReader tree (one per schema node)
+     → StringFieldReader (for string columns)
+     → StructFieldReader, ArrayFieldReader, ...
+   → Decoder (manages chunks, calls encoding APIs) +
+
Reads ALL rows sequentially. No filtering, no row selection, no predicate pushdown.
+
+ Decoder::next() +   → encoding_->materialize(n, buf)
+ Decoder::skip() +   → encoding_->skip(n) +
+
+ Uses materialize (virtual dispatch) — decode all n values into buffer sequentially. +
+
+
+

Selective Reader (nimble/velox/selective/)

+
+ NimbleRowReader
+   → SelectiveColumnReader tree (Velox base class)
+     → StringColumnReader (for string columns)
+     → StructColumnReader, FlatMapColumnReader, ...
+   → ChunkedDecoder (manages chunks, calls encoding APIs) +
+
Reads SELECTED rows. Has filters, struct nulls, lazy columns. Used by Prestissimo (Presto's C++ worker).
+
+ ChunkedDecoder::readWithVisitor() → encoding_->readWithVisitor(visitor)
+ ChunkedDecoder::skip()           → encoding_->skip(n) +
+
+ Uses readWithVisitor (template dispatch via switch + cast) for selective value reading. +
+
+
+ +
+ +

StringColumnReader — PreserveDictionaryEncoding

+ +
+
+ Velox Execution Engine +
TableScan operator → RowReader::next(batchSize, result)
+
+
↓ for each column in scanSpec
+
+ SelectiveColumnReader::readWithVisitor() +
Velox base class — dispatches to format-specific reader · computes rows (output row set) and incomingNulls (struct-level nulls)
+
+
↓ virtual dispatch
+
+ +
+
+

Other Columns (e.g. IntegerColumnReader)

+
Always flat — no decision needed, chunks load lazily.
+
+ decoder_.readWithVisitor(visitor)
+   → remainingValues_==0
+     → loadNextChunk(false) → read
+   → no pre-load, no inspection +
+
+
+

StringColumnReader

+
Must inspect the encoding before reading — pre-loads to decide dict vs flat.
+
+ 1. ensureLoaded(true)        ← load with flag FIRST
+ 2. dictionaryConvertible()?  ← inspect what we got
+ 3. yes → dict path           ← readDictionaryIndices
+    no  → flat path           ← readWithVisitor +
+
+
+ +
+
+ StringColumnReader StringColumnReader +
Decides dict vs flat path · manages output state (rawValues_, nulls, outputRows_) · applies filter post-hoc
+
+
+ ChunkedDecoder ChunkedDecoder +
Manages chunks · loads next chunk when exhausted · tracks remainingValues_ / rowPosition_
+
+
+ Encoding (Dictionary / RLE / Trivial / …) ENCODING +
Decodes one chunk's data · materializeIndices() for dict indices · readWithVisitor() for flat values
+
+
+ +
+ + +
+
+

How the Three Layers Interact for Dictionary Preservation

+ +
+ +
+
Session Property Check StringColumnReader
+
+ Only StringColumnReader checks these session properties. + Other column readers (Integer, Float, etc.) always read flat. +
+
+ StringColumnReader::readWithDictionary():
+   if (valueHook) return false;
+   if (!zeroCopy) return false;
+   if (!preserveDict) return false; +
+
+ +
+
Load Chunk ChunkedDecoder
+
+ Pass true optimistically — we don't know yet if the + column is dict-convertible (need the encoding loaded to check). +
+
+ decoder_.ensureLoaded(preserveDictionaryEncoding=true)
+   → ChunkedDecoder::loadNextChunk(true)
+     → options.preserveDictionaryEncoding = true
+     → EncodingFactory::create(pool, data, options) +
+
+ +
+
Flag Propagation ENCODING
+
+ EncodingFactory::create(options) calls each encoding's constructor + with the same options. Each encoding passes the same options to its children. + The flag reaches every node in the tree unconditionally. +
+
+ EncodingFactory::create(options{flag=true})
+   → NullableEncoding(options)     ← ignores
+     → factory.create(options)   ← same opts
+       → RLE<string_view>(opts) ← checks
+         → factory.create(opts) ← run values
+         → factory.create(opts) ← run lengths +
+
+ Only RLE checks the flag. All other encodings ignore it. + With isStringType fix, only RLE<string_view> acts on it. +
+ +
+
How RLE acts on the flag
+
+ // RLE<T> constructor
+ if (preserveDictionaryEncoding
+   && isStringType<physicalType>()
+   && valuesEncoding->dictionaryEnabled()) {
+   // dict mode: preserve inner dictionary
+   dictValues_ = BufferedDictEncoding(valuesEncoding);
+   values_ = nullptr;
+ } else {
+   // flat mode: consume inner dictionary
+   values_ = BufferedEncoding(valuesEncoding);
+   dictValues_ = nullptr;
+ } +
+
+ Dict mode: BufferedDictEncoding calls materializeIndices() → raw indices preserved, alphabet stored separately.
+ Flat mode: BufferedEncoding calls materialize() → indices resolved into values, dictionary consumed. +
+
+ +
+
DictionaryEncoding ignores the flag
+
+ DictionaryEncoding always supports both materialize() and + materializeIndices() regardless of the flag. It never hides its dictionary. + The flag only controls whether RLE on top preserves or flattens the dictionary + that DictionaryEncoding provides. +
+
+
+ +
+ ↑ Construction time (what mode to use?)   |   ↓ Read time (did it work?) +
+ +
+
Eligible Check ChunkedDecoder
+
+ Ask the root encoding: "can you produce raw dictionary indices?" + Per-chunk — different chunks can have different encodings. +
+
+ bool dictionaryConvertible() const {
+   return encoding_->dictionaryEnabled();
+ } +
+
+ +
+
Dict vs Flat Path StringColumnReader
+
+ Branch based on the eligible check result. +
+
+ if (!decoder_.dictionaryConvertible()) {
+   clearDictionaryState();
+   return false; // → flat path (readWithVisitor)
+ }
+ // → dict path (readDictionaryIndices → DictionaryVector) +
+
+ +
+
+
+ + +
+
+

3Encoding Catalog — Leaf vs Nested

+ +

+ Leaf encodings store data directly — no child encodings. +

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
EncodingInline Data
ConstantEncodingconstant value
TrivialEncoding<numeric>compression byte + raw packed values
TrivialEncoding<bool>compression byte + bit-packed bitmap
FixedBitWidthcompression byte + baseline + bit width + packed bits
PrefixEncodingrestart interval + offsets + prefix-compressed entries
Varintbaseline value + varint-encoded bytes
+ +

+ Nested encodings have child encodings created via EncodingFactory::create(). + The preserveDictionaryEncoding flag propagates through the same options object to all children. +

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
EncodingSame-Type ChildHelper Children
NullableEncodingnon-null valuesnulls (bool)
MainlyConstantEncodingotherValuesisCommon (bool)
SentinelEncodingsentineledData—
RLE<T> (non-bool)run valuesrun lengths (uint32_t)
RLE<bool>—run lengths (uint32_t)
DictionaryEncodingalphabetindices (uint32_t)
TrivialEncoding<string_view>—lengths (uint32_t)
SparseBool—indices (uint32_t)
Deltadeltas + restatementsisRestatements (bool)
+
+ +
+

45dictionaryEnabled() — Delegation Through Same-Type Children

+ +

+ dictionaryEnabled() only follows same-type children. + Helper children (uint32_t lengths, indices, bools) are never queried. +

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
EncodingdictionaryEnabled()
NESTEDConstantEncodingtrue (always, size-1 dict)
DictionaryEncodingtrue (always)
RLE<T> (non-bool)dictValues_ != nullptr
NullableEncodingdelegates to same-type child
MainlyConstantEncodingdelegates to same-type child
SentinelEncodingfalse (blocks) [dead code]
LEAFTrivialEncodingfalse
FixedBitWidthfalse
PrefixEncodingfalse
Varint, SparseBool, Delta, RLE<bool>false
+ +
+ +
+
Chain follows same-type children ✓
+ 
+dictionaryConvertible()?
+  │
+  ▼
+NullableEncoding::dictionaryEnabled()
+  → nonNulls()->dictionaryEnabled()  ← delegates
+  │
+  ▼
+RLE<string_view>::dictionaryEnabled()
+  → dictValues_ != nullptr → true
+
+  run lengths RLE<uint32_t>?  ← helper child
+    NEVER CHECKED
+
+ +
+
Chain stops — no same-type child ✗
+ 
+dictionaryConvertible()?
+  │
+  ▼
+TrivialEncoding::dictionaryEnabled()
+  → false (no override, no same-type child)
+  │
+  STOPS HERE
+
+  lengths RLE<uint32_t>?  ← helper child
+    NEVER CHECKED
+    (but received preserveDictionaryEncoding=true)
+
+ +
+
+ +
+

Two Chains — preserveDictionaryEncoding (Down) vs dictionaryEnabled() (Up)

+ +
+ +
+
+ Downward (construction) +
+
+ options{flag=true}
+   ↓ same options
+ NullableEncoding(options)    ← ignores
+   ↓ same options
+ RLE<string_view>(options)   ← checks → dict mode
+   ↓ same options
+ Dictionary<sv>(options)     ← ignores
+
+   ↓ run lengths
+ RLE<uint32_t>(options)      ← checks!
+
+ Reaches every node. +
+
+ +
+
+ Upward (read time) +
+
+ dictionaryConvertible()?
+   ↑
+ NullableEncoding → delegates
+   ↑
+ RLE<string_view> → true
+
+ RLE<uint32_t>?  never checked
+
+ Only follows same-type chain. +
+
+ +
+
+ +
+ +
+ +

Use Dictionary Encoding

+ +
+

Dict Path — dictionaryEnabled() and materializeIndices() Come in Pairs

+

+ Each encoding in the dict chain overrides materializeIndices() (or readIndicesWithVisitor()) + and knows how to delegate to its child. If an encoding says "I can produce indices" + via dictionaryEnabled(), it must also implement the API to actually produce them. + Otherwise the reader calls materializeIndices() and hits NIMBLE_UNREACHABLE in the base class. +

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
EncodingdictionaryEnabled()Indices API
ConstantEncodingtrue alwaysmaterializeIndices() — fills buffer with 0
DictionaryEncodingtrue alwaysmaterializeIndices() → indicesEncoding_->materialize()
RLE (dict mode)dictValues_ != nullptrmaterializeIndices() → dictValues_->nextIndex()
NullableEncodingdelegates to childreadIndicesWithVisitor() → handle nulls → delegate to child
MainlyConstantEncodingdelegates to childmaterializeIndices() → common rows get common index, others delegate to child
SentinelEncodingfalse (no override)none — base class NIMBLE_UNREACHABLE
Everything elsefalsenone
+ +

+ The dict path flows through the same-type value chain only. + Helper children (run lengths, null bitmaps, isCommon bools, indices uint32_t) are always + called via materialize(), never materializeIndices(): +

+ +
+ Reader: readDictionaryIndices()
+   │
+   ▼ callReadIndicesWithVisitor(*encoding_)
+ NullableEncoding::readIndicesWithVisitor()
+   → handle nulls (nulls child → materialize(), not materializeIndices)
+   → callReadIndicesWithVisitor(*nonNullValues_)    ← delegates indices to child
+   │
+   ▼
+ RLE<string_view>::readIndicesWithVisitor()
+   → run lengths → materialize()                   ← helper child, always flat
+   → dictValues_->nextIndex()                     ← value child, indices path
+   │
+   ▼
+ BufferedDictEncoding → DictionaryEncoding::materializeIndices()
+   → indicesEncoding_->materialize()              ← returns raw ints, which ARE the indices
+   → done — recursion stops here +
+
+ +

Dictionary Path Optimization

+ +
+

ReadWithVisitorParams — Encoding → Reader Bridge

+

+ The encoding layer decodes data but can't access the reader's null state directly. + ReadWithVisitorParams carries callbacks from the decoder to the encoding, + allowing the encoding to modify reader state at the right moment during decoding. +

+ +
+
+
StringColumnReader
+
owns anyNulls_, returnReaderNulls_,
resultNulls_, nullsInReadRange_
+
+
←
+
+
ChunkedDecoder
+
creates ReadWithVisitorParams
populates callbacks with lambdas
that capture visitor.reader()
+
+
→
+
+
Encoding
+
receives params
calls callbacks at the right
moment during decoding
+
+
+ +

Three Callbacks

+ + + + + + + + + + + + + + + + + + + + + + + + + +
CallbackWhat it does on the ReaderCalled by Encoding when
makeReaderNullsCreates null bitmap buffer, returns writable pointerNullableEncoding needs to write column-level nulls
setReturnNullsModeSets returnReaderNulls_ flag (which bitmap resultNulls() returns)After NullableEncoding finishes decoding nulls
prepareResultNullsResets anyNulls_ / returnReaderNulls_, allocates result nulls bufferBefore scattering indices for null gaps in readDenseMaterializedIndices
+ +
+ All three are lazy — guarded by a flag, execute at most once per batch even if called + per chunk. They exist because the encoding layer can't access the reader directly; + the callbacks bridge the gap via visitor.reader(). +
+
+ +

+ Example: DictionaryEncoding<string_view> with alphabet ["apple","banana","cherry","date"], + 8 rows with indices [0, 3, null, 1, 0, null, 2, 3], filter: name LIKE 'a%'. +

+ +
+
+

Flat Path (default)

+

Reads resolved string values. Filter applied per-row during read.

+
+ prepareRead<string_view>(offset, rows)
+ decoder_.readWithVisitor(visitor)
+   → DictionaryEncoding::readWithVisitor
+     → materialize(): index 0 → alphabet[0] → "apple"
+                    index 3 → alphabet[3] → "date"
+                    ... resolve each index → copy string
+     → visitor applies filter per row:
+       filter("apple") → PASS, copy "apple"
+       filter("date")  → FAIL, skip
+       filter("banana")→ FAIL, skip
+       filter("apple") → PASS, copy "apple" again
+       ... 8 filter calls, 2 string copies +
+
+ rawValues_ holds string_view values → FlatVector<StringView>. + Every row resolves index → string and copies it. +
+
+
+

Dictionary Path (optimization)

+

Reads raw int32 indices, filters post-hoc with SIMD. Returns DictionaryVector.

+
+ prepareRead<int32_t>(offset, rows)
+ suppress filter
+ decoder_.readDictionaryIndices(visitor)
+   → DictionaryEncoding::materializeIndices
+     → raw indices: [0, 3, —, 1, 0, —, 2, 3]
+       no string resolution, no copy
+ restore filter
+ filterDictionaryIndices → SIMD
+   → filter("apple") → cache[0]=PASS  ← once
+   → filter("date")  → cache[3]=FAIL  ← once
+   → filter("banana")→ cache[1]=FAIL  ← once
+   → filter("cherry")→ cache[2]=FAIL  ← once
+   → 4 filter calls total (one per unique entry)
+   → SIMD: gather cache → compress passing indices +
+
+ rawValues_ holds int32_t indices → DictionaryVector<StringView>. + Indices + alphabet pointer. Zero string copy. +
+
+
+ +
+ Why the dictionary path is faster:
+   • Filter calls: flat = N calls (one per row, "apple" tested twice). Dict = K calls (one per unique entry, cached).
+   • String copies: flat = copies every passing string into FlatVector. Dict = zero copies (DictionaryVector references alphabet by pointer).
+   • SIMD: dict path processes 8 indices per cycle (gather → compare → compress). Flat path is scalar per-row.
+   • For 1M rows with 100 unique strings: flat = 1M filter calls + 1M string copies. Dict = 100 filter calls + 1M int32 SIMD ops. +
+ + + + +

+ Example: SELECT name FROM t WHERE name LIKE 'a%' on a dictionary-encoded string column. + The optimization reads int32 indices instead of strings, filters with SIMD, and returns a + DictionaryVector (zero string copy). +

+ + +
+
+
1
+ Prepare — Position & Build Alphabet +
+
+
+
▸ prepareRead<int32_t>(offset, rows, incomingNulls)
+
Position decoder, read null bitmap, allocate rawValues_ for int32_t indices (not strings).
+
▸ ensureDictionaryState()
+
Cache the alphabet: pointer to the encoding's dictionary table.
+
+
+
Key insight
+
rawValues_ is sized for int32_t (4 bytes/slot), not string_view (16 bytes/slot). The optimization starts here — we never allocate space for strings.
+
alphabet ["apple","banana","cherry","date"]
+
+
+
+ + +
+
+
2
+ Suppress Filter & Bulk Read Indices +
+
+
+
▸ scanSpec_->setFilter(nullptr)
+
Suppress filter — bulk read gets ALL rows unfiltered. Filter applied post-hoc with SIMD.
+
▸ decoder_.readDictionaryIndices(visitor, onChunkBoundary)
+
└→ readDictionaryIndicesImpl → for each chunk:
+
    callReadIndicesWithVisitor → encoding.materializeIndices()
+
▸ scanSpec_->setFilter(std::move(savedFilter))
+
Restore filter for post-hoc application.
+
+
+
Result — raw indices, not strings
+
rawValues_
+
+
0
3
—
1
0
—
2
3
+
+
8 int32 indices instead of 8 string copies. Nulls at positions 2, 5.
+
Resolved: apple, date, null, banana, apple, null, cherry, date
+
+
+
+ + +
+
+
3
+ Post-Hoc SIMD Filter +
+
+
+
▸ filterDictionaryIndices(rows, filter)
+
└→ filterByCache → filterDictionaryRunSimd
+
   ├→ SIMD gather: cache[idx] for 8 indices at once
+
   ├→ SIMD compare: pass / fail / unknown masks
+
   ├→ scalar resolve unknowns: alphabet[idx] → filter.testBytes()
+
   └→ SIMD compress: compact passing indices + rows
+
Each dictionary entry tested at most once — cached for all subsequent rows with the same index.
+
+
+
Filter cache (filter: name LIKE 'a%')
+
+
✓
✗
✗
✗
+
+
[0]=apple ✓  [1]=banana ✗  [2]=cherry ✗  [3]=date ✗
+
Only 4 scalar filter calls for any number of rows — the rest is pure SIMD (gather → compare → compress).
+
+
+
+ + +
+
+
4
+ Output — DictionaryVector (Zero Copy) +
+
+
+
▸ getValues(rows, &result)
+
Wrap surviving indices + alphabet pointer into a DictionaryVector. No string copy at any point.
+
+
+
Final output
+
indices [0, 0]
+
rows [0, 4]
+
dictionary ["apple","banana","cherry","date"]
+
→ resolves to ["apple", "apple"] via pointer lookup, not string copy.
+
+
+
+ +
+

SIMD Filtering — filterDictionaryRunSimd (AVX2 — 8 lanes)

+

+ The SIMD kernel at velox/dwio/common/DecoderUtil.h:65 processes 8 dictionary indices per iteration. + The filter is called at most once per unique dictionary entry; results are cached. +

+ +

Cold cache iteration (first encounter)

+
+
① Load 8 indices
+
+
0
3
1
0
+
2
1
3
0
+
+
+
+
② SIMD gather — load 8 cache entries in 1 instruction
+
+
?
?
?
?
+
?
?
?
?
+
+
All kUnknown — cold cache
+
+
+
③ Resolve unknowns — scalar, one per unique entry
+
+ filter("apple") → PASS → cache[0] = ✓
+ filter("date") → FAIL → cache[3] = ✗
+ filter("banana") → FAIL → cache[1] = ✗
+ filter("cherry") → FAIL → cache[2] = ✗ +
+
4 scalar calls total — each entry tested once
+
+
+
④ Pass mask
+
+
1
0
0
1
+
0
0
0
1
+
+
+
+
⑤ SIMD compress — pack 3 passing indices + rows to output
+
+
0
+
0
+
0
+
+
+ +

Warm cache iteration (subsequent batch)

+
+ All cache entries are known → zero scalar calls. Pure SIMD: gather → compare → compress. + 8 indices processed in ~3 CPU instructions. +
+
+ +
+

Multi-Chunk Alphabet Merging

+

+ When a read spans multiple chunks, each chunk may have its own dictionary alphabet. + The onChunkBoundary callback merges alphabets and offsets indices so the reader + sees one unified dictionary across the entire read. +

+ +
+ // Chunk 0: alphabet = ["apple", "banana"]      indices = [0, 1, 0, 1, ...]
+ // Chunk 1: alphabet = ["cherry", "date"]       indices = [0, 1, 0, 0, ...]
+ // Chunk 2: alphabet = ["elderberry"]           indices = [0, 0, 0, 0, ...]
+ //
+ // Merged alphabet = ["apple", "banana", "cherry", "date", "elderberry"]
+ // Merged indices  = [0, 1, 0, 1, ...,  2, 3, 2, 2, ...,  4, 4, 4, 4, ...]
+ //                     ↑ chunk 0          ↑ chunk 1 (+2)    ↑ chunk 2 (+4) +
+ +
+
+
Chunk 0
+
["apple","banana"]
+
+
→ onChunkBoundary →
+
+
Chunk 1
+
["cherry","date"]
+
+
→ onChunkBoundary →
+
+
Chunk 2
+
["elderberry"]
+
+
+ +
↓ tryExtendDictionaryAtChunkBoundary()
+ +
+
Merged Alphabet
+
["apple", "banana", "cherry", "date", "elderberry"]
+
Each chunk's indices offset by cumulative alphabet size
+
+ +

When a chunk is NOT dict-convertible → Abandon

+ +
+
+
Chunk 0
+
Dict ✓
+
+
→
+
+
Chunk 1
+
Dict ✓
+
+
→
+
+
Chunk 2
+
Trivial ✗
+
+
+ +
↓ onChunkBoundary returns false
+ +
+
Abandon Dictionary
+
+ Dict indices (chunks 0-1) → filterDictionaryIndices → expandDictionaryToFlat
+ Remaining rows (chunk 2) → flat read via readWithVisitor +
+
+
+ +
+ + +
+
Reading Process — Stripes, Chunks, Batches & the Metadata Contract
+ +
+

Physical Layout

+

+ A Nimble file is divided into stripes. Each stripe contains multiple streams (one or more per column). + Each stream is divided into chunks. Each chunk is a self-contained encoded block. +

+ +
+ Nimble file:
+   Stripe 0 (stripe metadata: 10,000 rows):
+     Stream 0 (column A values): [Chunk 0: 600 values] [Chunk 1: 400 values]
+     Stream 1 (column B values): [Chunk 0: 1000 values]
+     ...
+   Stripe 1 (stripe metadata: 8,000 rows):
+     ... +
+
+ +
+

Reading a Chunk

+

+ Each stream has its own ChunkedDecoder. Loading a chunk creates an encoding + from the chunk's bytes and sets remainingValues_ to the encoding's row count. + Reading values decrements remainingValues_. When it reaches 0, the chunk is exhausted. +

+ +
+ loadNextChunk():
+   read chunk header (4 bytes: size + compression type)
+   decompress if needed
+   encoding_ = EncodingFactory::create(bytes)  ← e.g. ConstantEncoding("val"), rowCount=22
+   remainingValues_ = encoding_->rowCount()    ← e.g. 22
+
+ Reading:
+   callReadIndicesWithVisitor(*encoding_, visitor, params) → reads N values
+   advancePosition(N) → remainingValues_ -= N
+
+ When remainingValues_ == 0 → chunk exhausted → load next chunk or done +
+
+ +
+

Batch-Driven Loop

+

+ The reader processes rows in batches (e.g., 1000 rows per batch). The loop runs until + visitor.rowIndex() reaches numRows (the batch size). A single batch may + span multiple chunks. +

+ +
+ numRows = visitor.numRows()          ← batch size, e.g. 1000
+
+ while (visitor.rowIndex() < numRows):
+   if (remainingValues_ == 0):
+     loadNextChunk()                 ← cross chunk boundary
+   computeNextRowIndex()             ← how many rows to read from current chunk
+   read values
+   advancePosition(numNonNulls)
+   visitor.rowIndex() += rows processed +
+
+ +
+

Metadata Contract — Why Scalar Types Never Overrun

+

+ For scalar columns, the stripe metadata guarantees that total rows = total values in the stream. + The reader distributes the stripe's rows across batches. The loop exits when + visitor.rowIndex() reaches numRows — before ever calling + loadNextChunk() beyond the last chunk. +

+ +
+ Scalar column: Stream has Chunk 0 (600 values) + Chunk 1 (400 values) = 1000 total
+ Stripe metadata: 1000 rows. Batch size: 1000.
+
+ Iteration 1:
+   remainingValues_ = 600 (chunk 0)
+   computeNextRowIndex → numNonNulls = 600 (capped by remainingValues_)
+   read 600, advancePosition(600) → remainingValues_ = 0
+   visitor.rowIndex() = 600
+
+ Iteration 2:
+   visitor.rowIndex() = 600 < 1000 → enter loop
+   remainingValues_ == 0 → loadNextChunk() → chunk 1, remainingValues_ = 400
+   read 400, advancePosition(400) → remainingValues_ = 0
+   visitor.rowIndex() = 1000
+
+ Iteration 3:
+   visitor.rowIndex() = 1000 < 1000 → FALSE → EXIT
+   Loop exits before reaching remainingValues_ == 0 check again.
+
+ Supply (1000 values) = Demand (1000 rows from stripe metadata) → never overruns. +
+
+ +
+

Struct Children — The Contract Breaks

+

+ For struct children, the writer only writes non-null values to the child stream. + The struct's null bitmap is a separate stream. So the child stream has fewer values than the + stripe has rows — the null bitmap accounts for the difference. +

+ +
+
+
Scalar Column
+
+ Stripe: 1000 rows
+ Stream: 1000 values
+
+ Supply = Demand
+ Loop exits naturally
+ No overrun possible +
+
+
+
Struct Child Column
+
+ Stripe: 20,821 rows
+ Struct null stream: 20,821 bits (22 non-null)
+ Child value stream: 22 values
+
+ Supply (22) ≠ Demand (20,821 rows)
+ Null bitmap bridges the gap
+ Loop driven by bitmap, not stream size +
+
+
+ +

+ The read order: struct null stream is read first (producing the null bitmap), + then the child value stream is read using that bitmap. By the time the child's + readDictionaryIndicesImpl runs, the bitmap is fully available. +

+ +
+ StructColumnReader::read(batch of 1000 rows):
+   1. Read struct null stream → bitmap = [0,0,...,0,1,0,...] for this batch
+      (struct's own ChunkedDecoder, loaded before child reads)
+   2. Call StringColumnReader::read()
+      → readDictionaryIndicesImpl(visitor, nulls=bitmap, ...)
+      → bitmap already complete for this batch +
+ +
+ The fix: when remainingValues_ == 0, instead of blindly loading the next chunk + (which crashes when none exist), check the null bitmap: countNonNulls(nulls, numScanned, endRow). + If non-nulls remain → load chunk (if it fails, it's corruption). + If all remaining rows are null → skip load, process nulls. + This is metadata-driven (the bitmap), not stream-peek-driven (hasMoreChunks()). +
+
+
+ + + +
+ + diff --git a/velox/dwio/nimble/docs/develop.rst b/velox/dwio/nimble/docs/develop.rst new file mode 100644 index 00000000000..8fba7d979e2 --- /dev/null +++ b/velox/dwio/nimble/docs/develop.rst @@ -0,0 +1,20 @@ +*************** +Developer Guide +*************** + +This guide is intended for Nimble contributors and developers of Nimble-based +applications. + +.. toctree:: + :maxdepth: 1 + + develop/nimble_selective_reader + develop/nimble_writer + develop/nullable_encoding + develop/nimble_for_interactive + develop/velox_cache_and_nimble_datapath + develop/nimble_metadata_cache + develop/nimble_java_reader_blockBitPacking + develop/dwio_statistics_architecture + develop/encoding_cost_model + develop/encoding_options_flow diff --git a/velox/dwio/nimble/docs/develop/dwio_statistics_architecture.rst b/velox/dwio/nimble/docs/develop/dwio_statistics_architecture.rst new file mode 100644 index 00000000000..892f3962e8e --- /dev/null +++ b/velox/dwio/nimble/docs/develop/dwio_statistics_architecture.rst @@ -0,0 +1,431 @@ + +DWIO Statistics Architecture +============================ + +IoStatistics • ColumnReaderStatistics • Cache Stats • Writer Stats • Operator Stats • Full Aggregation Pipeline + +1 Reader Stats +-------------- + +The Velox DWIO reader stack has two independent statistics pipelines that never reference each other. They join only at the operator level, where both contribute keys to the RuntimeMetric map. + +Operator (TableScan) + +OperatorStats::runtimeStats — map < string, RuntimeMetric > + +↑ addIoStatsToRuntimeStats() + +↑ toRuntimeMetricMap() + +IoStatistics + +velox/common/io/IoStatistics.h + +Storage / cache layer — "how much did we read?" + + +Granularity: whole file (one per split) + +RuntimeStatistics + +velox/dwio/common/Statistics.h + +Format reader layer — "how much CPU did decoding cost?" + + +ColumnReaderStatistics + +flattenStringDictionaryValues, pageLoadTimeNs skippedStrides, processedStrides, numStripes + +↓ optional (per column) + +DecodingStatsSet → DecodingStats + +typeKind, decodeCPUTimeNanos, decompressCPUTimeNanos + +IoStatistics +^^^^^^^^^^^^ + +FileDataSource + +velox/connectors/hive/FileDataSource.cpp:381 + +dataIoStats\_ = make_shared < IoStatistics > () — for column data IO metadataIoStats\_ = make_shared < IoStatistics > () — for footer/metadata IO + + +FileSplitReader → ReaderOptions::setDataIoStats(shared_ptr) → BufferedInput constructor + +velox/common/io/IoStatistics.h:44 + +read\_ + +IoCounter + +ramHit\_ + +IoCounter + +ssdRead\_ + +IoCounter + +prefetch\_ + +IoCounter + +queryThreadIoLatencyUs\_ + +IoCounter + +storageReadLatencyUs\_ + +IoCounter + +ssdCacheReadLatencyUs\_ + +IoCounter + +cacheWaitLatencyUs\_ + +IoCounter + +rawBytesRead\_ + +atomic + +rawOverreadBytes\_ + +atomic + +map < string, OperationCounters > + +OperationCounters + +(per operation, e.g. "ws_pread") + +Not in RuntimeMetric — ODS export only. + + +stored as shared_ptr in CachedBufferedInput / DirectBufferedInput + +BufferedInput + +IoStatistics* + +enqueue() + + +CacheInputStream + +(cached path) + +read\_ + +ramHit\_ + +ssdRead\_ + +queryThreadIoLatencyUs\_ + + +DirectInputStream + +(no-cache path) + +read\_ + +prefetch\_ + +queryThreadIoLatencyUs\_ + +CoalesceIoStats + +— returned by coalesceIo() in groupRequests() , feeds back into IoStatistics: + +gaps + +readGap().merge() + +duplicateRegions + +incDuplicateRead() + +extraBytes + +incRawOverreadBytes() + + +FileDataSource::getRuntimeStats() calls addIoStatsToRuntimeStats() + +OperatorStats::runtimeStats + +map < string, RuntimeMetric > + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - IoStatistics counter + - RuntimeMetric key + - Unit + * - queryThreadIoLatencyUs\_ + - ioWaitWallNanos + - nanos (us*1000) + * - read\_ + - storageReadBytes + - bytes (sum/count/min/max) + * - ramHit\_ + - numRamRead + ramReadBytes + - count + bytes + * - ssdRead\_ + - numLocalRead + localReadBytes + - count + bytes + * - prefetch\_ + - numPrefetch + prefetchBytes + - count + bytes + * - readGap\_ + - readGapBytes + - bytes (sum/count/min/max) + * - rawOverreadBytes\_ + - overreadBytes + - bytes + +metadata. + +rawBytesRead\_ + +getCompletedBytes() + +ColumnReaderStatistics +^^^^^^^^^^^^^^^^^^^^^^ + +Format RowReader + +DwrfRowReader / SelectiveNimbleRowReader + +owns ColumnReaderStatistics columnReaderStats\_ as value member + + +initColumnStatsCollection(schema, options) — if collectColumnCpuMetrics() is on + +velox/dwio/common/Statistics.h:669 + +flattenStringDictionaryValues + +int64 + +pageLoadTimeNs + +IoCounter + +optional < DecodingMetricsSet > + +:596 + +folly::Synchronized < F14FastMap < nodeId, unique_ptr < DecodingMetrics > > > + +DecodingMetrics + +(per column, keyed by nodeId) :580 + +TypeKind + +IoCounter + + +passed by reference: FormatParams(pool, columnReaderStats\_ & ) + +FormatParams + +base class + +(NimbleParams / DwrfParams) + +runtimeStatistics() returns ColumnReaderStatistics & + + +columnMetricsSet- > getOrCreate(id) + +SelectiveColumnReader + +stores DecodingMetrics* + + +readWithTiming() wraps read() + +FORMAT-AGNOSTIC + +decodeCPUTimeNanos + +DeltaCpuWallTimer callback + + +toFormatData() extracts IoCounter* + +FormatData + +NimbleData / DwrfData + + +& metrics- > decompressCPUTimeNanos + +FORMAT-SPECIFIC + +decompressCPUTimeNanos + +encoding construction site + + +updateRuntimeStats() → stats.columnReaderStats.mergeFrom(columnReaderStats\_) + +RuntimeStatistics + +toRuntimeMetricMap() + +OperatorStats::runtimeStats + +Keys: column\_{nodeId}.{TypeName}.decodeCPUTimeNanos , ...decompressCPUTimeNanos + +Decompress Timing: DWRF vs Nimble +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +DWRF Path +^^^^^^^^^ + +StripeStreamsImpl + +ColumnReaderStatistics* + + +getDecompressCounter(nodeId) + +& metrics- > decompressCPUTimeNanos + + +PagedInputStream + +decompressCounter\_ + + +``withDecompressStats(counter, [ & ]{ /* zlib/zstd decompress */ })`` + +velox/dwio/dwrf/common/DecoderUtil.h • PagedInputStream.cpp:186 + +Nimble Path D108361387 +^^^^^^^^^^^^^^^^^^^^^^ + +NimbleParams::toFormatData() + +DecodingStats* + +decodingStatsSet + + +NimbleData + +decodingStats\_ + + +makeValuesDecoder() + +make*Decoder() + +ChunkedDecoder + + +loadNextChunk() + +options.recordDecompressNanos(decompressNanos) + +velox/dwio/nimble/velox/selective/NimbleData.cpp • ChunkedDecoder.cpp + +2 Cache Stats +------------- + +AsyncDataCache + +velox/common/caching/AsyncDataCache.h — process-wide singleton + +Aggregates from all CacheShard instances + + +each CacheShard increments its local counters, aggregated on stats() call + +velox/common/caching/AsyncDataCache.h:543 + +Snapshot: + +Cumulative: + +shared_ptr < SsdCacheStats > + +SsdCacheStats + +SsdFile.h:142 + + +reported periodically, NOT per-query + +PeriodicStatsReporter + +NOT exported to RuntimeMetric — completely separate from per-query stats + +SimpleLRUCacheStats + +SimpleLRUCache.h:30 + +Used by FileHandleFactory for file handle caching + +ScanTracker + +ScanTracker.h:87 + +Feeds adaptive prefetch decisions, not exported as metrics + +3 Writer Stats +-------------- + +Writer + +(Nimble) + +DwrfWriter + +(DWRF) + +Owns stats as member; populated during write/flush/close + + +incremented during write(), flush(), close() + +velox/dwio/nimble/writer/Writer.h:61 (Nimble) + +Timing: + +Sizes: + +RuntimeMetric + +Embedded: TabletWriter::Stats + +duplicateStreamCount duplicateStreamBytes + +RatioTracker + +(DWRF only) + +feeds flush policy decisions + + +Caller + +writer.runtimeStats() + +returned to TableWriter operator + + +Scuba + +MetricsLog + +StripeFlushMetrics, FileCloseMetrics diff --git a/velox/dwio/nimble/docs/develop/encoding_cost_model.rst b/velox/dwio/nimble/docs/develop/encoding_cost_model.rst new file mode 100644 index 00000000000..f87b17c5940 --- /dev/null +++ b/velox/dwio/nimble/docs/develop/encoding_cost_model.rst @@ -0,0 +1,137 @@ + +Nimble Encoding Cost Model +========================== + +How encoding selection estimates sizes with per-block statistics + +New + +Modified + +Unchanged + +0 + +Caller + +EncodingFactory::encode() +^^^^^^^^^^^^^^^^^^^^^^^^^ + +Called by FieldWriter when flushing a stream. + +.. code-block:: text + + auto statistics = Statistics::create(values); + // O(1) ← data_ + + auto result = selectorPolicy->select(values, statistics); + +Encoding selection + +EncodingSelectionPolicy::select(values, statistics) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. code-block:: text + + hasPerBlockEncoding = any_of(encodingReadFactors_, == BlockBitPacking); + + for (const auto& [encodingType, readFactor] : encodingReadFactors_) { + estimatedSize = EncodingSizeEstimation::estimateSize(encodingType, count, statistics, hasPerBlockEncoding); + // dispatches by type: + // numeric → estimateNumericSize() + // bool → estimateBoolSize() + // string → estimateStringSize() + cost = estimatedSize * readFactor; + if (cost < minCost) { + selectedEncoding = encodingType; + } + } + +EncodingSizeEstimation::estimateNumericSize() +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +switch (encodingType) → dispatches to each encoding's estimateSize() + +UNCHANGED + +no stats needed + +.. code-block:: text + + return prefix + count × sizeof(T) + +UNCHANGED + +no stats needed + +.. code-block:: text + + return prefix + sizeof(T) + +UNCHANGED + +.. code-block:: text + + estimateSize(rowCount, statistics, options) { + return estimateSize(rowCount, statistics.min(), statistics.max(), options); + } + + estimateSize(rowCount, minValue, maxValue, options) { + bitWidth = bitsRequired(maxValue - minValue); + if (!options.fixedBitWidthUseExactBits) bitWidth = roundUpToByte(bitWidth); + payloadSize = nbytes(bitWidth * rowCount); + return prefix + kPrefixSize + payloadSize; + } + +statistics.min() + +← scans data\_ once + +// sets both min\_ AND max\_ in one pass + +// via std::minmax_element(data\_) + +← cached after first call + +NEW + +.. code-block:: text + + estimateSize(const Statistics& statistics) { + const auto& blocks = statistics.minMaxBlocks(); + + for (const auto& b : blocks) { + range = b.max - b.min; + bw = bitsRequired(range); + packedSize = bufferSize(b.count, bw); + rawSize = b.count * sizeof(T); + if (packedSize < rawSize) { + dataSize += packedSize; + } else { + dataSize += rawSize; // skip encoding + } + } + + // per-block metadata stored as nested encodings: + // baselines → Trivial + // bitWidths → Trivial + // dataOffsets → Trivial + metadataSize = kMetadataHeaderSize + + Trivial::estimateSize(numBlocks) + + Trivial::estimateSize(numBlocks) + + Trivial::estimateSize(numBlocks); + + return prefix + metadataSize + dataSize; + } + +statistics.minMaxBlocks() + +← scans data\_ once + +// BlockStatsAccumulator: for each value, + +// track min/max, flush every 1024 rows + +// → vector < {count, min, max} > + +← cached after first call diff --git a/velox/dwio/nimble/docs/develop/encoding_options_flow.rst b/velox/dwio/nimble/docs/develop/encoding_options_flow.rst new file mode 100644 index 00000000000..6f1160aab10 --- /dev/null +++ b/velox/dwio/nimble/docs/develop/encoding_options_flow.rst @@ -0,0 +1,518 @@ + +Nimble Writer Configuration Flow +================================ + +Part 1: Encoding::Options Background + +What the struct is, where it's constructed, and which encodings consume which field + +1a + +velox/dwio/nimble/encodings/common/Encoding.h + +.. code-block:: text + + Encoding useVarintRowCount bufferPool preserveDict freqPartIdx bbpBlockSize + Trivial ● ○ + RLE ● ○ + Dictionary ● ○ + FixedBitWidth ● ○ + BlockBitPacking ● ○ ● + FreqPartition ● ○ ● ○ + Delta, Varint, + SparseBool, + MainlyConst, + Constant, + Sentinel, + Prefix, + Nullable, + ALP ● ○ + + +useVarintRowCount + +bufferPool + +Encoding + +1b + +Write Path + +Writer → file writing + +WriterOptions fields +^^^^^^^^^^^^^^^^^^^^^^^^^ + +(from serde params) + +Writer encode overview +^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +buildEncodingOptions() + +bbpBlockSize=N + +EncodingFactory::encode(options) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +options → select() + EncodingSelection + +Serializer / Deserializer + +Network/RPC vector transfer + +SerializerOptions +^^^^^^^^^^^^^^^^^ + +useVarint + +bbpBlockSize + +StreamData +^^^^^^^^^^ + +useVarintRowCount + +bufferPool + +EncodingFactory(opts).create() +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +options stored, passed to constructors + +Selective Reader + +ChunkedDecoder read path + +SelectiveNimbleRowReader read path +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +shared across all columns, default Options{} + +NimbleParams::toFormatData read overview +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +decodingStats* + +NimbleData read state +^^^^^^^^^^^^^^^^^^^^^ + +per-column: holds encodingFactory\_ + decodingStats\_ creates ChunkedDecoders for each stream + +ChunkedDecoder::loadChunk() +^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +copies factory opts and overrides per-create + +// copy, not reference + +options + +preserveDictionaryEncoding + +options + +decodingStats + +``encodingFactory_->create(pool, data, factory, options)`` + +virtual dispatch + +Part 2: Write Path Overview + +Encoding::Options + +2a + +Path A: DWIO FileWriter API + +TableService / Ingestion pipelines + +Hive Metastore +^^^^^^^^^^^^^^ + +HiveTableMetadata.sd().serdeInfo().parameters() + +FileWriterFactory::create() +^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +format = NIMBLE → copy serdeParams + +NimbleWriterOptionBuilder from FileWriter +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.withSerdeParams(schema, serdeParams).build() + +optionOverrides.nimbleOverrides(options) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +optional lambda-based overrides from caller + +dwio/api/FileWriter.cpp:484-544 + +Path B: Velox HiveDataSink + +Presto / Spark write queries + +HiveInsertTableHandle +^^^^^^^^^^^^^^^^^^^^^ + +serdeParameters\_ (from coordinator) + +createWriterOptions() +^^^^^^^^^^^^^^^^^^^^^ + +options → serdeParameters = table serde params + +processConfigs(connectorConfig, session) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +merges session & connector overrides via emplace() + +NimbleWriterOptionBuilder from HiveDataSink +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.withSerdeParams(schema, serdeParams).build() + +velox/connectors/hive/HiveDataSink.cpp + +OVERRIDE PRIORITY (highest → lowest) + +Table Serde Param + +Session Property + +Connector Config + +Default + +Enforced by std::map::emplace() — only inserts if key does not already exist + + +NimbleWriterOptionBuilder::withSerdeParams() + +map < string,string > + +WriterOptions + +2b + +WriterOptions in constructor +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Encoding + +encodingSelectionPolicyCreator + +readFactors + +(which encodings to try + weights) + +compressionOptions + +(codec, levels, accept ratio) + +blockBitPackingBlockSize + +buildEncodingOptions() + +→ Encoding::Options + +compressionOptions + +(duplicate for replay path) + +encodingLayoutTree + +Flush / Chunking + +(stripe size, memory thresholds) + +Schema + +flatMapColumns, dictionaryArrayColumns, deduplicatedMapColumns + +Index + +clusterIndexConfig + +Memory / Parallelism + +encodingExecutor, reclaimerFactory, spillConfig + +Stats / Misc + +enableStatsCollection, enableStreamDedup, metadata + +2c + +WriterOptions +^^^^^^^^^^^^^^^^^^ + +Writer constructor +^^^^^^^^^^^^^^^^^^^^^^^ + +distributes options to context, creates TabletWriter, invokes factories + +Writer + +enableChunking, chunk sizes stats collection schema columns encoding layout metadata + +FieldWriterContext + +buffer growth string buffers parallel encoding flat map nodes dict array nodes + +TabletWriter + +stream dedup chunk index metadata compression feature reorder + +FlushPolicy + +stripe raw size memory thresholds target storage sz compression factor + +Encoding::Options + +EncodingSelection + +(policy) + +blockBitPackingBlockSize + +(options) + +ClusterIndexWriter + +index columns sort orders key constraints encoding layout + +velox/dwio/nimble/writer/Writer.cpp · WriterOptions.h · FieldWriter.h · TabletWriter.h + +2d + +Writer encode implementation +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +buildEncodingOptions() + +(new) + +EncodingFactory::encode(policy, data, buffer, options ) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +options + +(new) + +Step 1 + +Estimation — pick encoding + +select(values, stats, options ) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +blockSize + +options + +(new) + +picks best encoding → returns EncodingSelectionResult{BBP, comprFactory} + +Step 2 + +Build selection — bundle result + options + +EncodingSelection < T > +^^^^^^^^^^^^^^^^^^^^^^^ + +options + +(new) + +Step 3 + +Encode — write data + +BBP::encode(selection, values, buffer, options ) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +options + +512 + +Estimation 512 = Encoding 512 ✓ + +Part 3: Read Path (Selective Reader) + +Encoding::Options + +3a + +SelectiveNimbleRowReader option flow +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +EncodingFactory() + +legacy::EncodingFactory() + +shared across all columns — default Options{} + +velox/dwio/nimble/velox/selective/SelectiveNimbleReader.cpp + +3b + +NimbleParams::toFormatData option extraction +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +per-column + +DecodingStats* + +runtimeStatistics().decodingStatsSet + +type- > id() + +type- > type()- > kind() + +NimbleData construction +^^^^^^^^^^^^^^^^^^^^^^^ + +per-column + +encodingFactory\_ + +decodingStats\_ + +all streams for the same column share the same decodingStats\_ + +velox/dwio/nimble/velox/selective/NimbleData.cpp + +3c + +ChunkedDecoder::loadNextChunk() +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +copies factory ’ s base options, overrides per-create: + +// Copy, not reference. + +preserveDictionaryEncoding + +decodingStats + +options + + +4-arg + +create() + +virtual + +zeroCopy + +velox/dwio/nimble/velox/selective/ChunkedDecoder.cpp + +3d + +``encodingFactory_->create(pool, data, sbf, options)`` + +virtual dispatch — depends on zeroCopy flag + +nimble::EncodingFactory + +zeroCopy = true + +options + +legacy::EncodingFactory + +(prod default) + +/*options*/ + +// drops options, delegates to 3-arg + +options + +Encoding tree (zeroCopy=true path only) + +NullableEncoding + +options + +options\_ + +``nimble::`` + +options + +// local, always base class + + +DictionaryEncoding + +options + +``nimble::`` + +options + + +TrivialEncoding + +options + +Compression::uncompress + +options + + +zeroCopy=true + +options + +decodingStats + +zeroCopy + + +All 12 production encodings use this pattern: ALP, Dictionary, Delta, Nullable, MainlyConstant, FreqPartition, ForEncoding, RLE, BBP, SubIntSplit, Trivial (string), DeltaEncoding. + +3e + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - Encoding + - What it decompresses + - Passes decompressCounter? + * - Trivial < T > + - values blob + - ● + * - Trivial < string > + - data blob + bitmap + - ● + * - FixedBitWidth + - packed bit array + - ● + * - BlockBitPacking + - packed blocks + - ● + * - ForEncoding + - packed data + - ● + + +encodings/legacy/ + +Compression::uncompress() + +/*decompressCounter=*/nullptr + +decompressCounter + +velox/dwio/nimble/compression/Compression.h · velox/dwio/common/Statistics.h + +Ke Wang · 2026-06-25 diff --git a/velox/dwio/nimble/docs/develop/nimble_for_interactive.rst b/velox/dwio/nimble/docs/develop/nimble_for_interactive.rst new file mode 100644 index 00000000000..3e93ace2e83 --- /dev/null +++ b/velox/dwio/nimble/docs/develop/nimble_for_interactive.rst @@ -0,0 +1,681 @@ +Nimble for Interactive +====================== + +Background • Architecture • Filter Pushdown + +Row Store vs Columnar Store +=========================== + +users (id INT, name STRING, age INT) + +Row Store + +CSV, JSON, TEXT + +1 + +alice + +25 + +2 + +bob + +30 + +3 + +charlie + +28 + +4 + +dave + +35 + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - id + - name + - age + * - 1 + - "alice" + - 25 + * - 2 + - "bob" + - 30 + * - 3 + - "charlie" + - 28 + * - 4 + - "dave" + - 35 + * - more + - rows + - omitted + +Columnar Store + +DWRF(ORC), Parquet, Nimble + +1 + +2 + +3 + +4 + +alice + +bob + +charlie + +dave + +25 + +30 + +28 + +35 + +Why Columnar Store? +=================== + +SELECT + +FROM + +Row Store + +CSV, JSON, TEXT + +1 + +alice + +25 + +2 + +bob + +30 + +3 + +charlie + +28 + +4 + +dave + +35 + +every row + +Columnar Store + +DWRF(ORC), Parquet, Nimble + +1 + +2 + +3 + +4 + +alice + +bob + +charlie + +dave + +25 + +30 + +28 + +35 + +only the age column + +* Column Pruning — read only needed columns, skip the rest +* Better Compression — same-type values compress more efficiently +* Less I/O — dramatically less data read from disk + +Columnar Store — Scaling with Stripes +===================================== + +users (id INT, name STRING, age INT) + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - id + - name + - age + * - 1 + - "alice" + - 25 + * - more + - rows + - omitted + * - 10000 + - "kate" + - 31 + * - 10001 + - "leo" + - 28 + * - more + - rows + - omitted + * - 20000 + - "zoe" + - 44 + +(rows 1–10,000) + +Worker A + +(rows 10,001–20,000) + +Worker B + +Why not store each column as one giant block? + +* Parallelism — each stripe is a self-contained unit, assigned to a different worker +* Reliability — process one stripe at a time (~256 MB), never load entire file +* I/O efficiency — per-stripe min/max stats enable predicate pushdown, skip irrelevant stripes entirely + +DWRF (ORC) vs Nimble +==================== + +File Layout • Metadata Path • Encoding Path + +Why Nimble? +----------- + +Why invent a new format when we already have DWRF? After a decade of evolution in hardware, query engines, and data scale, DWRF's core assumptions have become limiting. Nimble is a ground-up redesign that advances three areas: + +* Advanced file layout +* Metadata efficiency +* Encoding flexibility + +Area 1: Advanced File Layout +---------------------------- + +Richer metadata — Better data pruning + +* StripeGroup — targeted column reads +* ChunkIndex — sub-stripe predicate pushdown +* ClusterIndex — sort-key pruning + +Extensible structure — Optional section design allows for future flexibility and extensibility to enrich the metadata. + +Area 2: Metadata Path Efficiency +-------------------------------- + +DWRF (Protobuf): O(N) — must deserialize the entire blob to access any field. Nimble (FlatBuffers): O(1) — the buffer is the data structure; zero-copy, selective deserialization and random access. + +Must parse entire Footer + entire Stripe Footer to read 1 column + +O(1) access — only read what you need + +DWRF: + +Nimble: + +Area 3: Data Path Efficiency (Encoding) +--------------------------------------- + +Nimble's encoding system improves on DWRF across three layers — more algorithms, smarter selection, and finer granularity — resulting in smaller files on disk and less CPU spent on encoding/decoding. + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - Layer + - DWRF + - Nimble + * - Encoding Algorithms + - 4 kinds — DIRECT, DICTIONARY, V2 variants Flat structure — sub-encodings hardcoded + - 12 types — Constant, Trivial, FixedBitWidth, MainlyConst, SparseBool, Dictionary, RLE, Varint, Delta, Prefix, Nullable, Sentinel Cascading structure — each node picks its own best sub-encoding → Compounding savings at each tree level + * - Encoding Selection + - Threshold-based — binary cardinality check Only 2 choices: DIRECT or DICTIONARY + - Cost-based — estimates size for all 12 candidates, picks smallest → Optimal encoding per data shape + * - Encoding Granularity + - Per-file — locked after first stripe All subsequent stripes use the same encoding + - Per-stripe-group — re-evaluated as data changes → Adapts as data patterns shift + +Look Ahead — Nimble for Interactive + +Where we started + +ML workload rollout + +Next focus + +interactive workloads + +Step 1: Stats and Index + +Interactive queries are highly selective — they touch a small fraction of data. + +Encoded column chunks + + +Stripe N — Data Streams + +Stream offsets/sizes + +NEW Per-stride min/max/nullCount + +NEW Per-chunk stream positions + +NEW Sorted-column key lookup + +NEW File-level min/max/size per column + +Schema, stripe locations + +Footer size, checksum + +Stride Stats + +UNDER REVIEW + +Chunk Index + +MERGED + +Cluster Index + +MERGED + +VectorizedFileStats + +ROLLED OUT + +Levels of Pruning + +user_id BETWEEN 1000 AND 2000 + +NEW + +No overlap → ✗ skip entire file + +Overlaps → ✓ open file + +NEW + +File B + +Stripe 0 + +0–499 + +skip + +500–999 + +skip + +1000–1499 + +read + +2 of 3 skipped + +Stripe 1 + +2500–2999 + +skip + +3000–3499 + +skip + +3500–3999 + +skip + +All skipped — never decompressed + +Level 3: Residual Predicate (Row-Level) — Zooming into Stride 2 + +BETWEEN 1000 AND 2000 + +.. list-table:: + :widths: auto + + * - user_id + - 1000 ✓ + - 1001 ✓ + - … + - 1200 ✓ + - 1201 ✗ + - 1202 ✗ + - … + - 1499 ✗ + +201 of 500 rows pass → 60% filtered out at row level + +Result: + +Cross-Column: Skip decisions propagate to all projected columns + +.. list-table:: + :widths: auto + + * - + - S0 + - S1 + - S2 + * - user_id (filter) + - 0–499 + - 500–999 + - 1000–1499 + * - name (proj) + - skip + - inherit + - inherit + * - amount (proj) + - skip + - inherit + - inherit + +Gold + +Grey + +Dashed + +Multiplier effect: + +NEW + +When columns skip strides, they need to jump past skipped data. Without ChunkIndex, they must sequentially decompress every chunk in between. With ChunkIndex: O(1) seek. + +Without ChunkIndex + +Chunk 0 decompress & discard + +Chunk 1 decompress & discard + +Chunk 2 decompress & discard + +Chunk 3 target + +Must decompress 3 chunks to reach target + +With ChunkIndex + +Chunk 0 skipped + +Chunk 1 skipped + +Chunk 2 skipped + +Chunk 3 O(1) seek + +lookupChunk(row) → direct file seek, zero decompression + +PROTOTYPE + +Learning from Impulse + +FixedBitWidthEncoding + +Wins: + +Before: FixedBitWidth (global) + +One global min/max → same bit-width for all: + +Chunk 0 + +100 – 105 + +17b + +17 bits + +Chunk 1 + +5K – 5.2K + +17b + +17 bits + +Chunk 2 + +42 – 42 + +17b + +17 bits + +Chunk 3 + +0 – 100K + +17b + +17 bits + +Chunk 4 + +0 – 10 + +17b + +17 bits + +global range 0–100K forces 17b everywhere + +After: ChunkedBitPacking (per-chunk) + +Each 1024-row chunk uses its own local range: + +Chunk 0 + +100 – 105 + +3b + +3 bits + +Chunk 1 + +5K – 5.2K + +8b + +8 bits + +Chunk 2 + +42 – 42 + +0 bits ! + +Chunk 3 + +0 – 100K + +17b + +17 bits + +Chunk 4 + +0 – 10 + +4b + +4 bits + +only chunk 3 pays 17b — rest use 0–8b + +PROTOTYPE + +Float / double columns. + +Wins: + +Before: Trivial (raw float) + +Chunk 0 + +prices + +64b + +64 bits + +Chunk 1 + +whole #s + +64b + +64 bits + +Chunk 2 + +all 0.0 + +64b + +64 bits + +every float = 64 bits, no compression + +After: ChunkedALP (per-chunk) + +Chunk 0 + +e=2 (×100) + +8b + +8 bits + +Chunk 1 + +e=0 (×1) + +13b + +13 bits + +Chunk 2 + +constant + +0 bits ! + +64-bit floats → 0–13 bit integers, lossless + +Step 2: Better Encoding — TPCH Benchmark + +CPU time before and after ALP+CBP encodings + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - Query + - Before + - After (ALP+CBP) + - Speedup + * - q1 + - 15.58m + - 8.13m + - 1.9x + * - q5 + - 11.02m + - 4.09m + - 2.7x + * - q6 + - 3.62m + - 2.24m + - 1.6x + * - q7 + - 9.42m + - 6.49m + - 1.5x + * - q8 + - 2.95h + - 2.75h + - 1.1x + * - q9 + - 36.12m + - 35.20m + - 1.0x + * - q10 + - 14.95m + - 8.78m + - 1.7x + * - q11 + - 1.40m + - 31.45s + - 2.7x + * - q12 + - 14.47m + - 10.54m + - 1.4x + * - q13 + - 17.79m + - 12.32m + - 1.4x + * - q14 + - 5.84m + - 3.12m + - 1.9x + * - q15 + - 13.41m + - 6.42m + - 2.1x + * - q16 + - 6.00m + - 5.69m + - 1.1x + * - q18 + - 32.45m + - 31.28m + - 1.0x + * - q19 + - 7.89m + - 3.66m + - 2.2x + * - Median + - + - + - 1.6x + +≥1.4x significant speedup + +~1.0x minimal change + +11 of 15 queries improved by ≥1.4x + +Stride stats write: D103095055 • Stride stats read: D103095054 StrideIndex FBS: velox/dwio/nimble/tablet/StrideIndex.fbs • Writer: velox/dwio/nimble/tablet/StrideIndexWriter.cpp Reader: velox/dwio/nimble/velox/selective/NimbleData.cpp • Skip logic: SelectiveNimbleReader.cpp VectorizedFileStats: velox/dwio/nimble/velox/stats/VectorizedStatistics.h diff --git a/velox/dwio/nimble/docs/develop/nimble_java_reader_blockBitPacking.rst b/velox/dwio/nimble/docs/develop/nimble_java_reader_blockBitPacking.rst new file mode 100644 index 00000000000..d18e29d975f --- /dev/null +++ b/velox/dwio/nimble/docs/develop/nimble_java_reader_blockBitPacking.rst @@ -0,0 +1,249 @@ +Nimble Java Reader & Writer +=========================== + +Using BlockBitPacking (EncodingType=15) as a concrete example + +1 Write Path +------------ + +There is no pure Java Nimble writer . Spark writes Nimble files via JNI, delegating to the C++ Writer for encoding and file layout. + +Java AlphaPageWriter.writePage(Page) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Serializes Presto Page → ByteBuffer + + +JNI JniAlphaWriter.writePageJni(writerId, buffer) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +xldb/alpha_jni/cpp/AlphaWriterJNI.cpp + + +C++ Writer::write(RowVector) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Encoding selection → BlockBitPacking / Dictionary / RLE / ... → flush to disk + +Implication: + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - Layer + - Key Class + - File + * - Java + - AlphaPageWriter + - xldb/alpha_jni/alpha-common/.../writer/AlphaPageWriter.java + * - Java + - JniAlphaWriter + - xldb/alpha_jni/alpha-common/.../writer/JniAlphaWriter.java + * - JNI + - AlphaWriterJNI.cpp + - xldb/alpha_jni/cpp/AlphaWriterJNI.cpp + * - C++ + - Writer + - velox/dwio/nimble/writer/Writer.h + +2 Read Path +----------- + +Unlike the writer, the Java Nimble reader is a pure Java implementation in alpha.encodings . It is the production decoder used by Spark via xldb-orc → shaded AlphaRecordReader . + +Three Java Nimble Packages +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +There are three Java packages related to Nimble encoding. Only one is production: + +Production + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - Package + - Location + - Consumers + * - com.facebook.presto.alpha.encodings + - presto-facebook-alpha/.../alpha/encodings/ + - Spark reads via AlphaRecordReader + +Dead — prototype + its test JNI bridge + +.. list-table:: + :widths: auto + :header-rows: 1 + + * - Package + - Location + - Role + * - com.facebook.presto.nimble + - presto-facebook-alpha/.../nimble/ + - Selective-read prototype decoders — never shipped, zero consumers + * - com.facebook.xldb.nimble.jni + - xldb/alpha_jni/nimble-encodings/ + - JNI bridge that encodes test data via C++ for the prototype ’ s tests ( < scope > test < /scope > ) + +The prototype ’ s tests call NimbleEncodings.encode() via JNI to generate C++-encoded data, then verify the prototype ’ s Java decoders can read it. No alpha.* code imports xldb.nimble.jni . + +Important: + +alpha.encodings + +There are also two separate EncodingType enums: + +* alpha.common.EncodingType — production used by the read path; EncodingFactory calls fromValue(15) to dispatch decoding. Must be updated for new encodings. +* xldb.nimble.jni.EncodingType — test-only only imported by com.facebook.presto.nimble test files. Updated defensively to keep integer values in sync, but not required for production. + +Java AlphaRecordReader +^^^^^^^^^^^^^^^^^^^^^^ + +Spark → xldb-orc → AlphaRecordReader + + +EncodingFactory.deserializeEncoding() +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Reads encoding prefix → dispatches on EncodingType + + +BlockBitPackingEncoding.createEncoding(dataType, input, memCtx) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Dispatches to Byte / Short / Int / Long variant based on DataType + + +Constructor: parse nested metadata + decompress packed data +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Reads 3 sub-streams via recursive EncodingFactory calls + + +materialize(rowCount) → Block +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Per-block bit-unpack + baseline addition → ByteArrayBlock / IntArrayBlock / ... + +BlockBitPacking Binary Format +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +What the Java decoder reads + +.. code-block:: text + + [Encoding Prefix] encodingType=15, dataType, rowCount + + [compressionType] uint8 — Uncompressed / Zstd / Zstrong + [blockSize] uint16 — default 1024 + [numBlocks] uint16 — ceil(rowCount / blockSize) + + [size₁] [baselines encoding] physicalType[numBlocks] + [size₂] [bitWidths encoding] uint8[numBlocks], 255 = skip-encoding + [size₃] [offsets encoding] uint32[numBlocks], byte offset per block + + [packed data] all blocks contiguously, compressed with compressionType + +Type Specialization +^^^^^^^^^^^^^^^^^^^ + +ByteBlockBitPackingEncoding + +Int8 / Uint8 — baseline: byte , reads short for bit extraction + +ShortBlockBitPackingEncoding + +Int16 / Uint16 — baseline: short , reads int for bit extraction + +IntBlockBitPackingEncoding + +Int32 / Uint32 / Float — baseline: int , reads long for extraction + +LongBlockBitPackingEncoding + +Int64 / Uint64 / Double — baseline: long , two-word read for bw > 58 + +3 E2E Verification +------------------ + +The E2E test verifies that the pure Java decoder correctly reads what the C++ encoder writes. Artifacts are generated once by the C++ test generator and committed to the repo. + +Step 1 — Generate Artifacts C++ +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +test_generator ( //velox/dwio/nimble/encodings/tests:test_generator ) produces artifact files for each (encoding, dataType, rowCount) combination: + +For each encoding × type × rowCount: + +1. Generate random data — seeded RNG creates rowCount random values of the given type +2. Write .data file — raw bytes of each value in native byte order (no header, no encoding — just rowCount × sizeof(T) bytes) +3. Encode 3 × — calls the real C++ encoder ( E::encode() ) via a test wrapper ( nimble::test::Encoder < E > ::encode() ) with each compression type +4. Write .encoding files — the full encoded blob (prefix + nested metadata + compressed packed data) for Uncompressed, Zstd, and Zstrong + +Real encoder, controlled selection: + +BlockBitPackingEncoding::encode() + +TestTrivialEncodingSelectionPolicy + +.. code-block:: text + + # Generate only BlockBitPacking artifacts (use --encoding_filter) + buck2 run //velox/dwio/nimble/encodings/tests:test_generator \ + -- --output_dir=presto-facebook-alpha/src/test/resources/encodings \ + --encoding_filter=BlockBitPacking + + # Output per (dataType, rowCount): + # BlockBitPacking_Int32_256.data ← raw values (ground truth) + # BlockBitPacking_Int32_256_Uncompressed.encoding ← encoded blob, no compression + # BlockBitPacking_Int32_256_Zstd.encoding ← encoded blob, Zstd compressed + # BlockBitPacking_Int32_256_Zstrong.encoding ← encoded blob, Zstrong compressed + + # Omit --encoding_filter to regenerate ALL encoding artifacts + +Step 2 — Verify Java +^^^^^^^^^^^^^^^^^^^^ + +List data files in resources/encodings/ +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Parse filename → encoding type, data type, row count + + +Read .data file → expected values[] +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Raw source values written by C++ (ground truth) + + +Read .encoding file → EncodingFactory.deserialize() +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Recursive Nimble decoding → encoding.materialize(rowCount) → decoded values[] + + +assertEquals(expected[i], decoded[i]) ∀ i +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +3 compressions × 10 data types × 2 sizes = 60 test cases per encoding + +Running the Test +^^^^^^^^^^^^^^^^ + +.. code-block:: text + + # Generate artifacts (only needed when C++ format changes) + buck2 run //velox/dwio/nimble/encodings/tests:test_generator \ + -- --output_dir=presto-facebook-alpha/src/test/resources/encodings + + # Run the Java E2E test + cd github/presto-facebook-trunk + mvn test -pl presto-facebook-alpha \ + -Dtest=TestEncodingsE2E \ + -Dmaven.gitcommitid.skip=true \ + -Dcheckstyle.skip=true + +Coverage: + +60 test cases diff --git a/velox/dwio/nimble/docs/develop/nimble_metadata_cache.rst b/velox/dwio/nimble/docs/develop/nimble_metadata_cache.rst new file mode 100644 index 00000000000..e461b95cca0 --- /dev/null +++ b/velox/dwio/nimble/docs/develop/nimble_metadata_cache.rst @@ -0,0 +1,802 @@ + + + +Nimble Metadata Cache +===================== + +File layout, MetadataSection, CachedMetadataInput, and TabletReader init flow + +5 Why TabletReader Doesn't Use CachedBufferedInput +-------------------------------------------------- + +Contiguous vs Non-Contiguous +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +FlatBuffers (metadata, indexes) + +.. code-block:: text + + flatbuffers::GetRoot