You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "lyne7-sc (via GitHub)" <gi...@apache.org> on 2026/02/01 14:47:33 UTC
[PR] Fix `array_repeat` handling of null count values [datafusion]
lyne7-sc opened a new pull request, #20102:
URL: https://github.com/apache/datafusion/pull/20102
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
- Closes https://github.com/apache/datafusion/issues/20075.
## Rationale for this change
The previous implementation of `array_repeat` relied on Arrow defaults when handling null and negative count values. As a result, null counts were implicitly treated as zero and returned empty arrays, which is a correctness issue.
This PR makes the handling of these edge cases explicit and aligns the function with SQL null semantics.
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
## What changes are included in this PR?
- Explicit handling of null and negative count values
- Planner-time coercion of the count argument to `Int64`
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
## Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
Yes, SLTs added and pass.
## Are there any user-facing changes?
Yes. When the count value is null, `array_repeat` now returns a null array instead of an empty array.
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api change` label.
-->
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix `array_repeat` handling of null count values [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey commented on code in PR #20102:
URL: https://github.com/apache/datafusion/pull/20102#discussion_r2753000666
##########
datafusion/functions-nested/src/repeat.rs:
##########
@@ -222,7 +215,7 @@ fn general_repeat<O: OffsetSizeTrait>(
Arc::new(Field::new_list_field(array.data_type().to_owned(), true)),
OffsetBuffer::new(offsets.into()),
repeated_values,
- None,
+ Some(NullBuffer::new(nulls.finish())),
Review Comment:
We can copy the nulls buffer from the count array instead of using a builder
```suggestion
count_array.nulls().cloned(),
```
##########
datafusion/functions-nested/src/repeat.rs:
##########
@@ -193,21 +180,27 @@ fn array_repeat_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
/// ```
fn general_repeat<O: OffsetSizeTrait>(
array: &ArrayRef,
- count_array: &UInt64Array,
+ count_array: &Int64Array,
) -> Result<ArrayRef> {
- // Build offsets and take_indices
- let total_repeated_values: usize =
- count_array.values().iter().map(|&c| c as usize).sum();
+ let total_repeated_values: usize = count_array
+ .values()
+ .iter()
+ .map(|&c| if c > 0 { c as usize } else { 0 })
Review Comment:
nit: technically the spec allows null slots to have non-0 values, so this could overestimate
##########
datafusion/functions-nested/src/repeat.rs:
##########
@@ -298,8 +295,22 @@ fn general_list_repeat<O: OffsetSizeTrait>(
list_array.data_type().to_owned(),
true,
)),
- OffsetBuffer::<O>::from_lengths(counts.iter().map(|&c| c as usize)),
+ OffsetBuffer::<O>::from_lengths(
+ count_array
+ .iter()
+ .map(|c| c.map(|v| if v > 0 { v as usize } else { 0 }).unwrap_or(0)),
+ ),
Arc::new(inner_list),
- None,
+ Some(NullBuffer::new(outer_nulls.finish())),
Review Comment:
Same note here about reusing input null buffer
##########
datafusion/functions-nested/src/repeat.rs:
##########
@@ -238,42 +231,46 @@ fn general_repeat<O: OffsetSizeTrait>(
/// ```
fn general_list_repeat<O: OffsetSizeTrait>(
list_array: &GenericListArray<O>,
- count_array: &UInt64Array,
+ count_array: &Int64Array,
) -> Result<ArrayRef> {
let counts = count_array.values();
let list_offsets = list_array.value_offsets();
// calculate capacities for pre-allocation
- let outer_total = counts.iter().map(|&c| c as usize).sum();
- let inner_total = counts
- .iter()
- .enumerate()
- .filter(|&(i, _)| !list_array.is_null(i))
- .map(|(i, &c)| {
- let len = list_offsets[i + 1].to_usize().unwrap()
- - list_offsets[i].to_usize().unwrap();
- len * (c as usize)
- })
- .sum();
+ let mut outer_total = 0usize;
+ let mut inner_total = 0usize;
+ for (i, &c) in counts.iter().enumerate() {
+ if c > 0 {
+ outer_total += c as usize;
+ if !list_array.is_null(i) {
+ let len = list_offsets[i + 1].to_usize().unwrap()
+ - list_offsets[i].to_usize().unwrap();
+ inner_total += len * (c as usize);
+ }
+ }
+ }
// Build inner structures
let mut inner_offsets = Vec::with_capacity(outer_total + 1);
let mut take_indices = Vec::with_capacity(inner_total);
let mut inner_nulls = BooleanBufferBuilder::new(outer_total);
+ let mut outer_nulls = BooleanBufferBuilder::new(count_array.len());
let mut inner_running = 0usize;
inner_offsets.push(O::zero());
- for (row_idx, &count) in counts.iter().enumerate() {
- let is_valid = !list_array.is_null(row_idx);
+ for row_idx in 0..count_array.len() {
+ let (count, count_is_valid) = get_count_with_validity(count_array, row_idx);
+ outer_nulls.append(count_is_valid);
+ let list_is_valid = !list_array.is_null(row_idx);
Review Comment:
```suggestion
let list_is_valid = list_array.is_valid(row_idx);
```
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix `array_repeat` handling of null count values [datafusion]
Posted by "martin-g (via GitHub)" <gi...@apache.org>.
martin-g commented on code in PR #20102:
URL: https://github.com/apache/datafusion/pull/20102#discussion_r2754146646
##########
datafusion/functions-nested/src/repeat.rs:
##########
@@ -193,21 +180,27 @@ fn array_repeat_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
/// ```
fn general_repeat<O: OffsetSizeTrait>(
array: &ArrayRef,
- count_array: &UInt64Array,
+ count_array: &Int64Array,
) -> Result<ArrayRef> {
- // Build offsets and take_indices
- let total_repeated_values: usize =
- count_array.values().iter().map(|&c| c as usize).sum();
+ let total_repeated_values: usize = count_array
+ .values()
+ .iter()
+ .map(|&c| if c > 0 { c as usize } else { 0 })
+ .sum();
+
let mut take_indices = Vec::with_capacity(total_repeated_values);
+ let mut nulls = BooleanBufferBuilder::new(count_array.len());
let mut offsets = Vec::with_capacity(count_array.len() + 1);
offsets.push(O::zero());
let mut running_offset = 0usize;
- for (idx, &count) in count_array.values().iter().enumerate() {
- let count = count as usize;
+ for idx in 0..count_array.len() {
+ let (count, is_valid) = get_count_with_validity(count_array, idx);
+
running_offset += count;
offsets.push(O::from_usize(running_offset).unwrap());
Review Comment:
Extreme case:
1) if the input values is ListArray, then its offset type will be i32
2) and if the count value is bigger than i32::MAX
3) then i32::from_usize(i32::MAX + 1) will return None and it will panic
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix `array_repeat` handling of null count values [datafusion]
Posted by "martin-g (via GitHub)" <gi...@apache.org>.
martin-g commented on code in PR #20102:
URL: https://github.com/apache/datafusion/pull/20102#discussion_r2754035422
##########
datafusion/sqllogictest/test_files/array.slt:
##########
@@ -3256,24 +3256,97 @@ drop table array_repeat_table;
statement ok
drop table large_array_repeat_table;
-
+# array_repeat: arrays with NULL counts
statement ok
create table array_repeat_null_count_table
as values
(1, 2),
(2, null),
-(3, 1);
+(3, 1),
+(4, -1);
query I?
select column1, array_repeat(column1, column2) from array_repeat_null_count_table;
----
1 [1, 1]
-2 []
+2 NULL
3 [3]
+4 []
Review Comment:
```suggestion
4 []
NULL NULL
```
##########
datafusion/functions-nested/src/repeat.rs:
##########
@@ -238,42 +231,46 @@ fn general_repeat<O: OffsetSizeTrait>(
/// ```
fn general_list_repeat<O: OffsetSizeTrait>(
list_array: &GenericListArray<O>,
- count_array: &UInt64Array,
+ count_array: &Int64Array,
) -> Result<ArrayRef> {
let counts = count_array.values();
Review Comment:
This also does not check for NULLs in the count_array and may lead to overestimates. You need to use `get_count_with_validity()` at line 245 too
##########
datafusion/functions-nested/src/repeat.rs:
##########
@@ -193,21 +180,27 @@ fn array_repeat_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
/// ```
fn general_repeat<O: OffsetSizeTrait>(
array: &ArrayRef,
- count_array: &UInt64Array,
+ count_array: &Int64Array,
) -> Result<ArrayRef> {
- // Build offsets and take_indices
- let total_repeated_values: usize =
- count_array.values().iter().map(|&c| c as usize).sum();
+ let total_repeated_values: usize = count_array
+ .values()
+ .iter()
+ .map(|&c| if c > 0 { c as usize } else { 0 })
Review Comment:
Possible fix:
```rust
let total_repeated_values: usize = (0..count_array.len())
.map(|i| get_count_with_validity(count_array, i).0)
.sum();
```
##########
datafusion/sqllogictest/test_files/array.slt:
##########
@@ -3256,24 +3256,97 @@ drop table array_repeat_table;
statement ok
drop table large_array_repeat_table;
-
+# array_repeat: arrays with NULL counts
statement ok
create table array_repeat_null_count_table
as values
(1, 2),
(2, null),
-(3, 1);
+(3, 1),
+(4, -1);
Review Comment:
```suggestion
(4, -1),
(null, null);
```
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix `array_repeat` handling of null count values [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey commented on PR #20102:
URL: https://github.com/apache/datafusion/pull/20102#issuecomment-3850541499
Thanks @lyne7-sc & @martin-g
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix `array_repeat` handling of null count values [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey merged PR #20102:
URL: https://github.com/apache/datafusion/pull/20102
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix `array_repeat` handling of null count values [datafusion]
Posted by "lyne7-sc (via GitHub)" <gi...@apache.org>.
lyne7-sc commented on PR #20102:
URL: https://github.com/apache/datafusion/pull/20102#issuecomment-3835789317
Thanks for the review.
The code has been updated to handle overflow cases and and fix the capacity pre-calculation.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org