feat: Add temporal support

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):

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

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

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):