Add sanity checks for function return types

Algunenano · Algunenano · commit f79a402e6d92 · 2024-01-30T14:10:20.000+01:00
diff --git a/src/Columns/ColumnFunction.cpp b/src/Columns/ColumnFunction.cpp
@@ -308,6 +308,13 @@ ColumnWithTypeAndName ColumnFunction::reduce() const
         ProfileEvents::increment(ProfileEvents::CompiledFunctionExecute);
 
     res.column = function->execute(columns, res.type, elements_size);
+    if (res.column->getDataType() != res.type->getColumnType())
+        throw Exception(
+            ErrorCodes::LOGICAL_ERROR,
+            "Unexpected return type from {}. Expected {}. Got {}",
+            function->getName(),
+            res.type->getColumnType(),
+            res.column->getDataType());
     if (recursively_convert_result_to_full_column_if_low_cardinality)
     {
         res.column = recursiveRemoveLowCardinality(res.column);
diff --git a/src/DataTypes/DataTypeDate.h b/src/DataTypes/DataTypeDate.h
@@ -12,6 +12,7 @@ class DataTypeDate final : public DataTypeNumberBase<UInt16>
     static constexpr auto family_name = "Date";
 
     TypeIndex getTypeId() const override { return TypeIndex::Date; }
+    TypeIndex getColumnType() const override { return TypeIndex::UInt16; }
     const char * getFamilyName() const override { return family_name; }
 
     bool canBeUsedAsVersion() const override { return true; }
diff --git a/src/DataTypes/DataTypeDate32.h b/src/DataTypes/DataTypeDate32.h
@@ -12,6 +12,7 @@ class DataTypeDate32 final : public DataTypeNumberBase<Int32>
     static constexpr auto family_name = "Date32";
 
     TypeIndex getTypeId() const override { return TypeIndex::Date32; }
+    TypeIndex getColumnType() const override { return TypeIndex::Int32; }
     const char * getFamilyName() const override { return family_name; }
 
     Field getDefault() const override
diff --git a/src/DataTypes/DataTypeDateTime.h b/src/DataTypes/DataTypeDateTime.h
@@ -40,6 +40,7 @@ class DataTypeDateTime final : public DataTypeNumberBase<UInt32>, public Timezon
     const char * getFamilyName() const override { return family_name; }
     String doGetName() const override;
     TypeIndex getTypeId() const override { return TypeIndex::DateTime; }
+    TypeIndex getColumnType() const override { return TypeIndex::UInt32; }
 
     bool canBeUsedAsVersion() const override { return true; }
     bool canBeInsideNullable() const override { return true; }
diff --git a/src/DataTypes/DataTypeEnum.h b/src/DataTypes/DataTypeEnum.h
@@ -54,6 +54,7 @@ class DataTypeEnum final : public IDataTypeEnum, public EnumValues<Type>
     const char * getFamilyName() const override;
 
     TypeIndex getTypeId() const override { return type_id; }
+    TypeIndex getColumnType() const override { return sizeof(FieldType) == 1 ? TypeIndex::Int8 : TypeIndex::Int16; }
 
     FieldType readValue(ReadBuffer & istr) const
     {
diff --git a/src/DataTypes/DataTypeInterval.h b/src/DataTypes/DataTypeInterval.h
@@ -28,6 +28,7 @@ class DataTypeInterval final : public DataTypeNumberBase<Int64>
     std::string doGetName() const override { return fmt::format("Interval{}", kind.toString()); }
     const char * getFamilyName() const override { return "Interval"; }
     TypeIndex getTypeId() const override { return TypeIndex::Interval; }
+    TypeIndex getColumnType() const override { return TypeIndex::Int64; }
 
     bool equals(const IDataType & rhs) const override;
 
diff --git a/src/DataTypes/IDataType.h b/src/DataTypes/IDataType.h
@@ -86,6 +86,8 @@ class IDataType : private boost::noncopyable, public std::enable_shared_from_thi
 
     /// Data type id. It's used for runtime type checks.
     virtual TypeIndex getTypeId() const = 0;
+    /// Storage type (e.g. Int64 for Interval)
+    virtual TypeIndex getColumnType() const { return getTypeId(); }
 
     bool hasSubcolumn(std::string_view subcolumn_name) const;
 
diff --git a/src/Interpreters/ActionsDAG.cpp b/src/Interpreters/ActionsDAG.cpp
@@ -282,6 +282,13 @@ const ActionsDAG::Node & ActionsDAG::addFunctionImpl(
         {
             size_t num_rows = arguments.empty() ? 0 : arguments.front().column->size();
             column = node.function->execute(arguments, node.result_type, num_rows, true);
+            if (column->getDataType() != node.result_type->getColumnType())
+                throw Exception(
+                    ErrorCodes::LOGICAL_ERROR,
+                    "Unexpected return type from {}. Expected {}. Got {}",
+                    node.function->getName(),
+                    node.result_type->getColumnType(),
+                    column->getDataType());
         }
         else
         {
diff --git a/src/Interpreters/ExpressionActions.cpp b/src/Interpreters/ExpressionActions.cpp
@@ -611,6 +611,13 @@ static void executeAction(const ExpressionActions::Action & action, ExecutionCon
                     ProfileEvents::increment(ProfileEvents::CompiledFunctionExecute);
 
                 res_column.column = action.node->function->execute(arguments, res_column.type, num_rows, dry_run);
+                if (res_column.column->getDataType() != res_column.type->getColumnType())
+                    throw Exception(
+                        ErrorCodes::LOGICAL_ERROR,
+                        "Unexpected return type from {}. Expected {}. Got {}",
+                        action.node->function->getName(),
+                        res_column.type->getColumnType(),
+                        res_column.column->getDataType());
             }
             break;
         }

Original file line number	Diff line number	Diff line change
`@@ -54,6 +54,7 @@ class DataTypeEnum final : public IDataTypeEnum, public EnumValues<Type>`
`54`	`54`	`const char * getFamilyName() const override;`
`55`	`55`
`56`	`56`	`TypeIndex getTypeId() const override { return type_id; }`
	`57`	`+ TypeIndex getColumnType() const override { return sizeof(FieldType) == 1 ? TypeIndex::Int8 : TypeIndex::Int16; }`
`57`	`58`
`58`	`59`	`FieldType readValue(ReadBuffer & istr) const`
`59`	`60`	`{`
Original file line number	Diff line number	Diff line change
`@@ -282,6 +282,13 @@ const ActionsDAG::Node & ActionsDAG::addFunctionImpl(`
`282`	`282`	`{`
`283`	`283`	`size_t num_rows = arguments.empty() ? 0 : arguments.front().column->size();`
`284`	`284`	`column = node.function->execute(arguments, node.result_type, num_rows, true);`
	`285`	`+ if (column->getDataType() != node.result_type->getColumnType())`
	`286`	`+ throw Exception(`
	`287`	`+ ErrorCodes::LOGICAL_ERROR,`
	`288`	`+ "Unexpected return type from {}. Expected {}. Got {}",`
	`289`	`+ node.function->getName(),`
	`290`	`+ node.result_type->getColumnType(),`
	`291`	`+ column->getDataType());`
`285`	`292`	`}`
`286`	`293`	`else`
`287`	`294`	`{`
Original file line number	Diff line number	Diff line change
`@@ -611,6 +611,13 @@ static void executeAction(const ExpressionActions::Action & action, ExecutionCon`
`611`	`611`	`ProfileEvents::increment(ProfileEvents::CompiledFunctionExecute);`
`612`	`612`
`613`	`613`	`res_column.column = action.node->function->execute(arguments, res_column.type, num_rows, dry_run);`
	`614`	`+ if (res_column.column->getDataType() != res_column.type->getColumnType())`
	`615`	`+ throw Exception(`
	`616`	`+ ErrorCodes::LOGICAL_ERROR,`
	`617`	`+ "Unexpected return type from {}. Expected {}. Got {}",`
	`618`	`+ action.node->function->getName(),`
	`619`	`+ res_column.type->getColumnType(),`
	`620`	`+ res_column.column->getDataType());`
`614`	`621`	`}`
`615`	`622`	`break;`
`616`	`623`	`}`