This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 0d994d08 chore: remove some unwraps from shuffle module (#601)
0d994d08 is described below
commit 0d994d0830f5d116830ce8bd457911165d0c9030
Author: Andy Grove <[email protected]>
AuthorDate: Thu Jun 27 13:37:54 2024 -0700
chore: remove some unwraps from shuffle module (#601)
* remove some unwraps from shuffle module
* simplifiy
---
core/src/execution/shuffle/list.rs | 26 +++++++++++------------
core/src/execution/shuffle/row.rs | 43 ++++++++++++++++++++------------------
2 files changed, 36 insertions(+), 33 deletions(-)
diff --git a/core/src/execution/shuffle/list.rs
b/core/src/execution/shuffle/list.rs
index 53d155f8..d8bdcb19 100644
--- a/core/src/execution/shuffle/list.rs
+++ b/core/src/execution/shuffle/list.rs
@@ -192,7 +192,7 @@ pub fn append_list_element<T: ArrayBuilder>(
list_builder
.as_any_mut()
.downcast_mut::<ListBuilder<BooleanBuilder>>()
- .unwrap(),
+ .expect("ListBuilder<BooleanBuilder>"),
list,
idx,
),
@@ -200,7 +200,7 @@ pub fn append_list_element<T: ArrayBuilder>(
list_builder
.as_any_mut()
.downcast_mut::<ListBuilder<Int8Builder>>()
- .unwrap(),
+ .expect("ListBuilder<Int8Builder>"),
list,
idx,
),
@@ -208,7 +208,7 @@ pub fn append_list_element<T: ArrayBuilder>(
list_builder
.as_any_mut()
.downcast_mut::<ListBuilder<Int16Builder>>()
- .unwrap(),
+ .expect("ListBuilder<Int16Builder>"),
list,
idx,
),
@@ -216,7 +216,7 @@ pub fn append_list_element<T: ArrayBuilder>(
list_builder
.as_any_mut()
.downcast_mut::<ListBuilder<Int32Builder>>()
- .unwrap(),
+ .expect("ListBuilder<Int32Builder>"),
list,
idx,
),
@@ -224,7 +224,7 @@ pub fn append_list_element<T: ArrayBuilder>(
list_builder
.as_any_mut()
.downcast_mut::<ListBuilder<Int64Builder>>()
- .unwrap(),
+ .expect("ListBuilder<Int64Builder>"),
list,
idx,
),
@@ -232,7 +232,7 @@ pub fn append_list_element<T: ArrayBuilder>(
list_builder
.as_any_mut()
.downcast_mut::<ListBuilder<Float32Builder>>()
- .unwrap(),
+ .expect("ListBuilder<Float32Builder>"),
list,
idx,
),
@@ -240,7 +240,7 @@ pub fn append_list_element<T: ArrayBuilder>(
list_builder
.as_any_mut()
.downcast_mut::<ListBuilder<Float64Builder>>()
- .unwrap(),
+ .expect("ListBuilder<Float64Builder>"),
list,
idx,
),
@@ -248,7 +248,7 @@ pub fn append_list_element<T: ArrayBuilder>(
list_builder
.as_any_mut()
.downcast_mut::<ListBuilder<Date32Builder>>()
- .unwrap(),
+ .expect("ListBuilder<Date32Builder>"),
list,
idx,
),
@@ -256,7 +256,7 @@ pub fn append_list_element<T: ArrayBuilder>(
list_builder
.as_any_mut()
.downcast_mut::<ListBuilder<TimestampMicrosecondBuilder>>()
- .unwrap(),
+ .expect("ListBuilder<TimestampMicrosecondBuilder>"),
list,
idx,
),
@@ -264,7 +264,7 @@ pub fn append_list_element<T: ArrayBuilder>(
list_builder
.as_any_mut()
.downcast_mut::<ListBuilder<BinaryBuilder>>()
- .unwrap(),
+ .expect("ListBuilder<BinaryBuilder>"),
list,
idx,
),
@@ -272,7 +272,7 @@ pub fn append_list_element<T: ArrayBuilder>(
list_builder
.as_any_mut()
.downcast_mut::<ListBuilder<StringBuilder>>()
- .unwrap(),
+ .expect("ListBuilder<StringBuilder>"),
list,
idx,
),
@@ -281,7 +281,7 @@ pub fn append_list_element<T: ArrayBuilder>(
.values()
.as_any_mut()
.downcast_mut::<Decimal128Builder>()
- .unwrap();
+ .expect("ListBuilder<Decimal128Builder>");
let is_null = list.is_null_at(idx);
if is_null {
@@ -319,7 +319,7 @@ pub fn append_list_element<T: ArrayBuilder>(
.values()
.as_any_mut()
.downcast_mut::<StructBuilder>()
- .unwrap();
+ .expect("StructBuilder");
let is_null = list.is_null_at(idx);
let nested_row = if is_null {
diff --git a/core/src/execution/shuffle/row.rs
b/core/src/execution/shuffle/row.rs
index 2d1312c1..2aeb4881 100644
--- a/core/src/execution/shuffle/row.rs
+++ b/core/src/execution/shuffle/row.rs
@@ -39,7 +39,7 @@ use arrow_array::{
types::Int32Type,
Array, ArrayRef, RecordBatch, RecordBatchOptions,
};
-use arrow_schema::{DataType, Field, Schema, TimeUnit};
+use arrow_schema::{ArrowError, DataType, Field, Schema, TimeUnit};
use jni::sys::{jint, jlong};
use std::{
fs::OpenOptions,
@@ -275,7 +275,10 @@ impl SparkUnsafeRow {
macro_rules! downcast_builder {
($builder_type:ty, $builder:expr) => {
- $builder.into_box_any().downcast::<$builder_type>().unwrap()
+ $builder
+ .into_box_any()
+ .downcast::<$builder_type>()
+ .expect(stringify!($builder_type))
};
}
@@ -284,7 +287,7 @@ macro_rules! downcast_builder_ref {
$builder
.as_any_mut()
.downcast_mut::<$builder_type>()
- .unwrap()
+ .expect(stringify!($builder_type))
};
}
@@ -348,8 +351,7 @@ pub(crate) fn append_field(
$field,
field_builder,
&row.get_map(idx),
- )
- .unwrap();
+ )?;
}
}
}};
@@ -378,8 +380,7 @@ pub(crate) fn append_field(
$element_dt,
field_builder,
&row.get_array(idx),
- )
- .unwrap()
+ )?
}
}
}};
@@ -1057,7 +1058,7 @@ pub(crate) fn append_columns(
let element_builder = builder
.as_any_mut()
.downcast_mut::<$builder_type>()
- .unwrap();
+ .expect(stringify!($builder_type));
let mut row = SparkUnsafeRow::new(schema);
for i in row_start..row_end {
@@ -1084,7 +1085,7 @@ pub(crate) fn append_columns(
let list_builder = builder
.as_any_mut()
.downcast_mut::<ListBuilder<$builder_type>>()
- .unwrap();
+ .expect(stringify!($builder_type));
let mut row = SparkUnsafeRow::new(schema);
for i in row_start..row_end {
@@ -1103,8 +1104,7 @@ pub(crate) fn append_columns(
$element_dt,
list_builder,
&row.get_array(column_idx),
- )
- .unwrap()
+ )?
}
}
}};
@@ -1116,7 +1116,11 @@ pub(crate) fn append_columns(
let map_builder = builder
.as_any_mut()
.downcast_mut::<MapBuilder<$key_builder_type,
$value_builder_type>>()
- .unwrap();
+ .expect(&format!(
+ "MapBuilder<{},{}>",
+ stringify!($key_builder_type),
+ stringify!($value_builder_type)
+ ));
let mut row = SparkUnsafeRow::new(schema);
for i in row_start..row_end {
@@ -1135,8 +1139,7 @@ pub(crate) fn append_columns(
$field,
map_builder,
&row.get_map(column_idx),
- )
- .unwrap()
+ )?
}
}
}};
@@ -1148,7 +1151,7 @@ pub(crate) fn append_columns(
let struct_builder = builder
.as_any_mut()
.downcast_mut::<StructBuilder>()
- .unwrap();
+ .expect("StructBuilder");
let mut row = SparkUnsafeRow::new(schema);
for i in row_start..row_end {
@@ -3347,7 +3350,7 @@ pub fn process_sorted_row_partition(
.zip(schema.iter())
.map(|(builder, datatype)| builder_to_array(builder, datatype,
prefer_dictionary_ratio))
.collect();
- let batch = make_batch(array_refs?, n);
+ let batch = make_batch(array_refs?, n)?;
let mut frozen: Vec<u8> = vec![];
let mut cursor = Cursor::new(&mut frozen);
@@ -3382,7 +3385,7 @@ fn builder_to_array(
let builder = builder
.as_any_mut()
.downcast_mut::<StringDictionaryBuilder<Int32Type>>()
- .unwrap();
+ .expect("StringDictionaryBuilder<Int32Type>");
let dict_array = builder.finish();
let num_keys = dict_array.keys().len();
@@ -3401,7 +3404,7 @@ fn builder_to_array(
let builder = builder
.as_any_mut()
.downcast_mut::<BinaryDictionaryBuilder<Int32Type>>()
- .unwrap();
+ .expect("BinaryDictionaryBuilder<Int32Type>");
let dict_array = builder.finish();
let num_keys = dict_array.keys().len();
@@ -3420,7 +3423,7 @@ fn builder_to_array(
}
}
-fn make_batch(arrays: Vec<ArrayRef>, row_count: usize) -> RecordBatch {
+fn make_batch(arrays: Vec<ArrayRef>, row_count: usize) -> Result<RecordBatch,
ArrowError> {
let mut dict_id = 0;
let fields = arrays
.iter()
@@ -3442,5 +3445,5 @@ fn make_batch(arrays: Vec<ArrayRef>, row_count: usize) ->
RecordBatch {
.collect::<Vec<_>>();
let schema = Arc::new(Schema::new(fields));
let options =
RecordBatchOptions::new().with_row_count(Option::from(row_count));
- RecordBatch::try_new_with_options(schema, arrays, &options).unwrap()
+ RecordBatch::try_new_with_options(schema, arrays, &options)
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
