jayzhan211 commented on code in PR #25952: URL: https://github.com/apache/datafusion/pull/25952#discussion_r4207466603
########## datafusion/common/benches/create_hashes.rs: ########## @@ -0,0 +1,300 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{ + ArrayRef, ArrowPrimitiveType, DictionaryArray, FixedSizeListArray, Int32Array, + Int64Array, ListArray, ListViewArray, MapArray, NullArray, PrimitiveArray, RunArray, + StringArray, StringViewArray, StructArray, UnionArray, +}; +use arrow::buffer::{OffsetBuffer, ScalarBuffer}; +use arrow::datatypes::{DataType, Field, Fields, Int32Type, Int64Type, UnionFields}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::hash_utils::{RandomState, create_hashes}; +use rand::Rng; +use rand::SeedableRng; +use rand::distr::{Distribution, StandardUniform}; +use rand::prelude::StdRng; +use std::sync::Arc; + +const BATCH_SIZE: usize = 8192; +const ELEMENTS_PER_ROW: usize = 5; + +fn make_rng() -> StdRng { + StdRng::seed_from_u64(42) +} + +fn bench_create_hashes(c: &mut Criterion, name: &str, array: &ArrayRef) { + let state = RandomState::default(); + let mut hashes = vec![0u64; BATCH_SIZE]; + c.bench_function(name, |b| { + b.iter(|| { + create_hashes(std::slice::from_ref(array), &state, &mut hashes).unwrap(); + }) + }); +} + +fn primitive_array<T>() -> ArrayRef +where + T: ArrowPrimitiveType, + StandardUniform: Distribution<T::Native>, +{ + let mut rng = make_rng(); + Arc::new( + (0..BATCH_SIZE) + .map(|_| Some(rng.random::<T::Native>())) + .collect::<PrimitiveArray<T>>(), + ) +} + +fn null_array() -> ArrayRef { + Arc::new(NullArray::new(BATCH_SIZE)) +} + +fn utf8_array() -> ArrayRef { + let mut rng = make_rng(); + Arc::new(StringArray::from_iter_values((0..BATCH_SIZE).map(|_| { + let len = rng.random_range(1usize..=32); + (0..len).map(|_| rng.random::<char>()).collect::<String>() + }))) +} + +fn string_view_array() -> ArrayRef { + let mut rng = make_rng(); + Arc::new(StringViewArray::from_iter_values((0..BATCH_SIZE).map( + |_| { + let len = rng.random_range(1usize..=32); + (0..len).map(|_| rng.random::<char>()).collect::<String>() + }, + ))) +} + +fn dictionary_array() -> ArrayRef { + let pool: Vec<String> = (0..100).map(|i| format!("value_{i}")).collect(); + let mut rng = make_rng(); + Arc::new(DictionaryArray::<Int32Type>::from_iter( + (0..BATCH_SIZE).map(|_| Some(pool[rng.random_range(0..pool.len())].as_str())), + )) +} + +fn struct_array() -> ArrayRef { + Arc::new(StructArray::new( + Fields::from(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int32, false), + ]), + vec![ + primitive_array::<Int64Type>(), + primitive_array::<Int32Type>(), + ], + None, + )) +} + +fn list_array() -> ArrayRef { + let mut rng = make_rng(); + let total = BATCH_SIZE * ELEMENTS_PER_ROW; + let values: Int64Array = (0..total).map(|_| Some(rng.random::<i64>())).collect(); + let offsets: Vec<i32> = (0..=BATCH_SIZE) + .map(|i| (i * ELEMENTS_PER_ROW) as i32) + .collect(); + Arc::new(ListArray::new( + Arc::new(Field::new("item", DataType::Int64, true)), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + Arc::new(values), + None, + )) +} + +fn list_view_array() -> ArrayRef { + let mut rng = make_rng(); + let total = BATCH_SIZE * ELEMENTS_PER_ROW; + let values: Int64Array = (0..total).map(|_| Some(rng.random::<i64>())).collect(); + let offsets: ScalarBuffer<i32> = (0..BATCH_SIZE) + .map(|i| (i * ELEMENTS_PER_ROW) as i32) + .collect(); + let sizes: ScalarBuffer<i32> = + (0..BATCH_SIZE).map(|_| ELEMENTS_PER_ROW as i32).collect(); + Arc::new(ListViewArray::new( + Arc::new(Field::new("item", DataType::Int64, true)), + offsets, + sizes, + Arc::new(values), + None, + )) +} + +fn map_array() -> ArrayRef { + let mut rng = make_rng(); + let total = BATCH_SIZE * ELEMENTS_PER_ROW; + let keys: Int32Array = (0..total).map(|_| Some(rng.random::<i32>())).collect(); + let vals: Int64Array = (0..total).map(|_| Some(rng.random::<i64>())).collect(); + let offsets: Vec<i32> = (0..=BATCH_SIZE) + .map(|i| (i * ELEMENTS_PER_ROW) as i32) + .collect(); + let entries = StructArray::try_new( + Fields::from(vec![ + Field::new("keys", DataType::Int32, false), + Field::new("values", DataType::Int64, true), + ]), + vec![Arc::new(keys), Arc::new(vals)], + None, + ) + .unwrap(); + Arc::new(MapArray::new( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("keys", DataType::Int32, false), + Field::new("values", DataType::Int64, true), + ])), + false, + )), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + entries, + None, + false, + )) +} + +fn fixed_size_list_array() -> ArrayRef { Review Comment: `primitive_array()` returns 8192 values, and with `list_size = 4` that is a 2048-row `FixedSizeListArray` hashed into an 8192-slot buffer. This case therefore measures a quarter of the rows of every other case. It only runs because `hash_fixed_list_array` skips the buffer-length check that `hash_array_primitive` has (`hash_utils.rs:316`). Running the same shape through `with_hashes`'s length check fails with `left: 2048, right: 8192`. ```diff fn fixed_size_list_array() -> ArrayRef { - let list_size = 4i32; + let list_size = 4; + let mut rng = make_rng(); + let values: Int64Array = (0..BATCH_SIZE * list_size) + .map(|_| Some(rng.random::<i64>())) + .collect(); Arc::new(FixedSizeListArray::new( Arc::new(Field::new("item", DataType::Int64, true)), - list_size, - primitive_array::<Int64Type>(), + list_size as i32, + Arc::new(values), None, )) } ``` -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
