diff --git a/Makefile b/Makefile index c1767755..960b30fd 100644 --- a/Makefile +++ b/Makefile @@ -11,6 +11,7 @@ SRCS = src/scan/heap_reader.cpp \ src/scan/postgres_seq_scan.cpp \ src/utility/copy.cpp \ src/vendor/pg_explain.cpp \ + src/types/pgduckdb_enum.cpp \ src/pgduckdb_metadata_cache.cpp \ src/pgduckdb_detoast.cpp \ src/pgduckdb_duckdb.cpp \ @@ -26,6 +27,7 @@ SRCS = src/scan/heap_reader.cpp \ src/catalog/pgduckdb_storage.cpp \ src/catalog/pgduckdb_schema.cpp \ src/catalog/pgduckdb_table.cpp \ + src/catalog/pgduckdb_type.cpp \ src/catalog/pgduckdb_transaction.cpp \ src/catalog/pgduckdb_transaction_manager.cpp \ src/catalog/pgduckdb_catalog.cpp diff --git a/include/pgduckdb/catalog/pgduckdb_transaction.hpp b/include/pgduckdb/catalog/pgduckdb_transaction.hpp index c22362a0..b6bb46af 100644 --- a/include/pgduckdb/catalog/pgduckdb_transaction.hpp +++ b/include/pgduckdb/catalog/pgduckdb_transaction.hpp @@ -2,6 +2,7 @@ #include "duckdb/transaction/transaction.hpp" #include "pgduckdb/catalog/pgduckdb_table.hpp" +#include "pgduckdb/catalog/pgduckdb_type.hpp" #include "pgduckdb/catalog/pgduckdb_schema.hpp" namespace duckdb { @@ -15,11 +16,13 @@ class SchemaItems { public: optional_ptr GetTable(const string &name, PlannerInfo *planner_info); + optional_ptr GetType(const string &name); public: string name; unique_ptr schema; case_insensitive_map_t> tables; + case_insensitive_map_t> types; }; class PostgresTransaction : public Transaction { diff --git a/include/pgduckdb/catalog/pgduckdb_type.hpp b/include/pgduckdb/catalog/pgduckdb_type.hpp new file mode 100644 index 00000000..162d393f --- /dev/null +++ b/include/pgduckdb/catalog/pgduckdb_type.hpp @@ -0,0 +1,31 @@ +#pragma once + +#include "duckdb/catalog/catalog_entry/type_catalog_entry.hpp" +#include "duckdb/parser/parsed_data/create_type_info.hpp" +#include "duckdb/storage/table_storage_info.hpp" + +extern "C" { +#include "postgres.h" +#include "utils/snapshot.h" +#include "postgres.h" +#include "catalog/namespace.h" +#include "catalog/pg_class.h" +#include "optimizer/planmain.h" +#include "optimizer/planner.h" +#include "utils/builtins.h" +#include "utils/regproc.h" +#include "utils/snapmgr.h" +#include "utils/syscache.h" +#include "access/htup_details.h" +} + +namespace duckdb { + +class PostgresType : public TypeCatalogEntry { +public: + ~PostgresType() { + } + PostgresType(Catalog &catalog, SchemaCatalogEntry &schema, CreateTypeInfo &info); +}; + +} // namespace duckdb diff --git a/include/pgduckdb/duckdb_vendor/.clang-format b/include/pgduckdb/duckdb_vendor/.clang-format new file mode 100644 index 00000000..9d159247 --- /dev/null +++ b/include/pgduckdb/duckdb_vendor/.clang-format @@ -0,0 +1,2 @@ +DisableFormat: true +SortIncludes: false diff --git a/include/pgduckdb/duckdb_vendor/enum_type_info_templated.hpp b/include/pgduckdb/duckdb_vendor/enum_type_info_templated.hpp new file mode 100644 index 00000000..3f4b5ca9 --- /dev/null +++ b/include/pgduckdb/duckdb_vendor/enum_type_info_templated.hpp @@ -0,0 +1,57 @@ +#pragma once + +#include "duckdb/common/extra_type_info.hpp" +#include "duckdb/common/vector.hpp" +#include "duckdb/common/types/vector.hpp" +#include "duckdb/common/string_map_set.hpp" +#include "duckdb/common/serializer/deserializer.hpp" + +// This is copied directly without changes from 'duckdb/src/common/extra_type_info.cpp' +// The reason this is copied is to be able to inherit from it. + +namespace duckdb { +template +struct EnumTypeInfoTemplated : public EnumTypeInfo { + explicit EnumTypeInfoTemplated(Vector &values_insert_order_p, idx_t size_p) + : EnumTypeInfo(values_insert_order_p, size_p) { + D_ASSERT(values_insert_order_p.GetType().InternalType() == PhysicalType::VARCHAR); + + UnifiedVectorFormat vdata; + values_insert_order.ToUnifiedFormat(size_p, vdata); + + auto data = UnifiedVectorFormat::GetData(vdata); + for (idx_t i = 0; i < size_p; i++) { + auto idx = vdata.sel->get_index(i); + if (!vdata.validity.RowIsValid(idx)) { + throw InternalException("Attempted to create ENUM type with NULL value"); + } + if (values.count(data[idx]) > 0) { + throw InvalidInputException("Attempted to create ENUM type with duplicate value %s", + data[idx].GetString()); + } + values[data[idx]] = UnsafeNumericCast(i); + } + } + + static shared_ptr Deserialize(Deserializer &deserializer, uint32_t size) { + Vector values_insert_order(LogicalType::VARCHAR, size); + auto strings = FlatVector::GetData(values_insert_order); + + deserializer.ReadList(201, "values", [&](Deserializer::List &list, idx_t i) { + strings[i] = StringVector::AddStringOrBlob(values_insert_order, list.ReadElement()); + }); + return make_shared_ptr(values_insert_order, size); + } + + const string_map_t &GetValues() const { + return values; + } + + EnumTypeInfoTemplated(const EnumTypeInfoTemplated &) = delete; + EnumTypeInfoTemplated &operator=(const EnumTypeInfoTemplated &) = delete; + +private: + string_map_t values; +}; + +} // namespace duckdb diff --git a/include/pgduckdb/types/pgduckdb_enum.hpp b/include/pgduckdb/types/pgduckdb_enum.hpp new file mode 100644 index 00000000..2b903be0 --- /dev/null +++ b/include/pgduckdb/types/pgduckdb_enum.hpp @@ -0,0 +1,65 @@ +#pragma once + +#include "pgduckdb/duckdb_vendor/enum_type_info_templated.hpp" +#include "duckdb/common/types/vector.hpp" +#include "duckdb/common/types.hpp" + +extern "C" { +#include "postgres.h" +#include "catalog/pg_enum.h" +#include "catalog/pg_type.h" +#include "utils/syscache.h" +#include "access/htup_details.h" +} + +namespace pgduckdb { + +using duckdb::idx_t; +using duckdb::LogicalType; +using duckdb::Vector; + +// To store additional metadata for a type, DuckDB's LogicalType class contains a 'type_info_' field. +// The type of this is ExtraTypeInfo, ENUMs make use of this in the form of EnumTypeInfo (see +// duckdb/include/common/extra_type_info.hpp). We derive from this further to store the ENUMs connection with the oids +// of Postgres' enum members. + +template +class PGDuckDBEnumTypeInfo : public duckdb::EnumTypeInfoTemplated { +public: + PGDuckDBEnumTypeInfo(Vector &values_insert_order_p, idx_t dict_size_p, Vector &enum_member_oids) + : duckdb::EnumTypeInfoTemplated(values_insert_order_p, dict_size_p), enum_member_oids(enum_member_oids) { + } + +public: + const Vector & + GetMemberOids() const { + return enum_member_oids; + } + + duckdb::shared_ptr + Copy() const override { + auto &insert_order = this->GetValuesInsertOrder(); + Vector values_insert_order_copy(LogicalType::VARCHAR, false, false, 0); + values_insert_order_copy.Reference(insert_order); + + Vector enum_member_oids_copy(LogicalType::UINTEGER, false, false, 0); + enum_member_oids_copy.Reference(enum_member_oids); + + return duckdb::make_shared_ptr(values_insert_order_copy, this->GetDictSize(), + enum_member_oids_copy); + } + +private: + Vector enum_member_oids; +}; + +struct PGDuckDBEnum { + static LogicalType CreateEnumType(std::vector &enum_members); + static idx_t GetDuckDBEnumPosition(duckdb::Value &val); + static idx_t GetEnumPosition(Datum enum_member_oid, const duckdb::LogicalType &type); + static bool IsEnumType(Oid type_oid); + static Oid GetEnumTypeOid(const Vector &oids); + static const Vector &GetMemberOids(const duckdb::LogicalType &type); +}; + +} // namespace pgduckdb diff --git a/src/catalog/pgduckdb_transaction.cpp b/src/catalog/pgduckdb_transaction.cpp index 67258ac4..d0c3f018 100644 --- a/src/catalog/pgduckdb_transaction.cpp +++ b/src/catalog/pgduckdb_transaction.cpp @@ -1,8 +1,11 @@ #include "pgduckdb/catalog/pgduckdb_catalog.hpp" #include "pgduckdb/catalog/pgduckdb_transaction.hpp" #include "pgduckdb/catalog/pgduckdb_table.hpp" +#include "pgduckdb/catalog/pgduckdb_type.hpp" +#include "pgduckdb/pgduckdb_types.hpp" #include "pgduckdb/scan/postgres_scan.hpp" #include "duckdb/parser/parsed_data/create_table_info.hpp" +#include "duckdb/parser/parsed_data/create_type_info.hpp" #include "duckdb/parser/parsed_data/create_schema_info.hpp" #include "duckdb/catalog/catalog.hpp" @@ -10,6 +13,7 @@ extern "C" { #include "postgres.h" #include "catalog/namespace.h" #include "catalog/pg_class.h" +#include "catalog/pg_type.h" #include "optimizer/planmain.h" #include "optimizer/planner.h" #include "utils/builtins.h" @@ -23,6 +27,7 @@ extern "C" { #include "nodes/makefuncs.h" #include "nodes/nodeFuncs.h" #include "parser/parsetree.h" +#include "parser/parse_type.h" #include "utils/rel.h" } @@ -68,6 +73,56 @@ IsIndexScan(const Path *nodePath) { return false; } +optional_ptr +SchemaItems::GetType(const string &entry_name) { + auto it = types.find(entry_name); + if (it != types.end()) { + return it->second.get(); + } + + auto &catalog = schema->catalog; + + List *name_list = NIL; + name_list = lappend(name_list, makeString(pstrdup(name.c_str()))); + name_list = lappend(name_list, makeString(pstrdup(entry_name.c_str()))); + + TypeName *type_name = makeTypeNameFromNameList(name_list); + + // Try to look up the type by name + Oid type_oid = InvalidOid; + HeapTuple type_entry = LookupTypeName(NULL, type_name, NULL, false); + + if (HeapTupleIsValid(type_entry)) { + // Cast to the correct form to access the fields + Form_pg_type type_form = (Form_pg_type)GETSTRUCT(type_entry); + + // Extract the OID from the type entry + type_oid = type_form->oid; + + // Release the system cache entry + ReleaseSysCache(type_entry); + } + + // If the type could not be found, handle it here + if (type_oid == InvalidOid) { + return nullptr; + } + + FormData_pg_attribute dummy_attr; + dummy_attr.atttypid = type_oid; + dummy_attr.atttypmod = 0; + dummy_attr.attndims = 0; + Form_pg_attribute attribute = &dummy_attr; + + CreateTypeInfo info; + info.name = entry_name; + info.type = pgduckdb::ConvertPostgresToDuckColumnType(attribute); + + auto type = make_uniq(catalog, *schema, info); + types[entry_name] = std::move(type); + return types[entry_name].get(); +} + optional_ptr SchemaItems::GetTable(const string &entry_name, PlannerInfo *planner_info) { auto it = tables.find(entry_name); @@ -157,23 +212,27 @@ PostgresTransaction::GetCatalogEntry(CatalogType type, const string &schema, con throw InternalException("Could not find 'postgres_state' in 'PostgresTransaction::GetCatalogEntry'"); } auto planner_info = scan_data->m_query_planner_info; + if (type == CatalogType::SCHEMA_ENTRY) { + return GetSchema(schema); + } + + auto it = schemas.find(schema); + if (it == schemas.end()) { + return nullptr; + } + auto &schema_entry = it->second; + switch (type) { case CatalogType::TABLE_ENTRY: { - auto it = schemas.find(schema); - if (it == schemas.end()) { - return nullptr; - } - auto &schema_entry = it->second; return schema_entry.GetTable(name, planner_info); } - case CatalogType::SCHEMA_ENTRY: { - return GetSchema(schema); + case CatalogType::TYPE_ENTRY: { + return schema_entry.GetType(name); } default: + D_ASSERT(type != CatalogType::SCHEMA_ENTRY); return nullptr; } } } // namespace duckdb - -// namespace duckdb diff --git a/src/catalog/pgduckdb_type.cpp b/src/catalog/pgduckdb_type.cpp new file mode 100644 index 00000000..582e4b44 --- /dev/null +++ b/src/catalog/pgduckdb_type.cpp @@ -0,0 +1,28 @@ +#include "pgduckdb/pgduckdb_types.hpp" +#include "pgduckdb/catalog/pgduckdb_schema.hpp" +#include "pgduckdb/catalog/pgduckdb_type.hpp" + +extern "C" { +#include "postgres.h" +#include "access/tableam.h" +#include "access/heapam.h" +#include "storage/bufmgr.h" +#include "catalog/namespace.h" +#include "catalog/pg_class.h" +#include "optimizer/planmain.h" +#include "optimizer/planner.h" +#include "utils/builtins.h" +#include "utils/regproc.h" +#include "utils/snapmgr.h" +#include "utils/syscache.h" +#include "access/htup_details.h" +#include "parser/parsetree.h" +} + +namespace duckdb { + +PostgresType::PostgresType(Catalog &catalog, SchemaCatalogEntry &schema, CreateTypeInfo &info) + : TypeCatalogEntry(catalog, schema, info) { +} + +} // namespace duckdb diff --git a/src/pgduckdb_filter.cpp b/src/pgduckdb_filter.cpp index 081271d2..029c0956 100644 --- a/src/pgduckdb_filter.cpp +++ b/src/pgduckdb_filter.cpp @@ -10,6 +10,7 @@ extern "C" { } #include "pgduckdb/pgduckdb_filter.hpp" +#include "pgduckdb/types/pgduckdb_enum.hpp" #include "pgduckdb/pgduckdb_detoast.hpp" #include "pgduckdb/pgduckdb_types.hpp" @@ -71,8 +72,22 @@ FilterOperationSwitch(Datum &value, duckdb::Value &constant, Oid type_oid) { case VARCHAROID: return StringFilterOperation(value, constant); default: + if (PGDuckDBEnum::IsEnumType(type_oid)) { + auto position = PGDuckDBEnum::GetEnumPosition(value, constant.type()); + auto physical_type = duckdb::EnumType::GetPhysicalType(constant.type()); + switch (physical_type) { + case duckdb::PhysicalType::UINT8: + return TemplatedFilterOperation(static_cast(position), constant); + case duckdb::PhysicalType::UINT16: + return TemplatedFilterOperation(static_cast(position), constant); + case duckdb::PhysicalType::UINT32: + return TemplatedFilterOperation(static_cast(position), constant); + default: + throw duckdb::InternalException("Invalid Physical Type for ENUMs"); + } + } throw duckdb::InvalidTypeException( - duckdb::string("(DuckDB/FilterOperationSwitch) Unsupported duckdb type: %d", type_oid)); + duckdb::StringUtil::Format("(DuckDB/FilterOperationSwitch) Unsupported duckdb type: %d", type_oid)); } } diff --git a/src/pgduckdb_types.cpp b/src/pgduckdb_types.cpp index edd433ea..a74ff948 100644 --- a/src/pgduckdb_types.cpp +++ b/src/pgduckdb_types.cpp @@ -11,6 +11,7 @@ extern "C" { #include "executor/tuptable.h" #include "utils/builtins.h" #include "utils/numeric.h" +#include "utils/catcache.h" #include "utils/uuid.h" #include "utils/array.h" #include "fmgr.h" @@ -21,14 +22,20 @@ extern "C" { } #include "pgduckdb/pgduckdb.h" +#include "pgduckdb/pgduckdb_utils.hpp" #include "pgduckdb/scan/postgres_scan.hpp" #include "pgduckdb/types/decimal.hpp" #include "pgduckdb/pgduckdb_filter.hpp" #include "pgduckdb/pgduckdb_detoast.hpp" #include "pgduckdb/pgduckdb_types.hpp" +#include "pgduckdb/types/pgduckdb_enum.hpp" namespace pgduckdb { +using duckdb::EnumTypeInfo; +using duckdb::LogicalTypeId; +using duckdb::PhysicalType; + struct BoolArray { public: static ArrayType * @@ -324,6 +331,14 @@ ConvertDuckToPostgresArray(TupleTableSlot *slot, duckdb::Value &value, idx_t col slot->tts_values[col] = PointerGetDatum(arr); } +void +ConvertDuckToPostgresEnumValue(TupleTableSlot *slot, duckdb::Value &val, idx_t col) { + auto position = PGDuckDBEnum::GetDuckDBEnumPosition(val); + auto &enum_member_oids = PGDuckDBEnum::GetMemberOids(val.type()); + auto enum_member_oid = duckdb::FlatVector::GetData(enum_member_oids)[position]; + slot->tts_values[col] = ObjectIdGetDatum(enum_member_oid); +} + bool ConvertDuckToPostgresValue(TupleTableSlot *slot, duckdb::Value &value, idx_t col) { Oid oid = slot->tts_tupleDescriptor->attrs[col].atttypid; @@ -465,9 +480,18 @@ ConvertDuckToPostgresValue(TupleTableSlot *slot, duckdb::Value &value, idx_t col ConvertDuckToPostgresArray>>(slot, value, col); break; } - default: - elog(WARNING, "(PGDuckDB/ConvertDuckToPostgresValue) Unsuported pgduckdb type: %d", oid); - return false; + default: { + auto type_id = value.type().id(); + switch (type_id) { + case LogicalTypeId::ENUM: { + ConvertDuckToPostgresEnumValue(slot, value, col); + break; + default: + elog(WARNING, "(PGDuckDB/ConvertDuckToPostgresValue) Unsuported pgduckdb type: %d", oid); + return false; + } + } + } } return true; } @@ -498,6 +522,39 @@ ChildTypeFromArray(Oid array_type) { } } +duckdb::LogicalType +ConvertPostgresEnumToDuckEnum(Oid enum_type_oid) { + /* Get the list of existing members of the enum */ + auto list = + PostgresFunctionGuard([](int cacheId, Datum key) { return SearchSysCacheList1(cacheId, key); }, + ENUMTYPOIDNAME, ObjectIdGetDatum(enum_type_oid)); + auto nelems = list->n_members; + + /* Sort the existing members by enumsortorder */ + std::vector enum_members(nelems); + for (int i = 0; i < nelems; ++i) { + enum_members[i] = &list->members[i]->tuple; + } + + auto sort_order_cmp = [](const HeapTuple &v1, const HeapTuple &v2) { + Form_pg_enum en1 = (Form_pg_enum)GETSTRUCT(v1); + Form_pg_enum en2 = (Form_pg_enum)GETSTRUCT(v2); + + if (en1->enumsortorder < en2->enumsortorder) { + return true; // v1 < v2 + } else if (en1->enumsortorder > en2->enumsortorder) { + return false; // v1 > v2 + } else { + return false; // v1 == v2 + } + }; + + std::sort(enum_members.begin(), enum_members.end(), sort_order_cmp); + auto enum_type = PGDuckDBEnum::CreateEnumType(enum_members); + PostgresFunctionGuard(ReleaseCatCacheList, list); + return enum_type; +} + duckdb::LogicalType ConvertPostgresToDuckColumnType(Form_pg_attribute &attribute) { auto &type = attribute->atttypid; @@ -551,6 +608,9 @@ ConvertPostgresToDuckColumnType(Form_pg_attribute &attribute) { case REGCLASSOID: return duckdb::LogicalTypeId::UINTEGER; default: { + if (PGDuckDBEnum::IsEnumType(type)) { + return ConvertPostgresEnumToDuckEnum(type); + } std::string name = "UnsupportedPostgresType (Oid=" + std::to_string(type) + ")"; return duckdb::LogicalType::USER(name); } @@ -618,6 +678,10 @@ GetPostgresDuckDBType(duckdb::LogicalType type) { } } } + case duckdb::LogicalTypeId::ENUM: { + auto &member_oids = PGDuckDBEnum::GetMemberOids(type); + return PGDuckDBEnum::GetEnumTypeOid(member_oids); + } default: { elog(WARNING, "(PGDuckDB/GetPostgresDuckDBType) Could not convert DuckDB type: %s to Postgres type", type.ToString().c_str()); @@ -927,6 +991,24 @@ ConvertPostgresToDuckValue(Datum value, duckdb::Vector &result, idx_t offset) { } break; } + case LogicalTypeId::ENUM: { + auto enum_position = PGDuckDBEnum::GetEnumPosition(value, result.GetType()); + auto physical_type = result.GetType().InternalType(); + switch (physical_type) { + case PhysicalType::UINT8: + Append(result, static_cast(enum_position), offset); + break; + case PhysicalType::UINT16: + Append(result, static_cast(enum_position), offset); + break; + case PhysicalType::UINT32: + Append(result, static_cast(enum_position), offset); + break; + default: + throw duckdb::InternalException("Invalid Physical Type for ENUMs"); + } + break; + } default: throw duckdb::NotImplementedException("(DuckDB/ConvertPostgresToDuckValue) Unsupported pgduckdb type: %s", result.GetType().ToString().c_str()); diff --git a/src/types/pgduckdb_enum.cpp b/src/types/pgduckdb_enum.cpp new file mode 100644 index 00000000..ffd587ec --- /dev/null +++ b/src/types/pgduckdb_enum.cpp @@ -0,0 +1,140 @@ +#include "pgduckdb/types/pgduckdb_enum.hpp" +#include "pgduckdb/pgduckdb_utils.hpp" +#include "duckdb/common/extra_type_info.hpp" + +namespace pgduckdb { + +using duckdb::idx_t; +using duckdb::LogicalTypeId; +using duckdb::PhysicalType; + +// This is medium hacky, we create extra space in the Vector that would normally only hold the strings corresponding to +// the enum's members. In that extra space we store the enum member oids. + +// We do this so we can find the OID from the offset (which is what DuckDB stores in a Vector of type ENUM) and return +// that when converting the result from DuckDB -> Postgres + +LogicalType +PGDuckDBEnum::CreateEnumType(std::vector &enum_members) { + auto size = enum_members.size(); + + auto duck_enum_vec = duckdb::Vector(duckdb::LogicalType::VARCHAR, size); + auto enum_member_oid_vec = duckdb::Vector(duckdb::LogicalType::UINTEGER, size); + auto enum_vec_data = duckdb::FlatVector::GetData(duck_enum_vec); + auto enum_member_oid_data = duckdb::FlatVector::GetData(enum_member_oid_vec); + for (idx_t i = 0; i < size; i++) { + auto &member = enum_members[i]; + auto enum_data = (Form_pg_enum)GETSTRUCT(member); + enum_vec_data[i] = duckdb::StringVector::AddString(duck_enum_vec, enum_data->enumlabel.data); + enum_member_oid_data[i] = enum_data->oid; + } + + // Generate EnumTypeInfo + duckdb::shared_ptr info; + auto enum_internal_type = duckdb::EnumTypeInfo::DictType(size); + switch (enum_internal_type) { + case PhysicalType::UINT8: + info = duckdb::make_shared_ptr>(duck_enum_vec, size, enum_member_oid_vec); + break; + case PhysicalType::UINT16: + info = duckdb::make_shared_ptr>(duck_enum_vec, size, enum_member_oid_vec); + break; + case PhysicalType::UINT32: + info = duckdb::make_shared_ptr>(duck_enum_vec, size, enum_member_oid_vec); + break; + default: + throw duckdb::InternalException("Invalid Physical Type for ENUMs"); + } + // Generate Actual Enum Type + return LogicalType(LogicalTypeId::ENUM, info); +} + +idx_t +PGDuckDBEnum::GetDuckDBEnumPosition(duckdb::Value &val) { + D_ASSERT(val.type().id() == LogicalTypeId::ENUM); + auto physical_type = val.type().InternalType(); + switch (physical_type) { + case PhysicalType::UINT8: + return val.GetValue(); + case PhysicalType::UINT16: + return val.GetValue(); + case PhysicalType::UINT32: + return val.GetValue(); + default: + throw duckdb::InternalException("Invalid Physical Type for ENUMs"); + } +} + +idx_t +PGDuckDBEnum::GetEnumPosition(Datum enum_member_oid_p, const duckdb::LogicalType &type) { + auto enum_member_oid = DatumGetObjectId(enum_member_oid_p); + + auto &enum_type_info = type.AuxInfo()->Cast(); + auto dict_size = enum_type_info.GetDictSize(); + + auto &enum_member_oids = PGDuckDBEnum::GetMemberOids(type); + auto oids_data = duckdb::FlatVector::GetData(enum_member_oids); + + const uint32_t *begin = oids_data; + const uint32_t *end = oids_data + dict_size; + const uint32_t *result = std::find(begin, end, enum_member_oid); + + if (result == end) { + throw duckdb::InternalException("Could not find enum_member_oid: %d", enum_member_oid); + } + return (idx_t)(result - begin); +} + +bool +PGDuckDBEnum::IsEnumType(Oid type_oid) { + bool result = false; + auto type_tuple = PostgresFunctionGuard(SearchSysCache1, TYPEOID, ObjectIdGetDatum(type_oid)); + + if (HeapTupleIsValid(type_tuple)) { + auto type_form = (Form_pg_type)GETSTRUCT(type_tuple); + + // Check if the type is an enum + if (type_form->typtype == 'e') { + result = true; + } + PostgresFunctionGuard(ReleaseSysCache, type_tuple); + } + return result; +} + +Oid +PGDuckDBEnum::GetEnumTypeOid(const Vector &oids) { + /* Get the pg_type tuple for the enum type */ + auto member_oid = duckdb::FlatVector::GetData(oids)[0]; + auto tuple = PostgresFunctionGuard(SearchSysCache1, ENUMOID, ObjectIdGetDatum(member_oid)); + Oid result = InvalidOid; + if (!HeapTupleIsValid(tuple)) { + throw duckdb::InvalidInputException("Cache lookup failed for enum member with oid %d", member_oid); + } + + auto enum_form = (Form_pg_enum)GETSTRUCT(tuple); + result = enum_form->enumtypid; + + /* Release the cache tuple */ + PostgresFunctionGuard(ReleaseSysCache, tuple); + return result; +} + +const Vector & +PGDuckDBEnum::GetMemberOids(const duckdb::LogicalType &type) { + auto type_info = type.AuxInfo(); + auto &enum_type_info = type_info->Cast(); + + switch (enum_type_info.DictType(enum_type_info.GetDictSize())) { + case PhysicalType::UINT8: + return enum_type_info.Cast>().GetMemberOids(); + case PhysicalType::UINT16: + return enum_type_info.Cast>().GetMemberOids(); + case PhysicalType::UINT32: + return enum_type_info.Cast>().GetMemberOids(); + default: + throw duckdb::InternalException("Invalid Physical Type for ENUMs"); + } +} + +} // namespace pgduckdb diff --git a/test/regression/expected/enum.out b/test/regression/expected/enum.out new file mode 100644 index 00000000..9145cca0 --- /dev/null +++ b/test/regression/expected/enum.out @@ -0,0 +1,54 @@ +CREATE TYPE mood AS ENUM ( + 'sad', + 'neutral', + 'happy' +); +create table tbl (id serial primary key, a mood); +insert into tbl(a) select 'happy'; +insert into tbl(a) select 'neutral'; +insert into tbl(a) select 'sad'; +select * from tbl; + id | a +----+--------- + 1 | happy + 2 | neutral + 3 | sad +(3 rows) + +select * from tbl where a = 'sad'; -- returns `sad` + id | a +----+----- + 3 | sad +(1 row) + +select * from tbl where a > 'neutral'; --- returns `happy` + id | a +----+------- + 1 | happy +(1 row) + +ALTER TYPE mood ADD VALUE 'a-bit-happy' BEFORE 'happy'; +ALTER TYPE mood ADD VALUE 'very-happy' AFTER 'happy'; +ALTER TYPE mood ADD VALUE 'very-sad' BEFORE 'sad'; +insert into tbl(a) values ('very-sad'), ('very-happy'), ('a-bit-happy'); +select * from tbl where a > 'neutral'; --- returns `a-bit-happy`, `happy` and `very-happy`; + id | a +----+------------- + 1 | happy + 5 | very-happy + 6 | a-bit-happy +(3 rows) + +select * from tbl; + id | a +----+------------- + 1 | happy + 2 | neutral + 3 | sad + 4 | very-sad + 5 | very-happy + 6 | a-bit-happy +(6 rows) + +drop table tbl; +drop type mood; diff --git a/test/regression/schedule b/test/regression/schedule index 360ad831..428e136d 100644 --- a/test/regression/schedule +++ b/test/regression/schedule @@ -12,4 +12,5 @@ test: duckdb_only_functions test: cte test: create_table_as test: standard_conforming_strings +test: enum test: query_filter diff --git a/test/regression/sql/enum.sql b/test/regression/sql/enum.sql new file mode 100644 index 00000000..f445ec45 --- /dev/null +++ b/test/regression/sql/enum.sql @@ -0,0 +1,26 @@ +CREATE TYPE mood AS ENUM ( + 'sad', + 'neutral', + 'happy' +); +create table tbl (id serial primary key, a mood); + +insert into tbl(a) select 'happy'; +insert into tbl(a) select 'neutral'; +insert into tbl(a) select 'sad'; + +select * from tbl; + +select * from tbl where a = 'sad'; -- returns `sad` +select * from tbl where a > 'neutral'; --- returns `happy` + +ALTER TYPE mood ADD VALUE 'a-bit-happy' BEFORE 'happy'; +ALTER TYPE mood ADD VALUE 'very-happy' AFTER 'happy'; +ALTER TYPE mood ADD VALUE 'very-sad' BEFORE 'sad'; + +insert into tbl(a) values ('very-sad'), ('very-happy'), ('a-bit-happy'); +select * from tbl where a > 'neutral'; --- returns `a-bit-happy`, `happy` and `very-happy`; +select * from tbl; + +drop table tbl; +drop type mood;