diff --git a/datashader/core.py b/datashader/core.py index 7fb81fab0..0a21601ab 100644 --- a/datashader/core.py +++ b/datashader/core.py @@ -11,7 +11,7 @@ from xarray import DataArray, Dataset from .utils import Dispatcher, ngjit, calc_res, calc_bbox, orient_array, \ - dshape_from_xarray_dataset + dshape_from_xarray_dataset, _categorize_dask_columns from .utils import get_indices, dshape_from_pandas, dshape_from_dask from .utils import Expr # noqa (API import) from .resampling import resample_2d, resample_2d_distributed @@ -1309,36 +1309,7 @@ def _source_from_geopandas(self, source): else: return None - -def bypixel(source, canvas, glyph, agg, *, antialias=False): - """Compute an aggregate grouped by pixel sized bins. - - Aggregate input data ``source`` into a grid with shape and axis matching - ``canvas``, mapping data to bins by ``glyph``, and aggregating by reduction - ``agg``. - - Parameters - ---------- - source : pandas.DataFrame, dask.DataFrame - Input datasource - canvas : Canvas - glyph : Glyph - agg : Reduction - """ - source, dshape = _bypixel_sanitise(source, glyph, agg) - - schema = dshape.measure - glyph.validate(schema) - agg.validate(schema) - canvas.validate() - - # All-NaN objects (e.g. chunks of arrays with no data) are valid in Datashader - with warnings.catch_warnings(): - warnings.filterwarnings('ignore', r'All-NaN (slice|axis) encountered') - return bypixel.pipeline(source, schema, canvas, glyph, agg, antialias=antialias) - - -def _bypixel_sanitise(source, glyph, agg): +def _sanitize_xarray(source, canvas, glyph, agg): # Convert 1D xarray DataArrays and DataSets into Dask DataFrames if isinstance(source, DataArray) and source.ndim == 1: if not source.name: @@ -1352,9 +1323,10 @@ def _bypixel_sanitise(source, glyph, agg): source = source.to_dask_dataframe() else: source = source.to_dataframe() + return source - if (isinstance(source, pd.DataFrame) or - (cudf and isinstance(source, cudf.DataFrame))): +def _sanitize_dataframe(source, canvas, glyph, agg): + if isinstance(source, (pd.DataFrame, cudf.DataFrame) if cudf else pd.DataFrame): # Avoid datashape.Categorical instantiation bottleneck # by only retaining the necessary columns: # https://github.com/bokeh/datashader/issues/396 @@ -1372,16 +1344,107 @@ def _bypixel_sanitise(source, glyph, agg): if (sindex is not None and getattr(source[glyph.geometry].array, "_sindex", None) is None): source[glyph.geometry].array._sindex = sindex - dshape = dshape_from_pandas(source) elif dd and isinstance(source, dd.DataFrame): - dshape, source = dshape_from_dask(source) + # Categorize unknown categorical columns (memoized) + source = _categorize_dask_columns(source) + return source + + +def _patch_temporal(source, canvas, glyph, agg): + if not hasattr(glyph, "x") or not hasattr(glyph, "y"): + return source, None + x, y = str(glyph.x), str(glyph.y) + dtypes = dict(source.dtypes) + if x not in dtypes or y not in dtypes: + return source, None + xkind, ykind = dtypes[x].kind, dtypes[y].kind + is_temporal, original_ranges = {}, {} + + # Handle temporal types (datetime64='M', timedelta64='m') + for col, kind, range_attr in ((x, xkind, 'x_range'), (y, ykind, 'y_range')): + if kind in "Mm": + # Remove timezone if present and convert to int64 + source_col = source[col] + if getattr(dtypes[col], "tz", None): + source_col = source_col.dt.tz_localize(None) + dtypes[col] = source_col.dtype + source[col] = source_col.astype(np.int64, copy=False) + is_temporal[col] = True + + # Convert canvas range to int64, preserving original + canvas_range = getattr(canvas, range_attr) + if canvas_range: + original_ranges[range_attr] = canvas_range + fn = ( + (lambda x: pd.to_datetime(x).tz_localize(None)) + if kind == "M" else pd.to_timedelta + ) + setattr(canvas, range_attr, tuple( + fn(canvas_range).to_numpy().astype(dtypes[col]).astype(np.int64) + )) + + if not is_temporal: + return source, None + + def post_temporal(output): + # Restore temporal types in output (converted to float by compute_scale_and_translate) + for col, kind, range_attr in ((x, xkind, 'x_range'), (y, ykind, 'y_range')): + if kind in "Mm": + output[col] = output[col].data.astype(np.int64).view(dtypes[col]) + output.attrs[range_attr] = tuple( + np.asarray(output.attrs[range_attr]).astype(dtypes[col]) + ) + return output + + return source, post_temporal + + +def _get_schema(source): + """Extract schema from the data source.""" + if isinstance(source, (pd.DataFrame, cudf.DataFrame) if cudf else pd.DataFrame): + return dshape_from_pandas(source).measure + elif dd and isinstance(source, dd.DataFrame): + return dshape_from_dask(source).measure elif isinstance(source, Dataset): - # Multi-dimensional Dataset - dshape = dshape_from_xarray_dataset(source) + return dshape_from_xarray_dataset(source).measure else: - raise ValueError("source must be a pandas or dask DataFrame") + raise ValueError("source must be a pandas, dask, cuDF DataFrame or xarray Dataset") + + +def bypixel(source, canvas, glyph, agg, *, antialias=False): + """Compute an aggregate grouped by pixel sized bins. + + Aggregate input data ``source`` into a grid with shape and axis matching + ``canvas``, mapping data to bins by ``glyph``, and aggregating by reduction + ``agg``. + + Parameters + ---------- + source : pandas.DataFrame, dask.DataFrame + Input datasource + canvas : Canvas + glyph : Glyph + agg : Reduction + """ + # Pre functions, note order matters + source = _sanitize_xarray(source, canvas, glyph, agg) + source = _sanitize_dataframe(source, canvas, glyph, agg) + source, post_temporal = _patch_temporal(source, canvas, glyph, agg) + + schema = _get_schema(source) + glyph.validate(schema) + agg.validate(schema) + canvas.validate() + + # All-NaN objects (e.g. chunks of arrays with no data) are valid in Datashader + with warnings.catch_warnings(): + warnings.filterwarnings('ignore', r'All-NaN (slice|axis) encountered') + output = bypixel.pipeline(source, schema, canvas, glyph, agg, antialias=antialias) + + # Post functions + output = post_temporal(output) if post_temporal else output - return source, dshape + return output def _cols_to_keep(columns, glyph, agg): diff --git a/datashader/glyphs/glyph.py b/datashader/glyphs/glyph.py index 6ea11cd0a..98ec80334 100644 --- a/datashader/glyphs/glyph.py +++ b/datashader/glyphs/glyph.py @@ -63,9 +63,16 @@ def _compute_bounds(s): else: return Glyph._compute_bounds_numba(s) + @staticmethod - @ngjit def _compute_bounds_numba(arr): + if arr.dtype.kind == "i": + return Glyph._compute_bounds_numba_int(arr) + return Glyph._compute_bounds_numba_float(arr) + + @staticmethod + @ngjit + def _compute_bounds_numba_float(arr): minval = np.inf maxval = -np.inf for x in arr: @@ -74,7 +81,20 @@ def _compute_bounds_numba(arr): minval = x if x > maxval: maxval = x + return minval, maxval + @staticmethod + @ngjit + def _compute_bounds_numba_int(arr): + minval = 2 ** 64 // 2 - 1 + maxval = - 2 ** 64 // 2 + natval = - 2 ** 64 // 2 + for x in arr: + if x != natval: + if x < minval: + minval = x + if x > maxval: + maxval = x return minval, maxval @staticmethod diff --git a/datashader/tests/test_polygons.py b/datashader/tests/test_polygons.py index 9e0124afe..e012dddf9 100644 --- a/datashader/tests/test_polygons.py +++ b/datashader/tests/test_polygons.py @@ -320,7 +320,7 @@ def test_spatial_index_not_dropped(): glyph = ds.glyphs.polygon.PolygonGeom('some_geom') agg = ds.count() - df2, _ = ds.core._bypixel_sanitise(df, glyph, agg) + df2 = ds.core._sanitize_dataframe(df, None, glyph, agg) assert df2.columns == ['some_geom'] assert df2.some_geom.array._sindex == df.some_geom.array._sindex diff --git a/datashader/utils.py b/datashader/utils.py index 778113a4d..18ddf158f 100644 --- a/datashader/utils.py +++ b/datashader/utils.py @@ -457,21 +457,34 @@ def dshape_from_pandas(df): return len(df) * datashape.Record([(k, dshape_from_pandas_helper(df[k])) for k in df.columns]) +@memoize(key=lambda args, kwargs: tuple(args[0].__dask_keys__())) +def _categorize_dask_columns(df): + """Categorize unknown categorical columns in a dask DataFrame. + + Returns the categorized DataFrame. Memoized to avoid redundant computation. + """ + cat_columns = [ + col for col in df.columns + if ( + isinstance(df[col].dtype, ( + type(pd.Categorical.dtype), + pd.api.types.CategoricalDtype, + )) and not getattr(df[col].cat, 'known', True) + ) + ] + if cat_columns: + return df.categorize(cat_columns, index=False) + return df + @memoize(key=lambda args, kwargs: tuple(args[0].__dask_keys__())) def dshape_from_dask(df): """Return a datashape.DataShape object given a dask dataframe.""" - cat_columns = [ - col for col in df.columns - if (isinstance(df[col].dtype, type(pd.Categorical.dtype)) or - isinstance(df[col].dtype, pd.api.types.CategoricalDtype)) - and not getattr(df[col].cat, 'known', True)] - df = df.categorize(cat_columns, index=False) # get_partition(0) used below because categories are sometimes repeated # for dask-cudf DataFrames with multiple partitions return datashape.var * datashape.Record([ (k, dshape_from_pandas_helper(df[k].get_partition(0))) for k in df.columns - ]), df + ]) def dshape_from_xarray_dataset(xr_ds):