From b1c993b95769540677fa4b2dbea0abaf9d9cf541 Mon Sep 17 00:00:00 2001
From: Masha Basmanova <mbasmanova@meta.com>
Date: Wed, 24 Apr 2024 17:55:24 -0700
Subject: [PATCH] Add support for custom types to DuckParser

Summary: Enable parsing of expressions like "cast x as json".

Reviewed By: bikramSingh91

Differential Revision: D56544676

fbshipit-source-id: 74f3406bbef4c47d9c0f65ef75d53ebc1ebee66c
---
 velox/duckdb/conversion/DuckConversion.cpp       | 7 +++++++
 velox/duckdb/conversion/tests/DuckParserTest.cpp | 7 +++++++
 2 files changed, 14 insertions(+)

diff --git a/velox/duckdb/conversion/DuckConversion.cpp b/velox/duckdb/conversion/DuckConversion.cpp
index 7764b3b99b04..e777a2b1b503 100644
--- a/velox/duckdb/conversion/DuckConversion.cpp
+++ b/velox/duckdb/conversion/DuckConversion.cpp
@@ -175,6 +175,13 @@ TypePtr toVeloxType(LogicalType type, bool fileColumnNamesReadAsLowerCase) {
       }
       return ROW(std::move(names), std::move(types));
     }
+    case LogicalTypeId::USER: {
+      const auto name = ::duckdb::UserType::GetTypeName(type);
+      if (auto customType = getCustomType(name)) {
+        return customType;
+      }
+      [[fallthrough]];
+    }
     default:
       throw std::runtime_error(
           "unsupported type for duckdb -> velox conversion: " +
diff --git a/velox/duckdb/conversion/tests/DuckParserTest.cpp b/velox/duckdb/conversion/tests/DuckParserTest.cpp
index ab72297239a1..fb0d5ffc45c7 100644
--- a/velox/duckdb/conversion/tests/DuckParserTest.cpp
+++ b/velox/duckdb/conversion/tests/DuckParserTest.cpp
@@ -16,6 +16,7 @@
 #include "velox/duckdb/conversion/DuckParser.h"
 #include "velox/common/base/tests/GTestUtils.h"
 #include "velox/core/PlanNode.h"
+#include "velox/functions/prestosql/types/JsonType.h"
 #include "velox/parse/Expressions.h"
 
 using namespace facebook::velox;
@@ -367,6 +368,12 @@ TEST(DuckParserTest, cast) {
       parseExpr("cast(c0 as struct(a bigint, b real, c varchar))")->toString());
 }
 
+TEST(DuckParserTest, castToJson) {
+  registerJsonType();
+  EXPECT_EQ("cast(\"c0\", JSON)", parseExpr("cast(c0 as json)")->toString());
+  EXPECT_EQ("cast(\"c0\", JSON)", parseExpr("cast(c0 as JSON)")->toString());
+}
+
 TEST(DuckParserTest, ifCase) {
   EXPECT_EQ("if(99,1,0)", parseExpr("if(99, 1, 0)")->toString());
   EXPECT_EQ(