@@ -222,6 +222,23 @@ def _arrow_size(schema: pa.Schema, num_rows: int) -> int:
222222 return num_rows * sum (widths )
223223
224224
225+ def _shares_memory_with_arrow (values : np .ndarray ) -> bool :
226+ """Whether ``values``'s Arrow column, its rows in C order, can be its
227+ own memory rather than a copy.
228+
229+ Arrow lays out a number or a datetime as NumPy does, in the machine's
230+ byte order. It packs booleans into bits, and keeps strings in buffers
231+ of its own, so those are copies; so is an array that isn't
232+ C-contiguous, such as one in Fortran order or a transposed or strided
233+ view, whose rows in C order aren't one run of memory.
234+ """
235+ return (
236+ values .dtype .kind in "iufmM"
237+ and values .dtype .isnative
238+ and values .flags .c_contiguous
239+ )
240+
241+
225242def _in_memory_partitions (
226243 ds : xr .Dataset ,
227244 chunks : Chunks = None ,
@@ -241,23 +258,29 @@ def _in_memory_partitions(
241258
242259 The rows are in C order, and each partition is a contiguous range of
243260 them, as many partitions as ``chunks`` makes chunks, up to the number
244- of CPUs. So a data variable whose array is C-contiguous isn't copied,
245- whatever the chunks: its Arrow column shares the array's memory, as
246- does DataFusion's table, through Arrow's C data interface. A NaN adds a
247- validity bitmap beside it, one bit per value, made here: a NaN written
248- into the array later isn't read as NULL. What the table adds is its
249- dimension columns, one value per row for each dimension, which
250- long-form rows need. An array that isn't C-contiguous (Fortran order)
251- is copied, once.
252-
253- Returns ``None`` for a Dataset that isn't in memory (a Dask array or a
254- lazily loaded file), is larger than ``limit``, or has no rows; it is
255- read lazily instead.
261+ of CPUs. So no data variable is copied, whatever the chunks: its Arrow
262+ column shares the array's memory, as does DataFusion's table, through
263+ Arrow's C data interface, and a write into the array in place shows in
264+ the next query, as it does for a lazy table. A NaN adds a validity
265+ bitmap beside it, one bit per value, made here: a NaN written into the
266+ array later isn't read as NULL. What the table adds is its dimension
267+ columns, one value per row for each dimension, which long-form rows
268+ need.
269+
270+ Returns ``None``, for the Dataset to be read lazily instead, when it
271+ isn't in memory (a Dask array or a lazily loaded file), is larger than
272+ ``limit``, or has no rows; or when a data variable's Arrow column
273+ would have to be a copy of its array, which a write in place wouldn't
274+ reach (see `_shares_memory_with_arrow`).
256275 """
257276 # `_in_memory` is xarray's own test for a variable whose values are
258277 # loaded: a NumPy array, not a Dask array or a lazily indexed file.
259278 if not all (var ._in_memory for var in ds .variables .values ()):
260279 return None
280+ if not all (
281+ _shares_memory_with_arrow (var .values ) for var in ds .data_vars .values ()
282+ ):
283+ return None
261284 # TODO(#286): register a dimension that no data variable has (only a
262285 # coordinate does) as the lazy reader should, then drop this fallback.
263286 data_vars = list (ds .data_vars .values ())
0 commit comments