NaN not converted to Arrow null after PR #134, breaking SQL aggregations on missing data

[ghostiee-11]

Description

After PR #134 (Convert xarray partition to pyarrow record batch directly without going through pandas), NaN values in float arrays are no longer converted to Arrow nulls. This breaks standard SQL semantics for aggregations and filtering on missing data.

Before PR #134: Data went through pa.RecordBatch.from_pandas(df), which automatically converts NaN to Arrow null.

After PR #134: Data goes directly via pa.array(numpy_array, type=field.type), which preserves NaN as a regular float value with null_count=0.

Reproduction

import xarray as xr
import numpy as np
from xarray_sql import XarrayContext

ds = xr.Dataset({
    'temp': (['x'], [1.0, np.nan, 3.0, np.nan, 5.0]),
}, coords={'x': np.arange(5)})

ctx = XarrayContext()
ctx.from_dataset('temp', ds, chunks={'x': 5})

ctx.sql('SELECT AVG(temp) FROM temp').to_pandas()
# Returns NaN - expected 3.0

ctx.sql('SELECT COUNT(temp) FROM temp').to_pandas()
# Returns 5 - expected 3

ctx.sql('SELECT * FROM temp WHERE temp IS NULL').to_pandas()
# Returns 0 rows - expected 2

Root cause

pa.array() and pa.RecordBatch.from_pandas() handle NaN differently:

import pyarrow as pa
import numpy as np

arr = np.array([1.0, np.nan, 3.0])

pa.array(arr, type=pa.float64())
# [1, nan, 3] - null_count: 0

pa.RecordBatch.from_pandas(pd.DataFrame({'v': arr})).column('v')
# [1, null, 3] - null_count: 1

This affects four call sites in xarray_sql/df.py:

• dataset_to_record_batch(): lines 204, 207
• iter_record_batches(): lines 279, 282

Suggested fix

Pass a NaN mask for float arrays:

# For float data, convert NaN to Arrow null
if arr.dtype.kind == 'f':
    pa.array(arr, mask=np.isnan(arr), type=field.type)
else:
    pa.array(arr, type=field.type)

pa.array with mask= is zero-copy for the data portion, so there's no performance regression.

Impact

Scientific datasets almost always have missing values represented as NaN. Without this fix, AVG, SUM, COUNT, MIN, MAX all return incorrect results on real-world data.

[alxmrs]

Thanks, this is a helpful issue.

[ghostiee-11]

NP!!