Merge pull request #5114 from parca-dev/color-by

area/ui: Add a Color by functionality

pkg/query/columnquery.go

@@ -975,7 +975,7 @@ func getMappingFilesAndLabels(
 
 	labels, err := q.GetProfileMetadataLabels(ctx, query, startTime, endTime)
 	if err != nil {
-		return nil, nil, fmt.Errorf("failed to get mappings: %w", err)
+		return nil, nil, fmt.Errorf("failed to get labels: %w", err)
 	}
 
 	return mappingFiles, labels, nil

ui/packages/shared/profile/src/ProfileIcicleGraph/IcicleGraph/useColoredGraph.ts

@@ -25,8 +25,8 @@ import {
   setFeatures,
   useAppDispatch,
   useAppSelector,
-  type FeatureType,
-  type FeaturesMap,
+  type BinaryFeatureType,
+  type BinaryFeaturesMap,
 } from '@parca/store';
 import type {ColorProfileName} from '@parca/utilities';
 
@@ -52,7 +52,7 @@ const colorNodes = (
   mappings: Mapping[],
   locations: Location[],
   functions: ParcaFunction[],
-  features: {[key: string]: FeatureType}
+  features: {[key: string]: BinaryFeatureType}
 ): ColoredFlamegraphNode[] => {
   if (nodes === undefined) {
     return [];
@@ -85,11 +85,11 @@ const useColoredGraph = (graph: Flamegraph): ColoredFlamegraph => {
   );
   const isDarkMode = useAppSelector(selectDarkMode);
 
-  const [coloredGraph, features]: [ColoredFlamegraph, FeaturesMap] = useMemo(() => {
+  const [coloredGraph, features]: [ColoredFlamegraph, BinaryFeaturesMap] = useMemo(() => {
     if (graph.root == null) {
       return [graph as ColoredFlamegraph, {}];
     }
-    const features: FeaturesMap = {};
+    const features: BinaryFeaturesMap = {};
     const coloredGraph = {
       ...graph,
       root: {

ui/packages/shared/profile/src/ProfileIcicleGraph/IcicleGraph/utils.ts

@@ -17,7 +17,7 @@ import {
   Mapping,
   Function as ParcaFunction,
 } from '@parca/client/dist/parca/metastore/v1alpha1/metastore';
-import {EVERYTHING_ELSE, FEATURE_TYPES, type Feature} from '@parca/store';
+import {BINARY_FEATURE_TYPES, EVERYTHING_ELSE, type BinaryFeature} from '@parca/store';
 import {getLastItem} from '@parca/utilities';
 
 import {hexifyAddress} from '../../utils';
@@ -92,11 +92,11 @@ export const extractFeature = (
   mappings: Mapping[],
   locations: Location[],
   strings: string[]
-): Feature => {
+): BinaryFeature => {
   const binaryName = getBinaryName(data, mappings, locations, strings);
   if (binaryName != null) {
-    return {name: binaryName, type: FEATURE_TYPES.Binary};
+    return {name: binaryName, type: BINARY_FEATURE_TYPES.Binary};
   }
 
-  return {name: EVERYTHING_ELSE, type: FEATURE_TYPES.Misc};
+  return {name: EVERYTHING_ELSE, type: BINARY_FEATURE_TYPES.Misc};
 };

ui/packages/shared/profile/src/ProfileIcicleGraph/IcicleGraphArrow/ColorStackLegend.tsx

@@ -35,7 +35,12 @@ const ColorStackLegend = ({mappings, compareMode = false, loading}: Props): Reac
   const [colorProfileName] = useUserPreference<string>(
     USER_PREFERENCES.FLAMEGRAPH_COLOR_PROFILE.key
   );
-  const [currentSearchString, setSearchString] = useURLState<string[]>('binary_frame_filter', {
+
+  const [colorByValue, _] = useURLState('color_by');
+
+  const colorBy = colorByValue === 'binary' || colorByValue === undefined ? 'binary' : 'filename';
+
+  const [currentSearchString, setSearchString] = useURLState<string[]>(`binary_frame_filter`, {
     alwaysReturnArray: true,
     defaultValue: [],
   });
@@ -88,7 +93,7 @@ const ColorStackLegend = ({mappings, compareMode = false, loading}: Props): Reac
               }
             )}
             onClick={() => {
-              if (!filteringAllowed || isHighlighted) {
+              if (!filteringAllowed || isHighlighted || colorBy !== 'binary') {
                 return;
               }
 

ui/packages/shared/profile/src/ProfileIcicleGraph/IcicleGraphArrow/IcicleGraphNodes.tsx

@@ -27,6 +27,7 @@ import {
   FIELD_CHILDREN,
   FIELD_CUMULATIVE,
   FIELD_DIFF,
+  FIELD_FUNCTION_FILE_NAME,
   FIELD_FUNCTION_NAME,
   FIELD_MAPPING_FILE,
 } from './index';
@@ -38,7 +39,8 @@ export const RowHeight = 26;
 interface IcicleGraphNodesProps {
   table: Table<any>;
   row: number;
-  mappingColors: mappingColors;
+  colors: colorByColors;
+  colorBy: string;
   childRows: number[];
   x: number;
   y: number;
@@ -67,7 +69,8 @@ interface IcicleGraphNodesProps {
 export const IcicleGraphNodes = React.memo(function IcicleGraphNodesNoMemo({
   table,
   childRows,
-  mappingColors,
+  colors,
+  colorBy,
   x,
   y,
   xScale,
@@ -114,7 +117,8 @@ export const IcicleGraphNodes = React.memo(function IcicleGraphNodesNoMemo({
         key={`node-${level}-${i}`}
         table={table}
         row={child}
-        mappingColors={mappingColors}
+        colors={colors}
+        colorBy={colorBy}
         x={xStart}
         y={0}
         totalWidth={totalWidth}
@@ -145,7 +149,7 @@ export const IcicleGraphNodes = React.memo(function IcicleGraphNodesNoMemo({
   return <g transform={`translate(${x}, ${y})`}>{childrenElements}</g>;
 });
 
-export interface mappingColors {
+export interface colorByColors {
   [key: string]: string;
 }
 
@@ -158,7 +162,8 @@ interface IcicleNodeProps {
   level: number;
   table: Table<any>;
   row: number;
-  mappingColors: mappingColors;
+  colors: colorByColors;
+  colorBy: string;
   path: string[];
   total: bigint;
   setCurPath: (path: string[]) => void;
@@ -192,7 +197,8 @@ const fadedIcicleRectStyles = {
 export const IcicleNode = React.memo(function IcicleNodeNoMemo({
   table,
   row,
-  mappingColors,
+  colors,
+  colorBy,
   x,
   y,
   height,
@@ -223,12 +229,28 @@ export const IcicleNode = React.memo(function IcicleNodeNoMemo({
   const functionNameColumn = table.getChild(FIELD_FUNCTION_NAME);
   const cumulativeColumn = table.getChild(FIELD_CUMULATIVE);
   const diffColumn = table.getChild(FIELD_DIFF);
+  const filenameColumn = table.getChild(FIELD_FUNCTION_FILE_NAME);
   // get the actual values from the columns
   const mappingFile: string | null = arrowToString(mappingColumn?.get(row));
   const functionName: string | null = arrowToString(functionNameColumn?.get(row));
   const cumulative = cumulativeColumn?.get(row) !== null ? BigInt(cumulativeColumn?.get(row)) : 0n;
   const diff: bigint | null = diffColumn?.get(row) !== null ? BigInt(diffColumn?.get(row)) : null;
   const childRows: number[] = Array.from(table.getChild(FIELD_CHILDREN)?.get(row) ?? []);
+  const filename: string | null = arrowToString(filenameColumn?.get(row));
+
+  const colorAttribute: string | null = useMemo(() => {
+    let attr: string | null | undefined;
+
+    if (colorBy === 'filename') {
+      attr = filename;
+    } else if (colorBy === 'binary') {
+      attr = mappingFile;
+    }
+
+    return attr ?? null; // Provide a default value of null if attr is undefined
+  }, [colorBy, filename, mappingFile]);
+
+  const colorsMap = colors;
 
   const highlightedNodes = useMemo(() => {
     if (!highlightSimilarStacksPreference) {
@@ -280,7 +302,6 @@ export const IcicleNode = React.memo(function IcicleNodeNoMemo({
         if (aDiff !== null) {
           const cumulative: bigint =
             cumulativeColumn?.get(a) !== null ? BigInt(cumulativeColumn?.get(a)) : 0n;
-          console.log(typeof cumulative, typeof aDiff);
           const prev: bigint = cumulative - aDiff;
           aRatio = Number(aDiff) / Number(prev);
         }
@@ -315,8 +336,8 @@ export const IcicleNode = React.memo(function IcicleNodeNoMemo({
     compareMode,
     cumulative,
     diff,
-    mappingColors,
-    mappingFile,
+    colorsMap,
+    colorAttribute,
   });
   const name = useMemo(() => {
     return isRoot ? 'root' : nodeLabel(table, row, level, binaries.length > 1);
@@ -401,7 +422,8 @@ export const IcicleNode = React.memo(function IcicleNodeNoMemo({
         <IcicleGraphNodes
           table={table}
           row={row}
-          mappingColors={mappingColors}
+          colors={colors}
+          colorBy={colorBy}
           childRows={childRows}
           x={x}
           y={RowHeight}

ui/packages/shared/profile/src/ProfileIcicleGraph/IcicleGraphArrow/index.tsx

@@ -34,8 +34,9 @@ import GraphTooltipArrowContent from '../../GraphTooltipArrow/Content';
 import {DockedGraphTooltip} from '../../GraphTooltipArrow/DockedGraphTooltip';
 import {useProfileViewContext} from '../../ProfileView/ProfileViewContext';
 import ContextMenu from './ContextMenu';
-import {IcicleNode, RowHeight, mappingColors} from './IcicleGraphNodes';
-import {arrowToString, extractFeature} from './utils';
+import {IcicleNode, RowHeight, colorByColors} from './IcicleGraphNodes';
+import {useFilenamesList} from './useMappingList';
+import {arrowToString, extractFeature, extractFilenameFeature} from './utils';
 
 export const FIELD_LABELS_ONLY = 'labels_only';
 export const FIELD_MAPPING_FILE = 'mapping_file';
@@ -72,16 +73,30 @@ export const getMappingColors = (
   mappingsList: string[],
   isDarkMode: boolean,
   currentColorProfile: ColorConfig
-): mappingColors => {
+): colorByColors => {
   const mappingFeatures = mappingsList.map(mapping => extractFeature(mapping));
 
-  const colors: mappingColors = {};
+  const colors: colorByColors = {};
   Object.entries(mappingFeatures).forEach(([_, feature]) => {
     colors[feature.name] = getColorForFeature(feature.name, isDarkMode, currentColorProfile.colors);
   });
   return colors;
 };
 
+export const getFilenameColors = (
+  filenamesList: string[],
+  isDarkMode: boolean,
+  currentColorProfile: ColorConfig
+): colorByColors => {
+  const filenameFeatures = filenamesList.map(filename => extractFilenameFeature(filename));
+
+  const colors: colorByColors = {};
+  Object.entries(filenameFeatures).forEach(([_, feature]) => {
+    colors[feature.name] = getColorForFeature(feature.name, isDarkMode, currentColorProfile.colors);
+  });
+  return colors;
+};
+
 const noop = (): void => {};
 
 export const IcicleGraphArrow = memo(function IcicleGraphArrow({
@@ -123,6 +138,11 @@ export const IcicleGraphArrow = memo(function IcicleGraphArrow({
   const currentColorProfile = useCurrentColorProfile();
   const colorForSimilarNodes = currentColorProfile.colorForSimilarNodes;
 
+  const [colorBy, _] = useURLState('color_by');
+  const colorByValue = colorBy === undefined || colorBy === '' ? 'binary' : (colorBy as string);
+
+  const filenamesList = useFilenamesList(table);
+
   const mappingsList = useMemo(() => {
     // Read the mappings from the dictionary that contains all mapping strings.
     // This is great, as might only have a dozen or so mappings,
@@ -152,11 +172,25 @@ export const IcicleGraphArrow = memo(function IcicleGraphArrow({
     return mappings;
   }, [table]);
 
+  const filenameColors = useMemo(() => {
+    const colors = getFilenameColors(filenamesList, isDarkMode, currentColorProfile);
+    return colors;
+  }, [isDarkMode, filenamesList, currentColorProfile]);
+
   const mappingColors = useMemo(() => {
     const colors = getMappingColors(mappingsList, isDarkMode, currentColorProfile);
     return colors;
   }, [isDarkMode, mappingsList, currentColorProfile]);
 
+  const colorByList = {
+    filename: filenameColors,
+    binary: mappingColors,
+  };
+
+  type ColorByKey = keyof typeof colorByList;
+
+  const colorByColors: colorByColors = colorByList[colorByValue as ColorByKey];
+
   useEffect(() => {
     if (ref.current != null) {
       setHeight(ref?.current.getBoundingClientRect().height);
@@ -231,7 +265,8 @@ export const IcicleGraphArrow = memo(function IcicleGraphArrow({
             <IcicleNode
               table={table}
               row={0} // root is always row 0 in the arrow record
-              mappingColors={mappingColors}
+              colors={colorByColors}
+              colorBy={colorByValue}
               x={0}
               y={0}
               totalWidth={width ?? 1}
@@ -266,7 +301,8 @@ export const IcicleGraphArrow = memo(function IcicleGraphArrow({
     height,
     displayMenu,
     table,
-    mappingColors,
+    colorByColors,
+    colorByValue,
     setCurPath,
     curPath,
     total,

ui/packages/shared/profile/src/ProfileIcicleGraph/IcicleGraphArrow/useMappingList.ts

@@ -13,8 +13,13 @@
 
 import {useMemo} from 'react';
 
+import {Dictionary, Table, Vector} from 'apache-arrow';
+
 import {getLastItem} from '@parca/utilities';
 
+import {FIELD_FUNCTION_FILE_NAME} from './index';
+import {arrowToString} from './utils';
+
 const useMappingList = (mappings: string[] | undefined): string[] => {
   const mappingsList = useMemo(() => {
     if (mappings === undefined) {
@@ -39,4 +44,31 @@ const useMappingList = (mappings: string[] | undefined): string[] => {
   return mappingsList;
 };
 
+export const useFilenamesList = (table: Table | null): string[] => {
+  if (table === null) {
+    return [];
+  }
+  const filenamesDict: Vector<Dictionary> | null = table.getChild(FIELD_FUNCTION_FILE_NAME);
+  const filenames =
+    filenamesDict?.data
+      .map(file => {
+        if (file.dictionary == null) {
+          return [];
+        }
+        const len = file.dictionary.length;
+        const entries: string[] = [];
+        for (let i = 0; i < len; i++) {
+          const fn = arrowToString(file.dictionary.get(i));
+          entries.push(getLastItem(fn) ?? '');
+        }
+        return entries;
+      })
+      .flat() ?? [];
+
+  filenames.push('');
+
+  filenames.sort((a, b) => a.localeCompare(b));
+  return filenames;
+};
+
 export default useMappingList;