Skip to content

Commit bdfa8ce

Browse files
author
Abel Aoun
committed
WIP: naive implementation of enum
TODO: - fix to_netcdf It doesn't work if we have missing values because xarray doesn't use masked array so we try to assign a fill_value but netCDF4 forbid assigning values outside the enum valid range, thus crashing. - Add unit tests - Validate with xarray team if it's ok to add attribute to Variable and DataArray - Add implementation for other backends ?
1 parent f13da94 commit bdfa8ce

9 files changed

Lines changed: 106 additions & 69 deletions

File tree

‎xarray/backends/netCDF4_.py‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -409,9 +409,16 @@ def ds(self):
409409
return self._acquire()
410410

411411
def open_store_variable(self, name, var):
412+
import netCDF4
413+
412414
dimensions = var.dimensions
413415
data = indexing.LazilyIndexedArray(NetCDF4ArrayWrapper(name, self))
414416
attributes = {k: var.getncattr(k) for k in var.ncattrs()}
417+
enum_meaning = None
418+
enum_name = None
419+
if isinstance(var.datatype, netCDF4.EnumType):
420+
enum_meaning = var.datatype.enum_dict
421+
enum_name = var.datatype.name
415422
_ensure_fill_value_valid(data, attributes)
416423
# netCDF4 specific encoding; save _FillValue for later
417424
encoding = {}
@@ -434,8 +441,9 @@ def open_store_variable(self, name, var):
434441
encoding["source"] = self._filename
435442
encoding["original_shape"] = var.shape
436443
encoding["dtype"] = var.dtype
437-
438-
return Variable(dimensions, data, attributes, encoding)
444+
return Variable(dimensions, data, attributes, encoding,
445+
enum_meaning= enum_meaning,
446+
enum_name=enum_name)
439447

440448
def get_variables(self):
441449
return FrozenDict(
@@ -478,7 +486,7 @@ def encode_variable(self, variable):
478486
return variable
479487

480488
def prepare_variable(
481-
self, name, variable, check_encoding=False, unlimited_dims=None
489+
self, name, variable: Variable, check_encoding=False, unlimited_dims=None
482490
):
483491
_ensure_no_forward_slash_in_name(name)
484492

@@ -503,12 +511,19 @@ def prepare_variable(
503511
variable, raise_on_invalid=check_encoding, unlimited_dims=unlimited_dims
504512
)
505513

514+
enum = None
515+
if variable.enum_meaning:
516+
enum = self.ds.createEnumType(
517+
variable.dtype,
518+
variable.enum_name,
519+
variable.enum_meaning)
520+
506521
if name in self.ds.variables:
507522
nc4_var = self.ds.variables[name]
508523
else:
509524
nc4_var = self.ds.createVariable(
510525
varname=name,
511-
datatype=datatype,
526+
datatype=enum if enum else datatype,
512527
dimensions=variable.dims,
513528
zlib=encoding.get("zlib", False),
514529
complevel=encoding.get("complevel", 4),

‎xarray/backends/store.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ def guess_can_open(
2828

2929
def open_dataset( # type: ignore[override] # allow LSP violation, not supporting **kwargs
3030
self,
31-
filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
31+
filename_or_obj: AbstractDataStore,
3232
*,
3333
mask_and_scale=True,
3434
decode_times=True,

‎xarray/coding/strings.py‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,7 @@
1010
lazy_elemwise_func,
1111
pop_to,
1212
safe_setitem,
13-
unpack_for_decoding,
14-
unpack_for_encoding,
13+
unpack,
1514
)
1615
from xarray.core import indexing
1716
from xarray.core.parallelcompat import get_chunked_array_type, is_chunked_array
@@ -48,7 +47,7 @@ def __init__(self, allows_unicode=True):
4847
self.allows_unicode = allows_unicode
4948

5049
def encode(self, variable, name=None):
51-
dims, data, attrs, encoding = unpack_for_encoding(variable)
50+
dims, data, attrs, encoding, enum_meaning, enum_name = unpack(variable)
5251

5352
contains_unicode = is_unicode_dtype(data.dtype)
5453
encode_as_char = encoding.get("dtype") == "S1"
@@ -69,17 +68,17 @@ def encode(self, variable, name=None):
6968
# TODO: figure out how to handle this in a lazy way with dask
7069
data = encode_string_array(data, string_encoding)
7170

72-
return Variable(dims, data, attrs, encoding)
71+
return Variable(dims, data, attrs, encoding, enum_meaning=enum_meaning, enum_name=enum_name)
7372

7473
def decode(self, variable, name=None):
75-
dims, data, attrs, encoding = unpack_for_decoding(variable)
74+
dims, data, attrs, encoding, enum_meaning, enum_name = unpack(variable)
7675

7776
if "_Encoding" in attrs:
7877
string_encoding = pop_to(attrs, encoding, "_Encoding")
7978
func = partial(decode_bytes_array, encoding=string_encoding)
8079
data = lazy_elemwise_func(data, func, np.dtype(object))
8180

82-
return Variable(dims, data, attrs, encoding)
81+
return Variable(dims, data, attrs, encoding, enum_meaning=enum_meaning, enum_name=enum_name)
8382

8483

8584
def decode_bytes_array(bytes_array, encoding="utf-8"):
@@ -97,11 +96,11 @@ def encode_string_array(string_array, encoding="utf-8"):
9796

9897
def ensure_fixed_length_bytes(var):
9998
"""Ensure that a variable with vlen bytes is converted to fixed width."""
100-
dims, data, attrs, encoding = unpack_for_encoding(var)
99+
dims, data, attrs, encoding, enum_meaning, enum_name = unpack(var)
101100
if check_vlen_dtype(data.dtype) == bytes:
102101
# TODO: figure out how to handle this with dask
103102
data = np.asarray(data, dtype=np.string_)
104-
return Variable(dims, data, attrs, encoding)
103+
return Variable(dims, data, attrs, encoding, enum_meaning=enum_meaning, enum_name=enum_name)
105104

106105

107106
class CharacterArrayCoder(VariableCoder):
@@ -110,24 +109,24 @@ class CharacterArrayCoder(VariableCoder):
110109
def encode(self, variable, name=None):
111110
variable = ensure_fixed_length_bytes(variable)
112111

113-
dims, data, attrs, encoding = unpack_for_encoding(variable)
112+
dims, data, attrs, encoding, enum_meaning, enum_name = unpack(variable)
114113
if data.dtype.kind == "S" and encoding.get("dtype") is not str:
115114
data = bytes_to_char(data)
116115
if "char_dim_name" in encoding.keys():
117116
char_dim_name = encoding.pop("char_dim_name")
118117
else:
119118
char_dim_name = f"string{data.shape[-1]}"
120119
dims = dims + (char_dim_name,)
121-
return Variable(dims, data, attrs, encoding)
120+
return Variable(dims, data, attrs, encoding, enum_meaning=enum_meaning, enum_name=enum_name)
122121

123122
def decode(self, variable, name=None):
124-
dims, data, attrs, encoding = unpack_for_decoding(variable)
123+
dims, data, attrs, encoding, enum_meaning, enum_name = unpack(variable)
125124

126125
if data.dtype == "S1" and dims:
127126
encoding["char_dim_name"] = dims[-1]
128127
dims = dims[:-1]
129128
data = char_to_bytes(data)
130-
return Variable(dims, data, attrs, encoding)
129+
return Variable(dims, data, attrs, encoding, enum_meaning=enum_meaning, enum_name=enum_name)
131130

132131

133132
def bytes_to_char(arr):

‎xarray/coding/times.py‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,7 @@
1717
lazy_elemwise_func,
1818
pop_to,
1919
safe_setitem,
20-
unpack_for_decoding,
21-
unpack_for_encoding,
20+
unpack,
2221
)
2322
from xarray.core import indexing
2423
from xarray.core.common import contains_cftime_datetimes, is_np_datetime_like
@@ -702,22 +701,22 @@ def encode(self, variable: Variable, name: T_Name = None) -> Variable:
702701
if np.issubdtype(
703702
variable.data.dtype, np.datetime64
704703
) or contains_cftime_datetimes(variable):
705-
dims, data, attrs, encoding = unpack_for_encoding(variable)
704+
dims, data, attrs, encoding, enum_meaning, enum_name = unpack(variable)
706705

707706
(data, units, calendar) = encode_cf_datetime(
708707
data, encoding.pop("units", None), encoding.pop("calendar", None)
709708
)
710709
safe_setitem(attrs, "units", units, name=name)
711710
safe_setitem(attrs, "calendar", calendar, name=name)
712711

713-
return Variable(dims, data, attrs, encoding, fastpath=True)
712+
return Variable(dims, data, attrs, encoding, fastpath=True, enum_meaning=enum_meaning, enum_name=enum_name)
714713
else:
715714
return variable
716715

717716
def decode(self, variable: Variable, name: T_Name = None) -> Variable:
718717
units = variable.attrs.get("units", None)
719718
if isinstance(units, str) and "since" in units:
720-
dims, data, attrs, encoding = unpack_for_decoding(variable)
719+
dims, data, attrs, encoding, enum_meaning, enum_name = unpack(variable)
721720

722721
units = pop_to(attrs, encoding, "units")
723722
calendar = pop_to(attrs, encoding, "calendar")
@@ -730,33 +729,33 @@ def decode(self, variable: Variable, name: T_Name = None) -> Variable:
730729
)
731730
data = lazy_elemwise_func(data, transform, dtype)
732731

733-
return Variable(dims, data, attrs, encoding, fastpath=True)
732+
return Variable(dims, data, attrs, encoding, fastpath=True, enum_meaning=enum_meaning, enum_name=enum_name)
734733
else:
735734
return variable
736735

737736

738737
class CFTimedeltaCoder(VariableCoder):
739738
def encode(self, variable: Variable, name: T_Name = None) -> Variable:
740739
if np.issubdtype(variable.data.dtype, np.timedelta64):
741-
dims, data, attrs, encoding = unpack_for_encoding(variable)
740+
dims, data, attrs, encoding, enum_meaning, enum_name = unpack(variable)
742741

743742
data, units = encode_cf_timedelta(data, encoding.pop("units", None))
744743
safe_setitem(attrs, "units", units, name=name)
745744

746-
return Variable(dims, data, attrs, encoding, fastpath=True)
745+
return Variable(dims, data, attrs, encoding, fastpath=True, enum_meaning=enum_meaning, enum_name=enum_name)
747746
else:
748747
return variable
749748

750749
def decode(self, variable: Variable, name: T_Name = None) -> Variable:
751750
units = variable.attrs.get("units", None)
752751
if isinstance(units, str) and units in TIME_UNITS:
753-
dims, data, attrs, encoding = unpack_for_decoding(variable)
752+
dims, data, attrs, encoding, enum_meaning, enum_name = unpack(variable)
754753

755754
units = pop_to(attrs, encoding, "units")
756755
transform = partial(decode_cf_timedelta, units=units)
757756
dtype = np.dtype("timedelta64[ns]")
758757
data = lazy_elemwise_func(data, transform, dtype=dtype)
759758

760-
return Variable(dims, data, attrs, encoding, fastpath=True)
759+
return Variable(dims, data, attrs, encoding, fastpath=True, enum_meaning=enum_meaning, enum_name=enum_name)
761760
else:
762761
return variable

0 commit comments

Comments
 (0)