diff --git a/velox/functions/remote/client/tests/RemoteFunctionTest.cpp b/velox/functions/remote/client/tests/RemoteFunctionTest.cpp index 712358753895..1ca1d5e12820 100644 --- a/velox/functions/remote/client/tests/RemoteFunctionTest.cpp +++ b/velox/functions/remote/client/tests/RemoteFunctionTest.cpp @@ -140,9 +140,7 @@ class RemoteFunctionTest registerFunction>( {params.functionPrefix + ".remote_opaque"}); - registerOpaqueType("Foo"); - OpaqueType::registerSerialization( - "Foo", Foo::serialize, Foo::deserialize); + registerOpaqueType("Foo", Foo::serialize, Foo::deserialize); } }; diff --git a/velox/type/OpaqueCustomTypes.h b/velox/type/OpaqueCustomTypes.h index ef95c211f14a..6dc2fa85c1e9 100644 --- a/velox/type/OpaqueCustomTypes.h +++ b/velox/type/OpaqueCustomTypes.h @@ -34,12 +34,20 @@ class CastOperator; template class OpaqueCustomTypeRegister { public: - static bool registerType() { - return facebook::velox::registerCustomType( + static bool registerType( + OpaqueType::SerializeFunc serialize = nullptr, + OpaqueType::DeserializeFunc deserialize = nullptr) { + bool success = facebook::velox::registerCustomType( customTypeName, std::make_unique()); + if (success) { + OpaqueType::registerSerialization( + customTypeName, serialize, deserialize); + } + return success; } static bool unregisterType() { + OpaqueType::unregisterSerialization(customTypeName); return facebook::velox::unregisterCustomType(customTypeName); } @@ -77,6 +85,13 @@ class OpaqueCustomTypeRegister { std::string toString() const override { return customTypeName; } + + folly::dynamic serialize() const override { + folly::dynamic obj = folly::dynamic::object; + obj["name"] = "Type"; + obj["type"] = customTypeName; + return obj; + } }; static const TypePtr& singletonTypePtr() { diff --git a/velox/type/Type.cpp b/velox/type/Type.cpp index b7a703ee793d..326318461867 100644 --- a/velox/type/Type.cpp +++ b/velox/type/Type.cpp @@ -937,6 +937,14 @@ void OpaqueType::registerSerializationTypeErased( registry.reverse[persistentName] = type; } +void OpaqueType::unregisterSerializationTypeErased( + const std::shared_ptr& type, + const std::string& persistentName) { + auto& registry = OpaqueSerdeRegistry::get(); + registry.mapping.erase(type->typeIndex_); + registry.reverse.erase(persistentName); +} + ArrayTypePtr ARRAY(TypePtr elementType) { return TypeFactory::create(std::move(elementType)); } diff --git a/velox/type/Type.h b/velox/type/Type.h index 24a620809917..d865951b0f5e 100644 --- a/velox/type/Type.h +++ b/velox/type/Type.h @@ -1353,6 +1353,12 @@ class OpaqueType : public TypeBase { static void clearSerializationRegistry(); + template + FOLLY_NOINLINE static void unregisterSerialization( + const std::string& persistentName) { + unregisterSerializationTypeErased(OpaqueType::create(), persistentName); + } + protected: bool equals(const Type& other) const override; @@ -1364,6 +1370,10 @@ class OpaqueType : public TypeBase { const std::string& persistentName, SerializeFunc serialize = nullptr, DeserializeFunc deserialize = nullptr); + + static void unregisterSerializationTypeErased( + const std::shared_ptr& type, + const std::string& persistentName); }; using IntegerType = ScalarType; @@ -2342,11 +2352,22 @@ std::string getOpaqueAliasForTypeId(std::type_index typeIndex); /// might not be able to deserialize it in another process. To solve this /// problem, we require that both the serializing and deserializing processes /// register the opaque type using registerOpaqueType() with the same alias. +/// +/// This function also registers the serialization for the opaque type. If +/// custom serialize/deserialize functions are provided, they will be used; +/// otherwise, default (nullptr) will be registered. template -bool registerOpaqueType(const std::string& alias) { +bool registerOpaqueType( + const std::string& alias, + OpaqueType::SerializeFunc serialize = nullptr, + OpaqueType::DeserializeFunc deserialize = nullptr) { auto typeIndex = std::type_index(typeid(Class)); - return getTypeIndexByOpaqueAlias().emplace(alias, typeIndex).second && + bool success = getTypeIndexByOpaqueAlias().emplace(alias, typeIndex).second && getOpaqueAliasByTypeIndex().emplace(typeIndex, alias).second; + if (success) { + OpaqueType::registerSerialization(alias, serialize, deserialize); + } + return success; } /// Unregisters an opaque type. Returns true if the type was unregistered. @@ -2355,6 +2376,7 @@ bool registerOpaqueType(const std::string& alias) { template bool unregisterOpaqueType(const std::string& alias) { auto typeIndex = std::type_index(typeid(Class)); + OpaqueType::unregisterSerialization(alias); return getTypeIndexByOpaqueAlias().erase(alias) == 1 && getOpaqueAliasByTypeIndex().erase(typeIndex) == 1; }