vincev commented on code in PR #6903:
URL: https://github.com/apache/arrow-datafusion/pull/6903#discussion_r1260595224
##########
datafusion/core/src/physical_plan/unnest.rs:
##########
@@ -264,41 +267,195 @@ fn unnest_batch<T>(
where
T: ArrayAccessor<Item = ArrayRef>,
{
- let mut batches = Vec::new();
- let mut num_rows = 0;
-
- for row in 0..batch.num_rows() {
- let arrays = batch
- .columns()
- .iter()
- .enumerate()
- .map(|(col_idx, arr)| {
- if col_idx == column.index() {
- // Unnest the value at the given row.
- if list_array.value(row).is_empty() {
- // If nested array is empty add an array with 1 null.
- Ok(new_null_array(list_array.value(row).data_type(),
1))
- } else {
- Ok(list_array.value(row))
- }
- } else {
- // Number of elements to duplicate, use max(1) to handle
null.
- let nested_len = list_array.value(row).len().max(1);
- // Duplicate rows for each value in the nested array.
- if arr.is_null(row) {
- Ok(new_null_array(arr.data_type(), nested_len))
- } else {
- let scalar = ScalarValue::try_from_array(arr, row)?;
- Ok(scalar.to_array_of_size(nested_len))
- }
- }
- })
- .collect::<Result<Vec<_>>>()?;
+ // Create an array with the unnested values of the list array, given the
list
+ // array:
+ //
+ // [1], null, [2, 3, 4], null, [5, 6]
+ //
+ // the result array is:
+ //
+ // 1, null, 2, 3, 4, null, 5, 6
+ //
+ let unnested_array = unnest_array(list_array)?;
+
+ // Create an array with the lengths of each list value in the nested array.
+ // Given the nested array:
+ //
+ // [1], null, [2, 3, 4], null, [5, 6]
+ //
+ // the result array is:
+ //
+ // 1, null, 3, null, 2
+ //
+ // Depending on the list type the result may be Int32Array or Int64Array.
+ let list_lengths = kernels::length::length(list_array)?;
+
+ // Create the indices for the take kernel and then use those indices to
create
+ // the unnested record batch.
+ match list_lengths.data_type() {
+ DataType::Int32 => {
+ let list_lengths = as_primitive_array::<Int32Type>(&list_lengths)?;
+ let indices = create_take_indices(list_lengths,
unnested_array.len());
+ batch_from_indices(batch, schema, column.index(), &unnested_array,
&indices)
Review Comment:
We need this final step so that all the other columns have the same number
of values as the `unnest_array`. Let's say we have the following `batch`:
```
+----+-----------------+
| id | tags |
+----+-----------------+
| 1 | [t11] |
| 2 | [t21, t22] |
| 3 | [t31, t32, t33] |
+----+-----------------+
```
and we want to unnest the `tags` column, the `unnest_array` will have 6
values:
```
+-------+
| tags |
+-------+
| t11 |
| t21 |
| t22 |
| t31 |
| t32 |
| t33 |
+-------+
```
now we need to adjust the `id` column so that it has the same number of
values as the unnested `tags` column while maintaining the association with the
unnested list values, so that we have:
```
+----+-------+
| id | tags |
+----+-------+
| 1 | t11 |
| 2 | t21 |
| 2 | t22 |
| 3 | t31 |
| 3 | t32 |
| 3 | t33 |
+----+-------+
```
The `batch_from_indices` function expands all the other unnested columns
using the `take` kernel to replicate their values so that they match the
unnested array.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
