performance improvements

Author: shoyer

xarray/core/alignment.py

@@ -75,6 +75,11 @@ def align(*objects, **kwargs):
         raise TypeError('align() got unexpected keyword arguments: %s'
                         % list(kwargs))
 
+    if not indexes and len(objects) == 1:
+        # fast path for the trivial case
+        obj, = objects
+        return (obj.copy(deep=copy),)
+
     all_indexes = defaultdict(list)
     for obj in objects:
         for dim, index in iteritems(obj.indexes):
@@ -102,7 +107,12 @@ def align(*objects, **kwargs):
     for obj in objects:
         valid_indexers = dict((k, v) for k, v in joined_indexes.items()
                               if k in obj.dims)
-        result.append(obj.reindex(copy=copy, **valid_indexers))
+        if not valid_indexers:
+            # fast path for no reindexing necessary
+            new_obj = obj.copy(deep=copy)
+        else:
+            new_obj = obj.reindex(copy=copy, **valid_indexers)
+        result.append(new_obj)
 
     return tuple(result)
 

xarray/core/computation.py

@@ -11,7 +11,7 @@
 from .merge import _align_for_merge as deep_align
 from .merge import merge_coords_without_align
 from .utils import is_dict_like
-from .pycompat import dask_array_type, OrderedDict, basestring
+from .pycompat import dask_array_type, OrderedDict, basestring, suppress
 
 
 SLICE_NONE = slice(None)
@@ -138,22 +138,41 @@ def _default_result_attrs(attrs, func, signature):
 
 
 def build_output_coords(args, signature, new_coords=None):
+    coord_variables = []
+    for arg in args:
+        try:
+            coords = arg.coords
+        except AttributeError:
+            pass  # skip this argument
+        else:
+            coord_vars = getattr(coords, 'variables', coords)
+            coord_variables.append(coord_vars)
 
-    coord_variables = [getattr(getattr(arg, 'coords', {}), 'variables', {})
-                       for arg in args]
     if new_coords is not None:
         coord_variables.append(getattr(new_coords, 'variables', new_coords))
 
-    merged = merge_coords_without_align(coord_variables)
+    if len(args) == 1 and new_coords is None:
+        # we can skip the expensive merge
+        merged, = coord_variables
+    else:
+        merged = merge_coords_without_align(coord_variables)
+
+    missing_dims = signature.all_output_core_dims - set(merged)
+    if missing_dims:
+        raise ValueError('new output dimensions must have matching entries in '
+                         '`new_coords`: %r' % missing_dims)
 
-    output = []
+    output_coords = []
     for output_dims in signature.output_core_dims:
         dropped_dims = signature.all_input_core_dims - set(output_dims)
-        coords = OrderedDict((k, v) for k, v in merged.items()
-                             if set(v.dims).isdisjoint(dropped_dims))
-        output.append(coords)
+        if dropped_dims:
+            coords = OrderedDict((k, v) for k, v in merged.items()
+                                 if set(v.dims).isdisjoint(dropped_dims))
+        else:
+            coords = merged
+        output_coords.append(coords)
 
-    return output
+    return output_coords
 
 
 def apply_dataarray_ufunc(func, *args, **kwargs):
@@ -172,7 +191,8 @@ def apply_dataarray_ufunc(func, *args, **kwargs):
     if signature is None:
         signature = _default_signature(len(args))
 
-    args = deep_align(args, join=join, copy=False, raise_on_invalid=False)
+    if len(args) > 1:
+        args = deep_align(args, join=join, copy=False, raise_on_invalid=False)
 
     name = result_name(args)
     result_coords = build_output_coords(args, signature, new_coords)
@@ -181,16 +201,23 @@ def apply_dataarray_ufunc(func, *args, **kwargs):
     result_var = func(*data_vars)
 
     if signature.n_outputs > 1:
-        return tuple(DataArray(variable, coords, name=name)
+        return tuple(DataArray(variable, coords, name=name, fastpath=True)
                      for variable, coords in zip(result_var, result_coords))
     else:
         coords, = result_coords
-        return DataArray(result_var, coords, name=name)
+        return DataArray(result_var, coords, name=name, fastpath=True)
 
 
 def join_dict_keys(objects, how='inner'):
+    all_keys = [obj.keys() for obj in objects if hasattr(obj, 'keys')]
+
+    if len(all_keys) == 1:
+        # shortcut
+        result_keys, = all_keys
+        return result_keys
+
     joiner = _get_joiner(how)
-    all_keys = (obj.keys() for obj in objects if hasattr(obj, 'keys'))
+    # TODO: use a faster ordered set than a pandas.Index
     result_keys = joiner([pd.Index(keys) for keys in all_keys])
     return result_keys
 
@@ -203,6 +230,17 @@ def collect_dict_values(objects, keys, fill_value=None):
             for key in keys]
 
 
+def _fast_dataset(variables, coord_variables):
+    """Create a dataset as quickly as possible.
+
+    Variables are modified *inplace*.
+    """
+    from .dataset import Dataset
+    variables.update(coord_variables)
+    coord_names = set(coord_variables)
+    return Dataset._from_vars_and_coord_names(variables, coord_names)
+
+
 def apply_dataset_ufunc(func, *args, **kwargs):
     """
     def apply_dataset_ufunc(func, args, signature=None, join='inner',
@@ -221,7 +259,8 @@ def apply_dataset_ufunc(func, args, signature=None, join='inner',
     if signature is None:
         signature = _default_signature(len(args))
 
-    args = deep_align(args, join=join, copy=False, raise_on_invalid=False)
+    if len(args) > 1:
+        args = deep_align(args, join=join, copy=False, raise_on_invalid=False)
 
     list_of_coords = build_output_coords(args, signature, new_coords)
 
@@ -243,12 +282,11 @@ def apply_dataset_ufunc(func, args, signature=None, join='inner',
             for value, results_dict in zip(values, result_dict_list):
                 results_dict[name] = value
 
-        return tuple(Dataset(*args)
+        return tuple(_fast_dataset(*args)
                      for args in zip(result_dict_list, list_of_coords))
     else:
-        data_vars = result_vars
         coord_vars, = list_of_coords
-        return Dataset(data_vars, coord_vars)
+        return _fast_dataset(result_vars, coord_vars)
 
 
 def _iter_over_selections(obj, dim, values):

xarray/core/dataset.py

@@ -20,7 +20,8 @@
 from .merge import (dataset_update_method, dataset_merge_method,
                     merge_data_and_coords)
 from .utils import Frozen, SortedKeysDict, maybe_wrap_array, hashable
-from .variable import (Variable, as_variable, IndexVariable, broadcast_variables)
+from .variable import (Variable, as_variable, IndexVariable,
+                       broadcast_variables, default_index_coordinate)
 from .pycompat import (iteritems, basestring, OrderedDict,
                        dask_array_type)
 from .combine import concat
@@ -87,9 +88,12 @@ def calculate_dimensions(variables):
     return dims
 
 
-def _assert_empty(args, msg='%s'):
-    if args:
-        raise ValueError(msg % args)
+def add_default_dim_coords_inplace(variables, dims):
+    # type: (MutableMapping[object, Variable], Mapping[object, int]) -> None
+    """Add missing coordinates to variables inplace."""
+    for dim, size in iteritems(dims):
+        if dim not in variables:
+            variables[dim] = default_index_coordinate(dim, size)
 
 
 def as_dataset(obj):
@@ -206,13 +210,6 @@ def __init__(self, data_vars=None, coords=None, attrs=None,
             self.attrs = attrs
         self._initialized = True
 
-    def _add_missing_coords_inplace(self):
-        """Add missing coordinates to self._variables
-        """
-        for dim, size in iteritems(self.dims):
-            if dim not in self._variables:
-                self._variables[dim] = default_index_coordinate(dim, size)
-
     def _set_init_vars_and_dims(self, data_vars, coords, compat):
         """Set the initial value of Dataset variables and dimensions
         """
@@ -678,9 +675,11 @@ def reset_coords(self, names=None, drop=False, inplace=False):
             if isinstance(names, basestring):
                 names = [names]
             self._assert_all_in_dataset(names)
-            _assert_empty(
-                set(names) & set(self.dims),
-                'cannot remove index coordinates with reset_coords: %s')
+            bad_coords = set(names) & set(self.dims)
+            if bad_coords:
+                raise ValueError(
+                    'cannot remove index coordinates with reset_coords: %s'
+                    % bad_coords)
         obj = self if inplace else self.copy()
         obj._coord_names.difference_update(names)
         if drop:
@@ -1695,8 +1694,10 @@ def reduce(self, func, dim=None, keep_attrs=False, numeric_only=False,
         else:
             dims = set(dim)
 
-        _assert_empty([dim for dim in dims if dim not in self.dims],
-                      'Dataset does not contain the dimensions: %s')
+        missing_dimensions = [dim for dim in dims if dim not in self.dims]
+        if missing_dimensions:
+            raise ValueError('Dataset does not contain the dimensions: %s'
+                             % missing_dimensions)
 
         variables = OrderedDict()
         for name, var in iteritems(self._variables):

xarray/core/merge.py

@@ -2,7 +2,7 @@
 
 from .alignment import align
 from .utils import Frozen, is_dict_like
-from .variable import as_variable, default_index_coordinate
+from .variable import as_variable
 from .pycompat import (basestring, OrderedDict)
 
 
@@ -423,7 +423,7 @@ def merge_core(objs, compat='broadcast_equals', join='outer', priority_arg=None,
     ------
     MergeError if the merge cannot be done successfully.
     """
-    from .dataset import calculate_dimensions
+    from .dataset import calculate_dimensions, add_default_dim_coords_inplace
 
     _assert_compat_valid(compat)
 
@@ -440,10 +440,7 @@ def merge_core(objs, compat='broadcast_equals', join='outer', priority_arg=None,
     variables = merge_variables(expanded, priority_vars, compat=compat)
 
     dims = calculate_dimensions(variables)
-
-    for dim, size in dims.items():
-        if dim not in variables:
-            variables[dim] = default_index_coordinate(dim, size)
+    add_default_dim_coords_inplace(variables, dims)
 
     coord_names.update(dims)
 

xarray/core/utils.py

@@ -155,7 +155,7 @@ def remove_incompatible_items(first_dict, second_dict, compat=equivalent):
 
 
 def is_dict_like(value):
-    return hasattr(value, '__getitem__') and hasattr(value, 'keys')
+    return hasattr(value, 'keys') and hasattr(value, '__getitem__')
 
 
 def is_full_slice(value):

xarray/test/test_computation.py

@@ -217,13 +217,11 @@ def stack2(obj):
         # no new_coords
         return xr.apply_ufunc(func, obj, signature=sig)
 
-    actual = stack2(data_array)
-    expected_data_array.coords['sign'] = [0, 1]
-    assert_identical(actual, expected_data_array)
+    with pytest.raises(ValueError):
+        stack2(data_array)
 
-    actual = stack2(dataset)
-    expected_dataset.coords['sign'] = [0, 1]
-    assert_identical(actual, expected_dataset)
+    with pytest.raises(ValueError):
+        stack2(dataset)
 
 
 def test_broadcast_compat_data_1d():