Skip to content

Commit 23efb24

Browse files
authored
Merge branch 'main' into fix/5783-duplicate-fields
2 parents b00bc9a + 4eeb1f8 commit 23efb24

4 files changed

Lines changed: 509 additions & 26 deletions

File tree

‎native/core/Cargo.toml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,3 +127,7 @@ harness = false
127127
[[bench]]
128128
name = "explode"
129129
harness = false
130+
131+
[[bench]]
132+
name = "parquet_timestamp_conversion"
133+
harness = false
Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
//! Parquet timestamp conversion with nullable containers and unchanged array siblings.
19+
//! Run the same benchmark on the reviewed head and the candidate with Criterion baselines.
20+
21+
use std::{hint::black_box, sync::Arc};
22+
23+
use arrow::array::{
24+
Array, ArrayRef, Int32Array, ListArray, MapArray, StructArray, TimestampMillisecondArray,
25+
};
26+
use arrow::buffer::{NullBuffer, OffsetBuffer};
27+
use arrow::datatypes::{DataType, Field, TimeUnit};
28+
use comet::parquet::parquet_support::{spark_parquet_convert, SparkParquetOptions};
29+
use criterion::{criterion_group, criterion_main, Criterion};
30+
use datafusion::physical_plan::ColumnarValue;
31+
use datafusion_comet_spark_expr::EvalMode;
32+
33+
const ROWS: usize = 1024;
34+
35+
// Null rows have empty ranges, as in a decoded Parquet list/map. Their presence still
36+
// caused the original repeated_visibility implementation to scan all backing items.
37+
fn offsets(width: usize) -> (OffsetBuffer<i32>, NullBuffer, usize) {
38+
let valid: Vec<bool> = (0..ROWS).map(|row| row % 8 != 0).collect();
39+
let mut offsets = vec![0_i32];
40+
for &present in &valid {
41+
offsets.push(offsets.last().unwrap() + if present { width as i32 } else { 0 });
42+
}
43+
let len = *offsets.last().unwrap() as usize;
44+
(
45+
OffsetBuffer::new(offsets.into()),
46+
NullBuffer::from(valid),
47+
len,
48+
)
49+
}
50+
51+
fn timestamp_field(unit: TimeUnit) -> Arc<Field> {
52+
Arc::new(Field::new("ts", DataType::Timestamp(unit, None), true))
53+
}
54+
55+
fn array_sibling(width: usize) -> (ArrayRef, DataType) {
56+
let (offsets, nulls, len) = offsets(width);
57+
let ints: ArrayRef = Arc::new(ListArray::new(
58+
Arc::new(Field::new("item", DataType::Int32, false)),
59+
offsets,
60+
Arc::new(Int32Array::from(vec![7; len])),
61+
Some(nulls),
62+
));
63+
let field = Arc::new(Field::new("ints", ints.data_type().clone(), true));
64+
let target =
65+
DataType::Struct(vec![timestamp_field(TimeUnit::Microsecond), Arc::clone(&field)].into());
66+
let input = StructArray::new(
67+
vec![timestamp_field(TimeUnit::Millisecond), field].into(),
68+
vec![
69+
Arc::new(TimestampMillisecondArray::from(vec![7; ROWS])),
70+
ints,
71+
],
72+
None,
73+
);
74+
(Arc::new(input), target)
75+
}
76+
77+
fn timestamp_containers() -> [(ArrayRef, DataType); 2] {
78+
let (offsets, nulls, len) = offsets(32);
79+
let timestamps: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![7; len]));
80+
let list = ListArray::new(
81+
timestamp_field(TimeUnit::Millisecond),
82+
offsets.clone(),
83+
Arc::clone(&timestamps),
84+
Some(nulls.clone()),
85+
);
86+
let key = Arc::new(Field::new("key", DataType::Int32, false));
87+
let entries = StructArray::new(
88+
vec![Arc::clone(&key), timestamp_field(TimeUnit::Millisecond)].into(),
89+
vec![
90+
Arc::new(Int32Array::from(
91+
(0..len).map(|i| (i % 32) as i32).collect::<Vec<_>>(),
92+
)),
93+
timestamps,
94+
],
95+
None,
96+
);
97+
let map = MapArray::new(
98+
Arc::new(Field::new("entries", entries.data_type().clone(), false)),
99+
offsets,
100+
entries,
101+
Some(nulls),
102+
false,
103+
);
104+
let target_map = DataType::Map(
105+
Arc::new(Field::new(
106+
"entries",
107+
DataType::Struct(vec![key, timestamp_field(TimeUnit::Microsecond)].into()),
108+
false,
109+
)),
110+
false,
111+
);
112+
[
113+
(
114+
Arc::new(list),
115+
DataType::List(timestamp_field(TimeUnit::Microsecond)),
116+
),
117+
(Arc::new(map), target_map),
118+
]
119+
}
120+
121+
fn benchmark(c: &mut Criterion) {
122+
let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false);
123+
let mut group = c.benchmark_group("parquet_timestamp_conversion");
124+
for width in [8, 1024] {
125+
let (array, target) = array_sibling(width);
126+
group.bench_function(format!("array_sibling_width_{width}"), |b| {
127+
b.iter(|| {
128+
spark_parquet_convert(
129+
ColumnarValue::Array(Arc::clone(black_box(&array))),
130+
black_box(&target),
131+
black_box(&options),
132+
)
133+
.unwrap()
134+
})
135+
});
136+
}
137+
for (kind, (array, target)) in ["list", "map"].into_iter().zip(timestamp_containers()) {
138+
for sliced in [false, true] {
139+
let input = if sliced {
140+
array.slice(ROWS / 4, ROWS / 2)
141+
} else {
142+
Arc::clone(&array)
143+
};
144+
group.bench_function(format!("timestamp_{kind}_sliced_{sliced}"), |b| {
145+
b.iter(|| {
146+
spark_parquet_convert(
147+
ColumnarValue::Array(Arc::clone(black_box(&input))),
148+
black_box(&target),
149+
black_box(&options),
150+
)
151+
.unwrap()
152+
})
153+
});
154+
}
155+
}
156+
group.finish();
157+
}
158+
159+
criterion_group!(benches, benchmark);
160+
criterion_main!(benches);

0 commit comments

Comments
 (0)