[BUG] Register super extension on to_arrow

[colin-ho]

This is an issue where Daft Extension types were not getting converted to PyArrow properly. @jaychia discovered this while trying to write parquet with a tensor column, where the Extension metadata for tensor was getting dropped.

A simple test to reproduce the error:

import daft
import numpy as np
from daft import Series

# Create sample tensor data with some null values
tensor_data = [np.array([[1, 2], [3, 4]]), None, None]

# Uncomment this and it will work
# from daft.datatype import _ensure_registered_super_ext_type 
# _ensure_registered_super_ext_type()

df_original = daft.from_pydict({"tensor_col": Series.from_pylist(tensor_data)})
print(df_original.to_arrow().schema)

Output:

tensor_col: struct<data: large_list<item: int64>, shape: large_list<item: uint64>>
  child 0, data: large_list<item: int64>
      child 0, item: int64
  child 1, shape: large_list<item: uint64>
      child 0, item: uint64

It's not a tensor type! However if you uncomment the _ensure_registered_super_ext_type(), you will now see:

tensor_col: extension<daft.super_extension<DaftExtension>>

The issue here is that the class DaftExtension(pa.ExtensionType): is not imported during the FFI, as it is now a lazy import that must be called via _ensure_registered_super_ext_type().

This PR adds calls to this import in to_arrow for series and schema. However, I do not know if this is exhaustive, and I will give this more thought. @desmondcheongzx @samster25

[codspeed-hq]

CodSpeed Performance Report

Merging #3030 will not alter performance

_{Comparing colin/register-ext (225ea82) with main (f5d365b)}

Summary

✅ 17 untouched benchmarks

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 78.47%. Comparing base (9ae8122) to head (225ea82).
Report is 3 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3030      +/-   ##
==========================================
- Coverage   78.47%   78.47%   -0.01%     
==========================================
  Files         610      610              
  Lines       71740    71746       +6     
==========================================
+ Hits        56297    56301       +4     
- Misses      15443    15445       +2     

Files with missing lines	Coverage Δ
daft/logical/schema.py	92.24% <100.00%> (+0.06%)	⬆️
daft/series.py	89.75% <100.00%> (+0.01%)	⬆️

... and 3 files with indirect coverage changes

[desmondcheongzx]

Ah shucks yeah, I guess we missed these in the lazy import PR.

Do we have roundtrip tests that would catch this btw?

[jaychia]

Ah shucks yeah, I guess we missed these in the lazy import PR.

Do we have roundtrip tests that would catch this btw?

We do, but they will pass if the modules happen to be lazily imported by earlier test code.