diff --git a/dask_expr/_collection.py b/dask_expr/_collection.py index 370a3afc..b317efb9 100644 --- a/dask_expr/_collection.py +++ b/dask_expr/_collection.py @@ -121,6 +121,7 @@ ) from dask_expr._str_accessor import StringAccessor from dask_expr._util import ( + PANDAS_GE_300, _BackendData, _convert_to_list, _get_shuffle_preferring_order, @@ -400,7 +401,7 @@ def __getitem__(self, other): if ( self.ndim == 2 and is_integer_slice - and not is_float_dtype(self.index.dtype) + and (not is_float_dtype(self.index.dtype) or PANDAS_GE_300) ): return self.iloc[other] else: @@ -1577,6 +1578,10 @@ def std( if needs_time_conversion: numeric_dd = _convert_to_numeric(self, skipna) + units = None + if needs_time_conversion and time_cols is not None: + units = [getattr(self._meta[c].array, "unit", None) for c in time_cols] + if axis == 1: return numeric_dd.map_partitions( M.std if not needs_time_conversion else _sqrt_and_convert_to_timedelta, @@ -1598,6 +1603,8 @@ def std( "time_cols": time_cols, "axis": axis, "dtype": getattr(meta, "dtype", None), + "unit": getattr(meta, "unit", None), + "units": units, } sqrt_func = _sqrt_and_convert_to_timedelta else: diff --git a/dask_expr/_groupby.py b/dask_expr/_groupby.py index d3c2b64d..56297bac 100644 --- a/dask_expr/_groupby.py +++ b/dask_expr/_groupby.py @@ -1084,6 +1084,8 @@ def npartitions(self): def groupby_get_group(df, *by_key, get_key=None, columns=None): + if PANDAS_GE_300 and is_scalar(get_key): + get_key = (get_key,) return _groupby_get_group(df, list(by_key), get_key, columns) diff --git a/dask_expr/_indexing.py b/dask_expr/_indexing.py index 493263d8..8fb10c19 100644 --- a/dask_expr/_indexing.py +++ b/dask_expr/_indexing.py @@ -87,7 +87,9 @@ def _loc(self, iindexer, cindexer): return self._loc(iindexer(self.obj), cindexer) if self.obj.known_divisions: - iindexer = self._maybe_partial_time_string(iindexer) + idx = self.obj.index._meta + unit = idx.unit if hasattr(idx, "unit") else None + iindexer = self._maybe_partial_time_string(iindexer, unit=unit) if isinstance(iindexer, slice): return self._loc_slice(iindexer, cindexer) @@ -133,13 +135,13 @@ def _loc_array(self, iindexer, cindexer): ) return self._loc_series(iindexer_series, cindexer, check_alignment=False) - def _maybe_partial_time_string(self, iindexer): + def _maybe_partial_time_string(self, iindexer, unit): """ Convert index-indexer for partial time string slicing if obj.index is DatetimeIndex / PeriodIndex """ idx = meta_nonempty(self.obj._meta.index) - iindexer = _maybe_partial_time_string(idx, iindexer) + iindexer = _maybe_partial_time_string(idx, iindexer, unit) return iindexer def _loc_slice(self, iindexer, cindexer):