index-aware map method (#1559)

* index-aware map method

* fix type

* fix jsdocs for MapFunction and ReducerFunction

* copy edits

* test custom map methods

* test mapIndex, too

---------

Co-authored-by: Philippe Rivière <[email protected]>

Author: mbostock

docs/transforms/map.md

@@ -114,6 +114,7 @@ The following map methods are supported:
 * *rank* - the rank of each value in the sorted array
 * *quantile* - the rank, normalized between 0 and 1
 * a function to be passed an array of values, returning new values
+* a function to be passed an index and array of channel values, returning new values
 * an object that implements the *mapIndex* method
 
 If a function is used, it must return an array of the same length as the given input. If a *mapIndex* method is used, it is repeatedly passed the index for each series (an array of integers), the corresponding input channel’s array of values, and the output channel’s array of values; it must populate the slots specified by the index in the output array.

src/reducer.d.ts

@@ -52,8 +52,10 @@ export type ReducerName =
   | ReducerPercentile;
 
 /**
- * A shorthand functional reducer implementation: given an array of input
- * channel *values*, returns the corresponding reduced output value.
+ * A shorthand functional reducer implementation: given an *index* and the
+ * corresponding input channel *values* array, returns the corresponding reduced
+ * output value. If the function only takes a single argument, it is instead
+ * passed a subset of values from the input channel.
  */
 export type ReducerFunction<S = any, T = S> = ((index: number[], values: S[]) => T) | ((values: S[]) => T);
 

src/transforms/map.d.ts

@@ -2,11 +2,15 @@ import type {ChannelName, ChannelValue} from "../channel.js";
 import type {Transformed} from "./basic.js";
 
 /**
- * A shorthand functional map implementation: given an array of input channel
- * *values*, returns the corresponding array of mapped output channel values.
- * The returned array must have the same length as the given input.
+ * A shorthand functional map implementation: given an *index* and the
+ * corresponding input channel *values* array, returns the corresponding array
+ * of mapped output channel values. The returned array must have the same length
+ * as the given input index. If the function only takes a single argument, it is
+ * instead passed a subset of values from the input channel.
  */
-export type MapFunction<S = any, T = S> = (values: S[]) => T[];
+export type MapFunction<S = any, T = S> =
+  | ((index: number[], values: S[]) => ArrayLike<T>)
+  | ((values: S[]) => ArrayLike<T>);
 
 /**
  * The built-in map implementations; one of:

src/transforms/map.js

@@ -1,5 +1,5 @@
 import {count, group, rank} from "d3";
-import {column, identity, isObject, maybeInput, maybeZ, take, valueof} from "../options.js";
+import {column, identity, isObject, maybeInput, maybeZ, taker, valueof} from "../options.js";
 import {basic} from "./basic.js";
 
 export function mapX(mapper, options = {}) {
@@ -50,14 +50,14 @@ function maybeMap(map) {
   if (map == null) throw new Error("missing map");
   if (typeof map.mapIndex === "function") return map;
   if (typeof map.map === "function" && isObject(map)) return mapMap(map); // N.B. array.map
-  if (typeof map === "function") return mapFunction(map);
+  if (typeof map === "function") return mapFunction(taker(map));
   switch (`${map}`.toLowerCase()) {
     case "cumsum":
       return mapCumsum;
     case "rank":
-      return mapFunction(rank);
+      return mapFunction((I, V) => rank(I, (i) => V[i]));
     case "quantile":
-      return mapFunction(rankQuantile);
+      return mapFunction((I, V) => rankQuantile(I, (i) => V[i]));
   }
   throw new Error(`invalid map: ${map}`);
 }
@@ -67,15 +67,15 @@ function mapMap(map) {
   return {mapIndex: map.map.bind(map)};
 }
 
-function rankQuantile(V) {
-  const n = count(V) - 1;
-  return rank(V).map((r) => r / n);
+function rankQuantile(I, f) {
+  const n = count(I, f) - 1;
+  return rank(I, f).map((r) => r / n);
 }
 
 function mapFunction(f) {
   return {
     mapIndex(I, S, T) {
-      const M = f(take(S, I));
+      const M = f(I, S);
       if (M.length !== I.length) throw new Error("map function returned a mismatched length");
       for (let i = 0, n = I.length; i < n; ++i) T[I[i]] = M[i];
     }