Hi! A few days ago I saved a parquet dataset with

FixedShapeTensor

columns, but now when I check the schema, it looks like daft doesn't recognize its native datatype extensions anymore:

Extension[arrow.fixed_shape_tensor; FixedSizeList[UInt8; 3174400]]

Will this affect the performance or correctness of the column parsing? Is it related to the versions of daft/pyarrow/python I used to write the dataset?

Hi Sagi, thanks for reporting this. Did Daft recognize the data type before? And if so, what changes (versioning or otherwise) have been made in between this change of behavior?

It is possible that I wrote the daft dataframe with pyarrow since I had some issues with the

write_parquet

method. I assumed that all the needed information is kept in the arrow extension dtypes. If I read with daft a file that was written with pyarrow, the schema can't be inferred:

import daft
import tempfile
import numpy as np
import pyarrow.parquet as pq


df = daft.from_pydict({"tensor": [np.ones((2,2), dtype=np.uint8)]})

with tempfile.NamedTemporaryFile(suffix=".parquet") as f:
    pq.write_table(df.to_arrow(), f.name)
    
    df = daft.read_parquet(f.name)
    print(df.schema())
    
    df = daft.from_arrow(pq.read_table(f.name))
    print(df.schema())`

╭─────────────┬────────────────────────────────────────────────╮
│ Column Name ┆ Type                                           │
╞═════════════╪════════════════════════════════════════════════╡
│ tensor      ┆ Struct[data: List[UInt8], shape: List[UInt64]] │
╰─────────────┴────────────────────────────────────────────────╯

╭─────────────┬───────────────╮
│ Column Name ┆ Type          │
╞═════════════╪═══════════════╡
│ tensor      ┆ Tensor(UInt8) │

I see. It looks like if you write directly with Daft (e.g.

df.write_parquet(...)

) it seems to infer the schema just fine with

daft.read_parquet

. I'll take a look at why there is a discrepancy. Thanks for the info!