[Bug]: load_dataset: Some datasets raise a python error

[khanlou]

What happened?

I was able to get the docker container loaded and running, and while some Hugging Face datasets load fine (like AWeirdDev/airports), I have one (very large) dataset that I'm having trouble getting imported: foursquare/fsq-os-places.

When I try to import it, here's what I see:

postgres=# select ai.load_dataset('foursquare/fsq-os-places', table_name => 'hf_places');
ERROR:  AttributeError: 'dict' object has no attribute 'dtype'
CONTEXT:  Traceback (most recent call last):
  PL/Python function "load_dataset", line 28, in <module>
    return ai.load_dataset.load_dataset(
  PL/Python function "load_dataset", line 194, in load_dataset
  PL/Python function "load_dataset", line 39, in get_column_info
  PL/Python function "load_dataset", line 39, in <dictcomp>
PL/Python function "load_dataset"

It appears to be some kind of object/dictionary mixup?

I also tried using load_dataset_multi_txn, which results in the same error.

pgai extension affected

0.6.0

pgai library affected

No response

PostgreSQL version used

17.2

What operating system did you use?

Ubuntu 24.04.1

What installation method did you use?

Docker

What platform did you run on?

Other

Relevant log output and stack trace

No response

How can we reproduce the bug?

Run `select ai.load_dataset('foursquare/fsq-os-places', table_name => 'hf_places');`

Are you going to work on the bugfix?

🆘 No, could someone else please work on the bugfix?