Improved embedding for larger datasets

Author: slinjhu

smoosense-gui/src/components/emb/Umap2D.tsx

@@ -4,6 +4,7 @@ import { useState, useEffect, useCallback, useMemo } from 'react'
 import _ from 'lodash'
 import { useAppSelector } from '@/lib/hooks'
 import { useSingleColumnRenderType } from '@/lib/hooks/useRenderType'
+import { BUBBLE_SIZE_COLOR_VALUE } from '@/lib/features/bubblePlot/BubblePlotMoreControls'
 import Umap2DControls from './Umap2DControls'
 import Umap2DScatterPlot from './Umap2DScatterPlot'
 import TextPlaceHolder from '@/components/common/TextPlaceHolder'
@@ -14,6 +15,9 @@ export interface UmapResult {
   columnValues: Record<string, unknown[]>
   count: number
   runtime: number
+  sampled: boolean
+  totalRows: number
+  maxRows: number
   params: {
     nNeighbors: number
     minDist: number
@@ -37,7 +41,9 @@ export default function Umap2D({ onResultChange, onSelectionChange }: Umap2DProp
   const visualColumn = useAppSelector((state) => state.ui.columnForGalleryVisual)
   const captionColumn = useAppSelector((state) => state.ui.columnForGalleryCaption)
   const breakdownColumn = useAppSelector((state) => state.ui.bubblePlotBreakdownColumn)
-  const colorColumn = useAppSelector((state) => state.ui.bubblePlotColorColumn)
+  const rawColorColumn = useAppSelector((state) => state.ui.bubblePlotColorColumn)
+  // Filter out __bubble_size__ which is only valid for bubble plots
+  const colorColumn = rawColorColumn === BUBBLE_SIZE_COLOR_VALUE ? '' : rawColorColumn
   const nNeighbors = useAppSelector((state) => state.ui.umapNNeighbors)
   const minDist = useAppSelector((state) => state.ui.umapMinDist)
   const visualRenderType = useSingleColumnRenderType(visualColumn || '')
@@ -85,6 +91,9 @@ export default function Umap2D({ onResultChange, onSelectionChange }: Umap2DProp
         columnValues: data.columnValues || {},
         count: data.count,
         runtime: data.runtime,
+        sampled: data.sampled || false,
+        totalRows: data.totalRows || data.count,
+        maxRows: data.maxRows || 10000,
         params: data.params,
       }
       setResult(newResult)
@@ -175,8 +184,15 @@ export default function Umap2D({ onResultChange, onSelectionChange }: Umap2DProp
               visualRenderType={visualRenderType}
               onSelectionChange={handleSelectionChange}
             />
-            <div className="flex-shrink-0 px-4 py-2 text-xs text-muted-foreground border-t">
-              {result.count} points | n_neighbors={result.params?.nNeighbors ?? nNeighbors} | min_dist={result.params?.minDist ?? minDist} | {result.runtime.toFixed(2)}s
+            <div className="flex-shrink-0 px-4 py-2 text-xs text-muted-foreground border-t flex justify-between">
+              <span>
+                {result.count} points | n_neighbors={result.params?.nNeighbors ?? nNeighbors} | min_dist={result.params?.minDist ?? minDist}{result.runtime != null && ` | ${result.runtime.toFixed(2)}s`}
+              </span>
+              {result.sampled && (
+                <span className="text-amber-600 dark:text-amber-400">
+                  Sampled {result.count.toLocaleString()} of {result.totalRows.toLocaleString()} rows
+                </span>
+              )}
             </div>
           </div>
         ) : (

smoosense-gui/src/components/emb/Umap2DScatterPlot.tsx

@@ -68,6 +68,8 @@ const Umap2DScatterPlot = React.memo(function Umap2DScatterPlot({
     setSelectedIndices([])
   }, [x, y])
 
+  const hasColorValues = colorValues && colorValues.length > 0
+
   const plotData = useMemo((): Partial<PlotData>[] => {
     if (!x || !y || x.length === 0) {
       return []
@@ -89,11 +91,11 @@ const Umap2DScatterPlot = React.memo(function Umap2DScatterPlot({
       const traces: Partial<PlotData>[] = []
       const groupNames = Array.from(groups.keys()).sort()
 
-      groupNames.forEach((groupName) => {
+      groupNames.forEach((groupName, groupIndex) => {
         const indices = groups.get(groupName)!
         const groupX = indices.map(i => x[i])
         const groupY = indices.map(i => y[i])
-        const groupColorValues = colorValues ? indices.map(i => colorValues[i]) : undefined
+        const groupColorValues = hasColorValues ? indices.map(i => colorValues[i]) : undefined
         const customdata = indices.map(i => ({ index: i }))
 
         // Determine marker color
@@ -112,7 +114,12 @@ const Umap2DScatterPlot = React.memo(function Umap2DScatterPlot({
             size: markerSize,
             opacity: opacity,
             color: markerColor,
-            colorscale: colorValues ? colorScale : undefined,
+            colorscale: hasColorValues ? colorScale : undefined,
+            showscale: hasColorValues && groupIndex === 0, // Only show colorbar on first trace
+            colorbar: hasColorValues ? {
+              thickness: 15,
+              len: 0.5
+            } : undefined,
             line: {
               width: 0.5,
               color: colors.foreground
@@ -130,7 +137,7 @@ const Umap2DScatterPlot = React.memo(function Umap2DScatterPlot({
     // Single trace (no breakdown)
     // Determine marker color based on colorValues or selection
     let markerColor: string[] | number[]
-    if (colorValues && colorValues.length > 0) {
+    if (hasColorValues) {
       markerColor = colorValues
     } else {
       markerColor = x.map((_, i) =>
@@ -150,7 +157,12 @@ const Umap2DScatterPlot = React.memo(function Umap2DScatterPlot({
         size: markerSize,
         opacity: opacity,
         color: markerColor,
-        colorscale: colorValues ? colorScale : undefined,
+        colorscale: hasColorValues ? colorScale : undefined,
+        showscale: hasColorValues,
+        colorbar: hasColorValues ? {
+          thickness: 15,
+          len: 0.5
+        } : undefined,
         line: {
           width: 0.5,
           color: colors.foreground
@@ -160,7 +172,7 @@ const Umap2DScatterPlot = React.memo(function Umap2DScatterPlot({
       customdata: customdata as any,
       hoverinfo: 'none', // Disable default hover
     }]
-  }, [x, y, selectedIndices, colors.primary, colors.foreground, markerSize, opacity, colorScale, breakdownValues, colorValues])
+  }, [x, y, selectedIndices, colors.primary, colors.foreground, markerSize, opacity, colorScale, breakdownValues, colorValues, hasColorValues])
 
   const hasBreakdown = breakdownValues && breakdownValues.length > 0
 
@@ -172,6 +184,7 @@ const Umap2DScatterPlot = React.memo(function Umap2DScatterPlot({
     ...baseLayout,
     dragmode: 'lasso',
     hovermode: 'closest',
+    showlegend: hasBreakdown,
     margin: {
       l: 0,
       r: 0,

smoosense-gui/src/lib/features/ui/uiSlice.ts

@@ -89,7 +89,7 @@ const initialState: UiState = {
   bubblePlotMinMarkerSize: 7,
   bubblePlotOpacity: 0.7,
   bubblePlotMarkerSizeContrastRatio: 4.2,
-  bubblePlotColorColumn: '__bubble_size__',
+  bubblePlotColorColumn: '',
   bubblePlotColorScale: 'Jet',
   bubblePlotLogScaleX: false,
   bubblePlotLogScaleY: false,

smoosense-gui/src/lib/utils/cellRenderers/EmbeddingCellRenderer.tsx

@@ -92,7 +92,7 @@ function SimilarRowsGallery({ embedding, columnName }: SimilarRowsGalleryProps)
     <div className="flex flex-col h-full">
       <GalleryControls showRandom={false} />
       <div className="px-2 py-1 text-xs text-muted-foreground">
-        Cosine distance (0 = identical, 1 = orthogonal, 2 = opposite; may have small error due to indexing)
+        Cosine distance (0 = identical, 1 = orthogonal, 2 = opposite; may have small error due to vector indexing)
       </div>
       <div className="flex-1 overflow-auto p-2">
         {_.isEmpty(rows) ? (
@@ -114,7 +114,7 @@ function SimilarRowsGallery({ embedding, columnName }: SimilarRowsGalleryProps)
               return (
                 <div key={index} className="relative">
                   <div className="absolute top-1 left-1 z-10 bg-black/70 text-white text-xs px-1.5 py-0.5 rounded-full">
-                    {distance.toFixed(3)}
+                    distance: {distance.toFixed(3)}
                   </div>
                   <GalleryItem
                     row={row}

smoosense-py/pyproject.toml

@@ -100,6 +100,7 @@ build-backend = "setuptools.build_meta"
 [project.scripts]
 sense = "smoosense.cli:main"
 sense-ingest-images = "smoosense.images.ingest:main"
+parquet-to-lance = "smoosense.lance.parquet_to_lance:main"
 
 # Ruff configuration
 [tool.ruff]

smoosense-py/smoosense/handlers/umap.py

@@ -13,6 +13,8 @@
 logger = logging.getLogger(__name__)
 umap_bp = Blueprint("umap", __name__)
 
+# Maximum number of rows to compute UMAP on (random sample if exceeded)
+UMAP_MAX_ROWS = 1_000
 
 @umap_bp.post("/umap")
 @requires_auth_api
@@ -86,19 +88,34 @@ def compute_umap() -> Response:
     if len(embeddings) < 2:
         raise ValueError("Not enough embeddings to compute UMAP (need at least 2)")
 
+    # Random sample if exceeds max rows
+    sampled = False
+    total_rows = len(embeddings)
+    if total_rows > UMAP_MAX_ROWS:
+        logger.info(f"Random sampling UMAP input from {total_rows} to {UMAP_MAX_ROWS} rows")
+        rng = np.random.default_rng(seed=42)
+        indices = rng.choice(total_rows, size=UMAP_MAX_ROWS, replace=False)
+        indices.sort()  # Keep relative order for consistency
+        embeddings = [embeddings[i] for i in indices]
+        for col in extra_values:
+            extra_values[col] = [extra_values[col][i] for i in indices]
+        sampled = True
+
     # Convert to numpy array
     embeddings_array = np.array(embeddings, dtype=np.float32)
 
     # Adjust n_neighbors if larger than dataset
     actual_n_neighbors = min(n_neighbors, len(embeddings) - 1)
 
-    # Compute UMAP
+    # Compute UMAP with performance optimizations
     reducer = umap.UMAP(
         n_neighbors=actual_n_neighbors,
         min_dist=min_dist,
         n_components=2,
         metric="cosine",
         random_state=42,
+        low_memory=False,  # Trade memory for speed
+        n_jobs=-1,  # Use all CPU cores
     )
     projection = reducer.fit_transform(embeddings_array)
 
@@ -113,6 +130,9 @@ def compute_umap() -> Response:
             "y": y_coords,
             "columnValues": {col: serialize(vals) for col, vals in extra_values.items()},
             "count": len(x_coords),
+            "sampled": sampled,
+            "totalRows": total_rows,
+            "maxRows": UMAP_MAX_ROWS,
             "runtime": default_timer() - time_start,
             "params": {
                 "nNeighbors": actual_n_neighbors,