From 7df3645ee9126370cd03250df5d337f10075fc17 Mon Sep 17 00:00:00 2001 From: Angus Hollands Date: Wed, 29 Nov 2023 13:41:58 +0000 Subject: [PATCH] fix: merging --- src/awkward/_do/content.py | 13 +++++++++--- src/awkward/_do/meta.py | 22 ++++++--------------- src/awkward/_meta/numpymeta.py | 14 +++++++++---- src/awkward/_meta/recordmeta.py | 4 ++-- src/awkward/_nplikes/array_module.py | 5 ----- src/awkward/_nplikes/numpy_like.py | 7 ++----- src/awkward/_nplikes/typetracer.py | 5 ----- src/awkward/forms/numpyform.py | 6 ++++++ src/awkward/operations/ak_strings_astype.py | 4 ++-- src/awkward/operations/ak_values_astype.py | 3 ++- src/awkward/operations/str/__init__.py | 2 +- src/awkward/typetracer.py | 2 +- 12 files changed, 42 insertions(+), 45 deletions(-) diff --git a/src/awkward/_do/content.py b/src/awkward/_do/content.py index 271a175e56..7c63904c59 100644 --- a/src/awkward/_do/content.py +++ b/src/awkward/_do/content.py @@ -9,14 +9,16 @@ import awkward as ak from awkward._backends.backend import Backend from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._typing import Any, AxisMaybeNone, Literal -from awkward.contents.content import ActionType, Content +from awkward._typing import TYPE_CHECKING, Any, AxisMaybeNone, Literal from awkward.errors import AxisError from awkward.forms import form -from awkward.record import Record np = NumpyMetadata.instance() +if TYPE_CHECKING: + from awkward.contents.content import ActionType, Content + from awkward.record import Record + def recursively_apply( layout: Content | Record, @@ -32,6 +34,9 @@ def recursively_apply( function_name: str | None = None, regular_to_jagged=False, ) -> Content | Record | None: + from awkward.contents.content import Content + from awkward.record import Record + if isinstance(layout, Content): return layout._recursively_apply( action, @@ -201,6 +206,8 @@ def remove_structure( allow_records: bool = False, list_to_regular: bool = False, ): + from awkward.record import Record + if isinstance(layout, Record): return remove_structure( layout._array[layout._at : layout._at + 1], diff --git a/src/awkward/_do/meta.py b/src/awkward/_do/meta.py index 3176d649cf..09ddca9dfb 100644 --- a/src/awkward/_do/meta.py +++ b/src/awkward/_do/meta.py @@ -27,15 +27,11 @@ def is_option( meta: Meta -) -> TypeGuard[ - IndexedOptionMeta | BitMaskedMeta | ByteMaskedMeta | UnmaskedMeta -]: +) -> TypeGuard[IndexedOptionMeta | BitMaskedMeta | ByteMaskedMeta | UnmaskedMeta]: return meta.is_option -def is_list( - meta: Meta -) -> TypeGuard[RegularMeta | ListOffsetMeta | ListMeta]: +def is_list(meta: Meta) -> TypeGuard[RegularMeta | ListOffsetMeta | ListMeta]: return meta.is_list @@ -59,19 +55,13 @@ def is_indexed(meta: Meta) -> TypeGuard[IndexedOptionMeta, IndexedMeta]: return meta.is_indexed -class ImplementsTuple(RecordMeta): # Intersection - _fields: None - - -def is_record_tuple(meta: Meta) -> TypeGuard[ImplementsTuple]: +# FIXME: narrow this to have `is_tuple` be a const True +def is_record_tuple(meta: Meta) -> TypeGuard[RecordMeta]: return meta.is_record and meta.is_tuple -class ImplementsRecord(RecordMeta): - _fields: list[str] - - -def is_record_record(meta: Meta) -> TypeGuard[ImplementsRecord]: +# FIXME: narrow this to have `is_tuple` be a const False +def is_record_record(meta: Meta) -> TypeGuard[RecordMeta]: return meta.is_record and not meta.is_tuple diff --git a/src/awkward/_meta/numpymeta.py b/src/awkward/_meta/numpymeta.py index 7175cd7ac7..e81241f04e 100644 --- a/src/awkward/_meta/numpymeta.py +++ b/src/awkward/_meta/numpymeta.py @@ -4,10 +4,12 @@ from awkward._do.meta import is_indexed, is_numpy, is_option from awkward._meta.meta import Meta +from awkward._nplikes.numpy_like import NumpyMetadata from awkward._nplikes.shape import ShapeItem from awkward._parameters import type_parameters_equal -from awkward._typing import TYPE_CHECKING, JSONSerializable +from awkward._typing import TYPE_CHECKING, DType, JSONSerializable +np = NumpyMetadata.instance() if TYPE_CHECKING: from awkward._meta.regularmeta import RegularMeta @@ -17,6 +19,10 @@ class NumpyMeta(Meta): is_leaf = True inner_shape: tuple[ShapeItem, ...] + @property + def dtype(self) -> DType: + raise NotImplementedError + def purelist_parameters(self, *keys: str) -> JSONSerializable: if self._parameters is not None: for key in keys: @@ -104,9 +110,9 @@ def _mergeable_next(self, other: Meta, mergebool: bool) -> bool: # Default merging (can we cast one to the other) else: - return self.backend.nplike.can_cast( - self.dtype, other.dtype - ) or self.backend.nplike.can_cast(other.dtype, self.dtype) + return np.can_cast( + self.dtype, other.dtype, casting="same_kind" + ) or np.can_cast(other.dtype, self.dtype, casting="same_kind") else: return False diff --git a/src/awkward/_meta/recordmeta.py b/src/awkward/_meta/recordmeta.py index 7448227159..61cad64ed5 100644 --- a/src/awkward/_meta/recordmeta.py +++ b/src/awkward/_meta/recordmeta.py @@ -170,10 +170,10 @@ def _mergeable_next(self, other: Meta, mergebool: bool) -> bool: return False elif is_record_record(self) and is_record_record(other): - if set(self._fields) != set(other._fields): + if set(self._fields) != set(other._fields): # type: ignore[arg-type] return False - for i, field in enumerate(self._fields): + for i, field in enumerate(self._fields): # type: ignore[arg-type] x = self._contents[i] y = other.contents[other.field_to_index(field)] if not x._mergeable_next(y, mergebool): diff --git a/src/awkward/_nplikes/array_module.py b/src/awkward/_nplikes/array_module.py index 82ef7a11d1..7d5e4bb151 100644 --- a/src/awkward/_nplikes/array_module.py +++ b/src/awkward/_nplikes/array_module.py @@ -653,11 +653,6 @@ def astype( assert not isinstance(x, PlaceholderArray) return x.astype(dtype, copy=copy) # type: ignore[attr-defined] - def can_cast( - self, from_: DTypeLike | ArrayLikeT, to: DTypeLike | ArrayLikeT - ) -> bool: - return self._module.can_cast(from_, to, casting="same_kind") - @classmethod def is_own_array(cls, obj) -> bool: return cls.is_own_array_type(type(obj)) diff --git a/src/awkward/_nplikes/numpy_like.py b/src/awkward/_nplikes/numpy_like.py index 19119d754c..7cf550adbe 100644 --- a/src/awkward/_nplikes/numpy_like.py +++ b/src/awkward/_nplikes/numpy_like.py @@ -83,7 +83,8 @@ class NumpyMetadata(PublicSingleton): datetime_data = staticmethod(numpy.datetime_data) issubdtype = staticmethod(numpy.issubdtype) - AxisError = numpy.AxisError + AxisError = staticmethod(numpy.AxisError) + can_cast = staticmethod(numpy.can_cast) if hasattr(numpy, "float16"): @@ -537,10 +538,6 @@ def astype( ) -> ArrayLikeT: ... - @abstractmethod - def can_cast(self, from_: DType | ArrayLikeT, to: DType | ArrayLikeT) -> bool: - ... - @abstractmethod def is_c_contiguous(self, x: ArrayLikeT | PlaceholderArray) -> bool: ... diff --git a/src/awkward/_nplikes/typetracer.py b/src/awkward/_nplikes/typetracer.py index 236626ef02..423af36e7c 100644 --- a/src/awkward/_nplikes/typetracer.py +++ b/src/awkward/_nplikes/typetracer.py @@ -1512,11 +1512,6 @@ def astype( x.touch_data() return TypeTracerArray._new(np.dtype(dtype), x.shape) - def can_cast( - self, from_: DTypeLike | TypeTracerArray, to: DTypeLike | TypeTracerArray - ) -> bool: - return numpy.can_cast(from_, to, casting="same_kind") - @classmethod def is_own_array_type(cls, type_: type) -> bool: return issubclass(type_, TypeTracerArray) diff --git a/src/awkward/forms/numpyform.py b/src/awkward/forms/numpyform.py index 90b2139c0b..ec5344cd74 100644 --- a/src/awkward/forms/numpyform.py +++ b/src/awkward/forms/numpyform.py @@ -273,3 +273,9 @@ def _expected_from_buffers( def _to_regular_primitive(self) -> RegularForm | NumpyForm: return self.to_RegularForm() + + @property + def dtype(self) -> DType: + from awkward.types.numpytype import primitive_to_dtype + + return primitive_to_dtype(self.primitive) diff --git a/src/awkward/operations/ak_strings_astype.py b/src/awkward/operations/ak_strings_astype.py index 16ff6eeba0..3f799068d1 100644 --- a/src/awkward/operations/ak_strings_astype.py +++ b/src/awkward/operations/ak_strings_astype.py @@ -4,7 +4,7 @@ import awkward as ak from awkward._dispatch import high_level_function -from awkward._do.content import recursively_apply +from awkward._do.content import pad_none, recursively_apply from awkward._layout import HighLevelContext from awkward._nplikes.numpy import Numpy from awkward._nplikes.numpy_like import NumpyMetadata @@ -68,7 +68,7 @@ def action(layout, **kwargs): layout, highlevel=False, behavior=behavior ) max_length = ak.operations.max(ak.operations.num(layout, behavior=behavior)) - regulararray = ak._do.pad_none(layout, max_length, 1) + regulararray = pad_none(layout, max_length, 1) maskedarray = ak.operations.to_numpy(regulararray, allow_missing=True) npstrings = maskedarray.data if maskedarray.mask is not False: diff --git a/src/awkward/operations/ak_values_astype.py b/src/awkward/operations/ak_values_astype.py index 714a4320d9..49d1207ff9 100644 --- a/src/awkward/operations/ak_values_astype.py +++ b/src/awkward/operations/ak_values_astype.py @@ -4,6 +4,7 @@ import awkward as ak from awkward._dispatch import high_level_function +from awkward._do.content import numbers_to_type from awkward._layout import HighLevelContext from awkward._nplikes.numpy_like import NumpyMetadata @@ -73,5 +74,5 @@ def _impl(array, to, including_unknown, highlevel, behavior, attrs): with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") to_str = ak.types.numpytype.dtype_to_primitive(np.dtype(to)) - out = ak._do.numbers_to_type(layout, to_str, including_unknown) + out = numbers_to_type(layout, to_str, including_unknown) return ctx.wrap(out, highlevel=highlevel) diff --git a/src/awkward/operations/str/__init__.py b/src/awkward/operations/str/__init__.py index 3a6d776904..e362571a72 100644 --- a/src/awkward/operations/str/__init__.py +++ b/src/awkward/operations/str/__init__.py @@ -85,7 +85,7 @@ def _drop_option_preserving_form(layout, ensure_empty_mask: bool = False): - from awkward._do import recursively_apply + from awkward._do.content import recursively_apply from awkward.contents import UnmaskedArray, IndexedOptionArray, IndexedArray def action(_, continuation, **kwargs): diff --git a/src/awkward/typetracer.py b/src/awkward/typetracer.py index abaf542ceb..6cbcacf246 100644 --- a/src/awkward/typetracer.py +++ b/src/awkward/typetracer.py @@ -6,7 +6,7 @@ import awkward.forms from awkward._backends.typetracer import TypeTracerBackend -from awkward._do import touch_data as _touch_data +from awkward._do.content import touch_data as _touch_data from awkward._layout import HighLevelContext, wrap_layout from awkward._nplikes.numpy import NumpyMetadata from awkward._nplikes.placeholder import PlaceholderArray