[FEA] Support element_at function on GPU

[Dooyoung-Hwang]

I request a feature of element_at function on GPU.

The API link : https://spark.apache.org/docs/latest/api/sql/index.html#element_at

> SELECT element_at(array(1, 2, 3), 2);
 2
> SELECT element_at(map(1, 'a', 2, 'b'), 2);
 b

Driver log :
I!NOT_FOUND element_at cannot run on GPU because no GPU enabled version of expression class org.apache.spark.sql.catalyst.expressions.ElementAt could be found

[revans2]

We really need some more information to understand how you want to use this. Spark is very generic in what it supports compared to CUDF. So we can support some of this right now with just a little bit of work. But full support is going to take a lot more, so it would be good to understand exactly what the requirements are.

Do you want to use this for arrays or maps or both?
Is the key/index always a constant value or can it be a column of data?
What are they types that you need to support for the array (just the values) or map (both the key and the value)?

[Dooyoung-Hwang]

We just want to get an element of ArrayType of atomic DataType with index. We found a workaround for this case and applied and ran the job with GPU.

[Before]

    # prediction : Array of float
    features = ["c1", "c2", "c3"]
    for i in range(len(features)):
        inf_df = inf_df.withColumn(features[i], element_at("prediction", i + 1))

[After]

    # prediction : Array of float
    features = ["c1", "c2", "c3"]
    for i in range(len(features)):
        inf_df = inf_df.withColumn(features[i], col("prediction")[i])

[sameerz]

Tying element_at to what we already support for element indexing is the reason for the good first issue label. Supporting additional features, like columnar indexing (as opposed to indexing with a literal) will be more challenging.

[wjxiz1992]

We really need some more information to understand how you want to use this. Spark is very generic in what it supports compared to CUDF. So we can support some of this right now with just a little bit of work. But full support is going to take a lot more, so it would be good to understand exactly what the requirements are.

Do you want to use this for arrays or maps or both?
Is the key/index always a constant value or can it be a column of data?
What are they types that you need to support for the array (just the values) or map (both the key and the value)?

Hi @revans2 , since SKT's requirement is the simple case (not nested), can I take this feature issue? To support the simple case first.

[revans2]

can I take this feature issue? To support the simple case first.

Sure

[wjxiz1992]

Hi @revans2 , during the implementation, I found the 3.1.1 introduced ansiMode parameter to this expression. It causes different behaviors dealing with invalid index or key.

I want to confirm with you if my understanding is correct.

For select element_at(Array(1,2,3), 4):

Spark 3.0.2 + spark.sql.ansi.enabled=true:

+-----------------------------+
|element_at(array(1, 2, 3), 4)|
+-----------------------------+
|                         null|
+-----------------------------+

Spark 3.1.1 + spark.sql.ansi.enabled=true:

java.lang.ArrayIndexOutOfBoundsException: Invalid index: 4, numElements: 3

For select element_at(map(1, 'a', 2, 'b'), 3):

Spark 3.0.2 + spark.sql.ansi.enabled=true:

+------------------------------+
|element_at(map(1, a, 2, b), 3)|
+------------------------------+
|                          null|
+------------------------------+

Spark 3.1.1 + spark.sql.ansi.enabled=true:

java.util.NoSuchElementException: Key 3 does not exist.

I think I should then put corresponding codes into shim layers.
For shim layers, I saw the higher version inherits it's preceding version, and the older version's codes will be updated by the latter in the Exprs map.
So I could put one version in 3.0.0 and the other in 3.1.1.

Am I on the right way or there's should be a better solution? Thanks in advance!

[revans2]

Your analysis looks correct. At a minimum we need the rule in the shim layer so it can match the correct version of ElementAt. We probably don't need multiple implementations of GpuElementAt because everything prior to 3.1.1 is the same as setting failOnError to false. I don't even think we need to support failOnError initially. We could fall back to the CPU and have a follow on issue to make it work properly in ANSI mode.