[cherry-pick] compute: Add Vector data type and functions (#153)

* Add Vector data type (#141) Signed-off-by: Wish <[email protected]> * compute: Add vector functions (#146) Signed-off-by: Wish <[email protected]> * compute: Fix vector precision (#147) Signed-off-by: Wish <[email protected]> * compute: Fix vector distance over NULLs (#148) Signed-off-by: Wish <[email protected]> * [cherry-pick] Storages: Make DMFile ready for new column indexes/types (#149) Signed-off-by: Wish <[email protected]> --------- Signed-off-by: Wish <[email protected]> Co-authored-by: Wenxuan <[email protected]>
pingcap · Aug 6, 2024 · dc1eb52 · dc1eb52
1 parent 8ed9835
commit dc1eb52
Show file tree

Hide file tree

Showing 41 changed files with 2,740 additions and 568 deletions.
diff --git a/contrib/tipb b/contrib/tipb
diff --git a/dbms/src/Columns/ColumnArray.cpp b/dbms/src/Columns/ColumnArray.cpp
@@ -25,8 +25,13 @@
 #include <Common/SipHash.h>
 #include <Common/typeid_cast.h>
 #include <DataStreams/ColumnGathererStream.h>
+#include <Functions/FunctionHelpers.h>
+#include <IO/Endian.h>
+#include <IO/WriteHelpers.h>
 #include <string.h> // memcpy
 
+#include <memory>
+
 namespace DB
 {
 namespace ErrorCodes
@@ -761,10 +766,44 @@ void ColumnArray::getPermutation(bool reverse, size_t limit, int nan_direction_h
     }
 }
 
-ColumnPtr ColumnArray::replicateRange(size_t /*start_row*/, size_t /*end_row*/, const IColumn::Offsets & /*offsets*/)
+ColumnPtr ColumnArray::replicateRange(size_t start_row, size_t end_row, const IColumn::Offsets & replicate_offsets)
     const
 {
-    throw Exception("not implement.", ErrorCodes::NOT_IMPLEMENTED);
+    size_t col_size = size();
+    if (col_size != replicate_offsets.size())
+        throw Exception("Size of offsets doesn't match size of column.", ErrorCodes::SIZES_OF_COLUMNS_DOESNT_MATCH);
+
+    // We only support replicate to full column.
+    RUNTIME_CHECK(start_row == 0, start_row);
+    RUNTIME_CHECK(end_row == replicate_offsets.size(), end_row, replicate_offsets.size());
+
+    if (typeid_cast<const ColumnUInt8 *>(data.get()))
+        return replicateNumber<UInt8>(replicate_offsets);
+    if (typeid_cast<const ColumnUInt16 *>(data.get()))
+        return replicateNumber<UInt16>(replicate_offsets);
+    if (typeid_cast<const ColumnUInt32 *>(data.get()))
+        return replicateNumber<UInt32>(replicate_offsets);
+    if (typeid_cast<const ColumnUInt64 *>(data.get()))
+        return replicateNumber<UInt64>(replicate_offsets);
+    if (typeid_cast<const ColumnUInt128 *>(data.get()))
+        return replicateNumber<UInt128>(replicate_offsets);
+    if (typeid_cast<const ColumnInt8 *>(data.get()))
+        return replicateNumber<Int8>(replicate_offsets);
+    if (typeid_cast<const ColumnInt16 *>(data.get()))
+        return replicateNumber<Int16>(replicate_offsets);
+    if (typeid_cast<const ColumnInt32 *>(data.get()))
+        return replicateNumber<Int32>(replicate_offsets);
+    if (typeid_cast<const ColumnInt64 *>(data.get()))
+        return replicateNumber<Int64>(replicate_offsets);
+    if (typeid_cast<const ColumnFloat32 *>(data.get()))
+        return replicateNumber<Float32>(replicate_offsets);
+    if (typeid_cast<const ColumnFloat64 *>(data.get()))
+        return replicateNumber<Float64>(replicate_offsets);
+    if (typeid_cast<const ColumnConst *>(data.get()))
+        return replicateConst(replicate_offsets);
+    if (typeid_cast<const ColumnNullable *>(data.get()))
+        return replicateNullable(replicate_offsets);
+    return replicateGeneric(replicate_offsets);
 }
 
 
@@ -1011,4 +1050,50 @@ void ColumnArray::gather(ColumnGathererStream & gatherer)
     gatherer.gather(*this);
 }
 
+bool ColumnArray::decodeTiDBRowV2Datum(size_t cursor, const String & raw_value, size_t length, bool /* force_decode */)
+{
+    RUNTIME_CHECK(raw_value.size() >= cursor + length);
+    insertFromDatumData(raw_value.c_str() + cursor, length);
+    return true;
+}
+
+void ColumnArray::insertFromDatumData(const char * data, size_t length)
+{
+    RUNTIME_CHECK(boost::endian::order::native == boost::endian::order::little);
+
+    RUNTIME_CHECK(checkAndGetColumn<ColumnVector<Float32>>(&getData()));
+    RUNTIME_CHECK(getData().isFixedAndContiguous());
+
+    RUNTIME_CHECK(length >= sizeof(UInt32), length);
+    auto n = readLittleEndian<UInt32>(data);
+    data += sizeof(UInt32);
+
+    auto precise_data_size = n * sizeof(Float32);
+    RUNTIME_CHECK(length >= sizeof(UInt32) + precise_data_size, n, length);
+    insertData(data, precise_data_size);
+}
+
+size_t ColumnArray::encodeIntoDatumData(size_t element_idx, WriteBuffer & writer) const
+{
+    RUNTIME_CHECK(boost::endian::order::native == boost::endian::order::little);
+
+    size_t encoded_size = 0;
+
+    RUNTIME_CHECK(checkAndGetColumn<ColumnVector<Float32>>(&getData()));
+    RUNTIME_CHECK(getData().isFixedAndContiguous());
+
+    auto n = static_cast<UInt32>(sizeAt(element_idx));
+
+    writeIntBinary(n, writer);
+    encoded_size += sizeof(UInt32);
+
+    auto data = getDataAt(element_idx);
+    RUNTIME_CHECK(data.size == n * sizeof(Float32));
+    writer.write(data.data, data.size);
+    encoded_size += data.size;
+
+    return encoded_size;
+}
+
+
 } // namespace DB
diff --git a/dbms/src/Columns/ColumnArray.h b/dbms/src/Columns/ColumnArray.h
@@ -156,6 +156,15 @@ class ColumnArray final : public COWPtrHelper<IColumn, ColumnArray>
         callback(data);
     }
 
+    bool canBeInsideNullable() const override { return true; }
+
+    bool decodeTiDBRowV2Datum(size_t cursor, const String & raw_value, size_t /* length */, bool /* force_decode */)
+        override;
+
+    void insertFromDatumData(const char * data, size_t length) override;
+
+    size_t encodeIntoDatumData(size_t element_idx, WriteBuffer & writer) const;
+
 private:
     ColumnPtr data;
     ColumnPtr offsets;

diff --git a/dbms/src/Columns/ColumnNullable.cpp b/dbms/src/Columns/ColumnNullable.cpp
@@ -167,14 +167,29 @@ void ColumnNullable::get(size_t n, Field & res) const
         getNestedColumn().get(n, res);
 }
 
-StringRef ColumnNullable::getDataAt(size_t /*n*/) const
+StringRef ColumnNullable::getDataAt(size_t n) const
 {
-    throw Exception(fmt::format("Method getDataAt is not supported for {}", getName()), ErrorCodes::NOT_IMPLEMENTED);
+    if (!isNullAt(n))
+        return getNestedColumn().getDataAt(n);
+
+    throw Exception(
+        ErrorCodes::NOT_IMPLEMENTED,
+        "Method getDataAt is not supported for {} in case if value is NULL",
+        getName());
 }
 
-void ColumnNullable::insertData(const char * /*pos*/, size_t /*length*/)
+void ColumnNullable::insertData(const char * pos, size_t length)
 {
-    throw Exception(fmt::format("Method insertData is not supported for {}", getName()), ErrorCodes::NOT_IMPLEMENTED);
+    if (pos == nullptr)
+    {
+        getNestedColumn().insertDefault();
+        getNullMapData().push_back(1);
+    }
+    else
+    {
+        getNestedColumn().insertData(pos, length);
+        getNullMapData().push_back(0);
+    }
 }
 
 bool ColumnNullable::decodeTiDBRowV2Datum(size_t cursor, const String & raw_value, size_t length, bool force_decode)
@@ -185,6 +200,12 @@ bool ColumnNullable::decodeTiDBRowV2Datum(size_t cursor, const String & raw_valu
     return true;
 }
 
+void ColumnNullable::insertFromDatumData(const char * cursor, size_t len)
+{
+    getNestedColumn().insertFromDatumData(cursor, len);
+    getNullMapData().push_back(0);
+}
+
 StringRef ColumnNullable::serializeValueIntoArena(
     size_t n,
     Arena & arena,

diff --git a/dbms/src/Columns/ColumnNullable.h b/dbms/src/Columns/ColumnNullable.h
@@ -65,9 +65,11 @@ class ColumnNullable final : public COWPtrHelper<IColumn, ColumnNullable>
     Field operator[](size_t n) const override;
     void get(size_t n, Field & res) const override;
     UInt64 get64(size_t n) const override { return nested_column->get64(n); }
-    StringRef getDataAt(size_t n) const override;
+    StringRef getDataAt(size_t) const override;
+    /// Will insert null value if pos=nullptr
     void insertData(const char * pos, size_t length) override;
     bool decodeTiDBRowV2Datum(size_t cursor, const String & raw_value, size_t length, bool force_decode) override;
+    void insertFromDatumData(const char *, size_t) override;
     StringRef serializeValueIntoArena(
         size_t n,
         Arena & arena,

diff --git a/dbms/src/Columns/IColumn.h b/dbms/src/Columns/IColumn.h
@@ -162,6 +162,11 @@ class IColumn : public COWPtr<IColumn>
         throw Exception("Method decodeTiDBRowV2Datum is not supported for " + getName(), ErrorCodes::NOT_IMPLEMENTED);
     }
 
+    virtual void insertFromDatumData(const char *, size_t)
+    {
+        throw Exception("Method insertFromDatumData is not supported for " + getName(), ErrorCodes::NOT_IMPLEMENTED);
+    }
+
     /// Like getData, but has special behavior for columns that contain variable-length strings.
     /// In this special case inserting data should be zero-ending (i.e. length is 1 byte greater than real string size).
     virtual void insertDataWithTerminatingZero(const char * pos, size_t length) { insertData(pos, length); }

diff --git a/dbms/src/DataTypes/DataTypeArray.h b/dbms/src/DataTypes/DataTypeArray.h
@@ -34,7 +34,7 @@ class DataTypeArray final : public IDataType
 
     const char * getFamilyName() const override { return "Array"; }
 
-    bool canBeInsideNullable() const override { return false; }
+    bool canBeInsideNullable() const override { return true; }
 
     TypeIndex getTypeId() const override { return TypeIndex::Array; }
 

diff --git a/dbms/src/DataTypes/IDataType.cpp b/dbms/src/DataTypes/IDataType.cpp
@@ -91,6 +91,18 @@ bool IDataType::isNullMap(const IDataType::SubstreamPath & path)
     return false;
 }
 
+bool IDataType::isArraySizes(const SubstreamPath & path)
+{
+    for (const auto & elem : path)
+    {
+        if (elem.type == IDataType::Substream::ArraySizes)
+        {
+            return true;
+        }
+    }
+    return false;
+}
+
 String IDataType::getFileNameForStream(const String & column_name, const IDataType::SubstreamPath & path)
 {
     String nested_table_name = Nested::extractTableName(column_name);

diff --git a/dbms/src/DataTypes/IDataType.h b/dbms/src/DataTypes/IDataType.h
@@ -243,7 +243,8 @@ class IDataType : private boost::noncopyable
         const IColumn & column,
         size_t row_num,
         WriteBuffer & ostr,
-        const FormatSettingsJSON & settings) const = 0;
+        const FormatSettingsJSON & settings) const
+        = 0;
     virtual void deserializeTextJSON(IColumn & column, ReadBuffer & istr) const = 0;
 
     /** Text serialization for putting into the XML format.
@@ -432,6 +433,7 @@ class IDataType : private boost::noncopyable
     static String getFileNameForStream(const String & column_name, const SubstreamPath & path);
 
     static bool isNullMap(const SubstreamPath & path);
+    static bool isArraySizes(const SubstreamPath & path);
 };
 
 

diff --git a/dbms/src/Debug/MockExecutor/AstToPB.cpp b/dbms/src/Debug/MockExecutor/AstToPB.cpp
@@ -110,6 +110,13 @@ void literalFieldToTiPBExpr(const ColumnInfo & ci, const Field & val_field, tipb
             encodeDAGInt64(val, ss);
             break;
         }
+        case TiDB::TypeTiDBVectorFloat32:
+        {
+            expr->set_tp(tipb::ExprType::TiDBVectorFloat32);
+            const auto & val = val_field.safeGet<Array>();
+            encodeDAGVectorFloat32(val, ss);
+            break;
+        }
         default:
             throw Exception(fmt::format(
                 "Type {} does not support literal in function unit test",

diff --git a/dbms/src/Debug/dbgTools.cpp b/dbms/src/Debug/dbgTools.cpp
@@ -469,6 +469,10 @@ struct BatchCtrl
             throw Exception(
                 "Not implented yet: BatchCtrl::encodeDatum, TiDB::CodecFlagJson",
                 ErrorCodes::LOGICAL_ERROR);
+        case TiDB::CodecFlagVectorFloat32:
+            throw Exception(
+                "Not implented yet: BatchCtrl::encodeDatum, TiDB::CodecFlagVectorFloat32",
+                ErrorCodes::LOGICAL_ERROR);
         case TiDB::CodecFlagMax:
             throw Exception("Not implented yet: BatchCtrl::encodeDatum, TiDB::CodecFlagMax", ErrorCodes::LOGICAL_ERROR);
         case TiDB::CodecFlagDuration:

diff --git a/dbms/src/Flash/Coprocessor/ArrowColCodec.cpp b/dbms/src/Flash/Coprocessor/ArrowColCodec.cpp
@@ -12,11 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Columns/ColumnArray.h>
 #include <Columns/ColumnDecimal.h>
 #include <Columns/ColumnNullable.h>
 #include <Columns/ColumnString.h>
 #include <Columns/ColumnVector.h>
 #include <Common/TiFlashException.h>
+#include <DataTypes/DataTypeArray.h>
 #include <DataTypes/DataTypeDecimal.h>
 #include <DataTypes/DataTypeEnum.h>
 #include <DataTypes/DataTypeMyDate.h>
@@ -292,6 +294,34 @@ void flashStringColToArrowCol(
     }
 }
 
+template <bool is_nullable>
+void flashArrayFloat32ColToArrowCol(
+    TiDBColumn & dag_column,
+    const IColumn * flash_col_untyped,
+    size_t start_index,
+    size_t end_index)
+{
+    // We only unwrap the NULLABLE() part.
+    const IColumn * nested_col = getNestedCol(flash_col_untyped);
+    const auto * flash_col = checkAndGetColumn<ColumnArray>(nested_col);
+    for (size_t i = start_index; i < end_index; i++)
+    {
+        // todo check if we can convert flash_col to DAG col directly since the internal representation is almost the same
+        if constexpr (is_nullable)
+        {
+            if (flash_col_untyped->isNullAt(i))
+            {
+                dag_column.appendNull();
+                continue;
+            }
+        }
+
+        auto encoded_size = flash_col->encodeIntoDatumData(i, *dag_column.data);
+        RUNTIME_CHECK(encoded_size > 0);
+        dag_column.finishAppendVar(encoded_size);
+    }
+}
+
 template <bool is_nullable>
 void flashBitColToArrowCol(
     TiDBColumn & dag_column,
@@ -461,6 +491,20 @@ void flashColToArrowCol(
         else
             flashStringColToArrowCol<true>(dag_column, col, start_index, end_index);
         break;
+    case TiDB::TypeTiDBVectorFloat32:
+    {
+        const auto * data_type = checkAndGetDataType<DataTypeArray>(type);
+        if (!data_type || data_type->getNestedType()->getTypeId() != TypeIndex::Float32)
+            throw TiFlashException(
+                "Type un-matched during arrow encode, target col type is array<float32> and source column type is "
+                    + type->getName(),
+                Errors::Coprocessor::Internal);
+        if (tidb_column_info.hasNotNullFlag())
+            flashArrayFloat32ColToArrowCol<false>(dag_column, col, start_index, end_index);
+        else
+            flashArrayFloat32ColToArrowCol<true>(dag_column, col, start_index, end_index);
+        break;
+    }
     case TiDB::TypeBit:
         if (!checkDataType<DataTypeUInt64>(type))
             throw TiFlashException(
@@ -525,6 +569,35 @@ const char * arrowStringColToFlashCol(
     return pos + offsets[length];
 }
 
+const char * arrowArrayFloat32ColToFlashCol(
+    const char * pos,
+    UInt8,
+    UInt32 null_count,
+    const std::vector<UInt8> & null_bitmap,
+    const std::vector<UInt64> & offsets,
+    const ColumnWithTypeAndName & col,
+    const ColumnInfo &,
+    UInt32 length)
+{
+    const auto * data_type = checkAndGetDataType<DataTypeArray>(&*col.type);
+    if (!data_type || data_type->getNestedType()->getTypeId() != TypeIndex::Float32)
+        throw TiFlashException(
+            "Type un-matched during arrow decode, target col type is array<float32> and source column type is "
+                + col.type->getName(),
+            Errors::Coprocessor::Internal);
+
+    for (UInt32 i = 0; i < length; i++)
+    {
+        if (checkNull(i, null_count, null_bitmap, col))
+            continue;
+
+        auto arrow_data_size = offsets[i + 1] - offsets[i];
+        const auto * base_offset = pos + offsets[i];
+        col.column->assumeMutable()->insertFromDatumData(base_offset, arrow_data_size);
+    }
+    return pos + offsets[length];
+}
+
 const char * arrowEnumColToFlashCol(
     const char * pos,
     UInt8,
@@ -819,6 +892,16 @@ const char * arrowColToFlashCol(
             length);
     case TiDB::TypeBit:
         return arrowBitColToFlashCol(pos, field_length, null_count, null_bitmap, offsets, flash_col, col_info, length);
+    case TiDB::TypeTiDBVectorFloat32:
+        return arrowArrayFloat32ColToFlashCol(
+            pos,
+            field_length,
+            null_count,
+            null_bitmap,
+            offsets,
+            flash_col,
+            col_info,
+            length);
     case TiDB::TypeEnum:
         return arrowEnumColToFlashCol(pos, field_length, null_count, null_bitmap, offsets, flash_col, col_info, length);
     default: