Multi-dimensional GPU-based data filtering (#335)

kylebarron · web-flow · commit f5544e4cabb5 · 2024-01-29T12:20:00.000-05:00
Filter on up to 4 numeric attributes at the same time! This is working
locally, but I need to figure out a clean way to connect it to the
filter_value, since that's a list of lists.

The best solution (to enable jslink) might be to create our own custom
widget that can display 1-4 float sliders and stores its value as a list
of lists?
diff --git a/docs/api/experimental/traits.md b/docs/api/experimental/traits.md
@@ -1,3 +1,4 @@
 # lonboard.experimental.traits
 
 ::: lonboard.experimental.traits.PointAccessor
+::: lonboard.experimental.traits.GetFilterValueAccessor
diff --git a/lonboard/experimental/layer_extension.py b/lonboard/experimental/layer_extension.py
@@ -1,7 +1,7 @@
 import traitlets
 
 from lonboard._base import BaseExtension
-from lonboard.experimental.traits import PointAccessor
+from lonboard.experimental.traits import GetFilterValueAccessor, PointAccessor
 from lonboard.traits import FloatAccessor
 
 
@@ -160,12 +160,19 @@ class DataFilterExtension(BaseExtension):
 
     ## `filter_range`
 
-    The (min, max) bounds which defines whether an object should be rendered.
+    The bounds which defines whether an object should be rendered. If an object's
+    filtered value is within the bounds, the object will be rendered; otherwise it will
+    be hidden. This prop can be updated on user input or animation with very little
+    cost.
 
-    If an object's filtered value is within the bounds, the object will be rendered;
-    otherwise it will be hidden.
+    Format:
 
-    - Type: Tuple[float, float], optional
+    If `filter_size` is 1, provide a single tuple of `(min, max)`.
+
+    If `filter_size` is 2 to 4, provide a list of tuples: `[(min0, max0), (min1,
+    max1), ...]` for each filtered property, respectively.
+
+    - Type: either Tuple[float, float] or List[Tuple[float, float]], optional
     - Default: `(-1, 1)`
 
     ## `filter_soft_range`
@@ -199,8 +206,9 @@ class DataFilterExtension(BaseExtension):
 
     Accessor to retrieve the value for each object that it will be filtered by.
 
-    - Type: [FloatAccessor][lonboard.traits.FloatAccessor]
-        - If a number is provided, it is used as the value for all objects.
+    - Type:
+      [GetFilterValueAccessor][lonboard.experimental.traits.GetFilterValueAccessor]
+        - If a scalar value is provided, it is used as the value for all objects.
         - If an array is provided, each value in the array will be used as the value
           for the object at the same row index.
     """
@@ -209,22 +217,25 @@ class DataFilterExtension(BaseExtension):
 
     _layer_traits = {
         "filter_enabled": traitlets.Bool(True).tag(sync=True),
-        "filter_range": traitlets.Tuple(
-            traitlets.Float(), traitlets.Float(), default_value=(-1, 1)
+        "filter_range": traitlets.Union(
+            [
+                traitlets.List(traitlets.Float(), minlen=2, maxlen=2),
+                traitlets.List(
+                    traitlets.List(traitlets.Float(), minlen=2, maxlen=2),
+                    minlen=2,
+                    maxlen=4,
+                ),
+            ]
         ).tag(sync=True),
         "filter_soft_range": traitlets.Tuple(
             traitlets.Float(), traitlets.Float(), default_value=None, allow_none=True
         ).tag(sync=True),
         "filter_transform_size": traitlets.Bool(True).tag(sync=True),
         "filter_transform_color": traitlets.Bool(True).tag(sync=True),
-        "get_filter_value": FloatAccessor(None, allow_none=False),
+        "get_filter_value": GetFilterValueAccessor(None, allow_none=False),
     }
 
-    # TODO: support filterSize > 1
-    # In order to support filterSize > 1, we need to allow the get_filter_value accessor
-    # to be either a single float or a fixed size list of up to 4 floats.
-
-    # filter_size = traitlets.Int(1).tag(sync=True)
+    filter_size = traitlets.Int(1, min=1, max=4).tag(sync=True)
     """The size of the filter (number of columns to filter by).
 
     The data filter can show/hide data based on 1-4 numeric properties of each object.
diff --git a/lonboard/experimental/traits.py b/lonboard/experimental/traits.py
@@ -44,39 +44,37 @@ def validate(
     ) -> Union[Tuple[int, ...], List[int], pa.ChunkedArray, pa.FixedSizeListArray]:
         if isinstance(value, np.ndarray):
             if value.ndim != 2:
-                self.error(obj, value, info="Point array must have 2 dimensions.")
+                self.error(obj, value, info="Point array to have 2 dimensions")
 
             list_size = value.shape[1]
             if list_size not in (2, 3):
                 self.error(
                     obj,
                     value,
-                    info="Point array must have 2 or 3 as its second dimension.",
+                    info="Point array to have 2 or 3 as its second dimension",
                 )
 
             return pa.FixedSizeListArray.from_arrays(value.flatten("C"), list_size)
 
         if isinstance(value, (pa.ChunkedArray, pa.Array)):
             if not pa.types.is_fixed_size_list(value.type):
-                self.error(
-                    obj, value, info="Point pyarrow array must be a FixedSizeList."
-                )
+                self.error(obj, value, info="Point pyarrow array to be a FixedSizeList")
 
             if value.type.list_size not in (2, 3):
                 self.error(
                     obj,
                     value,
                     info=(
-                        "Color pyarrow array must have a FixedSizeList inner size of "
-                        "2 or 3."
+                        "Color pyarrow array to be a FixedSizeList with list size of "
+                        "2 or 3"
                     ),
                 )
 
             if not pa.types.is_floating(value.type.value_type):
                 self.error(
                     obj,
                     value,
-                    info="Point pyarrow array must have a floating point child.",
+                    info="Point pyarrow array to have a floating point child",
                 )
 
             return value
@@ -89,8 +87,8 @@ def validate(
                     obj,
                     value,
                     info=(
-                        "Color string must be a hex string interpretable by "
-                        "matplotlib.colors.to_rgba."
+                        "Color string to be a hex string interpretable by "
+                        "matplotlib.colors.to_rgba"
                     ),
                 )
                 return
@@ -99,3 +97,191 @@ def validate(
 
         self.error(obj, value)
         assert False
+
+
+class GetFilterValueAccessor(FixedErrorTraitType):
+    """
+    A trait to validate input for the `get_filter_value` accessor added by the
+    [`DataFilterExtension`][lonboard.experimental.DataFilterExtension], which can have
+    between 1 and 4 float values per row.
+
+    Various input is allowed:
+
+    - An `int` or `float`. This will be used as the value for all objects. The
+      `filter_size` of the
+      [`DataFilterExtension`][lonboard.experimental.DataFilterExtension] instance must
+      be 1.
+    - A one-dimensional numpy `ndarray` with a numeric data type. This will be casted to
+      an array of data type [`np.float32`][numpy.float32]. Each value in the array will
+      be used as the value for the object at the same row index. The `filter_size` of
+      the [`DataFilterExtension`][lonboard.experimental.DataFilterExtension] instance
+      must be 1.
+    - A two-dimensional numpy `ndarray` with a numeric data type. This will be casted to
+      an array of data type [`np.float32`][numpy.float32]. Each value in the array will
+      be used as the value for the object at the same row index. The `filter_size` of
+      the [`DataFilterExtension`][lonboard.experimental.DataFilterExtension] instance
+      must match the size of the second dimension of the array.
+    - A pandas `Series` with a numeric data type. This will be casted to an array of
+      data type [`np.float32`][numpy.float32]. Each value in the array will be used as
+      the value for the object at the same row index. The `filter_size` of
+      the [`DataFilterExtension`][lonboard.experimental.DataFilterExtension] instance
+      must be 1.
+    - A pyarrow [`FloatArray`][pyarrow.FloatArray], [`DoubleArray`][pyarrow.DoubleArray]
+      or [`ChunkedArray`][pyarrow.ChunkedArray] containing either a `FloatArray` or
+      `DoubleArray`. Each value in the array will be used as the value for the object at
+      the same row index. The `filter_size` of the
+      [`DataFilterExtension`][lonboard.experimental.DataFilterExtension] instance must
+      be 1.
+
+      Alternatively, you can pass any corresponding Arrow data structure from a library
+      that implements the [Arrow PyCapsule
+      Interface](https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html).
+    - A pyarrow [`FixedSizeListArray`][pyarrow.FixedSizeListArray] or
+      [`ChunkedArray`][pyarrow.ChunkedArray] containing `FixedSizeListArray`s. The child
+      array of the fixed size list must be of floating point type. The `filter_size` of
+      the [`DataFilterExtension`][lonboard.experimental.DataFilterExtension] instance
+      must match the list size.
+
+      Alternatively, you can pass any corresponding Arrow data structure from a library
+      that implements the [Arrow PyCapsule
+      Interface](https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html).
+    """
+
+    default_value = float(0)
+    info_text = (
+        "a float value or numpy ndarray or pyarrow array representing an array"
+        " of floats"
+    )
+
+    def __init__(
+        self: TraitType,
+        *args,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.tag(sync=True, **ACCESSOR_SERIALIZATION)
+
+    def validate(self, obj, value) -> Union[float, pa.ChunkedArray, pa.DoubleArray]:
+        # Find the data filter extension in the attributes of the parent object so we
+        # can validate against the filter size.
+        data_filter_extension = [
+            ext for ext in obj.extensions if ext._extension_type == "data-filter"
+        ]
+        assert len(data_filter_extension) == 1
+        filter_size = data_filter_extension[0].filter_size
+
+        if isinstance(value, (int, float)):
+            if filter_size != 1:
+                self.error(obj, value, info="filter_size==1 with scalar value")
+
+            return float(value)
+
+        if isinstance(value, (tuple, list)):
+            if filter_size != len(value):
+                self.error(
+                    obj,
+                    value,
+                    info=f"filter_size ({filter_size}) to match length of tuple/list",
+                )
+
+            if any(not isinstance(v, (int, float)) for v in value):
+                self.error(
+                    obj,
+                    value,
+                    info="all values in tuple or list to be numeric",
+                )
+
+            return value
+
+        # pandas Series
+        if (
+            value.__class__.__module__.startswith("pandas")
+            and value.__class__.__name__ == "Series"
+        ):
+            # Assert that filter_size == 1 for a pandas series.
+            # Pandas series can technically contain Python list objects inside them, but
+            # for simplicity we disallow that.
+            if filter_size != 1:
+                self.error(obj, value, info="filter_size==1 with pandas Series")
+
+            # Cast pandas Series to numpy ndarray
+            value = np.asarray(value)
+
+        if isinstance(value, np.ndarray):
+            if not np.issubdtype(value.dtype, np.number):
+                self.error(obj, value, info="numeric dtype")
+
+            # Cast to float32
+            value = value.astype(np.float32)
+
+            if len(value.shape) == 1:
+                if filter_size != 1:
+                    self.error(obj, value, info="filter_size==1 with 1-D numpy array")
+
+                return pa.array(value)
+
+            if len(value.shape) != 2:
+                self.error(obj, value, info="1-D or 2-D numpy array")
+
+            if value.shape[1] != filter_size:
+                self.error(
+                    obj,
+                    value,
+                    info=(
+                        f"filter_size ({filter_size}) to match 2nd dimension of "
+                        "numpy array"
+                    ),
+                )
+
+            return pa.FixedSizeListArray.from_arrays(value.flatten("C"), filter_size)
+
+        # Check for Arrow PyCapsule Interface
+        # https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
+        # TODO: with pyarrow v16 also import chunked array from stream
+        if not isinstance(value, (pa.ChunkedArray, pa.Array)):
+            if hasattr(value, "__arrow_c_array__"):
+                value = pa.array(value)
+
+        if isinstance(value, (pa.ChunkedArray, pa.Array)):
+            # Allowed inputs are either a FixedSizeListArray or numeric array.
+            # If not a fixed size list array, check for floating and cast to float32
+            if not pa.types.is_fixed_size_list(value.type):
+                if filter_size != 1:
+                    self.error(
+                        obj,
+                        value,
+                        info="filter_size==1 with non-FixedSizeList type arrow array",
+                    )
+
+                if not pa.types.is_floating(value.type):
+                    self.error(
+                        obj,
+                        value,
+                        info="arrow array to be a floating point type",
+                    )
+
+                return value.cast(pa.float32())
+
+            # We have a FixedSizeListArray
+            if filter_size != value.type.list_size:
+                self.error(
+                    obj,
+                    value,
+                    info=(
+                        f"filter_size ({filter_size}) to match list size of "
+                        "FixedSizeList arrow array"
+                    ),
+                )
+
+            if not pa.types.is_floating(value.type.value_type):
+                self.error(
+                    obj,
+                    value,
+                    info="arrow array to have floating point child type",
+                )
+
+            # Cast values to float32
+            return value.cast(pa.list_(pa.float32(), value.type.list_size))
+
+        self.error(obj, value)
+        assert False
diff --git a/src/model/extension.ts b/src/model/extension.ts
@@ -1,8 +1,5 @@
 import { LayerExtension } from "@deck.gl/core/typed";
 import {
-  BrushingExtensionProps,
-  CollisionFilterExtensionProps,
-  DataFilterExtensionProps,
   BrushingExtension as _BrushingExtension,
   CollisionFilterExtension as _CollisionFilterExtension,
   DataFilterExtension as _DataFilterExtension,
@@ -26,11 +23,6 @@ export class BrushingExtension extends BaseExtensionModel {
 
   extensionInstance: _BrushingExtension;
 
-  protected getBrushingTarget?: BrushingExtensionProps["getBrushingTarget"];
-  protected brushingEnabled?: BrushingExtensionProps["brushingEnabled"];
-  protected brushingTarget?: BrushingExtensionProps["brushingTarget"];
-  protected brushingRadius?: BrushingExtensionProps["brushingRadius"];
-
   constructor(
     model: WidgetModel,
     layerModel: BaseLayerModel,
@@ -65,11 +57,6 @@ export class CollisionFilterExtension extends BaseExtensionModel {
 
   extensionInstance: _CollisionFilterExtension;
 
-  protected collisionEnabled?: CollisionFilterExtensionProps["collisionEnabled"];
-  protected collisionGroup?: CollisionFilterExtensionProps["collisionGroup"];
-  protected collisionTestProps?: CollisionFilterExtensionProps["collisionTestProps"];
-  protected getCollisionPriority?: CollisionFilterExtensionProps["getCollisionPriority"];
-
   constructor(
     model: WidgetModel,
     layerModel: BaseLayerModel,
@@ -107,22 +94,17 @@ export class DataFilterExtension extends BaseExtensionModel {
 
   extensionInstance: _DataFilterExtension;
 
-  protected getFilterValue?: DataFilterExtensionProps["getFilterValue"];
-
-  protected filterEnabled?: DataFilterExtensionProps["filterEnabled"];
-  protected filterRange?: DataFilterExtensionProps["filterRange"];
-  protected filterSoftRange?: DataFilterExtensionProps["filterSoftRange"];
-  protected filterTransformSize?: DataFilterExtensionProps["filterTransformSize"];
-  protected filterTransformColor?: DataFilterExtensionProps["filterTransformColor"];
-
   constructor(
     model: WidgetModel,
     layerModel: BaseLayerModel,
     updateStateCallback: () => void,
   ) {
     super(model, updateStateCallback);
+
     // TODO: set filterSize, fp64, countItems in constructor
-    this.extensionInstance = new _DataFilterExtension();
+    // TODO: should filter_size automatically update from python?
+    const filterSize = this.model.get("filter_size");
+    this.extensionInstance = new _DataFilterExtension({ filterSize });
 
     // Properties added by the extension onto the layer
     layerModel.initRegularAttribute("filter_enabled", "filterEnabled");
diff --git a/tests/test_traits.py b/tests/test_traits.py

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
`1`	`1`	`# lonboard.experimental.traits`
`2`	`2`
`3`	`3`	`::: lonboard.experimental.traits.PointAccessor`
	`4`	`+::: lonboard.experimental.traits.GetFilterValueAccessor`