Pyspark pyspark.ml.stat.Correlation.corr failing with UnsupportedOperationException when number of features is greater than 16250. There is no official documentation available specifying the limitations of this library.

Code snippet

`
from pyspark.ml.stat import Correlation
from pyspark.ml.linalg import Vectors
from pyspark.sql.types import StringType, IntegerType, ArrayType, StructField, DoubleType, StructType
from pyspark.ml.linalg import VectorUDT, Vectors, SparseVector, ArrayType, DenseVector

data = spark.createDataFrame([('2020-08-03','10010076305','1','54516',1,
Vectors.sparse(20000,[0,15,4936,14925,19201,19258,19278,19340,19344],[1.0,1.0,1.0,1.0,3.0,1.0,1.0,1.0,1.0]),
Vectors.sparse(4754,[],[]), 8, 2020,
Vectors.sparse(16,[0,15],[1.0,1.0]),
Vectors.sparse(20000,[0,15,4936,14925,19201,19258,19278,19340,19344],[1.0,1.0,1.0,1.0,3.0,1.0,1.0,1.0,1.0]),
Vectors.sparse(20000,[0,15,4936,14925,19201,19258,19278,19340,19344],[3.3795359045759663,2.481239579312057,5.083862109711494,8.865372043017848,0.4381255479573206,17.415537953461378,11.045507632248706,1.2744463238709334,2.1180617031512665]),
Vectors.sparse(20000,[17841,17856],[3.3795359045759663,2.481239579312057])
), ('2020-08-04','10010076306','1','54516',1,
Vectors.sparse(20000,[0,15,4936,14925,19201,19258,19278,19340,19344],[1.0,1.0,1.0,1.0,3.0,1.0,1.0,1.0,1.0]),
Vectors.sparse(4754,[],[]), 8, 2020,
Vectors.sparse(16,[0,15],[1.0,1.0]),
Vectors.sparse(20000,[0,15,4936,14925,19201,19258,19278,19340,19344],[1.0,1.0,1.0,1.0,3.0,1.0,1.0,1.0,1.0]),
Vectors.sparse(20000,[0,15,4936,14925,19201,19258,19278,19340,19344],[3.3795359045759663,2.481239579312057,5.083862109711494,8.865372043017848,0.4381255479573206,17.415537953461378,11.045507632248706,1.2744463238709334,2.1180617031512665]),
Vectors.sparse(20000,[17841,17856],[3.3795359045759663,2.481239579312057])
)],StructType([StructField('pivot_date', StringType(), True), StructField('customer_id', StringType(), True), StructField('marketplace_id', StringType(), True), StructField('customer_hash_key', StringType(), True), StructField('retainSample', IntegerType(), True), StructField('feature_vector_with_lags', VectorUDT(), True), StructField('feature_vector', VectorUDT(), True), StructField('moy_cat', IntegerType(), True), StructField('year_cat', IntegerType(), True), StructField('encoded_categoricals', VectorUDT(), True), StructField('temp_vec_name', VectorUDT(), True), StructField('stdtemp_vec_name_w_zeroes', VectorUDT(), True), StructField('stdtemp_vec_name', VectorUDT(), True)]))

vec_name = 'temp_vec_name'
spark_corr_matrix = Correlation.corr(data, f"std{vec_name}")
`

Exception stacktrace when executing Correlation.corr:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/spark/python/pyspark/ml/stat.py", line 181, in corr
    return _java2py(sc, javaCorrObj.corr(*args))
  File "/usr/lib/spark/python/lib/py4j-0.10.9.7-src.zip/py4j/java_gateway.py", line 1322, in __call__
  File "/usr/lib/spark/python/pyspark/errors/exceptions/captured.py", line 169, in deco
    return f(*a, **kw)
  File "/usr/lib/spark/python/lib/py4j-0.10.9.7-src.zip/py4j/protocol.py", line 326, in get_return_value
py4j.protocol.Py4JJavaError: An error occurred while calling z:org.apache.spark.ml.stat.Correlation.corr.
: java.lang.UnsupportedOperationException: Cannot convert this array to unsafe format as it's too big.
        at org.apache.spark.sql.catalyst.expressions.UnsafeArrayData.fromPrimitiveArray(UnsafeArrayData.java:426)
        at org.apache.spark.sql.catalyst.expressions.UnsafeArrayData.fromPrimitiveArray(UnsafeArrayData.java:493)
        at org.apache.spark.ml.linalg.MatrixUDT.serialize(MatrixUDT.scala:66)
        at org.apache.spark.ml.linalg.MatrixUDT.serialize(MatrixUDT.scala:28)
        at org.apache.spark.sql.catalyst.CatalystTypeConverters$UDTConverter.toCatalystImpl(CatalystTypeConverters.scala:146)
        at org.apache.spark.sql.catalyst.CatalystTypeConverters$CatalystTypeConverter.toCatalyst(CatalystTypeConverters.scala:106)
        at org.apache.spark.sql.catalyst.CatalystTypeConverters$StructConverter.toCatalystImpl(CatalystTypeConverters.scala:251)
        at org.apache.spark.sql.catalyst.CatalystTypeConverters$StructConverter.toCatalystImpl(CatalystTypeConverters.scala:241)
        at org.apache.spark.sql.catalyst.CatalystTypeConverters$CatalystTypeConverter.toCatalyst(CatalystTypeConverters.scala:106)
        at org.apache.spark.sql.catalyst.CatalystTypeConverters$.$anonfun$createToCatalystConverter$2(CatalystTypeConverters.scala:477)
        at org.apache.spark.sql.catalyst.plans.logical.LocalRelation$.$anonfun$fromExternalRows$1(LocalRelation.scala:38)
        at scala.collection.TraversableLike.$anonfun$map$1(TraversableLike.scala:286)
        at scala.collection.Iterator.foreach(Iterator.scala:943)
        at scala.collection.Iterator.foreach$(Iterator.scala:943)
        at scala.collection.AbstractIterator.foreach(Iterator.scala:1431)
        at scala.collection.IterableLike.foreach(IterableLike.scala:74)
        at scala.collection.IterableLike.foreach$(IterableLike.scala:73)
        at scala.collection.AbstractIterable.foreach(Iterable.scala:56)
        at scala.collection.TraversableLike.map(TraversableLike.scala:286)
        at scala.collection.TraversableLike.map$(TraversableLike.scala:279)
        at scala.collection.AbstractTraversable.map(Traversable.scala:108)
        at org.apache.spark.sql.catalyst.plans.logical.LocalRelation$.fromExternalRows(LocalRelation.scala:38)
        at org.apache.spark.sql.SparkSession.$anonfun$createDataFrame$4(SparkSession.scala:393)
        at org.apache.spark.sql.SparkSession.withActive(SparkSession.scala:827)
        at org.apache.spark.sql.SparkSession.createDataFrame(SparkSession.scala:391)
        at org.apache.spark.ml.stat.Correlation$.corr(Correlation.scala:74)
        at org.apache.spark.ml.stat.Correlation.corr(Correlation.scala)
        at sun.reflect.NativeMethodAccessorImpl.invoke0(Native Method)
        at sun.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:62)        at sun.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)        at java.lang.reflect.Method.invoke(Method.java:498)        at py4j.reflection.MethodInvoker.invoke(MethodInvoker.java:244)
        at py4j.reflection.ReflectionEngine.invoke(ReflectionEngine.java:374)
        at py4j.Gateway.invoke(Gateway.java:282)
        at py4j.commands.AbstractCommand.invokeMethod(AbstractCommand.java:132)
        at py4j.commands.CallCommand.execute(CallCommand.java:79)
        at py4j.ClientServerConnection.waitForCommands(ClientServerConnection.java:182)
        at py4j.ClientServerConnection.run(ClientServerConnection.java:106)
        at java.lang.Thread.run(Thread.java:750)

I tried converting the sparse matrix datatype values to float/int, I reduced the decimal precision of the values as well. but seems like the values will always be double datatype as Vector.sparse seems to maintain it that way (link).

This seems to be a bug in the Correlation.corr library, it does not seem to scale even with larger ec2 instances where we have allocated 1000G driver memory and 1000G executor memory. The dataaframe in the code snippet above contains only 2 rows still the Correlation.corr is unable to scale.