Fix issue with returning incorrect datatype