Validate ColumnArray element types before appending

clickhouse/columns/array.cpp

@@ -6,6 +6,53 @@
 namespace clickhouse {
 
 namespace {
+bool CanAppendType(const TypeRef& destination_type, const TypeRef& source_type);
+
+bool CanAppendTupleType(const TypeRef& destination_type, const TypeRef& source_type) {
+    if (destination_type->GetCode() != Type::Tuple || source_type->GetCode() != Type::Tuple) {
+        return destination_type->IsEqual(source_type);
+    }
+
+    const auto destination_item_types = destination_type->As<TupleType>()->GetTupleType();
+    const auto source_item_types      = source_type->As<TupleType>()->GetTupleType();
+    if (destination_item_types.size() != source_item_types.size()) {
+        return false;
+    }
+
+    for (size_t i = 0; i < destination_item_types.size(); ++i) {
+        if (!CanAppendTupleType(destination_item_types[i], source_item_types[i])) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
+bool CanAppendType(const TypeRef& destination_type, const TypeRef& source_type) {
+    if (destination_type->IsEqual(source_type)) {
+        return true;
+    }
+
+    switch (destination_type->GetCode()) {
+        case Type::Array:
+            if (source_type->GetCode() != Type::Array) {
+                return false;
+            }
+            return CanAppendType(
+                destination_type->As<ArrayType>()->GetItemType(),
+                source_type->As<ArrayType>()->GetItemType());
+        case Type::Tuple:
+            return CanAppendTupleType(destination_type, source_type);
+        case Type::LowCardinality:
+            return source_type->GetCode() != Type::LowCardinality
+                && destination_type->As<LowCardinalityType>()->GetNestedType()->IsEqual(source_type);
+        case Type::Bool:
+            return source_type->GetCode() == Type::UInt8;
+        default:
+            return false;
+    }
+}
+
 std::shared_ptr<ColumnUInt64> make_single_offset(size_t value) {
     auto res = std::make_shared<ColumnUInt64>();
     if (value != 0) {
@@ -47,7 +94,12 @@ ColumnArray::ColumnArray(ColumnArray&& other)
 }
 
 void ColumnArray::AppendAsColumn(ColumnRef array) {
-    // appending data may throw (i.e. due to ype check failure), so do it first to avoid partly modified state.
+    if (!CanAppendType(data_->Type(), array->Type())) {
+        throw ValidationError(
+            "can't append column of type " + array->Type()->GetName() + " "
+            "to column type " + data_->Type()->GetName());
+    }
+
     data_->Append(array);
     AddOffset(array->Size());
 }

clickhouse/columns/date.cpp

@@ -254,15 +254,15 @@ void ColumnDateTime::Clear() {
 
 ColumnRef ColumnDateTime::Slice(size_t begin, size_t len) const {
     auto col = data_->Slice(begin, len)->As<ColumnUInt32>();
-    auto result = std::make_shared<ColumnDateTime>();
+    auto result = std::make_shared<ColumnDateTime>(Timezone());
 
     result->data_->Append(col);
 
     return result;
 }
 
 ColumnRef ColumnDateTime::CloneEmpty() const {
-    return std::make_shared<ColumnDateTime>();
+    return std::make_shared<ColumnDateTime>(Timezone());
 }
 
 void ColumnDateTime::Swap(Column& other) {

ut/column_array_ut.cpp

@@ -1,4 +1,5 @@
 #include <clickhouse/columns/array.h>
+#include <clickhouse/columns/bool.h>
 #include <clickhouse/columns/tuple.h>
 #include <clickhouse/columns/date.h>
 #include <clickhouse/columns/enum.h>
@@ -69,6 +70,153 @@ TEST(ColumnArray, Append) {
     ASSERT_EQ(col->As<ColumnUInt64>()->At(1), 3u);
 }
 
+TEST(ColumnArray, AppendPreservesDateTimeTimezoneCompatibility) {
+    auto source_data = std::make_shared<ColumnDateTime>("UTC");
+    source_data->AppendRaw(1);
+    auto source = std::make_shared<ColumnArray>(source_data);
+
+    auto destination = std::make_shared<ColumnArray>(std::make_shared<ColumnDateTime>("UTC"));
+
+    EXPECT_NO_THROW(destination->Append(source));
+    ASSERT_EQ(destination->Size(), 1u);
+    EXPECT_EQ(destination->GetSize(0), 1u);
+    EXPECT_EQ(destination->GetType().GetName(), source->GetType().GetName());
+
+    auto values = destination->GetAsColumn(0)->As<ColumnDateTime>();
+    ASSERT_NE(values, nullptr);
+    EXPECT_EQ(values->Timezone(), "UTC");
+    EXPECT_EQ(values->RawAt(0), 1u);
+}
+
+TEST(ColumnArray, AppendAsColumnRejectsDifferentDateTimeTimezone) {
+    auto data = std::make_shared<ColumnDateTime>("UTC");
+    auto array = std::make_shared<ColumnArray>(data);
+    auto source = std::make_shared<ColumnDateTime>("Europe/Berlin");
+    source->AppendRaw(1);
+
+    EXPECT_THROW(array->AppendAsColumn(source), ValidationError);
+    EXPECT_EQ(array->Size(), 0u);
+    EXPECT_EQ(data->Size(), 0u);
+    EXPECT_EQ(array->GetOffsets()->Size(), 0u);
+    EXPECT_EQ(source->Size(), 1u);
+}
+
+TEST(ColumnArray, AppendAsColumnAcceptsNestedArrayWithCompatibleElementType) {
+    auto nested_destination = std::make_shared<ColumnArray>(std::make_shared<ColumnBool>());
+    auto destination = std::make_shared<ColumnArray>(nested_destination);
+
+    auto nested_source = std::make_shared<ColumnArray>(std::make_shared<ColumnUInt8>());
+    auto values = std::make_shared<ColumnUInt8>();
+    values->Append(1);
+    values->Append(0);
+    nested_source->AppendAsColumn(values);
+
+    EXPECT_NO_THROW(destination->AppendAsColumn(nested_source));
+    ASSERT_EQ(destination->Size(), 1u);
+    EXPECT_EQ(destination->GetSize(0), 1u);
+    auto row = destination->GetAsColumn(0)->As<ColumnArray>();
+    ASSERT_NE(row, nullptr);
+    ASSERT_EQ(row->GetSize(0), 2u);
+    auto bool_values = row->GetAsColumn(0)->As<ColumnBool>();
+    ASSERT_NE(bool_values, nullptr);
+    EXPECT_TRUE(bool_values->At(0));
+    EXPECT_FALSE(bool_values->At(1));
+}
+
+TEST(ColumnArray, AppendAsColumnRejectsWrongTypeWithoutChangingState) {
+    auto data = std::make_shared<ColumnString>();
+    auto array = std::make_shared<ColumnArray>(data);
+
+    auto valid = std::make_shared<ColumnString>();
+    valid->Append("keep");
+    array->AppendAsColumn(valid);
+
+    auto wrong = std::make_shared<ColumnUInt64>();
+    wrong->Append(1);
+    wrong->Append(2);
+
+    EXPECT_THROW(array->AppendAsColumn(wrong), ValidationError);
+    EXPECT_EQ(array->Size(), 1u);
+    EXPECT_EQ(data->Size(), 1u);
+    EXPECT_EQ((*array->GetOffsets())[0], 1u);
+    EXPECT_EQ(data->At(0), "keep");
+    EXPECT_EQ(wrong->Size(), 2u);
+}
+
+TEST(ColumnArray, AppendAsColumnAcceptsEmptyColumn) {
+    auto array = std::make_shared<ColumnArray>(std::make_shared<ColumnString>());
+    auto empty = std::make_shared<ColumnString>();
+
+    EXPECT_NO_THROW(array->AppendAsColumn(empty));
+    ASSERT_EQ(array->Size(), 1u);
+    EXPECT_EQ(array->GetData()->Size(), 0u);
+    EXPECT_EQ(array->GetOffsets()->Size(), 1u);
+    EXPECT_EQ((*array->GetOffsets())[0], 0u);
+    EXPECT_EQ(array->GetSize(0), 0u);
+}
+
+TEST(ColumnArray, AppendAsColumnRejectsDifferentFixedStringSize) {
+    auto data = std::make_shared<ColumnFixedString>(2);
+    auto array = std::make_shared<ColumnArray>(data);
+    auto wrong = std::make_shared<ColumnFixedString>(4);
+    wrong->Append("abcd");
+
+    EXPECT_THROW(array->AppendAsColumn(wrong), ValidationError);
+    EXPECT_EQ(array->Size(), 0u);
+    EXPECT_EQ(data->Size(), 0u);
+    EXPECT_EQ(array->GetOffsets()->Size(), 0u);
+    EXPECT_EQ(wrong->Size(), 1u);
+}
+
+TEST(ColumnArray, AppendAsColumnAcceptsTupleWithDifferentFieldNames) {
+    auto data = std::make_shared<ColumnTuple>(
+        std::vector<ColumnRef>{std::make_shared<ColumnUInt64>(), std::make_shared<ColumnString>()},
+        std::vector<std::string>{"destination_id", "destination_name"});
+    auto array = std::make_shared<ColumnArray>(data);
+
+    auto source = std::make_shared<ColumnTuple>(
+        std::vector<ColumnRef>{std::make_shared<ColumnUInt64>(), std::make_shared<ColumnString>()},
+        std::vector<std::string>{"source_id", "source_name"});
+    (*source)[0]->As<ColumnUInt64>()->Append(7);
+    (*source)[1]->As<ColumnString>()->Append("value");
+
+    EXPECT_NO_THROW(array->AppendAsColumn(source));
+    ASSERT_EQ(array->Size(), 1u);
+    EXPECT_EQ(array->GetSize(0), 1u);
+    auto row = array->GetAsColumn(0)->As<ColumnTuple>();
+    ASSERT_NE(row, nullptr);
+    EXPECT_EQ((*row)[0]->As<ColumnUInt64>()->At(0), 7u);
+    EXPECT_EQ((*row)[1]->As<ColumnString>()->At(0), "value");
+}
+
+TEST(ColumnArray, AppendAsColumnAcceptsLowCardinalityDictionaryColumn) {
+    auto data = std::make_shared<ColumnLowCardinality>(std::make_shared<ColumnString>());
+    auto array = std::make_shared<ColumnArray>(data);
+    auto source = std::make_shared<ColumnString>();
+    source->Append("value");
+
+    EXPECT_NO_THROW(array->AppendAsColumn(source));
+    ASSERT_EQ(array->Size(), 1u);
+    EXPECT_EQ(array->GetSize(0), 1u);
+    EXPECT_EQ(array->GetData()->Size(), 1u);
+    EXPECT_EQ(array->GetAsColumn(0)->As<ColumnLowCardinality>()->GetItem(0).get<std::string_view>(), "value");
+}
+
+TEST(ColumnArray, AppendAsColumnAcceptsUInt8ForBool) {
+    auto array = std::make_shared<ColumnArray>(std::make_shared<ColumnBool>());
+    auto source = std::make_shared<ColumnUInt8>();
+    source->Append(1);
+    source->Append(0);
+
+    EXPECT_NO_THROW(array->AppendAsColumn(source));
+    ASSERT_EQ(array->Size(), 1u);
+    EXPECT_EQ(array->GetSize(0), 2u);
+    auto values = array->GetAsColumn(0)->As<ColumnBool>();
+    ASSERT_NE(values, nullptr);
+    EXPECT_TRUE(values->At(0));
+    EXPECT_FALSE(values->At(1));
+}
+
 TEST(ColumnArray, ArrayOfDecimal) {
     auto column = std::make_shared<clickhouse::ColumnDecimal>(18, 10);
     auto array = std::make_shared<clickhouse::ColumnArray>(column->CloneEmpty());

ut/columns_ut.cpp

@@ -598,6 +598,23 @@ TEST(ColumnsCase, DateTime_construct_from_rvalue_data) {
     EXPECT_TRUE(CompareRecursive(*col1, expected));
 }
 
+TEST(ColumnsCase, DateTimeTimezonePreservedBySliceAndCloneEmpty) {
+    auto column = std::make_shared<ColumnDateTime>("UTC");
+    column->AppendRaw(1);
+
+    auto slice = column->Slice(0, 1)->As<ColumnDateTime>();
+    ASSERT_NE(slice, nullptr);
+    EXPECT_EQ(slice->GetType().GetName(), column->GetType().GetName());
+    EXPECT_EQ(slice->Timezone(), "UTC");
+    EXPECT_EQ(slice->RawAt(0), 1u);
+
+    auto empty = column->CloneEmpty()->As<ColumnDateTime>();
+    ASSERT_NE(empty, nullptr);
+    EXPECT_EQ(empty->GetType().GetName(), column->GetType().GetName());
+    EXPECT_EQ(empty->Timezone(), "UTC");
+    EXPECT_EQ(empty->Size(), 0u);
+}
+
 TEST(ColumnsCase, DateTime64_0) {
     auto column = std::make_shared<ColumnDateTime64>(0ul);