In the process of understanding how Word2Vec in Spark differs from gensim one, I got very confused by the example presented in the Spark docs (reference link: https://spark.apache.org/docs/2.2.0/ml-features.html#word2vec) and I was wondering why instead of transforming single words, they transform entire sentences. Isn't word2vec purpose to embed single words into a vector space? Why they embed entire sentences? How can one properly train word2vec and then applying to single words in Spark?
Spark (naively) uses average of vectors for all words in the document as representation of the document. Check the API documentation little carefully.
"The Word2VecModel transforms each document into a vector using the average of all words in the document ; this vector can then be used as features for prediction, document similarity calculations, etc."
If you are interested specifically in vector for a word (and not document), you can check getVectors method which will return data-frame of word and vector. The API behaviour does give rise to confusion though as it assumes that everyone wants to use averaging by default.