Add chunks='auto' support for cftime datasets - #10527

Merged
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime
Oct 15, 2025
Merged

Add chunks='auto' support for cftime datasets#10527
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime

Conversation

@charles-turner-1

@charles-turner-1charles-turner-1 commented Jul 13, 2025

Copy link
Copy Markdown
Contributor

@welcome

welcomeBot commented Jul 13, 2025

Copy link
Copy Markdown

Thank you for opening this pull request! It may take us a few days to respond here, so thank you for being patient.
If you have questions, some answers may be found in our contributing guidelines.

@github-actionsgithub-actionsBot added topic-documentation topic-NamedArray Lightweight version of Variable labels Jul 13, 2025
@charles-turner-1charles-turner-1 changed the title All works, just need to satisfy mypy and whatnot nowAdd chunks='auto' support for cftime datasetsJul 13, 2025
@jemmajeffree

Copy link
Copy Markdown
Contributor

Would these changes also work for cf timedeltas or are they going to still cause problems?
I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

Would these changes also work for cf timedeltas or are they going to still cause problems? I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

If you can find something thats specifically a cftimedelta and run the _contains_cftime_datetimes function on it that'd be super helpful to know whether it returns True or False.

@charles-turner-1
charles-turner-1 marked this pull request as draft July 14, 2025 05:02
@jemmajeffree

Copy link
Copy Markdown
Contributor

TLDR: don't mind me, it's not going to cause any issues

Firstly, what I thought was a cftimedelta turned out to be a numpy timedelta hanging out with a cftime
Screenshot 2025-07-14 at 5 23 31 pm
When I did manage to coerce this timedelta into cftime conventions, it just contained a floating point number of days, so I can't see anything having issues with its size

coder=xr.coding.times.CFTimedeltaCoder()
result=coder.encode(oops.average_DT).load()
print(result.dtype)
result
Screenshot 2025-07-14 at 5 38 33 pm

Comment threadxarray/namedarray/daskmanager.py Outdated
Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

I did some prodding around yesterday and I realised this won't let us do something like

importxarrayasxrcftime_datafile="/path/to/file.nc"xr.open_dataset(cftime_datafile, chunks='auto')

yet, only stuff along the lines of

importxarrayasxrcftime_datafile="/path/to/file.nc"ds=xr.open_dataset(cftime_datafile, chunks=-1)
ds=ds.chunk('auto')

I think implementing the former is going to be a bit harder, but I'm starting to clock the code structure a bit more now so I'll have a decent crack.

@dcherian

Copy link
Copy Markdown
Contributor

Why so? Are we sending "auto" in to normalize_chunks first?

@charles-turner-1

charles-turner-1 commented Jul 23, 2025

Copy link
Copy Markdown
ContributorAuthor

Yup, this is the call stack:

---->3xr.open_dataset(
4"/Users/u1166368/xarray/tos_Omon_CESM2-WACCM_historical_r2i1p1f1_gr_185001-201412.nc", chunks="auto"/Users/u1166368/xarray/xarray/backends/api.py(721)open_dataset()
720 )
-->721ds=_dataset_from_backend_dataset(
722backend_ds,
/Users/u1166368/xarray/xarray/backends/api.py(418)_dataset_from_backend_dataset()
417ifchunksisnotNone:
-->418ds=_chunk_ds(
419ds,
/Users/u1166368/xarray/xarray/backends/api.py(368)_chunk_ds()
367forname, varinbackend_ds.variables.items():
-->368var_chunks=_get_chunk(var, chunks, chunkmanager)
369variables[name] =_maybe_chunk(
/Users/u1166368/xarray/xarray/structure/chunks.py(102)_get_chunk()
101-->102chunk_shape=chunkmanager.normalize_chunks(
103chunk_shape, shape=shape, dtype=var.dtype, previous_chunks=preferred_chunk_shape>/Users/u1166368/xarray/xarray/namedarray/daskmanager.py(60)normalize_chunks()

I've fixed it in the latest commit - but I think the implementation leaves a lot to be desired too.

Do I want to refactor to move the changes in xarray/structure/chunks.py into the daskmanager module if possible?

Once I've got the structure there cleaned up, I'll work on replacing the build_chunkspec function with something more sensible - I just need to work out how to extract the implementation in dask cleanly now I think - normalize_chunks also seems to calculate sensible chunk sizes.

Comment threadxarray/structure/chunks.py Outdated

from xarray.namedarray.utils import build_chunkspec

target_chunksize = parse_bytes(dask_config.get("array.chunk-size"))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about adding get_auto_chunk_size to the ChunkManager class; and put the dask-specific stuff in the DaskManager.

cc @TomNicholas

@dcherian

dcherian commented Jul 23, 2025

Copy link
Copy Markdown
Contributor

I guess one bit that's confusing here is that the code-path for backends and normal variables is different?

So let's add a test that reads form disk; and one that works iwth a DataArray constructed in memory.

Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

charles-turner-1 commented Sep 24, 2025

Copy link
Copy Markdown
ContributorAuthor

It looks like the failing test (same one I commented on above) might be flaky? Now only failing for windows & python3.11, not 3.13: https://github.com/pydata/xarray/actions/runs/17962030991/job/51087173288

raise NotImplementedError("Only chunks='auto' is supported at present.")
return dask.array.shuffle(x, indexer, axis, chunks="auto")

def get_auto_chunk_size(self) -> int:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tomwhite is there an equivalent for cubed? I didn't see it in the docs...

Comment on lines +232 to +244
if _contains_cftime_datetimes(data):
limit, dtype = fake_target_chunksize(data, chunkmanager.get_auto_chunk_size())
else:
limit = None
dtype = data.dtype

chunk_shape = chunkmanager.normalize_chunks(
chunk_shape,
shape=shape,
dtype=dtype,
limit=limit,
previous_chunks=preferred_chunk_shape,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does this seem fine to you @charles-turner-1 . I wanted to avoid calling get_auto_chunk_size as much as possible

@charles-turner-1charles-turner-1Oct 13, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, looks good! fake_target_chunksize also contains the same _contains_cf_datetimes check & early return if it's false, so we could remove the check in either fake_target_chunksize or here without causing issues if you think that's a good idea?

I'm guessing you meant calling fake_target_chunksize in your comment above, in which case we would probably want to either remove it in that function - or leave it in if we want to reuse fake_target_chunksize elsewhere?

@dcheriandcherian left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Phew, I think this is good to go. It would be good to clean up the types, but this PR has stalled for a long time.

Apologies for the delay (again). I was on vacation.

@dcheriandcherian added the plan to merge Final call for comments label Oct 13, 2025
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

No worries, thanks for all your help!

I'd love to keep getting my feet wet - do you happen to know if there are any other extant issues in roughly the same parts of the codebase off the top of your head? If not I'll go digging soon!

@dcherian

Copy link
Copy Markdown
Contributor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

@dcherian
dcherian merged commit 94798a0 into pydata:mainOct 15, 2025
45 of 47 checks passed
@welcome

welcomeBot commented Oct 15, 2025

Copy link
Copy Markdown

Congratulations on completing your first pull request! Welcome to Xarray! We are proud of you, and hope to see you again! celebration gif

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

🙏 I'll have a crack!

@spencerkclark

Copy link
Copy Markdown
Member

This is great—thanks @charles-turner-1 and @dcherian!

@Thomas-Moore-Creative

Copy link
Copy Markdown

This is great—thanks @charles-turner-1 and @dcherian!

Go Australia 🇦🇺 ( AKA @charles-turner-1 ), pulling our weight! 😉


output_dtype = np.dtype(np.float64)

nbytes_approx: int = sys.getsizeof(first_n_items(data, 1)) # type: ignore[no-untyped-call]

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm just came across this and I'm not quite sure it's the right size. I think sys.getsizeof is the in-memory size and and dtype.itemsize is the uncompressed disk size. Consider for instance:

importsysimportnumpyasnpimportcftimenp.dtype(np.float64).itemsize# 8sys.getsizeof(np.float64(1.0)) # 32sys.getsizeof(np.array([1.0], dtype=np.float64)) # 120sys.getsizeof(cftime.DatetimeGregorian.fromordinal(2450000)) #112

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm kind of wondering if setting the dtype to np.dtype(np.float64) would suffice

@charles-turner-1charles-turner-1Dec 16, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is for the assumed size as a float64 right? I think what you're saying is true but the array overhead rapidly becomes unimportant for reasonably large arrays? Very rough & ready analysis below:

# Does this still matter for decently sized arrays?importmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 8, num=50, dtype=int):
arr=np.zeros(n, dtype=np.float64)
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Exactly! Calling sys.getsizeof on an array containing a single cftime object is not going to be a good representation of the memory consumption of an array of these things. Even if you pop the object out of the array that is still not really a good representation of the memory consumption. I think you'd do better with just nbytes_approx: int = 8

I made that same plot you did but with cftimes inside the array:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 4, num=50, dtype=int):
arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My bad - I thought we'd popped the first cftime element out of the array and had a look at its size at that point.

It looks like the cftime elements in the array are 8bytes too - is that what we expect? I would have expected them to be a bit larger due to the extra overhead...

Assuming I'm wrong about that, it would be much simpler to just tell dask that a cftime is 8 bytes and leave the limit unadjusted - the ratio of two line should be pretty much 1 for all decently sized arrays.

On my phone right now but I'll have a proper look when I get to my computer.

@charles-turner-1charles-turner-1Dec 18, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So it does look like size per element in a numpy array is reliably 8 bytes, but I'm really unconvinced this can be correct tbh:

importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltcf_sizes : list[int] = []
f64_sizes : list[int] = []
num_elements :list[int] = []
forninnp.logspace(0, 4, num=50, dtype=int):
cf_arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
cf_sizes.append(sys.getsizeof(cf_arr)/n)
num_elements.append(n)
arr=np.zeros(n, dtype=np.float64)
f64_sizes.append(sys.getsizeof(arr)/n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
ratio= [s_cf/s_f64fors_cf, s_f64inzip(cf_sizes, f64_sizes)]
plt.plot(num_elements, cf_sizes, marker='o', label='cftime')
plt.plot(num_elements, f64_sizes, marker='o', label='float64')
plt.plot(num_elements, ratio, marker='o', label='cftime/float64 ratio')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte line, unit ratio lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.axhline(y=1, color='grey', linestyle='--', label='1')
plt.legend()
image

But if we look at the raw element (same as you did above):

>>>t=cftime.DatetimeGregorian.fromordinal(2450000)
>>>sys.getsizeof(t)
112

Since it's just not possible that the cftime objects are magically shrinking when we put them in a numpy array, I assume numpy is storing pointers to objects somewhere on the heap.

I've run a couple of more sophisticated (this does not necessarily mean more likely to be right!) tests here:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltimporttracemallocimportgccf_sizes : list[int] = []
numel :list[int] = []
forninnp.logspace(0, 5, num=50, dtype=int):
gc.collect()
tracemalloc.start()
snap1=tracemalloc.take_snapshot()
cf_arr=cftime.DatetimeGregorian.fromordinal(np.arange(2450000, 2450000+n))
snap2=tracemalloc.take_snapshot()
stats=snap2.compare_to(snap1, 'lineno')
tracemalloc.stop()
tot=sum(stat.size_diffforstatinstats)
cf_sizes.append(tot/n)
numel.append(n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot(numel, cf_sizes, marker='o', label='cftime')
plt.xscale('log')
plt.yscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.legend()
print(f"nbytes asymptotes to {cf_sizes[-1]:.2f} for large arrays")

nbytes asymptotes to 120.12 for large arrays
image

I'm still not convinced this is the right number, so I'm still digging. But it looks like we (accidentally/serendipitously) might have gotten in the right ballpark?


EDIT: I've more some more playing, and I reckon we're off by approximately a factor of 2-2.5 ish. No real justification yet, just empirical results.

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ioplan to mergeFinal call for commentstopic-backendstopic-documentationtopic-NamedArrayLightweight version of Variable

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Confusing error when use_cftime = True and chunks = 'auto' in xr.open_dataset()

6 participants

@charles-turner-1@jemmajeffree@dcherian@spencerkclark@Thomas-Moore-Creative@jsignell
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
Skip to content

Add chunks='auto' support for cftime datasets - #10527

Merged
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime
Oct 15, 2025
Merged

Add chunks='auto' support for cftime datasets#10527
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime

Conversation

@charles-turner-1

@charles-turner-1charles-turner-1 commented Jul 13, 2025

Copy link
Copy Markdown
Contributor

@welcome

welcomeBot commented Jul 13, 2025

Copy link
Copy Markdown

Thank you for opening this pull request! It may take us a few days to respond here, so thank you for being patient.
If you have questions, some answers may be found in our contributing guidelines.

@github-actionsgithub-actionsBot added topic-documentation topic-NamedArray Lightweight version of Variable labels Jul 13, 2025
@charles-turner-1charles-turner-1 changed the title All works, just need to satisfy mypy and whatnot nowAdd chunks='auto' support for cftime datasetsJul 13, 2025
@jemmajeffree

Copy link
Copy Markdown
Contributor

Would these changes also work for cf timedeltas or are they going to still cause problems?
I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

Would these changes also work for cf timedeltas or are they going to still cause problems? I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

If you can find something thats specifically a cftimedelta and run the _contains_cftime_datetimes function on it that'd be super helpful to know whether it returns True or False.

@charles-turner-1
charles-turner-1 marked this pull request as draft July 14, 2025 05:02
@jemmajeffree

Copy link
Copy Markdown
Contributor

TLDR: don't mind me, it's not going to cause any issues

Firstly, what I thought was a cftimedelta turned out to be a numpy timedelta hanging out with a cftime
Screenshot 2025-07-14 at 5 23 31 pm
When I did manage to coerce this timedelta into cftime conventions, it just contained a floating point number of days, so I can't see anything having issues with its size

coder=xr.coding.times.CFTimedeltaCoder()
result=coder.encode(oops.average_DT).load()
print(result.dtype)
result
Screenshot 2025-07-14 at 5 38 33 pm

Comment threadxarray/namedarray/daskmanager.py Outdated
Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

I did some prodding around yesterday and I realised this won't let us do something like

importxarrayasxrcftime_datafile="/path/to/file.nc"xr.open_dataset(cftime_datafile, chunks='auto')

yet, only stuff along the lines of

importxarrayasxrcftime_datafile="/path/to/file.nc"ds=xr.open_dataset(cftime_datafile, chunks=-1)
ds=ds.chunk('auto')

I think implementing the former is going to be a bit harder, but I'm starting to clock the code structure a bit more now so I'll have a decent crack.

@dcherian

Copy link
Copy Markdown
Contributor

Why so? Are we sending "auto" in to normalize_chunks first?

@charles-turner-1

charles-turner-1 commented Jul 23, 2025

Copy link
Copy Markdown
ContributorAuthor

Yup, this is the call stack:

---->3xr.open_dataset(
4"/Users/u1166368/xarray/tos_Omon_CESM2-WACCM_historical_r2i1p1f1_gr_185001-201412.nc", chunks="auto"/Users/u1166368/xarray/xarray/backends/api.py(721)open_dataset()
720 )
-->721ds=_dataset_from_backend_dataset(
722backend_ds,
/Users/u1166368/xarray/xarray/backends/api.py(418)_dataset_from_backend_dataset()
417ifchunksisnotNone:
-->418ds=_chunk_ds(
419ds,
/Users/u1166368/xarray/xarray/backends/api.py(368)_chunk_ds()
367forname, varinbackend_ds.variables.items():
-->368var_chunks=_get_chunk(var, chunks, chunkmanager)
369variables[name] =_maybe_chunk(
/Users/u1166368/xarray/xarray/structure/chunks.py(102)_get_chunk()
101-->102chunk_shape=chunkmanager.normalize_chunks(
103chunk_shape, shape=shape, dtype=var.dtype, previous_chunks=preferred_chunk_shape>/Users/u1166368/xarray/xarray/namedarray/daskmanager.py(60)normalize_chunks()

I've fixed it in the latest commit - but I think the implementation leaves a lot to be desired too.

Do I want to refactor to move the changes in xarray/structure/chunks.py into the daskmanager module if possible?

Once I've got the structure there cleaned up, I'll work on replacing the build_chunkspec function with something more sensible - I just need to work out how to extract the implementation in dask cleanly now I think - normalize_chunks also seems to calculate sensible chunk sizes.

Comment threadxarray/structure/chunks.py Outdated

from xarray.namedarray.utils import build_chunkspec

target_chunksize = parse_bytes(dask_config.get("array.chunk-size"))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about adding get_auto_chunk_size to the ChunkManager class; and put the dask-specific stuff in the DaskManager.

cc @TomNicholas

@dcherian

dcherian commented Jul 23, 2025

Copy link
Copy Markdown
Contributor

I guess one bit that's confusing here is that the code-path for backends and normal variables is different?

So let's add a test that reads form disk; and one that works iwth a DataArray constructed in memory.

Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

charles-turner-1 commented Sep 24, 2025

Copy link
Copy Markdown
ContributorAuthor

It looks like the failing test (same one I commented on above) might be flaky? Now only failing for windows & python3.11, not 3.13: https://github.com/pydata/xarray/actions/runs/17962030991/job/51087173288

raise NotImplementedError("Only chunks='auto' is supported at present.")
return dask.array.shuffle(x, indexer, axis, chunks="auto")

def get_auto_chunk_size(self) -> int:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tomwhite is there an equivalent for cubed? I didn't see it in the docs...

Comment on lines +232 to +244
if _contains_cftime_datetimes(data):
limit, dtype = fake_target_chunksize(data, chunkmanager.get_auto_chunk_size())
else:
limit = None
dtype = data.dtype

chunk_shape = chunkmanager.normalize_chunks(
chunk_shape,
shape=shape,
dtype=dtype,
limit=limit,
previous_chunks=preferred_chunk_shape,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does this seem fine to you @charles-turner-1 . I wanted to avoid calling get_auto_chunk_size as much as possible

@charles-turner-1charles-turner-1Oct 13, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, looks good! fake_target_chunksize also contains the same _contains_cf_datetimes check & early return if it's false, so we could remove the check in either fake_target_chunksize or here without causing issues if you think that's a good idea?

I'm guessing you meant calling fake_target_chunksize in your comment above, in which case we would probably want to either remove it in that function - or leave it in if we want to reuse fake_target_chunksize elsewhere?

@dcheriandcherian left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Phew, I think this is good to go. It would be good to clean up the types, but this PR has stalled for a long time.

Apologies for the delay (again). I was on vacation.

@dcheriandcherian added the plan to merge Final call for comments label Oct 13, 2025
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

No worries, thanks for all your help!

I'd love to keep getting my feet wet - do you happen to know if there are any other extant issues in roughly the same parts of the codebase off the top of your head? If not I'll go digging soon!

@dcherian

Copy link
Copy Markdown
Contributor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

@dcherian
dcherian merged commit 94798a0 into pydata:mainOct 15, 2025
45 of 47 checks passed
@welcome

welcomeBot commented Oct 15, 2025

Copy link
Copy Markdown

Congratulations on completing your first pull request! Welcome to Xarray! We are proud of you, and hope to see you again! celebration gif

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

🙏 I'll have a crack!

@spencerkclark

Copy link
Copy Markdown
Member

This is great—thanks @charles-turner-1 and @dcherian!

@Thomas-Moore-Creative

Copy link
Copy Markdown

This is great—thanks @charles-turner-1 and @dcherian!

Go Australia 🇦🇺 ( AKA @charles-turner-1 ), pulling our weight! 😉


output_dtype = np.dtype(np.float64)

nbytes_approx: int = sys.getsizeof(first_n_items(data, 1)) # type: ignore[no-untyped-call]

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm just came across this and I'm not quite sure it's the right size. I think sys.getsizeof is the in-memory size and and dtype.itemsize is the uncompressed disk size. Consider for instance:

importsysimportnumpyasnpimportcftimenp.dtype(np.float64).itemsize# 8sys.getsizeof(np.float64(1.0)) # 32sys.getsizeof(np.array([1.0], dtype=np.float64)) # 120sys.getsizeof(cftime.DatetimeGregorian.fromordinal(2450000)) #112

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm kind of wondering if setting the dtype to np.dtype(np.float64) would suffice

@charles-turner-1charles-turner-1Dec 16, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is for the assumed size as a float64 right? I think what you're saying is true but the array overhead rapidly becomes unimportant for reasonably large arrays? Very rough & ready analysis below:

# Does this still matter for decently sized arrays?importmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 8, num=50, dtype=int):
arr=np.zeros(n, dtype=np.float64)
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Exactly! Calling sys.getsizeof on an array containing a single cftime object is not going to be a good representation of the memory consumption of an array of these things. Even if you pop the object out of the array that is still not really a good representation of the memory consumption. I think you'd do better with just nbytes_approx: int = 8

I made that same plot you did but with cftimes inside the array:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 4, num=50, dtype=int):
arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My bad - I thought we'd popped the first cftime element out of the array and had a look at its size at that point.

It looks like the cftime elements in the array are 8bytes too - is that what we expect? I would have expected them to be a bit larger due to the extra overhead...

Assuming I'm wrong about that, it would be much simpler to just tell dask that a cftime is 8 bytes and leave the limit unadjusted - the ratio of two line should be pretty much 1 for all decently sized arrays.

On my phone right now but I'll have a proper look when I get to my computer.

@charles-turner-1charles-turner-1Dec 18, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So it does look like size per element in a numpy array is reliably 8 bytes, but I'm really unconvinced this can be correct tbh:

importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltcf_sizes : list[int] = []
f64_sizes : list[int] = []
num_elements :list[int] = []
forninnp.logspace(0, 4, num=50, dtype=int):
cf_arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
cf_sizes.append(sys.getsizeof(cf_arr)/n)
num_elements.append(n)
arr=np.zeros(n, dtype=np.float64)
f64_sizes.append(sys.getsizeof(arr)/n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
ratio= [s_cf/s_f64fors_cf, s_f64inzip(cf_sizes, f64_sizes)]
plt.plot(num_elements, cf_sizes, marker='o', label='cftime')
plt.plot(num_elements, f64_sizes, marker='o', label='float64')
plt.plot(num_elements, ratio, marker='o', label='cftime/float64 ratio')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte line, unit ratio lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.axhline(y=1, color='grey', linestyle='--', label='1')
plt.legend()
image

But if we look at the raw element (same as you did above):

>>>t=cftime.DatetimeGregorian.fromordinal(2450000)
>>>sys.getsizeof(t)
112

Since it's just not possible that the cftime objects are magically shrinking when we put them in a numpy array, I assume numpy is storing pointers to objects somewhere on the heap.

I've run a couple of more sophisticated (this does not necessarily mean more likely to be right!) tests here:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltimporttracemallocimportgccf_sizes : list[int] = []
numel :list[int] = []
forninnp.logspace(0, 5, num=50, dtype=int):
gc.collect()
tracemalloc.start()
snap1=tracemalloc.take_snapshot()
cf_arr=cftime.DatetimeGregorian.fromordinal(np.arange(2450000, 2450000+n))
snap2=tracemalloc.take_snapshot()
stats=snap2.compare_to(snap1, 'lineno')
tracemalloc.stop()
tot=sum(stat.size_diffforstatinstats)
cf_sizes.append(tot/n)
numel.append(n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot(numel, cf_sizes, marker='o', label='cftime')
plt.xscale('log')
plt.yscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.legend()
print(f"nbytes asymptotes to {cf_sizes[-1]:.2f} for large arrays")

nbytes asymptotes to 120.12 for large arrays
image

I'm still not convinced this is the right number, so I'm still digging. But it looks like we (accidentally/serendipitously) might have gotten in the right ballpark?


EDIT: I've more some more playing, and I reckon we're off by approximately a factor of 2-2.5 ish. No real justification yet, just empirical results.

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ioplan to mergeFinal call for commentstopic-backendstopic-documentationtopic-NamedArrayLightweight version of Variable

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Confusing error when use_cftime = True and chunks = 'auto' in xr.open_dataset()

6 participants

@charles-turner-1@jemmajeffree@dcherian@spencerkclark@Thomas-Moore-Creative@jsignell
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix"); } } catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

Add chunks='auto' support for cftime datasets - #10527

Merged
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime
Oct 15, 2025
Merged

Add chunks='auto' support for cftime datasets#10527
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime

Conversation

@charles-turner-1

@charles-turner-1charles-turner-1 commented Jul 13, 2025

Copy link
Copy Markdown
Contributor

@welcome

welcomeBot commented Jul 13, 2025

Copy link
Copy Markdown

Thank you for opening this pull request! It may take us a few days to respond here, so thank you for being patient.
If you have questions, some answers may be found in our contributing guidelines.

@github-actionsgithub-actionsBot added topic-documentation topic-NamedArray Lightweight version of Variable labels Jul 13, 2025
@charles-turner-1charles-turner-1 changed the title All works, just need to satisfy mypy and whatnot nowAdd chunks='auto' support for cftime datasetsJul 13, 2025
@jemmajeffree

Copy link
Copy Markdown
Contributor

Would these changes also work for cf timedeltas or are they going to still cause problems?
I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

Would these changes also work for cf timedeltas or are they going to still cause problems? I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

If you can find something thats specifically a cftimedelta and run the _contains_cftime_datetimes function on it that'd be super helpful to know whether it returns True or False.

@charles-turner-1
charles-turner-1 marked this pull request as draft July 14, 2025 05:02
@jemmajeffree

Copy link
Copy Markdown
Contributor

TLDR: don't mind me, it's not going to cause any issues

Firstly, what I thought was a cftimedelta turned out to be a numpy timedelta hanging out with a cftime
Screenshot 2025-07-14 at 5 23 31 pm
When I did manage to coerce this timedelta into cftime conventions, it just contained a floating point number of days, so I can't see anything having issues with its size

coder=xr.coding.times.CFTimedeltaCoder()
result=coder.encode(oops.average_DT).load()
print(result.dtype)
result
Screenshot 2025-07-14 at 5 38 33 pm

Comment threadxarray/namedarray/daskmanager.py Outdated
Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

I did some prodding around yesterday and I realised this won't let us do something like

importxarrayasxrcftime_datafile="/path/to/file.nc"xr.open_dataset(cftime_datafile, chunks='auto')

yet, only stuff along the lines of

importxarrayasxrcftime_datafile="/path/to/file.nc"ds=xr.open_dataset(cftime_datafile, chunks=-1)
ds=ds.chunk('auto')

I think implementing the former is going to be a bit harder, but I'm starting to clock the code structure a bit more now so I'll have a decent crack.

@dcherian

Copy link
Copy Markdown
Contributor

Why so? Are we sending "auto" in to normalize_chunks first?

@charles-turner-1

charles-turner-1 commented Jul 23, 2025

Copy link
Copy Markdown
ContributorAuthor

Yup, this is the call stack:

---->3xr.open_dataset(
4"/Users/u1166368/xarray/tos_Omon_CESM2-WACCM_historical_r2i1p1f1_gr_185001-201412.nc", chunks="auto"/Users/u1166368/xarray/xarray/backends/api.py(721)open_dataset()
720 )
-->721ds=_dataset_from_backend_dataset(
722backend_ds,
/Users/u1166368/xarray/xarray/backends/api.py(418)_dataset_from_backend_dataset()
417ifchunksisnotNone:
-->418ds=_chunk_ds(
419ds,
/Users/u1166368/xarray/xarray/backends/api.py(368)_chunk_ds()
367forname, varinbackend_ds.variables.items():
-->368var_chunks=_get_chunk(var, chunks, chunkmanager)
369variables[name] =_maybe_chunk(
/Users/u1166368/xarray/xarray/structure/chunks.py(102)_get_chunk()
101-->102chunk_shape=chunkmanager.normalize_chunks(
103chunk_shape, shape=shape, dtype=var.dtype, previous_chunks=preferred_chunk_shape>/Users/u1166368/xarray/xarray/namedarray/daskmanager.py(60)normalize_chunks()

I've fixed it in the latest commit - but I think the implementation leaves a lot to be desired too.

Do I want to refactor to move the changes in xarray/structure/chunks.py into the daskmanager module if possible?

Once I've got the structure there cleaned up, I'll work on replacing the build_chunkspec function with something more sensible - I just need to work out how to extract the implementation in dask cleanly now I think - normalize_chunks also seems to calculate sensible chunk sizes.

Comment threadxarray/structure/chunks.py Outdated

from xarray.namedarray.utils import build_chunkspec

target_chunksize = parse_bytes(dask_config.get("array.chunk-size"))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about adding get_auto_chunk_size to the ChunkManager class; and put the dask-specific stuff in the DaskManager.

cc @TomNicholas

@dcherian

dcherian commented Jul 23, 2025

Copy link
Copy Markdown
Contributor

I guess one bit that's confusing here is that the code-path for backends and normal variables is different?

So let's add a test that reads form disk; and one that works iwth a DataArray constructed in memory.

Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

charles-turner-1 commented Sep 24, 2025

Copy link
Copy Markdown
ContributorAuthor

It looks like the failing test (same one I commented on above) might be flaky? Now only failing for windows & python3.11, not 3.13: https://github.com/pydata/xarray/actions/runs/17962030991/job/51087173288

raise NotImplementedError("Only chunks='auto' is supported at present.")
return dask.array.shuffle(x, indexer, axis, chunks="auto")

def get_auto_chunk_size(self) -> int:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tomwhite is there an equivalent for cubed? I didn't see it in the docs...

Comment on lines +232 to +244
if _contains_cftime_datetimes(data):
limit, dtype = fake_target_chunksize(data, chunkmanager.get_auto_chunk_size())
else:
limit = None
dtype = data.dtype

chunk_shape = chunkmanager.normalize_chunks(
chunk_shape,
shape=shape,
dtype=dtype,
limit=limit,
previous_chunks=preferred_chunk_shape,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does this seem fine to you @charles-turner-1 . I wanted to avoid calling get_auto_chunk_size as much as possible

@charles-turner-1charles-turner-1Oct 13, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, looks good! fake_target_chunksize also contains the same _contains_cf_datetimes check & early return if it's false, so we could remove the check in either fake_target_chunksize or here without causing issues if you think that's a good idea?

I'm guessing you meant calling fake_target_chunksize in your comment above, in which case we would probably want to either remove it in that function - or leave it in if we want to reuse fake_target_chunksize elsewhere?

@dcheriandcherian left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Phew, I think this is good to go. It would be good to clean up the types, but this PR has stalled for a long time.

Apologies for the delay (again). I was on vacation.

@dcheriandcherian added the plan to merge Final call for comments label Oct 13, 2025
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

No worries, thanks for all your help!

I'd love to keep getting my feet wet - do you happen to know if there are any other extant issues in roughly the same parts of the codebase off the top of your head? If not I'll go digging soon!

@dcherian

Copy link
Copy Markdown
Contributor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

@dcherian
dcherian merged commit 94798a0 into pydata:mainOct 15, 2025
45 of 47 checks passed
@welcome

welcomeBot commented Oct 15, 2025

Copy link
Copy Markdown

Congratulations on completing your first pull request! Welcome to Xarray! We are proud of you, and hope to see you again! celebration gif

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

🙏 I'll have a crack!

@spencerkclark

Copy link
Copy Markdown
Member

This is great—thanks @charles-turner-1 and @dcherian!

@Thomas-Moore-Creative

Copy link
Copy Markdown

This is great—thanks @charles-turner-1 and @dcherian!

Go Australia 🇦🇺 ( AKA @charles-turner-1 ), pulling our weight! 😉


output_dtype = np.dtype(np.float64)

nbytes_approx: int = sys.getsizeof(first_n_items(data, 1)) # type: ignore[no-untyped-call]

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm just came across this and I'm not quite sure it's the right size. I think sys.getsizeof is the in-memory size and and dtype.itemsize is the uncompressed disk size. Consider for instance:

importsysimportnumpyasnpimportcftimenp.dtype(np.float64).itemsize# 8sys.getsizeof(np.float64(1.0)) # 32sys.getsizeof(np.array([1.0], dtype=np.float64)) # 120sys.getsizeof(cftime.DatetimeGregorian.fromordinal(2450000)) #112

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm kind of wondering if setting the dtype to np.dtype(np.float64) would suffice

@charles-turner-1charles-turner-1Dec 16, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is for the assumed size as a float64 right? I think what you're saying is true but the array overhead rapidly becomes unimportant for reasonably large arrays? Very rough & ready analysis below:

# Does this still matter for decently sized arrays?importmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 8, num=50, dtype=int):
arr=np.zeros(n, dtype=np.float64)
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Exactly! Calling sys.getsizeof on an array containing a single cftime object is not going to be a good representation of the memory consumption of an array of these things. Even if you pop the object out of the array that is still not really a good representation of the memory consumption. I think you'd do better with just nbytes_approx: int = 8

I made that same plot you did but with cftimes inside the array:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 4, num=50, dtype=int):
arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My bad - I thought we'd popped the first cftime element out of the array and had a look at its size at that point.

It looks like the cftime elements in the array are 8bytes too - is that what we expect? I would have expected them to be a bit larger due to the extra overhead...

Assuming I'm wrong about that, it would be much simpler to just tell dask that a cftime is 8 bytes and leave the limit unadjusted - the ratio of two line should be pretty much 1 for all decently sized arrays.

On my phone right now but I'll have a proper look when I get to my computer.

@charles-turner-1charles-turner-1Dec 18, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So it does look like size per element in a numpy array is reliably 8 bytes, but I'm really unconvinced this can be correct tbh:

importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltcf_sizes : list[int] = []
f64_sizes : list[int] = []
num_elements :list[int] = []
forninnp.logspace(0, 4, num=50, dtype=int):
cf_arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
cf_sizes.append(sys.getsizeof(cf_arr)/n)
num_elements.append(n)
arr=np.zeros(n, dtype=np.float64)
f64_sizes.append(sys.getsizeof(arr)/n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
ratio= [s_cf/s_f64fors_cf, s_f64inzip(cf_sizes, f64_sizes)]
plt.plot(num_elements, cf_sizes, marker='o', label='cftime')
plt.plot(num_elements, f64_sizes, marker='o', label='float64')
plt.plot(num_elements, ratio, marker='o', label='cftime/float64 ratio')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte line, unit ratio lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.axhline(y=1, color='grey', linestyle='--', label='1')
plt.legend()
image

But if we look at the raw element (same as you did above):

>>>t=cftime.DatetimeGregorian.fromordinal(2450000)
>>>sys.getsizeof(t)
112

Since it's just not possible that the cftime objects are magically shrinking when we put them in a numpy array, I assume numpy is storing pointers to objects somewhere on the heap.

I've run a couple of more sophisticated (this does not necessarily mean more likely to be right!) tests here:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltimporttracemallocimportgccf_sizes : list[int] = []
numel :list[int] = []
forninnp.logspace(0, 5, num=50, dtype=int):
gc.collect()
tracemalloc.start()
snap1=tracemalloc.take_snapshot()
cf_arr=cftime.DatetimeGregorian.fromordinal(np.arange(2450000, 2450000+n))
snap2=tracemalloc.take_snapshot()
stats=snap2.compare_to(snap1, 'lineno')
tracemalloc.stop()
tot=sum(stat.size_diffforstatinstats)
cf_sizes.append(tot/n)
numel.append(n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot(numel, cf_sizes, marker='o', label='cftime')
plt.xscale('log')
plt.yscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.legend()
print(f"nbytes asymptotes to {cf_sizes[-1]:.2f} for large arrays")

nbytes asymptotes to 120.12 for large arrays
image

I'm still not convinced this is the right number, so I'm still digging. But it looks like we (accidentally/serendipitously) might have gotten in the right ballpark?


EDIT: I've more some more playing, and I reckon we're off by approximately a factor of 2-2.5 ish. No real justification yet, just empirical results.

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ioplan to mergeFinal call for commentstopic-backendstopic-documentationtopic-NamedArrayLightweight version of Variable

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Confusing error when use_cftime = True and chunks = 'auto' in xr.open_dataset()

6 participants

@charles-turner-1@jemmajeffree@dcherian@spencerkclark@Thomas-Moore-Creative@jsignell
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms"); } } catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

Add chunks='auto' support for cftime datasets - #10527

Merged
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime
Oct 15, 2025
Merged

Add chunks='auto' support for cftime datasets#10527
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime

Conversation

@charles-turner-1

@charles-turner-1charles-turner-1 commented Jul 13, 2025

Copy link
Copy Markdown
Contributor

@welcome

welcomeBot commented Jul 13, 2025

Copy link
Copy Markdown

Thank you for opening this pull request! It may take us a few days to respond here, so thank you for being patient.
If you have questions, some answers may be found in our contributing guidelines.

@github-actionsgithub-actionsBot added topic-documentation topic-NamedArray Lightweight version of Variable labels Jul 13, 2025
@charles-turner-1charles-turner-1 changed the title All works, just need to satisfy mypy and whatnot nowAdd chunks='auto' support for cftime datasetsJul 13, 2025
@jemmajeffree

Copy link
Copy Markdown
Contributor

Would these changes also work for cf timedeltas or are they going to still cause problems?
I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

Would these changes also work for cf timedeltas or are they going to still cause problems? I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

If you can find something thats specifically a cftimedelta and run the _contains_cftime_datetimes function on it that'd be super helpful to know whether it returns True or False.

@charles-turner-1
charles-turner-1 marked this pull request as draft July 14, 2025 05:02
@jemmajeffree

Copy link
Copy Markdown
Contributor

TLDR: don't mind me, it's not going to cause any issues

Firstly, what I thought was a cftimedelta turned out to be a numpy timedelta hanging out with a cftime
Screenshot 2025-07-14 at 5 23 31 pm
When I did manage to coerce this timedelta into cftime conventions, it just contained a floating point number of days, so I can't see anything having issues with its size

coder=xr.coding.times.CFTimedeltaCoder()
result=coder.encode(oops.average_DT).load()
print(result.dtype)
result
Screenshot 2025-07-14 at 5 38 33 pm

Comment threadxarray/namedarray/daskmanager.py Outdated
Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

I did some prodding around yesterday and I realised this won't let us do something like

importxarrayasxrcftime_datafile="/path/to/file.nc"xr.open_dataset(cftime_datafile, chunks='auto')

yet, only stuff along the lines of

importxarrayasxrcftime_datafile="/path/to/file.nc"ds=xr.open_dataset(cftime_datafile, chunks=-1)
ds=ds.chunk('auto')

I think implementing the former is going to be a bit harder, but I'm starting to clock the code structure a bit more now so I'll have a decent crack.

@dcherian

Copy link
Copy Markdown
Contributor

Why so? Are we sending "auto" in to normalize_chunks first?

@charles-turner-1

charles-turner-1 commented Jul 23, 2025

Copy link
Copy Markdown
ContributorAuthor

Yup, this is the call stack:

---->3xr.open_dataset(
4"/Users/u1166368/xarray/tos_Omon_CESM2-WACCM_historical_r2i1p1f1_gr_185001-201412.nc", chunks="auto"/Users/u1166368/xarray/xarray/backends/api.py(721)open_dataset()
720 )
-->721ds=_dataset_from_backend_dataset(
722backend_ds,
/Users/u1166368/xarray/xarray/backends/api.py(418)_dataset_from_backend_dataset()
417ifchunksisnotNone:
-->418ds=_chunk_ds(
419ds,
/Users/u1166368/xarray/xarray/backends/api.py(368)_chunk_ds()
367forname, varinbackend_ds.variables.items():
-->368var_chunks=_get_chunk(var, chunks, chunkmanager)
369variables[name] =_maybe_chunk(
/Users/u1166368/xarray/xarray/structure/chunks.py(102)_get_chunk()
101-->102chunk_shape=chunkmanager.normalize_chunks(
103chunk_shape, shape=shape, dtype=var.dtype, previous_chunks=preferred_chunk_shape>/Users/u1166368/xarray/xarray/namedarray/daskmanager.py(60)normalize_chunks()

I've fixed it in the latest commit - but I think the implementation leaves a lot to be desired too.

Do I want to refactor to move the changes in xarray/structure/chunks.py into the daskmanager module if possible?

Once I've got the structure there cleaned up, I'll work on replacing the build_chunkspec function with something more sensible - I just need to work out how to extract the implementation in dask cleanly now I think - normalize_chunks also seems to calculate sensible chunk sizes.

Comment threadxarray/structure/chunks.py Outdated

from xarray.namedarray.utils import build_chunkspec

target_chunksize = parse_bytes(dask_config.get("array.chunk-size"))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about adding get_auto_chunk_size to the ChunkManager class; and put the dask-specific stuff in the DaskManager.

cc @TomNicholas

@dcherian

dcherian commented Jul 23, 2025

Copy link
Copy Markdown
Contributor

I guess one bit that's confusing here is that the code-path for backends and normal variables is different?

So let's add a test that reads form disk; and one that works iwth a DataArray constructed in memory.

Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

charles-turner-1 commented Sep 24, 2025

Copy link
Copy Markdown
ContributorAuthor

It looks like the failing test (same one I commented on above) might be flaky? Now only failing for windows & python3.11, not 3.13: https://github.com/pydata/xarray/actions/runs/17962030991/job/51087173288

raise NotImplementedError("Only chunks='auto' is supported at present.")
return dask.array.shuffle(x, indexer, axis, chunks="auto")

def get_auto_chunk_size(self) -> int:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tomwhite is there an equivalent for cubed? I didn't see it in the docs...

Comment on lines +232 to +244
if _contains_cftime_datetimes(data):
limit, dtype = fake_target_chunksize(data, chunkmanager.get_auto_chunk_size())
else:
limit = None
dtype = data.dtype

chunk_shape = chunkmanager.normalize_chunks(
chunk_shape,
shape=shape,
dtype=dtype,
limit=limit,
previous_chunks=preferred_chunk_shape,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does this seem fine to you @charles-turner-1 . I wanted to avoid calling get_auto_chunk_size as much as possible

@charles-turner-1charles-turner-1Oct 13, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, looks good! fake_target_chunksize also contains the same _contains_cf_datetimes check & early return if it's false, so we could remove the check in either fake_target_chunksize or here without causing issues if you think that's a good idea?

I'm guessing you meant calling fake_target_chunksize in your comment above, in which case we would probably want to either remove it in that function - or leave it in if we want to reuse fake_target_chunksize elsewhere?

@dcheriandcherian left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Phew, I think this is good to go. It would be good to clean up the types, but this PR has stalled for a long time.

Apologies for the delay (again). I was on vacation.

@dcheriandcherian added the plan to merge Final call for comments label Oct 13, 2025
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

No worries, thanks for all your help!

I'd love to keep getting my feet wet - do you happen to know if there are any other extant issues in roughly the same parts of the codebase off the top of your head? If not I'll go digging soon!

@dcherian

Copy link
Copy Markdown
Contributor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

@dcherian
dcherian merged commit 94798a0 into pydata:mainOct 15, 2025
45 of 47 checks passed
@welcome

welcomeBot commented Oct 15, 2025

Copy link
Copy Markdown

Congratulations on completing your first pull request! Welcome to Xarray! We are proud of you, and hope to see you again! celebration gif

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

🙏 I'll have a crack!

@spencerkclark

Copy link
Copy Markdown
Member

This is great—thanks @charles-turner-1 and @dcherian!

@Thomas-Moore-Creative

Copy link
Copy Markdown

This is great—thanks @charles-turner-1 and @dcherian!

Go Australia 🇦🇺 ( AKA @charles-turner-1 ), pulling our weight! 😉


output_dtype = np.dtype(np.float64)

nbytes_approx: int = sys.getsizeof(first_n_items(data, 1)) # type: ignore[no-untyped-call]

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm just came across this and I'm not quite sure it's the right size. I think sys.getsizeof is the in-memory size and and dtype.itemsize is the uncompressed disk size. Consider for instance:

importsysimportnumpyasnpimportcftimenp.dtype(np.float64).itemsize# 8sys.getsizeof(np.float64(1.0)) # 32sys.getsizeof(np.array([1.0], dtype=np.float64)) # 120sys.getsizeof(cftime.DatetimeGregorian.fromordinal(2450000)) #112

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm kind of wondering if setting the dtype to np.dtype(np.float64) would suffice

@charles-turner-1charles-turner-1Dec 16, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is for the assumed size as a float64 right? I think what you're saying is true but the array overhead rapidly becomes unimportant for reasonably large arrays? Very rough & ready analysis below:

# Does this still matter for decently sized arrays?importmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 8, num=50, dtype=int):
arr=np.zeros(n, dtype=np.float64)
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Exactly! Calling sys.getsizeof on an array containing a single cftime object is not going to be a good representation of the memory consumption of an array of these things. Even if you pop the object out of the array that is still not really a good representation of the memory consumption. I think you'd do better with just nbytes_approx: int = 8

I made that same plot you did but with cftimes inside the array:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 4, num=50, dtype=int):
arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My bad - I thought we'd popped the first cftime element out of the array and had a look at its size at that point.

It looks like the cftime elements in the array are 8bytes too - is that what we expect? I would have expected them to be a bit larger due to the extra overhead...

Assuming I'm wrong about that, it would be much simpler to just tell dask that a cftime is 8 bytes and leave the limit unadjusted - the ratio of two line should be pretty much 1 for all decently sized arrays.

On my phone right now but I'll have a proper look when I get to my computer.

@charles-turner-1charles-turner-1Dec 18, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So it does look like size per element in a numpy array is reliably 8 bytes, but I'm really unconvinced this can be correct tbh:

importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltcf_sizes : list[int] = []
f64_sizes : list[int] = []
num_elements :list[int] = []
forninnp.logspace(0, 4, num=50, dtype=int):
cf_arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
cf_sizes.append(sys.getsizeof(cf_arr)/n)
num_elements.append(n)
arr=np.zeros(n, dtype=np.float64)
f64_sizes.append(sys.getsizeof(arr)/n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
ratio= [s_cf/s_f64fors_cf, s_f64inzip(cf_sizes, f64_sizes)]
plt.plot(num_elements, cf_sizes, marker='o', label='cftime')
plt.plot(num_elements, f64_sizes, marker='o', label='float64')
plt.plot(num_elements, ratio, marker='o', label='cftime/float64 ratio')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte line, unit ratio lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.axhline(y=1, color='grey', linestyle='--', label='1')
plt.legend()
image

But if we look at the raw element (same as you did above):

>>>t=cftime.DatetimeGregorian.fromordinal(2450000)
>>>sys.getsizeof(t)
112

Since it's just not possible that the cftime objects are magically shrinking when we put them in a numpy array, I assume numpy is storing pointers to objects somewhere on the heap.

I've run a couple of more sophisticated (this does not necessarily mean more likely to be right!) tests here:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltimporttracemallocimportgccf_sizes : list[int] = []
numel :list[int] = []
forninnp.logspace(0, 5, num=50, dtype=int):
gc.collect()
tracemalloc.start()
snap1=tracemalloc.take_snapshot()
cf_arr=cftime.DatetimeGregorian.fromordinal(np.arange(2450000, 2450000+n))
snap2=tracemalloc.take_snapshot()
stats=snap2.compare_to(snap1, 'lineno')
tracemalloc.stop()
tot=sum(stat.size_diffforstatinstats)
cf_sizes.append(tot/n)
numel.append(n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot(numel, cf_sizes, marker='o', label='cftime')
plt.xscale('log')
plt.yscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.legend()
print(f"nbytes asymptotes to {cf_sizes[-1]:.2f} for large arrays")

nbytes asymptotes to 120.12 for large arrays
image

I'm still not convinced this is the right number, so I'm still digging. But it looks like we (accidentally/serendipitously) might have gotten in the right ballpark?


EDIT: I've more some more playing, and I reckon we're off by approximately a factor of 2-2.5 ish. No real justification yet, just empirical results.

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ioplan to mergeFinal call for commentstopic-backendstopic-documentationtopic-NamedArrayLightweight version of Variable

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Confusing error when use_cftime = True and chunks = 'auto' in xr.open_dataset()

6 participants

@charles-turner-1@jemmajeffree@dcherian@spencerkclark@Thomas-Moore-Creative@jsignell
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links"); } } catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); } })(); (function(){ try { var __m = "youtube.com"; var __re = new RegExp('^' + "youtube\\.com" + '
Skip to content

Add chunks='auto' support for cftime datasets - #10527

Merged
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime
Oct 15, 2025
Merged

Add chunks='auto' support for cftime datasets#10527
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime

Conversation

@charles-turner-1

@charles-turner-1charles-turner-1 commented Jul 13, 2025

Copy link
Copy Markdown
Contributor

@welcome

welcomeBot commented Jul 13, 2025

Copy link
Copy Markdown

Thank you for opening this pull request! It may take us a few days to respond here, so thank you for being patient.
If you have questions, some answers may be found in our contributing guidelines.

@github-actionsgithub-actionsBot added topic-documentation topic-NamedArray Lightweight version of Variable labels Jul 13, 2025
@charles-turner-1charles-turner-1 changed the title All works, just need to satisfy mypy and whatnot nowAdd chunks='auto' support for cftime datasetsJul 13, 2025
@jemmajeffree

Copy link
Copy Markdown
Contributor

Would these changes also work for cf timedeltas or are they going to still cause problems?
I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

Would these changes also work for cf timedeltas or are they going to still cause problems? I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

If you can find something thats specifically a cftimedelta and run the _contains_cftime_datetimes function on it that'd be super helpful to know whether it returns True or False.

@charles-turner-1
charles-turner-1 marked this pull request as draft July 14, 2025 05:02
@jemmajeffree

Copy link
Copy Markdown
Contributor

TLDR: don't mind me, it's not going to cause any issues

Firstly, what I thought was a cftimedelta turned out to be a numpy timedelta hanging out with a cftime
Screenshot 2025-07-14 at 5 23 31 pm
When I did manage to coerce this timedelta into cftime conventions, it just contained a floating point number of days, so I can't see anything having issues with its size

coder=xr.coding.times.CFTimedeltaCoder()
result=coder.encode(oops.average_DT).load()
print(result.dtype)
result
Screenshot 2025-07-14 at 5 38 33 pm

Comment threadxarray/namedarray/daskmanager.py Outdated
Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

I did some prodding around yesterday and I realised this won't let us do something like

importxarrayasxrcftime_datafile="/path/to/file.nc"xr.open_dataset(cftime_datafile, chunks='auto')

yet, only stuff along the lines of

importxarrayasxrcftime_datafile="/path/to/file.nc"ds=xr.open_dataset(cftime_datafile, chunks=-1)
ds=ds.chunk('auto')

I think implementing the former is going to be a bit harder, but I'm starting to clock the code structure a bit more now so I'll have a decent crack.

@dcherian

Copy link
Copy Markdown
Contributor

Why so? Are we sending "auto" in to normalize_chunks first?

@charles-turner-1

charles-turner-1 commented Jul 23, 2025

Copy link
Copy Markdown
ContributorAuthor

Yup, this is the call stack:

---->3xr.open_dataset(
4"/Users/u1166368/xarray/tos_Omon_CESM2-WACCM_historical_r2i1p1f1_gr_185001-201412.nc", chunks="auto"/Users/u1166368/xarray/xarray/backends/api.py(721)open_dataset()
720 )
-->721ds=_dataset_from_backend_dataset(
722backend_ds,
/Users/u1166368/xarray/xarray/backends/api.py(418)_dataset_from_backend_dataset()
417ifchunksisnotNone:
-->418ds=_chunk_ds(
419ds,
/Users/u1166368/xarray/xarray/backends/api.py(368)_chunk_ds()
367forname, varinbackend_ds.variables.items():
-->368var_chunks=_get_chunk(var, chunks, chunkmanager)
369variables[name] =_maybe_chunk(
/Users/u1166368/xarray/xarray/structure/chunks.py(102)_get_chunk()
101-->102chunk_shape=chunkmanager.normalize_chunks(
103chunk_shape, shape=shape, dtype=var.dtype, previous_chunks=preferred_chunk_shape>/Users/u1166368/xarray/xarray/namedarray/daskmanager.py(60)normalize_chunks()

I've fixed it in the latest commit - but I think the implementation leaves a lot to be desired too.

Do I want to refactor to move the changes in xarray/structure/chunks.py into the daskmanager module if possible?

Once I've got the structure there cleaned up, I'll work on replacing the build_chunkspec function with something more sensible - I just need to work out how to extract the implementation in dask cleanly now I think - normalize_chunks also seems to calculate sensible chunk sizes.

Comment threadxarray/structure/chunks.py Outdated

from xarray.namedarray.utils import build_chunkspec

target_chunksize = parse_bytes(dask_config.get("array.chunk-size"))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about adding get_auto_chunk_size to the ChunkManager class; and put the dask-specific stuff in the DaskManager.

cc @TomNicholas

@dcherian

dcherian commented Jul 23, 2025

Copy link
Copy Markdown
Contributor

I guess one bit that's confusing here is that the code-path for backends and normal variables is different?

So let's add a test that reads form disk; and one that works iwth a DataArray constructed in memory.

Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

charles-turner-1 commented Sep 24, 2025

Copy link
Copy Markdown
ContributorAuthor

It looks like the failing test (same one I commented on above) might be flaky? Now only failing for windows & python3.11, not 3.13: https://github.com/pydata/xarray/actions/runs/17962030991/job/51087173288

raise NotImplementedError("Only chunks='auto' is supported at present.")
return dask.array.shuffle(x, indexer, axis, chunks="auto")

def get_auto_chunk_size(self) -> int:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tomwhite is there an equivalent for cubed? I didn't see it in the docs...

Comment on lines +232 to +244
if _contains_cftime_datetimes(data):
limit, dtype = fake_target_chunksize(data, chunkmanager.get_auto_chunk_size())
else:
limit = None
dtype = data.dtype

chunk_shape = chunkmanager.normalize_chunks(
chunk_shape,
shape=shape,
dtype=dtype,
limit=limit,
previous_chunks=preferred_chunk_shape,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does this seem fine to you @charles-turner-1 . I wanted to avoid calling get_auto_chunk_size as much as possible

@charles-turner-1charles-turner-1Oct 13, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, looks good! fake_target_chunksize also contains the same _contains_cf_datetimes check & early return if it's false, so we could remove the check in either fake_target_chunksize or here without causing issues if you think that's a good idea?

I'm guessing you meant calling fake_target_chunksize in your comment above, in which case we would probably want to either remove it in that function - or leave it in if we want to reuse fake_target_chunksize elsewhere?

@dcheriandcherian left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Phew, I think this is good to go. It would be good to clean up the types, but this PR has stalled for a long time.

Apologies for the delay (again). I was on vacation.

@dcheriandcherian added the plan to merge Final call for comments label Oct 13, 2025
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

No worries, thanks for all your help!

I'd love to keep getting my feet wet - do you happen to know if there are any other extant issues in roughly the same parts of the codebase off the top of your head? If not I'll go digging soon!

@dcherian

Copy link
Copy Markdown
Contributor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

@dcherian
dcherian merged commit 94798a0 into pydata:mainOct 15, 2025
45 of 47 checks passed
@welcome

welcomeBot commented Oct 15, 2025

Copy link
Copy Markdown

Congratulations on completing your first pull request! Welcome to Xarray! We are proud of you, and hope to see you again! celebration gif

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

🙏 I'll have a crack!

@spencerkclark

Copy link
Copy Markdown
Member

This is great—thanks @charles-turner-1 and @dcherian!

@Thomas-Moore-Creative

Copy link
Copy Markdown

This is great—thanks @charles-turner-1 and @dcherian!

Go Australia 🇦🇺 ( AKA @charles-turner-1 ), pulling our weight! 😉


output_dtype = np.dtype(np.float64)

nbytes_approx: int = sys.getsizeof(first_n_items(data, 1)) # type: ignore[no-untyped-call]

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm just came across this and I'm not quite sure it's the right size. I think sys.getsizeof is the in-memory size and and dtype.itemsize is the uncompressed disk size. Consider for instance:

importsysimportnumpyasnpimportcftimenp.dtype(np.float64).itemsize# 8sys.getsizeof(np.float64(1.0)) # 32sys.getsizeof(np.array([1.0], dtype=np.float64)) # 120sys.getsizeof(cftime.DatetimeGregorian.fromordinal(2450000)) #112

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm kind of wondering if setting the dtype to np.dtype(np.float64) would suffice

@charles-turner-1charles-turner-1Dec 16, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is for the assumed size as a float64 right? I think what you're saying is true but the array overhead rapidly becomes unimportant for reasonably large arrays? Very rough & ready analysis below:

# Does this still matter for decently sized arrays?importmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 8, num=50, dtype=int):
arr=np.zeros(n, dtype=np.float64)
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Exactly! Calling sys.getsizeof on an array containing a single cftime object is not going to be a good representation of the memory consumption of an array of these things. Even if you pop the object out of the array that is still not really a good representation of the memory consumption. I think you'd do better with just nbytes_approx: int = 8

I made that same plot you did but with cftimes inside the array:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 4, num=50, dtype=int):
arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My bad - I thought we'd popped the first cftime element out of the array and had a look at its size at that point.

It looks like the cftime elements in the array are 8bytes too - is that what we expect? I would have expected them to be a bit larger due to the extra overhead...

Assuming I'm wrong about that, it would be much simpler to just tell dask that a cftime is 8 bytes and leave the limit unadjusted - the ratio of two line should be pretty much 1 for all decently sized arrays.

On my phone right now but I'll have a proper look when I get to my computer.

@charles-turner-1charles-turner-1Dec 18, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So it does look like size per element in a numpy array is reliably 8 bytes, but I'm really unconvinced this can be correct tbh:

importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltcf_sizes : list[int] = []
f64_sizes : list[int] = []
num_elements :list[int] = []
forninnp.logspace(0, 4, num=50, dtype=int):
cf_arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
cf_sizes.append(sys.getsizeof(cf_arr)/n)
num_elements.append(n)
arr=np.zeros(n, dtype=np.float64)
f64_sizes.append(sys.getsizeof(arr)/n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
ratio= [s_cf/s_f64fors_cf, s_f64inzip(cf_sizes, f64_sizes)]
plt.plot(num_elements, cf_sizes, marker='o', label='cftime')
plt.plot(num_elements, f64_sizes, marker='o', label='float64')
plt.plot(num_elements, ratio, marker='o', label='cftime/float64 ratio')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte line, unit ratio lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.axhline(y=1, color='grey', linestyle='--', label='1')
plt.legend()
image

But if we look at the raw element (same as you did above):

>>>t=cftime.DatetimeGregorian.fromordinal(2450000)
>>>sys.getsizeof(t)
112

Since it's just not possible that the cftime objects are magically shrinking when we put them in a numpy array, I assume numpy is storing pointers to objects somewhere on the heap.

I've run a couple of more sophisticated (this does not necessarily mean more likely to be right!) tests here:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltimporttracemallocimportgccf_sizes : list[int] = []
numel :list[int] = []
forninnp.logspace(0, 5, num=50, dtype=int):
gc.collect()
tracemalloc.start()
snap1=tracemalloc.take_snapshot()
cf_arr=cftime.DatetimeGregorian.fromordinal(np.arange(2450000, 2450000+n))
snap2=tracemalloc.take_snapshot()
stats=snap2.compare_to(snap1, 'lineno')
tracemalloc.stop()
tot=sum(stat.size_diffforstatinstats)
cf_sizes.append(tot/n)
numel.append(n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot(numel, cf_sizes, marker='o', label='cftime')
plt.xscale('log')
plt.yscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.legend()
print(f"nbytes asymptotes to {cf_sizes[-1]:.2f} for large arrays")

nbytes asymptotes to 120.12 for large arrays
image

I'm still not convinced this is the right number, so I'm still digging. But it looks like we (accidentally/serendipitously) might have gotten in the right ballpark?


EDIT: I've more some more playing, and I reckon we're off by approximately a factor of 2-2.5 ish. No real justification yet, just empirical results.

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ioplan to mergeFinal call for commentstopic-backendstopic-documentationtopic-NamedArrayLightweight version of Variable

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Confusing error when use_cftime = True and chunks = 'auto' in xr.open_dataset()

6 participants

@charles-turner-1@jemmajeffree@dcherian@spencerkclark@Thomas-Moore-Creative@jsignell
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default"); } } catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

Add chunks='auto' support for cftime datasets - #10527

Merged
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime
Oct 15, 2025
Merged

Add chunks='auto' support for cftime datasets#10527
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime

Conversation

@charles-turner-1

@charles-turner-1charles-turner-1 commented Jul 13, 2025

Copy link
Copy Markdown
Contributor

@welcome

welcomeBot commented Jul 13, 2025

Copy link
Copy Markdown

Thank you for opening this pull request! It may take us a few days to respond here, so thank you for being patient.
If you have questions, some answers may be found in our contributing guidelines.

@github-actionsgithub-actionsBot added topic-documentation topic-NamedArray Lightweight version of Variable labels Jul 13, 2025
@charles-turner-1charles-turner-1 changed the title All works, just need to satisfy mypy and whatnot nowAdd chunks='auto' support for cftime datasetsJul 13, 2025
@jemmajeffree

Copy link
Copy Markdown
Contributor

Would these changes also work for cf timedeltas or are they going to still cause problems?
I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

Would these changes also work for cf timedeltas or are they going to still cause problems? I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

If you can find something thats specifically a cftimedelta and run the _contains_cftime_datetimes function on it that'd be super helpful to know whether it returns True or False.

@charles-turner-1
charles-turner-1 marked this pull request as draft July 14, 2025 05:02
@jemmajeffree

Copy link
Copy Markdown
Contributor

TLDR: don't mind me, it's not going to cause any issues

Firstly, what I thought was a cftimedelta turned out to be a numpy timedelta hanging out with a cftime
Screenshot 2025-07-14 at 5 23 31 pm
When I did manage to coerce this timedelta into cftime conventions, it just contained a floating point number of days, so I can't see anything having issues with its size

coder=xr.coding.times.CFTimedeltaCoder()
result=coder.encode(oops.average_DT).load()
print(result.dtype)
result
Screenshot 2025-07-14 at 5 38 33 pm

Comment threadxarray/namedarray/daskmanager.py Outdated
Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

I did some prodding around yesterday and I realised this won't let us do something like

importxarrayasxrcftime_datafile="/path/to/file.nc"xr.open_dataset(cftime_datafile, chunks='auto')

yet, only stuff along the lines of

importxarrayasxrcftime_datafile="/path/to/file.nc"ds=xr.open_dataset(cftime_datafile, chunks=-1)
ds=ds.chunk('auto')

I think implementing the former is going to be a bit harder, but I'm starting to clock the code structure a bit more now so I'll have a decent crack.

@dcherian

Copy link
Copy Markdown
Contributor

Why so? Are we sending "auto" in to normalize_chunks first?

@charles-turner-1

charles-turner-1 commented Jul 23, 2025

Copy link
Copy Markdown
ContributorAuthor

Yup, this is the call stack:

---->3xr.open_dataset(
4"/Users/u1166368/xarray/tos_Omon_CESM2-WACCM_historical_r2i1p1f1_gr_185001-201412.nc", chunks="auto"/Users/u1166368/xarray/xarray/backends/api.py(721)open_dataset()
720 )
-->721ds=_dataset_from_backend_dataset(
722backend_ds,
/Users/u1166368/xarray/xarray/backends/api.py(418)_dataset_from_backend_dataset()
417ifchunksisnotNone:
-->418ds=_chunk_ds(
419ds,
/Users/u1166368/xarray/xarray/backends/api.py(368)_chunk_ds()
367forname, varinbackend_ds.variables.items():
-->368var_chunks=_get_chunk(var, chunks, chunkmanager)
369variables[name] =_maybe_chunk(
/Users/u1166368/xarray/xarray/structure/chunks.py(102)_get_chunk()
101-->102chunk_shape=chunkmanager.normalize_chunks(
103chunk_shape, shape=shape, dtype=var.dtype, previous_chunks=preferred_chunk_shape>/Users/u1166368/xarray/xarray/namedarray/daskmanager.py(60)normalize_chunks()

I've fixed it in the latest commit - but I think the implementation leaves a lot to be desired too.

Do I want to refactor to move the changes in xarray/structure/chunks.py into the daskmanager module if possible?

Once I've got the structure there cleaned up, I'll work on replacing the build_chunkspec function with something more sensible - I just need to work out how to extract the implementation in dask cleanly now I think - normalize_chunks also seems to calculate sensible chunk sizes.

Comment threadxarray/structure/chunks.py Outdated

from xarray.namedarray.utils import build_chunkspec

target_chunksize = parse_bytes(dask_config.get("array.chunk-size"))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about adding get_auto_chunk_size to the ChunkManager class; and put the dask-specific stuff in the DaskManager.

cc @TomNicholas

@dcherian

dcherian commented Jul 23, 2025

Copy link
Copy Markdown
Contributor

I guess one bit that's confusing here is that the code-path for backends and normal variables is different?

So let's add a test that reads form disk; and one that works iwth a DataArray constructed in memory.

Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

charles-turner-1 commented Sep 24, 2025

Copy link
Copy Markdown
ContributorAuthor

It looks like the failing test (same one I commented on above) might be flaky? Now only failing for windows & python3.11, not 3.13: https://github.com/pydata/xarray/actions/runs/17962030991/job/51087173288

raise NotImplementedError("Only chunks='auto' is supported at present.")
return dask.array.shuffle(x, indexer, axis, chunks="auto")

def get_auto_chunk_size(self) -> int:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tomwhite is there an equivalent for cubed? I didn't see it in the docs...

Comment on lines +232 to +244
if _contains_cftime_datetimes(data):
limit, dtype = fake_target_chunksize(data, chunkmanager.get_auto_chunk_size())
else:
limit = None
dtype = data.dtype

chunk_shape = chunkmanager.normalize_chunks(
chunk_shape,
shape=shape,
dtype=dtype,
limit=limit,
previous_chunks=preferred_chunk_shape,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does this seem fine to you @charles-turner-1 . I wanted to avoid calling get_auto_chunk_size as much as possible

@charles-turner-1charles-turner-1Oct 13, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, looks good! fake_target_chunksize also contains the same _contains_cf_datetimes check & early return if it's false, so we could remove the check in either fake_target_chunksize or here without causing issues if you think that's a good idea?

I'm guessing you meant calling fake_target_chunksize in your comment above, in which case we would probably want to either remove it in that function - or leave it in if we want to reuse fake_target_chunksize elsewhere?

@dcheriandcherian left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Phew, I think this is good to go. It would be good to clean up the types, but this PR has stalled for a long time.

Apologies for the delay (again). I was on vacation.

@dcheriandcherian added the plan to merge Final call for comments label Oct 13, 2025
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

No worries, thanks for all your help!

I'd love to keep getting my feet wet - do you happen to know if there are any other extant issues in roughly the same parts of the codebase off the top of your head? If not I'll go digging soon!

@dcherian

Copy link
Copy Markdown
Contributor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

@dcherian
dcherian merged commit 94798a0 into pydata:mainOct 15, 2025
45 of 47 checks passed
@welcome

welcomeBot commented Oct 15, 2025

Copy link
Copy Markdown

Congratulations on completing your first pull request! Welcome to Xarray! We are proud of you, and hope to see you again! celebration gif

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

🙏 I'll have a crack!

@spencerkclark

Copy link
Copy Markdown
Member

This is great—thanks @charles-turner-1 and @dcherian!

@Thomas-Moore-Creative

Copy link
Copy Markdown

This is great—thanks @charles-turner-1 and @dcherian!

Go Australia 🇦🇺 ( AKA @charles-turner-1 ), pulling our weight! 😉


output_dtype = np.dtype(np.float64)

nbytes_approx: int = sys.getsizeof(first_n_items(data, 1)) # type: ignore[no-untyped-call]

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm just came across this and I'm not quite sure it's the right size. I think sys.getsizeof is the in-memory size and and dtype.itemsize is the uncompressed disk size. Consider for instance:

importsysimportnumpyasnpimportcftimenp.dtype(np.float64).itemsize# 8sys.getsizeof(np.float64(1.0)) # 32sys.getsizeof(np.array([1.0], dtype=np.float64)) # 120sys.getsizeof(cftime.DatetimeGregorian.fromordinal(2450000)) #112

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm kind of wondering if setting the dtype to np.dtype(np.float64) would suffice

@charles-turner-1charles-turner-1Dec 16, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is for the assumed size as a float64 right? I think what you're saying is true but the array overhead rapidly becomes unimportant for reasonably large arrays? Very rough & ready analysis below:

# Does this still matter for decently sized arrays?importmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 8, num=50, dtype=int):
arr=np.zeros(n, dtype=np.float64)
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Exactly! Calling sys.getsizeof on an array containing a single cftime object is not going to be a good representation of the memory consumption of an array of these things. Even if you pop the object out of the array that is still not really a good representation of the memory consumption. I think you'd do better with just nbytes_approx: int = 8

I made that same plot you did but with cftimes inside the array:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 4, num=50, dtype=int):
arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My bad - I thought we'd popped the first cftime element out of the array and had a look at its size at that point.

It looks like the cftime elements in the array are 8bytes too - is that what we expect? I would have expected them to be a bit larger due to the extra overhead...

Assuming I'm wrong about that, it would be much simpler to just tell dask that a cftime is 8 bytes and leave the limit unadjusted - the ratio of two line should be pretty much 1 for all decently sized arrays.

On my phone right now but I'll have a proper look when I get to my computer.

@charles-turner-1charles-turner-1Dec 18, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So it does look like size per element in a numpy array is reliably 8 bytes, but I'm really unconvinced this can be correct tbh:

importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltcf_sizes : list[int] = []
f64_sizes : list[int] = []
num_elements :list[int] = []
forninnp.logspace(0, 4, num=50, dtype=int):
cf_arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
cf_sizes.append(sys.getsizeof(cf_arr)/n)
num_elements.append(n)
arr=np.zeros(n, dtype=np.float64)
f64_sizes.append(sys.getsizeof(arr)/n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
ratio= [s_cf/s_f64fors_cf, s_f64inzip(cf_sizes, f64_sizes)]
plt.plot(num_elements, cf_sizes, marker='o', label='cftime')
plt.plot(num_elements, f64_sizes, marker='o', label='float64')
plt.plot(num_elements, ratio, marker='o', label='cftime/float64 ratio')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte line, unit ratio lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.axhline(y=1, color='grey', linestyle='--', label='1')
plt.legend()
image

But if we look at the raw element (same as you did above):

>>>t=cftime.DatetimeGregorian.fromordinal(2450000)
>>>sys.getsizeof(t)
112

Since it's just not possible that the cftime objects are magically shrinking when we put them in a numpy array, I assume numpy is storing pointers to objects somewhere on the heap.

I've run a couple of more sophisticated (this does not necessarily mean more likely to be right!) tests here:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltimporttracemallocimportgccf_sizes : list[int] = []
numel :list[int] = []
forninnp.logspace(0, 5, num=50, dtype=int):
gc.collect()
tracemalloc.start()
snap1=tracemalloc.take_snapshot()
cf_arr=cftime.DatetimeGregorian.fromordinal(np.arange(2450000, 2450000+n))
snap2=tracemalloc.take_snapshot()
stats=snap2.compare_to(snap1, 'lineno')
tracemalloc.stop()
tot=sum(stat.size_diffforstatinstats)
cf_sizes.append(tot/n)
numel.append(n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot(numel, cf_sizes, marker='o', label='cftime')
plt.xscale('log')
plt.yscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.legend()
print(f"nbytes asymptotes to {cf_sizes[-1]:.2f} for large arrays")

nbytes asymptotes to 120.12 for large arrays
image

I'm still not convinced this is the right number, so I'm still digging. But it looks like we (accidentally/serendipitously) might have gotten in the right ballpark?


EDIT: I've more some more playing, and I reckon we're off by approximately a factor of 2-2.5 ish. No real justification yet, just empirical results.

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ioplan to mergeFinal call for commentstopic-backendstopic-documentationtopic-NamedArrayLightweight version of Variable

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Confusing error when use_cftime = True and chunks = 'auto' in xr.open_dataset()

6 participants

@charles-turner-1@jemmajeffree@dcherian@spencerkclark@Thomas-Moore-Creative@jsignell
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers"); } } catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

Add chunks='auto' support for cftime datasets - #10527

Merged
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime
Oct 15, 2025
Merged

Add chunks='auto' support for cftime datasets#10527
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime

Conversation

@charles-turner-1

@charles-turner-1charles-turner-1 commented Jul 13, 2025

Copy link
Copy Markdown
Contributor

@welcome

welcomeBot commented Jul 13, 2025

Copy link
Copy Markdown

Thank you for opening this pull request! It may take us a few days to respond here, so thank you for being patient.
If you have questions, some answers may be found in our contributing guidelines.

@github-actionsgithub-actionsBot added topic-documentation topic-NamedArray Lightweight version of Variable labels Jul 13, 2025
@charles-turner-1charles-turner-1 changed the title All works, just need to satisfy mypy and whatnot nowAdd chunks='auto' support for cftime datasetsJul 13, 2025
@jemmajeffree

Copy link
Copy Markdown
Contributor

Would these changes also work for cf timedeltas or are they going to still cause problems?
I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

Would these changes also work for cf timedeltas or are they going to still cause problems? I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

If you can find something thats specifically a cftimedelta and run the _contains_cftime_datetimes function on it that'd be super helpful to know whether it returns True or False.

@charles-turner-1
charles-turner-1 marked this pull request as draft July 14, 2025 05:02
@jemmajeffree

Copy link
Copy Markdown
Contributor

TLDR: don't mind me, it's not going to cause any issues

Firstly, what I thought was a cftimedelta turned out to be a numpy timedelta hanging out with a cftime
Screenshot 2025-07-14 at 5 23 31 pm
When I did manage to coerce this timedelta into cftime conventions, it just contained a floating point number of days, so I can't see anything having issues with its size

coder=xr.coding.times.CFTimedeltaCoder()
result=coder.encode(oops.average_DT).load()
print(result.dtype)
result
Screenshot 2025-07-14 at 5 38 33 pm

Comment threadxarray/namedarray/daskmanager.py Outdated
Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

I did some prodding around yesterday and I realised this won't let us do something like

importxarrayasxrcftime_datafile="/path/to/file.nc"xr.open_dataset(cftime_datafile, chunks='auto')

yet, only stuff along the lines of

importxarrayasxrcftime_datafile="/path/to/file.nc"ds=xr.open_dataset(cftime_datafile, chunks=-1)
ds=ds.chunk('auto')

I think implementing the former is going to be a bit harder, but I'm starting to clock the code structure a bit more now so I'll have a decent crack.

@dcherian

Copy link
Copy Markdown
Contributor

Why so? Are we sending "auto" in to normalize_chunks first?

@charles-turner-1

charles-turner-1 commented Jul 23, 2025

Copy link
Copy Markdown
ContributorAuthor

Yup, this is the call stack:

---->3xr.open_dataset(
4"/Users/u1166368/xarray/tos_Omon_CESM2-WACCM_historical_r2i1p1f1_gr_185001-201412.nc", chunks="auto"/Users/u1166368/xarray/xarray/backends/api.py(721)open_dataset()
720 )
-->721ds=_dataset_from_backend_dataset(
722backend_ds,
/Users/u1166368/xarray/xarray/backends/api.py(418)_dataset_from_backend_dataset()
417ifchunksisnotNone:
-->418ds=_chunk_ds(
419ds,
/Users/u1166368/xarray/xarray/backends/api.py(368)_chunk_ds()
367forname, varinbackend_ds.variables.items():
-->368var_chunks=_get_chunk(var, chunks, chunkmanager)
369variables[name] =_maybe_chunk(
/Users/u1166368/xarray/xarray/structure/chunks.py(102)_get_chunk()
101-->102chunk_shape=chunkmanager.normalize_chunks(
103chunk_shape, shape=shape, dtype=var.dtype, previous_chunks=preferred_chunk_shape>/Users/u1166368/xarray/xarray/namedarray/daskmanager.py(60)normalize_chunks()

I've fixed it in the latest commit - but I think the implementation leaves a lot to be desired too.

Do I want to refactor to move the changes in xarray/structure/chunks.py into the daskmanager module if possible?

Once I've got the structure there cleaned up, I'll work on replacing the build_chunkspec function with something more sensible - I just need to work out how to extract the implementation in dask cleanly now I think - normalize_chunks also seems to calculate sensible chunk sizes.

Comment threadxarray/structure/chunks.py Outdated

from xarray.namedarray.utils import build_chunkspec

target_chunksize = parse_bytes(dask_config.get("array.chunk-size"))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about adding get_auto_chunk_size to the ChunkManager class; and put the dask-specific stuff in the DaskManager.

cc @TomNicholas

@dcherian

dcherian commented Jul 23, 2025

Copy link
Copy Markdown
Contributor

I guess one bit that's confusing here is that the code-path for backends and normal variables is different?

So let's add a test that reads form disk; and one that works iwth a DataArray constructed in memory.

Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

charles-turner-1 commented Sep 24, 2025

Copy link
Copy Markdown
ContributorAuthor

It looks like the failing test (same one I commented on above) might be flaky? Now only failing for windows & python3.11, not 3.13: https://github.com/pydata/xarray/actions/runs/17962030991/job/51087173288

raise NotImplementedError("Only chunks='auto' is supported at present.")
return dask.array.shuffle(x, indexer, axis, chunks="auto")

def get_auto_chunk_size(self) -> int:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tomwhite is there an equivalent for cubed? I didn't see it in the docs...

Comment on lines +232 to +244
if _contains_cftime_datetimes(data):
limit, dtype = fake_target_chunksize(data, chunkmanager.get_auto_chunk_size())
else:
limit = None
dtype = data.dtype

chunk_shape = chunkmanager.normalize_chunks(
chunk_shape,
shape=shape,
dtype=dtype,
limit=limit,
previous_chunks=preferred_chunk_shape,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does this seem fine to you @charles-turner-1 . I wanted to avoid calling get_auto_chunk_size as much as possible

@charles-turner-1charles-turner-1Oct 13, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, looks good! fake_target_chunksize also contains the same _contains_cf_datetimes check & early return if it's false, so we could remove the check in either fake_target_chunksize or here without causing issues if you think that's a good idea?

I'm guessing you meant calling fake_target_chunksize in your comment above, in which case we would probably want to either remove it in that function - or leave it in if we want to reuse fake_target_chunksize elsewhere?

@dcheriandcherian left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Phew, I think this is good to go. It would be good to clean up the types, but this PR has stalled for a long time.

Apologies for the delay (again). I was on vacation.

@dcheriandcherian added the plan to merge Final call for comments label Oct 13, 2025
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

No worries, thanks for all your help!

I'd love to keep getting my feet wet - do you happen to know if there are any other extant issues in roughly the same parts of the codebase off the top of your head? If not I'll go digging soon!

@dcherian

Copy link
Copy Markdown
Contributor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

@dcherian
dcherian merged commit 94798a0 into pydata:mainOct 15, 2025
45 of 47 checks passed
@welcome

welcomeBot commented Oct 15, 2025

Copy link
Copy Markdown

Congratulations on completing your first pull request! Welcome to Xarray! We are proud of you, and hope to see you again! celebration gif

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

🙏 I'll have a crack!

@spencerkclark

Copy link
Copy Markdown
Member

This is great—thanks @charles-turner-1 and @dcherian!

@Thomas-Moore-Creative

Copy link
Copy Markdown

This is great—thanks @charles-turner-1 and @dcherian!

Go Australia 🇦🇺 ( AKA @charles-turner-1 ), pulling our weight! 😉


output_dtype = np.dtype(np.float64)

nbytes_approx: int = sys.getsizeof(first_n_items(data, 1)) # type: ignore[no-untyped-call]

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm just came across this and I'm not quite sure it's the right size. I think sys.getsizeof is the in-memory size and and dtype.itemsize is the uncompressed disk size. Consider for instance:

importsysimportnumpyasnpimportcftimenp.dtype(np.float64).itemsize# 8sys.getsizeof(np.float64(1.0)) # 32sys.getsizeof(np.array([1.0], dtype=np.float64)) # 120sys.getsizeof(cftime.DatetimeGregorian.fromordinal(2450000)) #112

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm kind of wondering if setting the dtype to np.dtype(np.float64) would suffice

@charles-turner-1charles-turner-1Dec 16, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is for the assumed size as a float64 right? I think what you're saying is true but the array overhead rapidly becomes unimportant for reasonably large arrays? Very rough & ready analysis below:

# Does this still matter for decently sized arrays?importmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 8, num=50, dtype=int):
arr=np.zeros(n, dtype=np.float64)
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Exactly! Calling sys.getsizeof on an array containing a single cftime object is not going to be a good representation of the memory consumption of an array of these things. Even if you pop the object out of the array that is still not really a good representation of the memory consumption. I think you'd do better with just nbytes_approx: int = 8

I made that same plot you did but with cftimes inside the array:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 4, num=50, dtype=int):
arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My bad - I thought we'd popped the first cftime element out of the array and had a look at its size at that point.

It looks like the cftime elements in the array are 8bytes too - is that what we expect? I would have expected them to be a bit larger due to the extra overhead...

Assuming I'm wrong about that, it would be much simpler to just tell dask that a cftime is 8 bytes and leave the limit unadjusted - the ratio of two line should be pretty much 1 for all decently sized arrays.

On my phone right now but I'll have a proper look when I get to my computer.

@charles-turner-1charles-turner-1Dec 18, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So it does look like size per element in a numpy array is reliably 8 bytes, but I'm really unconvinced this can be correct tbh:

importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltcf_sizes : list[int] = []
f64_sizes : list[int] = []
num_elements :list[int] = []
forninnp.logspace(0, 4, num=50, dtype=int):
cf_arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
cf_sizes.append(sys.getsizeof(cf_arr)/n)
num_elements.append(n)
arr=np.zeros(n, dtype=np.float64)
f64_sizes.append(sys.getsizeof(arr)/n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
ratio= [s_cf/s_f64fors_cf, s_f64inzip(cf_sizes, f64_sizes)]
plt.plot(num_elements, cf_sizes, marker='o', label='cftime')
plt.plot(num_elements, f64_sizes, marker='o', label='float64')
plt.plot(num_elements, ratio, marker='o', label='cftime/float64 ratio')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte line, unit ratio lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.axhline(y=1, color='grey', linestyle='--', label='1')
plt.legend()
image

But if we look at the raw element (same as you did above):

>>>t=cftime.DatetimeGregorian.fromordinal(2450000)
>>>sys.getsizeof(t)
112

Since it's just not possible that the cftime objects are magically shrinking when we put them in a numpy array, I assume numpy is storing pointers to objects somewhere on the heap.

I've run a couple of more sophisticated (this does not necessarily mean more likely to be right!) tests here:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltimporttracemallocimportgccf_sizes : list[int] = []
numel :list[int] = []
forninnp.logspace(0, 5, num=50, dtype=int):
gc.collect()
tracemalloc.start()
snap1=tracemalloc.take_snapshot()
cf_arr=cftime.DatetimeGregorian.fromordinal(np.arange(2450000, 2450000+n))
snap2=tracemalloc.take_snapshot()
stats=snap2.compare_to(snap1, 'lineno')
tracemalloc.stop()
tot=sum(stat.size_diffforstatinstats)
cf_sizes.append(tot/n)
numel.append(n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot(numel, cf_sizes, marker='o', label='cftime')
plt.xscale('log')
plt.yscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.legend()
print(f"nbytes asymptotes to {cf_sizes[-1]:.2f} for large arrays")

nbytes asymptotes to 120.12 for large arrays
image

I'm still not convinced this is the right number, so I'm still digging. But it looks like we (accidentally/serendipitously) might have gotten in the right ballpark?


EDIT: I've more some more playing, and I reckon we're off by approximately a factor of 2-2.5 ish. No real justification yet, just empirical results.

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ioplan to mergeFinal call for commentstopic-backendstopic-documentationtopic-NamedArrayLightweight version of Variable

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Confusing error when use_cftime = True and chunks = 'auto' in xr.open_dataset()

6 participants

@charles-turner-1@jemmajeffree@dcherian@spencerkclark@Thomas-Moore-Creative@jsignell
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode"); } } catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); } })(); })();
Skip to content

Add chunks='auto' support for cftime datasets - #10527

Merged
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime
Oct 15, 2025
Merged

Add chunks='auto' support for cftime datasets#10527
dcherian merged 63 commits into
pydata:mainfrom
charles-turner-1:autochunk-cftime

Conversation

@charles-turner-1

@charles-turner-1charles-turner-1 commented Jul 13, 2025

Copy link
Copy Markdown
Contributor

@welcome

welcomeBot commented Jul 13, 2025

Copy link
Copy Markdown

Thank you for opening this pull request! It may take us a few days to respond here, so thank you for being patient.
If you have questions, some answers may be found in our contributing guidelines.

@github-actionsgithub-actionsBot added topic-documentation topic-NamedArray Lightweight version of Variable labels Jul 13, 2025
@charles-turner-1charles-turner-1 changed the title All works, just need to satisfy mypy and whatnot nowAdd chunks='auto' support for cftime datasetsJul 13, 2025
@jemmajeffree

Copy link
Copy Markdown
Contributor

Would these changes also work for cf timedeltas or are they going to still cause problems?
I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

Would these changes also work for cf timedeltas or are they going to still cause problems? I'm tempted to write a script to bash through all the ACCESS-NRI intake datastores and see if there's anything else in there that's dtype object — let me know if this would be useful, or if we should just wait for it to break later

If you can find something thats specifically a cftimedelta and run the _contains_cftime_datetimes function on it that'd be super helpful to know whether it returns True or False.

@charles-turner-1
charles-turner-1 marked this pull request as draft July 14, 2025 05:02
@jemmajeffree

Copy link
Copy Markdown
Contributor

TLDR: don't mind me, it's not going to cause any issues

Firstly, what I thought was a cftimedelta turned out to be a numpy timedelta hanging out with a cftime
Screenshot 2025-07-14 at 5 23 31 pm
When I did manage to coerce this timedelta into cftime conventions, it just contained a floating point number of days, so I can't see anything having issues with its size

coder=xr.coding.times.CFTimedeltaCoder()
result=coder.encode(oops.average_DT).load()
print(result.dtype)
result
Screenshot 2025-07-14 at 5 38 33 pm

Comment threadxarray/namedarray/daskmanager.py Outdated
Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

I did some prodding around yesterday and I realised this won't let us do something like

importxarrayasxrcftime_datafile="/path/to/file.nc"xr.open_dataset(cftime_datafile, chunks='auto')

yet, only stuff along the lines of

importxarrayasxrcftime_datafile="/path/to/file.nc"ds=xr.open_dataset(cftime_datafile, chunks=-1)
ds=ds.chunk('auto')

I think implementing the former is going to be a bit harder, but I'm starting to clock the code structure a bit more now so I'll have a decent crack.

@dcherian

Copy link
Copy Markdown
Contributor

Why so? Are we sending "auto" in to normalize_chunks first?

@charles-turner-1

charles-turner-1 commented Jul 23, 2025

Copy link
Copy Markdown
ContributorAuthor

Yup, this is the call stack:

---->3xr.open_dataset(
4"/Users/u1166368/xarray/tos_Omon_CESM2-WACCM_historical_r2i1p1f1_gr_185001-201412.nc", chunks="auto"/Users/u1166368/xarray/xarray/backends/api.py(721)open_dataset()
720 )
-->721ds=_dataset_from_backend_dataset(
722backend_ds,
/Users/u1166368/xarray/xarray/backends/api.py(418)_dataset_from_backend_dataset()
417ifchunksisnotNone:
-->418ds=_chunk_ds(
419ds,
/Users/u1166368/xarray/xarray/backends/api.py(368)_chunk_ds()
367forname, varinbackend_ds.variables.items():
-->368var_chunks=_get_chunk(var, chunks, chunkmanager)
369variables[name] =_maybe_chunk(
/Users/u1166368/xarray/xarray/structure/chunks.py(102)_get_chunk()
101-->102chunk_shape=chunkmanager.normalize_chunks(
103chunk_shape, shape=shape, dtype=var.dtype, previous_chunks=preferred_chunk_shape>/Users/u1166368/xarray/xarray/namedarray/daskmanager.py(60)normalize_chunks()

I've fixed it in the latest commit - but I think the implementation leaves a lot to be desired too.

Do I want to refactor to move the changes in xarray/structure/chunks.py into the daskmanager module if possible?

Once I've got the structure there cleaned up, I'll work on replacing the build_chunkspec function with something more sensible - I just need to work out how to extract the implementation in dask cleanly now I think - normalize_chunks also seems to calculate sensible chunk sizes.

Comment threadxarray/structure/chunks.py Outdated

from xarray.namedarray.utils import build_chunkspec

target_chunksize = parse_bytes(dask_config.get("array.chunk-size"))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about adding get_auto_chunk_size to the ChunkManager class; and put the dask-specific stuff in the DaskManager.

cc @TomNicholas

@dcherian

dcherian commented Jul 23, 2025

Copy link
Copy Markdown
Contributor

I guess one bit that's confusing here is that the code-path for backends and normal variables is different?

So let's add a test that reads form disk; and one that works iwth a DataArray constructed in memory.

Comment threadxarray/namedarray/daskmanager.py Outdated
@charles-turner-1

charles-turner-1 commented Sep 24, 2025

Copy link
Copy Markdown
ContributorAuthor

It looks like the failing test (same one I commented on above) might be flaky? Now only failing for windows & python3.11, not 3.13: https://github.com/pydata/xarray/actions/runs/17962030991/job/51087173288

raise NotImplementedError("Only chunks='auto' is supported at present.")
return dask.array.shuffle(x, indexer, axis, chunks="auto")

def get_auto_chunk_size(self) -> int:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tomwhite is there an equivalent for cubed? I didn't see it in the docs...

Comment on lines +232 to +244
if _contains_cftime_datetimes(data):
limit, dtype = fake_target_chunksize(data, chunkmanager.get_auto_chunk_size())
else:
limit = None
dtype = data.dtype

chunk_shape = chunkmanager.normalize_chunks(
chunk_shape,
shape=shape,
dtype=dtype,
limit=limit,
previous_chunks=preferred_chunk_shape,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does this seem fine to you @charles-turner-1 . I wanted to avoid calling get_auto_chunk_size as much as possible

@charles-turner-1charles-turner-1Oct 13, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, looks good! fake_target_chunksize also contains the same _contains_cf_datetimes check & early return if it's false, so we could remove the check in either fake_target_chunksize or here without causing issues if you think that's a good idea?

I'm guessing you meant calling fake_target_chunksize in your comment above, in which case we would probably want to either remove it in that function - or leave it in if we want to reuse fake_target_chunksize elsewhere?

@dcheriandcherian left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Phew, I think this is good to go. It would be good to clean up the types, but this PR has stalled for a long time.

Apologies for the delay (again). I was on vacation.

@dcheriandcherian added the plan to merge Final call for comments label Oct 13, 2025
@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

No worries, thanks for all your help!

I'd love to keep getting my feet wet - do you happen to know if there are any other extant issues in roughly the same parts of the codebase off the top of your head? If not I'll go digging soon!

@dcherian

Copy link
Copy Markdown
Contributor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

@dcherian
dcherian merged commit 94798a0 into pydata:mainOct 15, 2025
45 of 47 checks passed
@welcome

welcomeBot commented Oct 15, 2025

Copy link
Copy Markdown

Congratulations on completing your first pull request! Welcome to Xarray! We are proud of you, and hope to see you again! celebration gif

@charles-turner-1

Copy link
Copy Markdown
ContributorAuthor

There's this one #9897 ;) but it's a bit gnarly, high impact though.

🙏 I'll have a crack!

@spencerkclark

Copy link
Copy Markdown
Member

This is great—thanks @charles-turner-1 and @dcherian!

@Thomas-Moore-Creative

Copy link
Copy Markdown

This is great—thanks @charles-turner-1 and @dcherian!

Go Australia 🇦🇺 ( AKA @charles-turner-1 ), pulling our weight! 😉


output_dtype = np.dtype(np.float64)

nbytes_approx: int = sys.getsizeof(first_n_items(data, 1)) # type: ignore[no-untyped-call]

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm just came across this and I'm not quite sure it's the right size. I think sys.getsizeof is the in-memory size and and dtype.itemsize is the uncompressed disk size. Consider for instance:

importsysimportnumpyasnpimportcftimenp.dtype(np.float64).itemsize# 8sys.getsizeof(np.float64(1.0)) # 32sys.getsizeof(np.array([1.0], dtype=np.float64)) # 120sys.getsizeof(cftime.DatetimeGregorian.fromordinal(2450000)) #112

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm kind of wondering if setting the dtype to np.dtype(np.float64) would suffice

@charles-turner-1charles-turner-1Dec 16, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is for the assumed size as a float64 right? I think what you're saying is true but the array overhead rapidly becomes unimportant for reasonably large arrays? Very rough & ready analysis below:

# Does this still matter for decently sized arrays?importmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 8, num=50, dtype=int):
arr=np.zeros(n, dtype=np.float64)
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Exactly! Calling sys.getsizeof on an array containing a single cftime object is not going to be a good representation of the memory consumption of an array of these things. Even if you pop the object out of the array that is still not really a good representation of the memory consumption. I think you'd do better with just nbytes_approx: int = 8

I made that same plot you did but with cftimes inside the array:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportmatplotlib.pyplotaspltsizes : list[tuple[int,int]] = []
forninnp.logspace(0, 4, num=50, dtype=int):
arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
sizes.append((n, sys.getsizeof(arr)/n))
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot([nforn,_insizes], [sizefor_,sizeinsizes], marker='o')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
image

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My bad - I thought we'd popped the first cftime element out of the array and had a look at its size at that point.

It looks like the cftime elements in the array are 8bytes too - is that what we expect? I would have expected them to be a bit larger due to the extra overhead...

Assuming I'm wrong about that, it would be much simpler to just tell dask that a cftime is 8 bytes and leave the limit unadjusted - the ratio of two line should be pretty much 1 for all decently sized arrays.

On my phone right now but I'll have a proper look when I get to my computer.

@charles-turner-1charles-turner-1Dec 18, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So it does look like size per element in a numpy array is reliably 8 bytes, but I'm really unconvinced this can be correct tbh:

importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltcf_sizes : list[int] = []
f64_sizes : list[int] = []
num_elements :list[int] = []
forninnp.logspace(0, 4, num=50, dtype=int):
cf_arr=np.array([cftime.DatetimeGregorian.fromordinal(2450000+i) foriinrange(n)])
cf_sizes.append(sys.getsizeof(cf_arr)/n)
num_elements.append(n)
arr=np.zeros(n, dtype=np.float64)
f64_sizes.append(sys.getsizeof(arr)/n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
ratio= [s_cf/s_f64fors_cf, s_f64inzip(cf_sizes, f64_sizes)]
plt.plot(num_elements, cf_sizes, marker='o', label='cftime')
plt.plot(num_elements, f64_sizes, marker='o', label='float64')
plt.plot(num_elements, ratio, marker='o', label='cftime/float64 ratio')
plt.xscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte line, unit ratio lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.axhline(y=1, color='grey', linestyle='--', label='1')
plt.legend()
image

But if we look at the raw element (same as you did above):

>>>t=cftime.DatetimeGregorian.fromordinal(2450000)
>>>sys.getsizeof(t)
112

Since it's just not possible that the cftime objects are magically shrinking when we put them in a numpy array, I assume numpy is storing pointers to objects somewhere on the heap.

I've run a couple of more sophisticated (this does not necessarily mean more likely to be right!) tests here:

# Does this still matter for decently sized arrays?importsysimportnumpyasnpimportcftimeimportmatplotlib.pyplotaspltimporttracemallocimportgccf_sizes : list[int] = []
numel :list[int] = []
forninnp.logspace(0, 5, num=50, dtype=int):
gc.collect()
tracemalloc.start()
snap1=tracemalloc.take_snapshot()
cf_arr=cftime.DatetimeGregorian.fromordinal(np.arange(2450000, 2450000+n))
snap2=tracemalloc.take_snapshot()
stats=snap2.compare_to(snap1, 'lineno')
tracemalloc.stop()
tot=sum(stat.size_diffforstatinstats)
cf_sizes.append(tot/n)
numel.append(n)
# Plot size per element vs number of elementsplt.figure(figsize=(10,6))
plt.plot(numel, cf_sizes, marker='o', label='cftime')
plt.xscale('log')
plt.yscale('log')
plt.xlabel('Number of elements in array (log scale)')
plt.ylabel('Size per element (bytes)')
# Add 8 byte lineplt.axhline(y=8, color='r', linestyle='--', label='8')
plt.legend()
print(f"nbytes asymptotes to {cf_sizes[-1]:.2f} for large arrays")

nbytes asymptotes to 120.12 for large arrays
image

I'm still not convinced this is the right number, so I'm still digging. But it looks like we (accidentally/serendipitously) might have gotten in the right ballpark?


EDIT: I've more some more playing, and I reckon we're off by approximately a factor of 2-2.5 ish. No real justification yet, just empirical results.

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ioplan to mergeFinal call for commentstopic-backendstopic-documentationtopic-NamedArrayLightweight version of Variable

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Confusing error when use_cftime = True and chunks = 'auto' in xr.open_dataset()

6 participants

@charles-turner-1@jemmajeffree@dcherian@spencerkclark@Thomas-Moore-Creative@jsignell