Add support for `MapSort` expression in Spark 4.0.0

[andygrove]

What is the problem the feature request solves?

Spark 4.0.0 adds a MapSort expression, which can be used in hash partitioning expressions. Comet will currently fall back to Spark if these expressions exist.

Describe the potential solution

No response

Additional context

No response

[rishvin]

I would like to work on this.

[rishvin]

Some Updates:
I have a simple test to start with, which will produce _groupingmapsort.

val data = Seq(
     |   Map("a" -> 1, "b" -> 2),
     |   Map("a" -> 3, "b" -> 4),
     |   Map("b" -> 2, "a" -> 1)
     | )

val df = data.toDF("map")
df.groupBy("map").count().show(false)

So far based on my understanding, it looks like - this will require some plumbing at the Parquet reader utils in Scala, because currently we only support primitive types but MapType is a complex type. I made some hack to bypass some type-checking, however, the Native.initColumnReader() expects primitiveTypeId. I have to understand how the MapType would translate to primitiveTypeId. The closest one at the moment seems to be parquet's BINARY physical type. So, thinking of passing BINARY type and annotating with MapKeyValueTypeAnnotation. However, I'm still trying to understand this piece of code, so I might be wrong. I will do some more experiments to have more clarity on this.

[rishvin]

Some Updates: I have a simple test to start with, which will produce _groupingmapsort.

val data = Seq(
     |   Map("a" -> 1, "b" -> 2),
     |   Map("a" -> 3, "b" -> 4),
     |   Map("b" -> 2, "a" -> 1)
     | )

val df = data.toDF("map")
df.groupBy("map").count().show(false)

So far based on my understanding, it looks like - this will require some plumbing at the Parquet reader utils in Scala, because currently we only support primitive types but MapType is a complex type. I made some hack to bypass some type-checking, however, the Native.initColumnReader() expects primitiveTypeId. I have to understand how the MapType would translate to primitiveTypeId. The closest one at the moment seems to be parquet's BINARY physical type. So, thinking of passing BINARY type and annotating with MapKeyValueTypeAnnotation. However, I'm still trying to understand this piece of code, so I might be wrong. I will do some more experiments to have more clarity on this.

I managed to scan the map-type by setting CometConf.COMET_NATIVE_SCAN_IMPL.key -> native_datafusion . Added map_sort UDF with return type as Map.

However, it looks like doing a group-by on the map-type is not possible at the moment. Native code throws error message like - Not yet implemented: not yet implemented: Map(...).

This error seems to be coming from Arrow. So, presently, we cannot use MapType for grouping, we might have to translate map to a more simple type. Will explore more on this.

[parthchandra]

managed to scan the map-type by setting CometConf.COMET_NATIVE_SCAN_IMPL.key -> native_datafusion . Added map_sort UDF with return type as Map.

Right. native_comet does not support complex types. You can also try with native_iceberg_compat.
For map_sort you could try an order by instead of a group by in your test query

[rishvin]

managed to scan the map-type by setting CometConf.COMET_NATIVE_SCAN_IMPL.key -> native_datafusion . Added map_sort UDF with return type as Map.

Right. native_comet does not support complex types. You can also try with native_iceberg_compat. For map_sort you could try an order by instead of a group by in your test query

Thanks @parthchandra. We cannot use order by because spark applies MapSort either for groupby or repartitioning using the following rules,
InsertMapSortInRepartitioningExpression
InsertMapSortInGroupingExpression

I doubled-checked applying orderby on mapType in Spark-4.0 and got,

The `sortorder` does not support ordering on type "MAP<STRING, INT>". SQLSTATE: 42K09;

I'm attempting if we can do some changes to the plan to use string-representation of map instead of map directly during groupby/shuffling. If this doesn't work correctly or turned out to be too convoluted, I will wait for the complex type support.

[parthchandra]

Oh well.
For Maps, you might want to look at the mapping of Spark's MapType to a Parquet Schema here:

datafusion-comet/common/src/main/scala/org/apache/spark/sql/comet/parquet/CometSparkToParquetSchemaConverter.scala

Line 254 in 320ce55

case MapType(keyType, valueType, valueContainsNull) if writeLegacyParquetFormat =>

datafusion-comet/common/src/main/scala/org/apache/spark/sql/comet/parquet/CometSparkToParquetSchemaConverter.scala

Line 287 in 320ce55

case MapType(keyType, valueType, valueContainsNull) =>

Essentially a Map type is a repeated group of key and value fields. Note however that both key and value can themselves be complex types as in this test file:

datafusion-comet/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala

Line 719 in 320ce55

def makeParquetFileComplexTypes(

[rishvin]

I have something working for HashAgg on map-type integrated with MapSort. Will open a PR soon.

[rishvin]

Opened #2221