chore: Cleanup code to use repeat_n in a few places (#20527)

## Which issue does this PR close?

N/A

## Rationale for this change

Using `repeat_n` is more readable and slightly faster than
`(0..n).map(|_| ...)`.

## What changes are included in this PR?

## Are these changes tested?

Yes.

## Are there any user-facing changes?

No.
This commit is contained in:
Neil Conway
2026-02-24 10:46:58 -05:00
committed by GitHub
parent 670dbf481c
commit e71e7a39bf
6 changed files with 18 additions and 16 deletions
+2 -3
View File
@@ -66,10 +66,9 @@ fn generate_spm_for_round_robin_tie_breaker(
RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap()
};
let rbs = (0..batch_count).map(|_| rb.clone()).collect::<Vec<_>>();
let partitions = vec![rbs.clone(); partition_count];
let schema = rb.schema();
let rbs = std::iter::repeat_n(rb, batch_count).collect::<Vec<_>>();
let partitions = vec![rbs.clone(); partition_count];
let sort = [
PhysicalSortExpr {
expr: col("b", &schema).unwrap(),
@@ -278,9 +278,11 @@ async fn run_sort_test_with_limited_memory(
let string_item_size =
record_batch_memory_size / record_batch_size as usize;
let string_array = Arc::new(StringArray::from_iter_values(
(0..record_batch_size).map(|_| "a".repeat(string_item_size)),
));
let string_array =
Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
"a".repeat(string_item_size),
record_batch_size as usize,
)));
RecordBatch::try_new(
Arc::clone(&schema),
@@ -536,9 +538,11 @@ async fn run_test_aggregate_with_high_cardinality(
let string_item_size =
record_batch_memory_size / record_batch_size as usize;
let string_array = Arc::new(StringArray::from_iter_values(
(0..record_batch_size).map(|_| "a".repeat(string_item_size)),
));
let string_array =
Arc::new(StringArray::from_iter_values(std::iter::repeat_n(
"a".repeat(string_item_size),
record_batch_size as usize,
)));
RecordBatch::try_new(
Arc::clone(&schema),
+1 -1
View File
@@ -392,7 +392,7 @@ mod tests {
vec![Arc::new(UInt32Array::from(vec![1, 2, 3, 4, 5, 6]))],
)?;
let batches: Vec<RecordBatch> = (0..50).map(|_| batch.clone()).collect();
let batches: Vec<RecordBatch> = std::iter::repeat_n(batch, 50).collect();
let session_config = SessionConfig::new().with_batch_size(200);
let task_ctx = TaskContext::default().with_session_config(session_config);
@@ -2483,7 +2483,7 @@ mod tests {
/// Create vector batches
fn create_vec_batches(n: usize) -> Vec<RecordBatch> {
let batch = create_batch();
(0..n).map(|_| batch.clone()).collect()
std::iter::repeat_n(batch, n).collect()
}
/// Create batch
@@ -475,10 +475,9 @@ mod tests {
let b: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("a"); row_size]));
let c: ArrayRef = Arc::new(Int64Array::from_iter(vec![0; row_size]));
let rb = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)])?;
let rbs = (0..1024).map(|_| rb.clone()).collect::<Vec<_>>();
let schema = rb.schema();
let rbs = std::iter::repeat_n(rb, 1024).collect::<Vec<_>>();
let sort = [
PhysicalSortExpr {
expr: col("b", &schema)?,
+2 -2
View File
@@ -181,8 +181,8 @@ fn map_deduplicate_keys(
let num_keys_entries = *next_keys_offset as usize - cur_keys_offset;
let num_values_entries = *next_values_offset as usize - cur_values_offset;
let mut keys_mask_one = [false].repeat(num_keys_entries);
let mut values_mask_one = [false].repeat(num_values_entries);
let mut keys_mask_one = vec![false; num_keys_entries];
let mut values_mask_one = vec![false; num_values_entries];
let key_is_valid = keys_nulls.is_none_or(|buf| buf.is_valid(row_idx));
let value_is_valid = values_nulls.is_none_or(|buf| buf.is_valid(row_idx));