Function bitPositionsToArray added support for big integers

ClickHouse · Jun 17, 2021 · ca672b0 · ca672b0
1 parent 6d44aa3
commit ca672b0
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 5 deletions.
diff --git a/src/Functions/FunctionsCoding.h b/src/Functions/FunctionsCoding.h
@@ -1523,7 +1523,7 @@ class FunctionBitPositionsToArray : public IFunction
 
     DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
     {
-        if (!isNativeInteger(arguments[0]))
+        if (!isInteger(arguments[0]))
             throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
                 "Illegal type {} of argument of function {}",
                 getName(),
@@ -1553,14 +1553,31 @@ class FunctionBitPositionsToArray : public IFunction
         result_array_values_data.reserve(size * 2);
 
         using UnsignedType = make_unsigned_t<T>;
+
         for (size_t row = 0; row < size; ++row)
         {
             UnsignedType x = static_cast<UnsignedType>(vec_from[row]);
 
-            while (x)
+            if constexpr (is_big_int_v<UnsignedType>)
+            {
+                size_t position = 0;
+
+                while (x)
+                {
+                    if (x & 1)
+                        result_array_values_data.push_back(position);
+
+                    x >>= 1;
+                    ++position;
+                }
+            }
+            else
             {
-                result_array_values_data.push_back(getTrailingZeroBitsUnsafe(x));
-                x &= (x - 1);
+                while (x)
+                {
+                    result_array_values_data.push_back(getTrailingZeroBitsUnsafe(x));
+                    x &= (x - 1);
+                }
             }
 
             result_array_offsets_data[row] = result_array_values_data.size();
@@ -1581,10 +1598,14 @@ class FunctionBitPositionsToArray : public IFunction
             || (result_column = executeType<UInt32>(in_column))
             || (result_column = executeType<UInt32>(in_column))
             || (result_column = executeType<UInt64>(in_column))
+            || (result_column = executeType<UInt128>(in_column))
+            || (result_column = executeType<UInt256>(in_column))
             || (result_column = executeType<Int8>(in_column))
             || (result_column = executeType<Int16>(in_column))
             || (result_column = executeType<Int32>(in_column))
-            || (result_column = executeType<Int64>(in_column))))
+            || (result_column = executeType<Int64>(in_column))
+            || (result_column = executeType<Int128>(in_column))
+            || (result_column = executeType<Int256>(in_column))))
         {
             throw Exception(ErrorCodes::ILLEGAL_COLUMN,
                "Illegal column {} of first argument of function {}",

diff --git a/tests/queries/0_stateless/01866_bit_positions_to_array.reference b/tests/queries/0_stateless/01866_bit_positions_to_array.reference
@@ -16,6 +16,12 @@ Int32
 Int64
 0	[]
 1	[0]
+Int128
+0	[]
+1	[0]
+Int256
+0	[]
+1	[0]
 UInt8
 0	[]
 1	[0]
@@ -29,3 +35,10 @@ UInt32
 UInt64
 0	[]
 1	[0]
+UInt128
+0	[]
+1	[0]
+340282366920938463463374607431768211455	[0]
+UInt256
+0	[]
+1	[0]
diff --git a/tests/queries/0_stateless/01866_bit_positions_to_array.sql b/tests/queries/0_stateless/01866_bit_positions_to_array.sql
@@ -20,6 +20,14 @@ SELECT 'Int64';
 SELECT toInt64(0), bitPositionsToArray(toInt64(0));
 SELECT toInt64(1), bitPositionsToArray(toInt64(1));
 
+SELECT 'Int128';
+SELECT toInt128(0), bitPositionsToArray(toInt128(0));
+SELECT toInt128(1), bitPositionsToArray(toInt128(1));
+
+SELECT 'Int256';
+SELECT toInt256(0), bitPositionsToArray(toInt256(0));
+SELECT toInt256(1), bitPositionsToArray(toInt256(1));
+
 SELECT 'UInt8';
 SELECT toUInt8(0), bitPositionsToArray(toUInt8(0));
 SELECT toUInt8(1), bitPositionsToArray(toUInt8(1));
@@ -36,3 +44,12 @@ SELECT toUInt32(1), bitPositionsToArray(toUInt32(1));
 SELECT 'UInt64';
 SELECT toUInt64(0), bitPositionsToArray(toUInt64(0));
 SELECT toUInt64(1), bitPositionsToArray(toUInt64(1));
+
+SELECT 'UInt128';
+SELECT toUInt128(0), bitPositionsToArray(toUInt128(0));
+SELECT toUInt128(1), bitPositionsToArray(toUInt128(1));
+SELECT toUInt128(-1), bitPositionsToArray(toUInt128(1));
+
+SELECT 'UInt256';
+SELECT toUInt256(0), bitPositionsToArray(toUInt256(0));
+SELECT toUInt256(1), bitPositionsToArray(toUInt256(1));