add Morton Coding (ZCurve)

.gitmodules

@@ -287,3 +287,6 @@
 [submodule "contrib/corrosion"]
 	path = contrib/corrosion
 	url = https://github.com/corrosion-rs/corrosion.git
+[submodule "contrib/morton-nd"]
+	path = contrib/morton-nd
+	url = https://github.com/morton-nd/morton-nd

cmake/cpu_features.cmake

@@ -75,6 +75,7 @@ elseif (ARCH_AMD64)
     option (ENABLE_AVX512 "Use AVX512 instructions on x86_64" 0)
     option (ENABLE_AVX512_VBMI "Use AVX512_VBMI instruction on x86_64 (depends on ENABLE_AVX512)" 0)
     option (ENABLE_BMI "Use BMI instructions on x86_64" 0)
+    option (ENABLE_BMI2 "Use BMI2 instructions on x86_64 (depends on ENABLE_AVX2)" 0)
     option (ENABLE_AVX2_FOR_SPEC_OP "Use avx2 instructions for specific operations on x86_64" 0)
     option (ENABLE_AVX512_FOR_SPEC_OP "Use avx512 instructions for specific operations on x86_64" 0)
 
@@ -90,6 +91,7 @@ elseif (ARCH_AMD64)
         SET(ENABLE_AVX512 0)
         SET(ENABLE_AVX512_VBMI 0)
         SET(ENABLE_BMI 0)
+        SET(ENABLE_BMI2 0)
         SET(ENABLE_AVX2_FOR_SPEC_OP 0)
         SET(ENABLE_AVX512_FOR_SPEC_OP 0)
     endif()
@@ -237,6 +239,20 @@ elseif (ARCH_AMD64)
         set (COMPILER_FLAGS "${COMPILER_FLAGS} ${TEST_FLAG}")
     endif ()
 
+    set (TEST_FLAG "-mbmi2")
+    set (CMAKE_REQUIRED_FLAGS "${TEST_FLAG} -O0")
+    check_cxx_source_compiles("
+        #include <immintrin.h>
+        int main() {
+            auto a = _pdep_u64(0, 0);
+            (void)a;
+            return 0;
+        }
+    " HAVE_BMI2)
+    if (HAVE_BMI2 AND HAVE_AVX2 AND ENABLE_AVX2 AND ENABLE_BMI2)
+        set (COMPILER_FLAGS "${COMPILER_FLAGS} ${TEST_FLAG}")
+    endif ()
+
     # Limit avx2/avx512 flag for specific source build
     set (X86_INTRINSICS_FLAGS "")
     if (ENABLE_AVX2_FOR_SPEC_OP)

contrib/CMakeLists.txt

@@ -164,6 +164,7 @@ add_contrib (sqlite-cmake sqlite-amalgamation)
 add_contrib (s2geometry-cmake s2geometry)
 add_contrib (c-ares-cmake c-ares)
 add_contrib (qpl-cmake qpl)
+add_contrib (morton-nd-cmake morton-nd)
 
 add_contrib(annoy-cmake annoy)
 

contrib/morton-nd

@@ -0,0 +1 @@
+Subproject commit 3795491a4aa3cdc916c8583094683f0d68df5bc0

contrib/morton-nd-cmake/CMakeLists.txt

@@ -0,0 +1,3 @@
+add_library(_morton_nd INTERFACE)
+target_include_directories(_morton_nd SYSTEM BEFORE INTERFACE "${ClickHouse_SOURCE_DIR}/contrib/morton-nd/include/")
+add_library(ch_contrib::morton_nd ALIAS _morton_nd)

docker/test/fasttest/run.sh

@@ -136,6 +136,7 @@ function clone_submodules
             contrib/wyhash
             contrib/hashidsxx
             contrib/c-ares
+            contrib/morton-nd
         )
 
         git submodule sync

src/Functions/CMakeLists.txt

@@ -22,6 +22,7 @@ list (APPEND PUBLIC_LIBS
         ch_contrib::metrohash
         ch_contrib::murmurhash
         ch_contrib::hashidsxx
+        ch_contrib::morton_nd
 )
 
 list (APPEND PRIVATE_LIBS

src/Functions/mortonDecode.cpp

@@ -0,0 +1,433 @@
+#include <Functions/IFunction.h>
+#include <Functions/FunctionFactory.h>
+#include <DataTypes/DataTypesNumber.h>
+#include <DataTypes/DataTypeTuple.h>
+#include <Columns/ColumnsNumber.h>
+#include <Functions/FunctionHelpers.h>
+#include <Columns/ColumnTuple.h>
+#include <Functions/PerformanceAdaptors.h>
+
+#include <morton-nd/mortonND_LUT.h>
+#if USE_MULTITARGET_CODE && defined(__BMI2__)
+#include <morton-nd/mortonND_BMI2.h>
+#endif
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
+    extern const int ILLEGAL_COLUMN;
+    extern const int ARGUMENT_OUT_OF_BOUND;
+}
+
+#define EXTRACT_VECTOR(INDEX) \
+        auto col##INDEX = ColumnUInt64::create(); \
+        auto & vec##INDEX = col##INDEX->getData(); \
+        vec##INDEX.resize(input_rows_count);
+
+#define DECODE(ND, ...) \
+        if (nd == (ND)) \
+        { \
+            for (size_t i = 0; i < input_rows_count; i++) \
+            { \
+                auto res = MortonND_##ND##D_Dec.Decode(col_code->getUInt(i)); \
+                __VA_ARGS__ \
+            } \
+        }
+
+#define MASK(IDX, ...) \
+        ((mask) ? shrink(mask->getColumn((IDX)).getUInt(0), std::get<IDX>(__VA_ARGS__)) : std::get<IDX>(__VA_ARGS__))
+
+#define EXECUTE() \
+    size_t nd; \
+    const auto * col_const = typeid_cast<const ColumnConst *>(arguments[0].column.get()); \
+    const auto * mask = typeid_cast<const ColumnTuple *>(col_const->getDataColumnPtr().get()); \
+    if (mask) \
+        nd = mask->tupleSize(); \
+    else \
+        nd = col_const->getUInt(0); \
+    auto non_const_arguments = arguments; \
+    non_const_arguments[1].column = non_const_arguments[1].column->convertToFullColumnIfConst(); \
+    const ColumnPtr & col_code = non_const_arguments[1].column; \
+    Columns tuple_columns(nd); \
+    EXTRACT_VECTOR(0) \
+    if (nd == 1) \
+    { \
+        if (mask) \
+        { \
+            for (size_t i = 0; i < input_rows_count; i++) \
+            { \
+                vec0[i] = shrink(mask->getColumn(0).getUInt(0), col_code->getUInt(i)); \
+            } \
+            tuple_columns[0] = std::move(col0); \
+        } \
+        else \
+        { \
+            for (size_t i = 0; i < input_rows_count; i++) \
+            { \
+                vec0[i] = col_code->getUInt(i); \
+            } \
+            tuple_columns[0] = std::move(col0); \
+        } \
+        return ColumnTuple::create(tuple_columns); \
+    } \
+    EXTRACT_VECTOR(1) \
+    DECODE(2, \
+       vec0[i] = MASK(0, res); \
+       vec1[i] = MASK(1, res);) \
+    EXTRACT_VECTOR(2) \
+    DECODE(3, \
+       vec0[i] = MASK(0, res); \
+       vec1[i] = MASK(1, res); \
+       vec2[i] = MASK(2, res);) \
+    EXTRACT_VECTOR(3) \
+    DECODE(4, \
+       vec0[i] = MASK(0, res); \
+       vec1[i] = MASK(1, res); \
+       vec2[i] = MASK(2, res); \
+       vec3[i] = MASK(3, res);) \
+    EXTRACT_VECTOR(4) \
+    DECODE(5, \
+       vec0[i] = MASK(0, res); \
+       vec1[i] = MASK(1, res); \
+       vec2[i] = MASK(2, res); \
+       vec3[i] = MASK(3, res); \
+       vec4[i] = MASK(4, res);) \
+    EXTRACT_VECTOR(5) \
+    DECODE(6, \
+       vec0[i] = MASK(0, res); \
+       vec1[i] = MASK(1, res); \
+       vec2[i] = MASK(2, res); \
+       vec3[i] = MASK(3, res); \
+       vec4[i] = MASK(4, res); \
+       vec5[i] = MASK(5, res);) \
+    EXTRACT_VECTOR(6) \
+    DECODE(7, \
+       vec0[i] = MASK(0, res); \
+       vec1[i] = MASK(1, res); \
+       vec2[i] = MASK(2, res); \
+       vec3[i] = MASK(3, res); \
+       vec4[i] = MASK(4, res); \
+       vec5[i] = MASK(5, res); \
+       vec6[i] = MASK(6, res);) \
+    EXTRACT_VECTOR(7) \
+    DECODE(8, \
+       vec0[i] = MASK(0, res); \
+       vec1[i] = MASK(1, res); \
+       vec2[i] = MASK(2, res); \
+       vec3[i] = MASK(3, res); \
+       vec4[i] = MASK(4, res); \
+       vec5[i] = MASK(5, res); \
+       vec6[i] = MASK(6, res); \
+       vec7[i] = MASK(7, res);) \
+    switch (nd) \
+    { \
+        case 2: \
+            tuple_columns[0] = std::move(col0); \
+            tuple_columns[1] = std::move(col1); \
+            break; \
+        case 3: \
+            tuple_columns[0] = std::move(col0); \
+            tuple_columns[1] = std::move(col1); \
+            tuple_columns[2] = std::move(col2); \
+            return ColumnTuple::create(tuple_columns); \
+        case 4: \
+            tuple_columns[0] = std::move(col0); \
+            tuple_columns[1] = std::move(col1); \
+            tuple_columns[2] = std::move(col2); \
+            tuple_columns[3] = std::move(col3); \
+            return ColumnTuple::create(tuple_columns); \
+        case 5: \
+            tuple_columns[0] = std::move(col0); \
+            tuple_columns[1] = std::move(col1); \
+            tuple_columns[2] = std::move(col2); \
+            tuple_columns[3] = std::move(col3); \
+            tuple_columns[4] = std::move(col4); \
+            return ColumnTuple::create(tuple_columns); \
+        case 6: \
+            tuple_columns[0] = std::move(col0); \
+            tuple_columns[1] = std::move(col1); \
+            tuple_columns[2] = std::move(col2); \
+            tuple_columns[3] = std::move(col3); \
+            tuple_columns[4] = std::move(col4); \
+            tuple_columns[5] = std::move(col5); \
+            return ColumnTuple::create(tuple_columns); \
+        case 7: \
+            tuple_columns[0] = std::move(col0); \
+            tuple_columns[1] = std::move(col1); \
+            tuple_columns[2] = std::move(col2); \
+            tuple_columns[3] = std::move(col3); \
+            tuple_columns[4] = std::move(col4); \
+            tuple_columns[5] = std::move(col5); \
+            tuple_columns[6] = std::move(col6); \
+            return ColumnTuple::create(tuple_columns); \
+        case 8: \
+            tuple_columns[0] = std::move(col0); \
+            tuple_columns[1] = std::move(col1); \
+            tuple_columns[2] = std::move(col2); \
+            tuple_columns[3] = std::move(col3); \
+            tuple_columns[4] = std::move(col4); \
+            tuple_columns[5] = std::move(col5); \
+            tuple_columns[6] = std::move(col6); \
+            tuple_columns[7] = std::move(col7); \
+            return ColumnTuple::create(tuple_columns); \
+    } \
+    return ColumnTuple::create(tuple_columns);
+
+DECLARE_DEFAULT_CODE(
+constexpr auto MortonND_2D_Dec = mortonnd::MortonNDLutDecoder<2, 32, 8>();
+constexpr auto MortonND_3D_Dec = mortonnd::MortonNDLutDecoder<3, 21, 8>();
+constexpr auto MortonND_4D_Dec = mortonnd::MortonNDLutDecoder<4, 16, 8>();
+constexpr auto MortonND_5D_Dec = mortonnd::MortonNDLutDecoder<5, 12, 8>();
+constexpr auto MortonND_6D_Dec = mortonnd::MortonNDLutDecoder<6, 10, 8>();
+constexpr auto MortonND_7D_Dec = mortonnd::MortonNDLutDecoder<7, 9, 8>();
+constexpr auto MortonND_8D_Dec = mortonnd::MortonNDLutDecoder<8, 8, 8>();
+class FunctionMortonDecode : public IFunction
+{
+public:
+    static constexpr auto name = "mortonDecode";
+    static FunctionPtr create(ContextPtr)
+    {
+        return std::make_shared<FunctionMortonDecode>();
+    }
+
+    String getName() const override
+    {
+        return name;
+    }
+
+    size_t getNumberOfArguments() const override
+    {
+        return 2;
+    }
+
+    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
+
+    ColumnNumbers getArgumentsThatAreAlwaysConstant() const override { return {0}; }
+
+    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
+    {
+        UInt64 tuple_size = 0;
+        const auto * col_const = typeid_cast<const ColumnConst *>(arguments[0].column.get());
+        if (!col_const)
+            throw Exception(ErrorCodes::ILLEGAL_COLUMN,
+                            "Illegal column type {} of function {}, should be a constant (UInt or Tuple)",
+                            arguments[0].type->getName(), getName());
+        if (!WhichDataType(arguments[1].type).isNativeUInt())
+            throw Exception(ErrorCodes::ILLEGAL_COLUMN,
+                            "Illegal column type {} of function {}, should be a native UInt",
+                            arguments[1].type->getName(), getName());
+        const auto * mask = typeid_cast<const ColumnTuple *>(col_const->getDataColumnPtr().get());
+        if (mask)
+        {
+            tuple_size = mask->tupleSize();
+        }
+        else if (WhichDataType(arguments[0].type).isNativeUInt())
+        {
+            tuple_size = col_const->getUInt(0);
+        }
+        else
+            throw Exception(ErrorCodes::ILLEGAL_COLUMN,
+                            "Illegal column type {} of function {}, should be UInt or Tuple",
+                            arguments[0].type->getName(), getName());
+        if (tuple_size > 8 || tuple_size < 1)
+            throw Exception(ErrorCodes::ARGUMENT_OUT_OF_BOUND,
+                            "Illegal first argument for function {}, should be a number in range 1-8 or a Tuple of such size",
+                            getName());
+        if (mask)
+        {
+            const auto * type_tuple = typeid_cast<const DataTypeTuple *>(arguments[0].type.get());
+            for (size_t i = 0; i < tuple_size; i++)
+            {
+                if (!WhichDataType(type_tuple->getElement(i)).isNativeUInt())
+                    throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
+                                    "Illegal type {} of argument in tuple for function {}, should be a native UInt",
+                                    type_tuple->getElement(i)->getName(), getName());
+                auto ratio = mask->getColumn(i).getUInt(0);
+                if (ratio > 8 || ratio < 1)
+                    throw Exception(ErrorCodes::ARGUMENT_OUT_OF_BOUND,
+                                    "Illegal argument {} in tuple for function {}, should be a number in range 1-8",
+                                    ratio, getName());
+            }
+        }
+        DataTypes types(tuple_size);
+        for (size_t i = 0; i < tuple_size; i++)
+        {
+            types[i] = std::make_shared<DataTypeUInt64>();
+        }
+        return std::make_shared<DataTypeTuple>(types);
+    }
+
+    static UInt64 shrink(UInt64 ratio, UInt64 value)
+    {
+        switch (ratio)
+        {
+            case 1:
+                return value;
+            case 2:
+                return std::get<1>(MortonND_2D_Dec.Decode(value));
+            case 3:
+                return std::get<2>(MortonND_3D_Dec.Decode(value));
+            case 4:
+                return std::get<3>(MortonND_4D_Dec.Decode(value));
+            case 5:
+                return std::get<4>(MortonND_5D_Dec.Decode(value));
+            case 6:
+                return std::get<5>(MortonND_6D_Dec.Decode(value));
+            case 7:
+                return std::get<6>(MortonND_7D_Dec.Decode(value));
+            case 8:
+                return std::get<7>(MortonND_8D_Dec.Decode(value));
+        }
+        return value;
+    }
+
+    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
+    {
+        EXECUTE()
+    }
+};
+) // DECLARE_DEFAULT_CODE
+
+#if defined(MORTON_ND_BMI2_ENABLED)
+#undef DECODE
+#define DECODE(ND, ...) \
+        if (nd == (ND)) \
+        { \
+            for (size_t i = 0; i < input_rows_count; i++) \
+            { \
+                auto res = MortonND_##ND##D::Decode(col_code->getUInt(i)); \
+                __VA_ARGS__ \
+            } \
+        }
+
+DECLARE_AVX2_SPECIFIC_CODE(
+using MortonND_2D = mortonnd::MortonNDBmi<2, uint64_t>;
+using MortonND_3D = mortonnd::MortonNDBmi<3, uint64_t>;
+using MortonND_4D = mortonnd::MortonNDBmi<4, uint64_t>;
+using MortonND_5D = mortonnd::MortonNDBmi<5, uint64_t>;
+using MortonND_6D = mortonnd::MortonNDBmi<6, uint64_t>;
+using MortonND_7D = mortonnd::MortonNDBmi<7, uint64_t>;
+using MortonND_8D = mortonnd::MortonNDBmi<8, uint64_t>;
+class FunctionMortonDecode: public TargetSpecific::Default::FunctionMortonDecode
+{
+    static UInt64 shrink(UInt64 ratio, UInt64 value)
+    {
+        switch (ratio)
+        {
+            case 1:
+                return value;
+            case 2:
+                return std::get<1>(MortonND_2D::Decode(value));
+            case 3:
+                return std::get<2>(MortonND_3D::Decode(value));
+            case 4:
+                return std::get<3>(MortonND_4D::Decode(value));
+            case 5:
+                return std::get<4>(MortonND_5D::Decode(value));
+            case 6:
+                return std::get<5>(MortonND_6D::Decode(value));
+            case 7:
+                return std::get<6>(MortonND_7D::Decode(value));
+            case 8:
+                return std::get<7>(MortonND_8D::Decode(value));
+        }
+        return value;
+    }
+
+    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
+    {
+        EXECUTE()
+    }
+};
+)
+#endif // MORTON_ND_BMI2_ENABLED
+
+#undef DECODE
+#undef MASK
+#undef EXTRACT_VECTOR
+#undef EXECUTE
+
+class FunctionMortonDecode: public TargetSpecific::Default::FunctionMortonDecode
+{
+public:
+    explicit FunctionMortonDecode(ContextPtr context) : selector(context)
+    {
+        selector.registerImplementation<TargetArch::Default,
+                                        TargetSpecific::Default::FunctionMortonDecode>();
+
+#if USE_MULTITARGET_CODE && defined(MORTON_ND_BMI2_ENABLED)
+        selector.registerImplementation<TargetArch::AVX2,
+                                        TargetSpecific::AVX2::FunctionMortonDecode>();
+#endif
+    }
+
+    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count) const override
+    {
+        return selector.selectAndExecute(arguments, result_type, input_rows_count);
+    }
+
+    static FunctionPtr create(ContextPtr context)
+    {
+        return std::make_shared<FunctionMortonDecode>(context);
+    }
+
+private:
+    ImplementationSelector<IFunction> selector;
+};
+
+REGISTER_FUNCTION(MortonDecode)
+{
+    factory.registerFunction<FunctionMortonDecode>({
+        R"(
+Decodes a Morton encoding (ZCurve) into the corresponding unsigned integer tuple
+
+The function has two modes of operation:
+- Simple
+- Expanded
+
+Simple: accepts a resulting tuple size as a first argument and the code as a second argument.
+[example:simple]
+Will decode into: `(1,2,3,4)`
+The resulting tuple size cannot be more than 8
+
+Expanded: accepts a range mask (tuple) as a first argument and the code as a second argument.
+Each number in mask configures the amount of range shrink
+1 - no shrink
+2 - 2x shrink
+3 - 3x shrink
+....
+Up to 8x shrink.
+[example:range_shrank]
+Note: see mortonEncode() docs on why range change might be beneficial.
+Still limited to 8 numbers at most.
+
+Morton code for one argument is always the argument itself (as a tuple).
+[example:identity]
+Produces: `(1)`
+
+You can shrink one argument too:
+[example:identity_shrank]
+Produces: `(128)`
+
+The function accepts a column of codes as a second argument:
+[example:from_table]
+
+The range tuple must be a constant:
+[example:from_table_range]
+)",
+        Documentation::Examples{
+            {"simple", "SELECT mortonDecode(4, 2149)"},
+            {"range_shrank", "SELECT mortonDecode((1,2), 1572864)"},
+            {"identity", "SELECT mortonDecode(1, 1)"},
+            {"identity_shrank", "SELECT mortonDecode(tuple(2), 32768)"},
+            {"from_table", "SELECT mortonDecode(2, code) FROM table"},
+            {"from_table_range", "SELECT mortonDecode((1,2), code) FROM table"},
+            },
+        Documentation::Categories {"ZCurve", "Morton coding"}
+    });
+}
+
+}

src/Functions/mortonEncode.cpp

@@ -0,0 +1,393 @@
+#include <Functions/IFunction.h>
+#include <Functions/FunctionFactory.h>
+#include <DataTypes/DataTypesNumber.h>
+#include <DataTypes/DataTypeTuple.h>
+#include <Columns/ColumnsNumber.h>
+#include <Columns/ColumnConst.h>
+#include <Columns/ColumnTuple.h>
+#include <Functions/PerformanceAdaptors.h>
+
+#include <morton-nd/mortonND_LUT.h>
+#if USE_MULTITARGET_CODE && defined(__BMI2__)
+#include <morton-nd/mortonND_BMI2.h>
+#endif
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
+    extern const int ARGUMENT_OUT_OF_BOUND;
+    extern const int TOO_FEW_ARGUMENTS_FOR_FUNCTION;
+}
+
+#define EXTRACT_VECTOR(INDEX) \
+    const ColumnPtr & col##INDEX = non_const_arguments[(INDEX) + vectorStartIndex].column;
+
+#define ENCODE(ND, ...) \
+    if (nd == (ND)) \
+    { \
+        for (size_t i = 0; i < input_rows_count; i++) \
+        {               \
+            vec_res[i] = MortonND_##ND##D_Enc.Encode(__VA_ARGS__); \
+        } \
+        return col_res; \
+    }
+
+#define EXPAND(IDX, ...) \
+    (mask) ? expand(mask->getColumn(IDX).getUInt(0), __VA_ARGS__) : __VA_ARGS__
+
+#define MASK(ND, IDX, ...) \
+    (EXPAND(IDX, __VA_ARGS__) & MortonND_##ND##D_Enc.InputMask())
+
+#define EXECUTE() \
+    size_t nd = arguments.size(); \
+    size_t vectorStartIndex = 0; \
+    const auto * const_col = typeid_cast<const ColumnConst *>(arguments[0].column.get()); \
+    const ColumnTuple * mask; \
+    if (const_col) \
+        mask = typeid_cast<const ColumnTuple *>(const_col->getDataColumnPtr().get()); \
+    else \
+        mask = typeid_cast<const ColumnTuple *>(arguments[0].column.get()); \
+    if (mask) \
+    { \
+        nd = mask->tupleSize(); \
+        vectorStartIndex = 1; \
+        for (size_t i = 0; i < nd; i++) \
+        { \
+            auto ratio = mask->getColumn(i).getUInt(0); \
+            if (ratio > 8 || ratio < 1) \
+                throw Exception(ErrorCodes::ARGUMENT_OUT_OF_BOUND, \
+                                "Illegal argument {} of function {}, should be a number in range 1-8", \
+                                arguments[0].column->getName(), getName()); \
+        } \
+    } \
+     \
+    auto non_const_arguments = arguments; \
+    for (auto & argument : non_const_arguments) \
+        argument.column = argument.column->convertToFullColumnIfConst(); \
+     \
+    auto col_res = ColumnUInt64::create(); \
+    ColumnUInt64::Container & vec_res = col_res->getData(); \
+    vec_res.resize(input_rows_count); \
+     \
+    EXTRACT_VECTOR(0) \
+    if (nd == 1) \
+    { \
+        for (size_t i = 0; i < input_rows_count; i++) \
+        { \
+            vec_res[i] = EXPAND(0, col0->getUInt(i)); \
+        } \
+        return col_res; \
+    } \
+     \
+    EXTRACT_VECTOR(1) \
+    ENCODE(2, \
+           MASK(2, 0, col0->getUInt(i)), \
+           MASK(2, 1, col1->getUInt(i))) \
+    EXTRACT_VECTOR(2) \
+    ENCODE(3, \
+           MASK(3, 0, col0->getUInt(i)), \
+           MASK(3, 1, col1->getUInt(i)), \
+           MASK(3, 2, col2->getUInt(i))) \
+    EXTRACT_VECTOR(3) \
+    ENCODE(4, \
+           MASK(4, 0, col0->getUInt(i)), \
+           MASK(4, 1, col1->getUInt(i)), \
+           MASK(4, 2, col2->getUInt(i)), \
+           MASK(4, 3, col3->getUInt(i))) \
+    EXTRACT_VECTOR(4) \
+    ENCODE(5, \
+           MASK(5, 0, col0->getUInt(i)), \
+           MASK(5, 1, col1->getUInt(i)), \
+           MASK(5, 2, col2->getUInt(i)), \
+           MASK(5, 3, col3->getUInt(i)), \
+           MASK(5, 4, col4->getUInt(i))) \
+    EXTRACT_VECTOR(5) \
+    ENCODE(6, \
+           MASK(6, 0, col0->getUInt(i)), \
+           MASK(6, 1, col1->getUInt(i)), \
+           MASK(6, 2, col2->getUInt(i)), \
+           MASK(6, 3, col3->getUInt(i)), \
+           MASK(6, 4, col4->getUInt(i)), \
+           MASK(6, 5, col5->getUInt(i))) \
+    EXTRACT_VECTOR(6) \
+    ENCODE(7, \
+           MASK(7, 0, col0->getUInt(i)), \
+           MASK(7, 1, col1->getUInt(i)), \
+           MASK(7, 2, col2->getUInt(i)), \
+           MASK(7, 3, col3->getUInt(i)), \
+           MASK(7, 4, col4->getUInt(i)), \
+           MASK(7, 5, col5->getUInt(i)), \
+           MASK(7, 6, col6->getUInt(i))) \
+    EXTRACT_VECTOR(7) \
+    ENCODE(8, \
+           MASK(8, 0, col0->getUInt(i)), \
+           MASK(8, 1, col1->getUInt(i)), \
+           MASK(8, 2, col2->getUInt(i)), \
+           MASK(8, 3, col3->getUInt(i)), \
+           MASK(8, 4, col4->getUInt(i)), \
+           MASK(8, 5, col5->getUInt(i)), \
+           MASK(8, 6, col6->getUInt(i)), \
+           MASK(8, 7, col7->getUInt(i))) \
+     \
+    throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, \
+                    "Illegal number of UInt arguments of function {}, max: 8", \
+                    getName()); \
+
+DECLARE_DEFAULT_CODE(
+constexpr auto MortonND_2D_Enc = mortonnd::MortonNDLutEncoder<2, 32, 8>();
+constexpr auto MortonND_3D_Enc = mortonnd::MortonNDLutEncoder<3, 21, 8>();
+constexpr auto MortonND_4D_Enc = mortonnd::MortonNDLutEncoder<4, 16, 8>();
+constexpr auto MortonND_5D_Enc = mortonnd::MortonNDLutEncoder<5, 12, 8>();
+constexpr auto MortonND_6D_Enc = mortonnd::MortonNDLutEncoder<6, 10, 8>();
+constexpr auto MortonND_7D_Enc = mortonnd::MortonNDLutEncoder<7, 9, 8>();
+constexpr auto MortonND_8D_Enc = mortonnd::MortonNDLutEncoder<8, 8, 8>();
+class FunctionMortonEncode : public IFunction
+{
+public:
+    static constexpr auto name = "mortonEncode";
+    static FunctionPtr create(ContextPtr)
+    {
+        return std::make_shared<FunctionMortonEncode>();
+    }
+
+    String getName() const override
+    {
+        return name;
+    }
+
+    bool isVariadic() const override
+    {
+        return true;
+    }
+
+    size_t getNumberOfArguments() const override
+    {
+        return 0;
+    }
+
+    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
+
+    bool useDefaultImplementationForConstants() const override { return true; }
+
+    DataTypePtr getReturnTypeImpl(const DB::DataTypes & arguments) const override
+    {
+        size_t vectorStartIndex = 0;
+        if (arguments.empty())
+            throw Exception(ErrorCodes::TOO_FEW_ARGUMENTS_FOR_FUNCTION,
+                            "At least one UInt argument is required for function {}",
+                            getName());
+        if (WhichDataType(arguments[0]).isTuple())
+        {
+            vectorStartIndex = 1;
+            const auto * type_tuple = typeid_cast<const DataTypeTuple *>(arguments[0].get());
+            auto tuple_size = type_tuple->getElements().size();
+            if (tuple_size != (arguments.size() - 1))
+                throw Exception(ErrorCodes::ARGUMENT_OUT_OF_BOUND,
+                                "Illegal argument {} for function {}, tuple size should be equal to number of UInt arguments",
+                                arguments[0]->getName(), getName());
+            for (size_t i = 0; i < tuple_size; i++)
+            {
+                if (!WhichDataType(type_tuple->getElement(i)).isNativeUInt())
+                    throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
+                                    "Illegal type {} of argument in tuple for function {}, should be a native UInt",
+                                    type_tuple->getElement(i)->getName(), getName());
+            }
+        }
+
+        for (size_t i = vectorStartIndex; i < arguments.size(); i++)
+        {
+            const auto & arg = arguments[i];
+            if (!WhichDataType(arg).isNativeUInt())
+                throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
+                                "Illegal type {} of argument of function {}, should be a native UInt",
+                                arg->getName(), getName());
+        }
+        return std::make_shared<DataTypeUInt64>();
+    }
+
+    static UInt64 expand(UInt64 ratio, UInt64 value)
+    {
+        switch (ratio)
+        {
+            case 1:
+                return value;
+            case 2:
+                return MortonND_2D_Enc.Encode(0, value & MortonND_2D_Enc.InputMask());
+            case 3:
+                return MortonND_3D_Enc.Encode(0, 0, value & MortonND_3D_Enc.InputMask());
+            case 4:
+                return MortonND_4D_Enc.Encode(0, 0, 0, value & MortonND_4D_Enc.InputMask());
+            case 5:
+                return MortonND_5D_Enc.Encode(0, 0, 0, 0, value & MortonND_5D_Enc.InputMask());
+            case 6:
+                return MortonND_6D_Enc.Encode(0, 0, 0, 0, 0, value & MortonND_6D_Enc.InputMask());
+            case 7:
+                return MortonND_7D_Enc.Encode(0, 0, 0, 0, 0, 0, value & MortonND_7D_Enc.InputMask());
+            case 8:
+                return MortonND_8D_Enc.Encode(0, 0, 0, 0, 0, 0, 0, value & MortonND_8D_Enc.InputMask());
+        }
+        return value;
+    }
+
+    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
+    {
+        EXECUTE()
+    }
+};
+) // DECLARE_DEFAULT_CODE
+
+#if defined(MORTON_ND_BMI2_ENABLED)
+#undef ENCODE
+#define ENCODE(ND, ...) \
+    if (nd == (ND)) \
+    { \
+        for (size_t i = 0; i < input_rows_count; i++) \
+        {               \
+            vec_res[i] = MortonND_##ND##D::Encode(__VA_ARGS__); \
+        } \
+        return col_res; \
+    }
+
+#undef MASK
+#define MASK(ND, IDX, ...) \
+    (EXPAND(IDX, __VA_ARGS__))
+
+DECLARE_AVX2_SPECIFIC_CODE(
+using MortonND_2D = mortonnd::MortonNDBmi<2, uint64_t>;
+using MortonND_3D = mortonnd::MortonNDBmi<3, uint64_t>;
+using MortonND_4D = mortonnd::MortonNDBmi<4, uint64_t>;
+using MortonND_5D = mortonnd::MortonNDBmi<5, uint64_t>;
+using MortonND_6D = mortonnd::MortonNDBmi<6, uint64_t>;
+using MortonND_7D = mortonnd::MortonNDBmi<7, uint64_t>;
+using MortonND_8D = mortonnd::MortonNDBmi<8, uint64_t>;
+
+class FunctionMortonEncode : public TargetSpecific::Default::FunctionMortonEncode
+{
+public:
+    static UInt64 expand(UInt64 ratio, UInt64 value)
+    {
+        switch (ratio)
+        {
+            case 1:
+                return value;
+            case 2:
+                return MortonND_2D::Encode(0, value);
+            case 3:
+                return MortonND_3D::Encode(0, 0, value);
+            case 4:
+                return MortonND_4D::Encode(0, 0, 0, value);
+            case 5:
+                return MortonND_5D::Encode(0, 0, 0, 0, value);
+            case 6:
+                return MortonND_6D::Encode(0, 0, 0, 0, 0, value);
+            case 7:
+                return MortonND_7D::Encode(0, 0, 0, 0, 0, 0, value);
+            case 8:
+                return MortonND_8D::Encode(0, 0, 0, 0, 0, 0, 0, value);
+        }
+        return value;
+    }
+
+    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
+    {
+        EXECUTE()
+    }
+};
+) // DECLARE_AVX2_SPECIFIC_CODE
+#endif // MORTON_ND_BMI2_ENABLED
+
+#undef ENCODE
+#undef MASK
+#undef EXTRACT_VECTOR
+#undef EXPAND
+#undef EXECUTE
+
+class FunctionMortonEncode: public TargetSpecific::Default::FunctionMortonEncode
+{
+public:
+    explicit FunctionMortonEncode(ContextPtr context) : selector(context)
+    {
+        selector.registerImplementation<TargetArch::Default,
+                                        TargetSpecific::Default::FunctionMortonEncode>();
+
+#if USE_MULTITARGET_CODE && defined(MORTON_ND_BMI2_ENABLED)
+        selector.registerImplementation<TargetArch::AVX2,
+                                        TargetSpecific::AVX2::FunctionMortonEncode>();
+#endif
+    }
+
+    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count) const override
+    {
+        return selector.selectAndExecute(arguments, result_type, input_rows_count);
+    }
+
+    static FunctionPtr create(ContextPtr context)
+    {
+        return std::make_shared<FunctionMortonEncode>(context);
+    }
+
+private:
+    ImplementationSelector<IFunction> selector;
+};
+
+REGISTER_FUNCTION(MortonEncode)
+{
+    factory.registerFunction<FunctionMortonEncode>({
+    R"(
+Calculates Morton encoding (ZCurve) for a list of unsigned integers
+
+The function has two modes of operation:
+- Simple
+- Expanded
+
+Simple: accepts up to 8 unsigned integers as arguments and produces a UInt64 code.
+[example:simple]
+
+Expanded: accepts a range mask (tuple) as a first argument and up to 8 unsigned integers as other arguments.
+Each number in mask configures the amount of range expansion
+1 - no expansion
+2 - 2x expansion
+3 - 3x expansion
+....
+Up to 8x expansion.
+[example:range_expanded]
+Note: tuple size must be equal to the number of the other arguments
+
+Range expansion can be beneficial when you need a similar distribution for arguments with wildly different ranges (or cardinality)
+For example: 'IP Address' (0...FFFFFFFF) and 'Country code' (0...FF)
+
+Morton encoding for one argument is always the argument itself.
+[example:identity]
+Produces: `1`
+
+You can expand one argument too:
+[example:identity_expanded]
+Produces: `32768`
+
+The function also accepts columns as arguments:
+[example:from_table]
+
+But the range tuple must still be a constant:
+[example:from_table_range]
+
+Please note that you can fit only so much bits of information into Morton code as UInt64 has.
+Two arguments will have a range of maximum 2^32 (64/2) each
+Three arguments: range of max 2^21 (64/3) each
+And so on, all overflow will be clamped to zero
+)",
+        Documentation::Examples{
+            {"simple", "SELECT mortonEncode(1, 2, 3)"},
+            {"range_expanded", "SELECT mortonEncode((1,2), 1024, 16)"},
+            {"identity", "SELECT mortonEncode(1)"},
+            {"identity_expanded", "SELECT mortonEncode(tuple(2), 128)"},
+            {"from_table", "SELECT mortonEncode(n1, n2) FROM table"},
+            {"from_table_range", "SELECT mortonEncode((1,2), n1, n2) FROM table"},
+            },
+        Documentation::Categories {"ZCurve", "Morton coding"}
+    });
+}
+
+}

tests/queries/0_stateless/02457_morton_coding.reference

@@ -0,0 +1,12 @@
+----- START -----
+----- CONST -----
+2149
+(1,2,3,4)
+4294967286
+(65534,65533)
+4294967286
+(4294967286)
+----- 256, 8 -----
+----- 65536, 4 -----
+----- 4294967296, 2 -----
+----- END -----

tests/queries/0_stateless/02457_morton_coding.sql

@@ -0,0 +1,137 @@
+SELECT '----- START -----';
+drop table if exists morton_numbers_02457;
+create table morton_numbers_02457(
+    n1 UInt32,
+    n2 UInt32,
+    n3 UInt16,
+    n4 UInt16,
+    n5 UInt8,
+    n6 UInt8,
+    n7 UInt8,
+    n8 UInt8
+)
+    Engine=MergeTree()
+    ORDER BY n1;
+
+SELECT '----- CONST -----';
+select mortonEncode(1,2,3,4);
+select mortonDecode(4, 2149);
+select mortonEncode(65534, 65533);
+select mortonDecode(2, 4294967286);
+select mortonEncode(4294967286);
+select mortonDecode(1, 4294967286);
+
+SELECT '----- 256, 8 -----';
+insert into morton_numbers_02457
+select n1.number, n2.number, n3.number, n4.number, n5.number, n6.number, n7.number, n8.number
+from numbers(256-4, 4) n1
+    cross join numbers(256-4, 4) n2
+    cross join numbers(256-4, 4) n3
+    cross join numbers(256-4, 4) n4
+    cross join numbers(256-4, 4) n5
+    cross join numbers(256-4, 4) n6
+    cross join numbers(256-4, 4) n7
+    cross join numbers(256-4, 4) n8
+;
+drop table if exists morton_numbers_1_02457;
+create table morton_numbers_1_02457(
+    n1 UInt64,
+    n2 UInt64,
+    n3 UInt64,
+    n4 UInt64,
+    n5 UInt64,
+    n6 UInt64,
+    n7 UInt64,
+    n8 UInt64
+)
+    Engine=MergeTree()
+    ORDER BY n1;
+
+insert into morton_numbers_1_02457
+select untuple(mortonDecode(8, mortonEncode(n1, n2, n3, n4, n5, n6, n7, n8)))
+from morton_numbers_02457;
+
+(
+    select * from morton_numbers_02457
+    union distinct
+    select * from morton_numbers_1_02457
+)
+except
+(
+    select * from morton_numbers_02457
+    intersect
+    select * from morton_numbers_1_02457
+);
+drop table if exists morton_numbers_1_02457;
+
+SELECT '----- 65536, 4 -----';
+insert into morton_numbers_02457
+select n1.number, n2.number, n3.number, n4.number, 0, 0, 0, 0
+from numbers(pow(2, 16)-8,8) n1
+    cross join numbers(pow(2, 16)-8, 8) n2
+    cross join numbers(pow(2, 16)-8, 8) n3
+    cross join numbers(pow(2, 16)-8, 8) n4
+;
+
+create table morton_numbers_2_02457(
+    n1 UInt64,
+    n2 UInt64,
+    n3 UInt64,
+    n4 UInt64
+)
+    Engine=MergeTree()
+    ORDER BY n1;
+
+insert into morton_numbers_2_02457
+select untuple(mortonDecode(4, mortonEncode(n1, n2, n3, n4)))
+from morton_numbers_02457;
+
+(
+    select n1, n2, n3, n4 from morton_numbers_02457
+    union distinct
+    select n1, n2, n3, n4 from morton_numbers_2_02457
+)
+except
+(
+    select n1, n2, n3, n4 from morton_numbers_02457
+    intersect
+    select n1, n2, n3, n4 from morton_numbers_2_02457
+);
+drop table if exists morton_numbers_2_02457;
+
+SELECT '----- 4294967296, 2 -----';
+insert into morton_numbers_02457
+select n1.number, n2.number, 0, 0, 0, 0, 0, 0
+from numbers(pow(2, 32)-8,8) n1
+    cross join numbers(pow(2, 32)-8, 8) n2
+    cross join numbers(pow(2, 32)-8, 8) n3
+    cross join numbers(pow(2, 32)-8, 8) n4
+;
+
+drop table if exists morton_numbers_3_02457;
+create table morton_numbers_3_02457(
+    n1 UInt64,
+    n2 UInt64
+)
+    Engine=MergeTree()
+    ORDER BY n1;
+
+insert into morton_numbers_3_02457
+select untuple(mortonDecode(2, mortonEncode(n1, n2)))
+from morton_numbers_02457;
+
+(
+    select n1, n2 from morton_numbers_3_02457
+    union distinct
+    select n1, n2 from morton_numbers_3_02457
+)
+except
+(
+    select n1, n2 from morton_numbers_3_02457
+    intersect
+    select n1, n2 from morton_numbers_3_02457
+);
+drop table if exists morton_numbers_3_02457;
+
+SELECT '----- END -----';
+drop table if exists morton_numbers_02457;

tests/queries/0_stateless/02457_morton_coding_with_mask.reference

@@ -0,0 +1,15 @@
+----- START -----
+----- CONST -----
+4205569
+(1,2,3,4)
+4294967286
+(65534,65533)
+4294967286
+(4294967286)
+2147483648
+(128)
+0
+----- (1,2,1,2) -----
+----- (1,4) -----
+----- (1,1,2) -----
+----- END -----

tests/queries/0_stateless/02457_morton_coding_with_mask.sql

@@ -0,0 +1,143 @@
+SELECT '----- START -----';
+
+SELECT '----- CONST -----';
+select mortonEncode((1,2,3,1), 1,2,3,4);
+select mortonDecode((1, 2, 3, 1), 4205569);
+select mortonEncode((1,1), 65534, 65533);
+select mortonDecode((1,1), 4294967286);
+select mortonEncode(tuple(1), 4294967286);
+select mortonDecode(tuple(1), 4294967286);
+select mortonEncode(tuple(4), 128);
+select mortonDecode(tuple(4), 2147483648);
+select mortonEncode((4,4,4,4), 128, 128, 128, 128);
+
+SELECT '----- (1,2,1,2) -----';
+drop table if exists morton_numbers_mask_02457;
+create table morton_numbers_mask_02457(
+    n1 UInt8,
+    n2 UInt8,
+    n3 UInt8,
+    n4 UInt8
+)
+    Engine=MergeTree()
+    ORDER BY n1;
+
+insert into morton_numbers_mask_02457
+select n1.number, n2.number, n3.number, n4.number
+from           numbers(256-16, 16) n1
+    cross join numbers(256-16, 16) n2
+    cross join numbers(256-16, 16) n3
+    cross join numbers(256-16, 16) n4
+;
+drop table if exists morton_numbers_mask_1_02457;
+create table morton_numbers_mask_1_02457(
+    n1 UInt64,
+    n2 UInt64,
+    n3 UInt64,
+    n4 UInt64
+)
+    Engine=MergeTree()
+    ORDER BY n1;
+
+insert into morton_numbers_mask_1_02457
+select untuple(mortonDecode((1,2,1,2), mortonEncode((1,2,1,2), n1, n2, n3, n4)))
+from morton_numbers_mask_02457;
+
+(
+    select * from morton_numbers_mask_02457
+    union distinct
+    select * from morton_numbers_mask_1_02457
+)
+except
+(
+    select * from morton_numbers_mask_02457
+    intersect
+    select * from morton_numbers_mask_1_02457
+);
+drop table if exists morton_numbers_mask_02457;
+drop table if exists morton_numbers_mask_1_02457;
+
+SELECT '----- (1,4) -----';
+drop table if exists morton_numbers_mask_02457;
+create table morton_numbers_mask_02457(
+    n1 UInt32,
+    n2 UInt8
+)
+    Engine=MergeTree()
+    ORDER BY n1;
+
+insert into morton_numbers_mask_02457
+select n1.number, n2.number
+from           numbers(pow(2, 32)-64, 64) n1
+    cross join numbers(pow(2, 8)-64, 64) n2
+;
+drop table if exists morton_numbers_mask_2_02457;
+create table morton_numbers_mask_2_02457(
+    n1 UInt64,
+    n2 UInt64
+)
+    Engine=MergeTree()
+    ORDER BY n1;
+
+insert into morton_numbers_mask_2_02457
+select untuple(mortonDecode((1,4), mortonEncode((1,4), n1, n2)))
+from morton_numbers_mask_02457;
+
+(
+    select * from morton_numbers_mask_02457
+    union distinct
+    select * from morton_numbers_mask_2_02457
+)
+except
+(
+    select * from morton_numbers_mask_02457
+                      intersect
+        select * from morton_numbers_mask_2_02457
+);
+drop table if exists morton_numbers_mask_02457;
+drop table if exists morton_numbers_mask_2_02457;
+
+SELECT '----- (1,1,2) -----';
+drop table if exists morton_numbers_mask_02457;
+create table morton_numbers_mask_02457(
+    n1 UInt16,
+    n2 UInt16,
+    n3 UInt8,
+)
+    Engine=MergeTree()
+    ORDER BY n1;
+
+insert into morton_numbers_mask_02457
+select n1.number, n2.number, n3.number
+from           numbers(pow(2, 16)-64, 64) n1
+                   cross join numbers(pow(2, 16)-64, 64) n2
+    cross join numbers(pow(2, 8)-64, 64) n3
+;
+drop table if exists morton_numbers_mask_3_02457;
+create table morton_numbers_mask_3_02457(
+    n1 UInt64,
+    n2 UInt64,
+    n3 UInt64
+)
+    Engine=MergeTree()
+    ORDER BY n1;
+
+insert into morton_numbers_mask_3_02457
+select untuple(mortonDecode((1,1,2), mortonEncode((1,1,2), n1, n2, n3)))
+from morton_numbers_mask_02457;
+
+(
+    select * from morton_numbers_mask_02457
+    union distinct
+    select * from morton_numbers_mask_3_02457
+)
+except
+(
+    select * from morton_numbers_mask_02457
+    intersect
+    select * from morton_numbers_mask_3_02457
+);
+drop table if exists morton_numbers_mask_02457;
+drop table if exists morton_numbers_mask_3_02457;
+
+SELECT '----- END -----';