In-memory tables: an in-place write that moves a NaN reads wrongly (the validity bitmap is computed at registration) #284

🌀 Claude (xarray-sql)

Summary

A small in-memory Dataset is registered as a native in-memory table that shares its data variables' memory (#275). from_dataset's docstring says that, as with a lazy table, "changing the Dataset's values in place changes what queries return". That holds for numbers, but not for NaN. xarray-sql reads a NaN as SQL NULL, and the in-memory table works out which values are NULL once, when it is registered, as a validity bitmap beside the shared values. After an in-place write:

• a NaN written where there was a number reads as NaN, not NULL;
• a number written where there was a NaN reads as NULL.

A lazy table, or the same Dataset registered again, reads the new values as expected.

Reproduction

On #275's branch (in-memory-tables), DataFusion 55.0.0:

import numpy as np
import xarray as xr
import xarray_sql as xql

x = np.array([1.0, np.nan, 3.0])
ctx = xql.Context()
ctx.from_dataset("t", xr.Dataset({"x": ("i", x)}, coords={"i": np.arange(3)}))
query = "SELECT i, x, x IS NULL AS null FROM t ORDER BY i"
print(ctx.sql(query).to_pydict())
# {'i': [0, 1, 2], 'x': [1.0, None, 3.0], 'null': [False, True, False]}

x[0] = np.nan
x[1] = 2.0
print(ctx.sql(query).to_pydict())
# {'i': [0, 1, 2], 'x': [nan, None, 3.0], 'null': [False, True, False]}

again = xql.Context()
again.from_dataset("t", xr.Dataset({"x": ("i", x)}, coords={"i": np.arange(3)}))
print(again.sql(query).to_pydict())
# {'i': [0, 1, 2], 'x': [None, 2.0, 3.0], 'null': [True, False, False]}

Options

• Say so in the docstring: an in-place write must keep its NaNs where they were, or the Dataset must be registered again.
• Keep NaN as a value in in-memory tables, with no bitmap (Optional mode to preserve floating NaNs instead of constructing Arrow nulls #251 asks for that mode in general). That would differ from lazy tables until Optional mode to preserve floating NaNs instead of constructing Arrow nulls #251 lands for both.

The same limit applies to Context.prepare, which runs a planned query again on values written in place. It checks that a write keeps its NaNs where they were.