From 2b25d86371cbdc4dbe93057d96e31a1fd0758e7e Mon Sep 17 00:00:00 2001
From: Kevin Wilfong <kevinwilfong@meta.com>
Date: Fri, 4 Oct 2024 14:57:20 -0700
Subject: [PATCH] Fix casting "Generic" Vectors to FlatVector<UnknownValue>
 (#11170)

Summary:
Pull Request resolved: https://github.com/facebookincubator/velox/pull/11170

In simple functions and simple aggregations, there's a bunch of places where we use
TypeToFlatVector to decide how to cast a VectorPtr to its concrete type.  Notably in
the VectorWriters for complex types which we use in simple functions, we use this to
cast the child Vectors of the complex Vectors we'll write to.

Today, if the type is Generic, we cast to a FlatVector<UnknownValue>.  This is wrong
as the Vector could be of any flat-like Vector (FlatVector, ArrayVector, MapVector,
etc.).  Ultimately this just sort of works AFAICT because the GenericWriter ignores
whatever type it was cast to and treats it as a generic BaseVector. It's generally
unsafe though, and UBSan is catching this in our Fuzzer tests.

To fix this, I updated the TypeKind of a Generic in SimpleTypeTrait to be INVALID
rather than UNKNOWN.  UNKNOWN is intended for Vectors of nulls where
the type can't be determined.  In this case the type can be determined in the code
designed to handle Generics, and if you're trying to use the TypeKind of a Generic
in SimpleTypeTrait you're probably doing something wrong, so INVALID seems like a
better fit.  The TypeTraits for INVALID are also not primitive and not fixed width unlike
those for UNKNOWN which matches the other fields in SimpleTypeTrait for Generic.

I also updated TypeToFlatVector to add a template specialization that sets type to
BaseVector, rather than using KindToFlatVector. This is consistent with where these
Vectors are used (e.g. GenericWriter) and makes the casts safe.

Reviewed By: bikramSingh91

Differential Revision: D63905959

fbshipit-source-id: 70be0fa0e539bc30e6418ec31688b1e1c7abfdac
---
 velox/type/SimpleFunctionApi.h | 2 +-
 velox/vector/VectorTypeUtils.h | 5 +++++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/velox/type/SimpleFunctionApi.h b/velox/type/SimpleFunctionApi.h
index 3b78e21f9ffe..6cd754170631 100644
--- a/velox/type/SimpleFunctionApi.h
+++ b/velox/type/SimpleFunctionApi.h
@@ -332,7 +332,7 @@ struct SimpleTypeTrait<IntervalYearMonth> : public SimpleTypeTrait<int32_t> {
 
 template <typename T, bool comparable, bool orderable>
 struct SimpleTypeTrait<Generic<T, comparable, orderable>> {
-  static constexpr TypeKind typeKind = TypeKind::UNKNOWN;
+  static constexpr TypeKind typeKind = TypeKind::INVALID;
   static constexpr bool isPrimitiveType = false;
   static constexpr bool isFixedWidth = false;
 };
diff --git a/velox/vector/VectorTypeUtils.h b/velox/vector/VectorTypeUtils.h
index 370ae9f9b2f3..ca6a6b3cef13 100644
--- a/velox/vector/VectorTypeUtils.h
+++ b/velox/vector/VectorTypeUtils.h
@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include "velox/type/SimpleFunctionApi.h"
 #include "velox/type/Type.h"
 #include "velox/vector/ComplexVector.h"
 
@@ -84,5 +85,9 @@ struct TypeToFlatVector {
   using type = typename KindToFlatVector<SimpleTypeTrait<T>::typeKind>::type;
 };
 
+template <typename T, bool comparable, bool orderable>
+struct TypeToFlatVector<Generic<T, comparable, orderable>> {
+  using type = BaseVector;
+};
 } // namespace velox
 } // namespace facebook