Scaling to 20M is not working

[Zsclarx]

---

⚡️ Performance Bottleneck on Large-scale HNSW Querying (20M x 384D vectors)

I'm working on scaling approximate nearest neighbor search using HNSW in Spark on a large dataset of sentence embeddings.

✅ Setup Summary

• Data: 20 million records, each with a 384-dimensional embedding (from all-MiniLM-L6-v2)

• Cluster Config:

  • 12 Executors, each with:

    • 8 cores
    • 64 GB RAM

  • 1 Driver with same config

---

⚙️ Code Snippet

val hnsw = new HnswSimilarity()
  .setIdentifierCol("row_id")
  .setFeaturesCol("normFeatures")  // Manually normalized to unit vectors
  .setDistanceFunction("inner-product")
  .setNumPartitions(224)
  .setM(16)
  .setEfConstruction(100)
  .setEf(30)
  .setK(10)
  .setPredictionCol("neighbors")

val model = hnsw.fit(data)            // ⚡ Fast: ~15 minutes for 20M records
val knnDf = model.transform(data)     // 🐢 Very slow: ~10 hours!
knnDf.cache()
knnDf.count()

---

❗ The Problem

• Index build time is excellent (~15 mins) ✅
• But the querying phase (model.transform) is extremely slow (~10 hours) ❌
• Reducing partitions causes executor OOM/failures
• Dataset is read in advance (no lazy transformations)
• Vectors are pre-normalized to avoid on-the-fly overhead

---

💡 Questions

1. Why is model.transform() so slow despite fast index construction?
2. Could this be due to internal broadcast or shuffle mechanisms used by the library?
3. Are there any best practices for:
   • Optimal numPartitions (e.g. per executor/core)?
   • Avoiding serialization/deserialization bottlenecks?

---

🙏 Any ideas, tuning tips, or architectural suggestions are highly appreciated!

---