import duckdb
import pyarrow as pa
import uuid

###
### Create an extension type (using the example from the pyarrow docs)
###
class UuidType(pa.ExtensionType):

    def __init__(self):
        super().__init__(pa.binary(16), "my_package.uuid")

    def __arrow_ext_serialize__(self):
        # Since we don't have a parameterized type, we don't need extra                                                                                                                                            
        # metadata to be deserialized                                                                                                                                                                              
        return b''

    @classmethod
    def __arrow_ext_deserialize__(cls, storage_type, serialized):
        # Sanity checks, not required but illustrate the method signature.                                                                                                                                         
        assert storage_type == pa.binary(16)
        assert serialized == b''
        # Return an instance of this subclass given the serialized                                                                                                                                                 
        # metadata.                                                                                                                                                                                                
        return UuidType()

storage_array = pa.array([uuid.uuid4().bytes for _ in range(4)], pa.binary(16))
arr = pa.ExtensionArray.from_storage(UuidType(), storage_array)

###
### Create a table with all columns (duckdb can't query this) ###
###
tab = pa.table({"uuids": arr, "normal_type": [1, 2, 3, 4]})

###
### Prune my table to just include the columns duckdb can handle ###
###
pruned_tab = tab.select(["normal_type"])

# This query succeeds
print(duckdb.sql("SELECT normal_type from pruned_tab").fetchall())
# This query fails with the error
#   duckdb.duckdb.NotImplementedException: Not implemented Error: Arrow Type with extension name: my_package.uuid and format: w:16, is not currently supported in DuckDB
print(duckdb.sql("SELECT normal_type FROM tab").fetchall())