[SPARK-23821][SQL] Collection function: flatten
## What changes were proposed in this pull request?
This PR adds a new collection function that transforms an array of arrays into a single array. The PR comprises:
- An expression for flattening array structure
- Flatten function
- A wrapper for PySpark
## How was this patch tested?
New tests added into:
- CollectionExpressionsSuite
- DataFrameFunctionsSuite
## Codegen examples
### Primitive type
```
val df = Seq(
Seq(Seq(1, 2), Seq(4, 5)),
Seq(null, Seq(1))
).toDF("i")
df.filter($"i".isNotNull || $"i".isNull).select(flatten($"i")).debugCodegen
```
Result:
```
/* 033 */ boolean inputadapter_isNull = inputadapter_row.isNullAt(0);
/* 034 */ ArrayData inputadapter_value = inputadapter_isNull ?
/* 035 */ null : (inputadapter_row.getArray(0));
/* 036 */
/* 037 */ boolean filter_value = true;
/* 038 */
/* 039 */ if (!(!inputadapter_isNull)) {
/* 040 */ filter_value = inputadapter_isNull;
/* 041 */ }
/* 042 */ if (!filter_value) continue;
/* 043 */
/* 044 */ ((org.apache.spark.sql.execution.metric.SQLMetric) references[0] /* numOutputRows */).add(1);
/* 045 */
/* 046 */ boolean project_isNull = inputadapter_isNull;
/* 047 */ ArrayData project_value = null;
/* 048 */
/* 049 */ if (!inputadapter_isNull) {
/* 050 */ for (int z = 0; !project_isNull && z < inputadapter_value.numElements(); z++) {
/* 051 */ project_isNull |= inputadapter_value.isNullAt(z);
/* 052 */ }
/* 053 */ if (!project_isNull) {
/* 054 */ long project_numElements = 0;
/* 055 */ for (int z = 0; z < inputadapter_value.numElements(); z++) {
/* 056 */ project_numElements += inputadapter_value.getArray(z).numElements();
/* 057 */ }
/* 058 */ if (project_numElements > 2147483632) {
/* 059 */ throw new RuntimeException("Unsuccessful try to flatten an array of arrays with " +
/* 060 */ project_numElements + " elements due to exceeding the array size limit 2147483632.");
/* 061 */ }
/* 062 */
/* 063 */ long project_size = UnsafeArrayData.calculateSizeOfUnderlyingByteArray(
/* 064 */ project_numElements,
/* 065 */ 4);
/* 066 */ if (project_size > 2147483632) {
/* 067 */ throw new RuntimeException("Unsuccessful try to flatten an array of arrays with " +
/* 068 */ project_size + " bytes of data due to exceeding the limit 2147483632" +
/* 069 */ " bytes for UnsafeArrayData.");
/* 070 */ }
/* 071 */
/* 072 */ byte[] project_array = new byte[(int)project_size];
/* 073 */ UnsafeArrayData project_tempArrayData = new UnsafeArrayData();
/* 074 */ Platform.putLong(project_array, 16, project_numElements);
/* 075 */ project_tempArrayData.pointTo(project_array, 16, (int)project_size);
/* 076 */ int project_counter = 0;
/* 077 */ for (int k = 0; k < inputadapter_value.numElements(); k++) {
/* 078 */ ArrayData arr = inputadapter_value.getArray(k);
/* 079 */ for (int l = 0; l < arr.numElements(); l++) {
/* 080 */ if (arr.isNullAt(l)) {
/* 081 */ project_tempArrayData.setNullAt(project_counter);
/* 082 */ } else {
/* 083 */ project_tempArrayData.setInt(
/* 084 */ project_counter,
/* 085 */ arr.getInt(l)
/* 086 */ );
/* 087 */ }
/* 088 */ project_counter++;
/* 089 */ }
/* 090 */ }
/* 091 */ project_value = project_tempArrayData;
/* 092 */
/* 093 */ }
/* 094 */
/* 095 */ }
```
### Non-primitive type
```
val df = Seq(
Seq(Seq("a", "b"), Seq(null, "d")),
Seq(null, Seq("a"))
).toDF("s")
df.filter($"s".isNotNull || $"s".isNull).select(flatten($"s")).debugCodegen
```
Result:
```
/* 033 */ boolean inputadapter_isNull = inputadapter_row.isNullAt(0);
/* 034 */ ArrayData inputadapter_value = inputadapter_isNull ?
/* 035 */ null : (inputadapter_row.getArray(0));
/* 036 */
/* 037 */ boolean filter_value = true;
/* 038 */
/* 039 */ if (!(!inputadapter_isNull)) {
/* 040 */ filter_value = inputadapter_isNull;
/* 041 */ }
/* 042 */ if (!filter_value) continue;
/* 043 */
/* 044 */ ((org.apache.spark.sql.execution.metric.SQLMetric) references[0] /* numOutputRows */).add(1);
/* 045 */
/* 046 */ boolean project_isNull = inputadapter_isNull;
/* 047 */ ArrayData project_value = null;
/* 048 */
/* 049 */ if (!inputadapter_isNull) {
/* 050 */ for (int z = 0; !project_isNull && z < inputadapter_value.numElements(); z++) {
/* 051 */ project_isNull |= inputadapter_value.isNullAt(z);
/* 052 */ }
/* 053 */ if (!project_isNull) {
/* 054 */ long project_numElements = 0;
/* 055 */ for (int z = 0; z < inputadapter_value.numElements(); z++) {
/* 056 */ project_numElements += inputadapter_value.getArray(z).numElements();
/* 057 */ }
/* 058 */ if (project_numElements > 2147483632) {
/* 059 */ throw new RuntimeException("Unsuccessful try to flatten an array of arrays with " +
/* 060 */ project_numElements + " elements due to exceeding the array size limit 2147483632.");
/* 061 */ }
/* 062 */
/* 063 */ Object[] project_arrayObject = new Object[(int)project_numElements];
/* 064 */ int project_counter = 0;
/* 065 */ for (int k = 0; k < inputadapter_value.numElements(); k++) {
/* 066 */ ArrayData arr = inputadapter_value.getArray(k);
/* 067 */ for (int l = 0; l < arr.numElements(); l++) {
/* 068 */ project_arrayObject[project_counter] = arr.getUTF8String(l);
/* 069 */ project_counter++;
/* 070 */ }
/* 071 */ }
/* 072 */ project_value = new org.apache.spark.sql.catalyst.util.GenericArrayData(project_arrayObject);
/* 073 */
/* 074 */ }
/* 075 */
/* 076 */ }
```
Author: mn-mikke <mrkAha12346github>
Closes #20938 from mn-mikke/feature/array-api-flatten-to-master.
5fea17b3be