Skip to content

Commit 0d994d0

Browse files
andygroveandygrove
authored
chore: remove some unwraps from shuffle module (#601)
* remove some unwraps from shuffle module * simplifiy
1 parent a1641ab commit 0d994d0

File tree

2 files changed

+36
-33
lines changed

2 files changed

+36
-33
lines changed

‎core/src/execution/shuffle/list.rs‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -192,87 +192,87 @@ pub fn append_list_element<T: ArrayBuilder>(
192192
list_builder
193193
.as_any_mut()
194194
.downcast_mut::<ListBuilder<BooleanBuilder>>()
195-
.unwrap(),
195+
.expect("ListBuilder<BooleanBuilder>"),
196196
list,
197197
idx,
198198
),
199199
DataType::Int8 => append_int8_element(
200200
list_builder
201201
.as_any_mut()
202202
.downcast_mut::<ListBuilder<Int8Builder>>()
203-
.unwrap(),
203+
.expect("ListBuilder<Int8Builder>"),
204204
list,
205205
idx,
206206
),
207207
DataType::Int16 => append_int16_element(
208208
list_builder
209209
.as_any_mut()
210210
.downcast_mut::<ListBuilder<Int16Builder>>()
211-
.unwrap(),
211+
.expect("ListBuilder<Int16Builder>"),
212212
list,
213213
idx,
214214
),
215215
DataType::Int32 => append_int32_element(
216216
list_builder
217217
.as_any_mut()
218218
.downcast_mut::<ListBuilder<Int32Builder>>()
219-
.unwrap(),
219+
.expect("ListBuilder<Int32Builder>"),
220220
list,
221221
idx,
222222
),
223223
DataType::Int64 => append_int64_element(
224224
list_builder
225225
.as_any_mut()
226226
.downcast_mut::<ListBuilder<Int64Builder>>()
227-
.unwrap(),
227+
.expect("ListBuilder<Int64Builder>"),
228228
list,
229229
idx,
230230
),
231231
DataType::Float32 => append_float32_element(
232232
list_builder
233233
.as_any_mut()
234234
.downcast_mut::<ListBuilder<Float32Builder>>()
235-
.unwrap(),
235+
.expect("ListBuilder<Float32Builder>"),
236236
list,
237237
idx,
238238
),
239239
DataType::Float64 => append_float64_element(
240240
list_builder
241241
.as_any_mut()
242242
.downcast_mut::<ListBuilder<Float64Builder>>()
243-
.unwrap(),
243+
.expect("ListBuilder<Float64Builder>"),
244244
list,
245245
idx,
246246
),
247247
DataType::Date32 => append_date32_element(
248248
list_builder
249249
.as_any_mut()
250250
.downcast_mut::<ListBuilder<Date32Builder>>()
251-
.unwrap(),
251+
.expect("ListBuilder<Date32Builder>"),
252252
list,
253253
idx,
254254
),
255255
DataType::Timestamp(TimeUnit::Microsecond, _) => append_timestamp_element(
256256
list_builder
257257
.as_any_mut()
258258
.downcast_mut::<ListBuilder<TimestampMicrosecondBuilder>>()
259-
.unwrap(),
259+
.expect("ListBuilder<TimestampMicrosecondBuilder>"),
260260
list,
261261
idx,
262262
),
263263
DataType::Binary => append_binary_element(
264264
list_builder
265265
.as_any_mut()
266266
.downcast_mut::<ListBuilder<BinaryBuilder>>()
267-
.unwrap(),
267+
.expect("ListBuilder<BinaryBuilder>"),
268268
list,
269269
idx,
270270
),
271271
DataType::Utf8 => append_string_element(
272272
list_builder
273273
.as_any_mut()
274274
.downcast_mut::<ListBuilder<StringBuilder>>()
275-
.unwrap(),
275+
.expect("ListBuilder<StringBuilder>"),
276276
list,
277277
idx,
278278
),
@@ -281,7 +281,7 @@ pub fn append_list_element<T: ArrayBuilder>(
281281
.values()
282282
.as_any_mut()
283283
.downcast_mut::<Decimal128Builder>()
284-
.unwrap();
284+
.expect("ListBuilder<Decimal128Builder>");
285285
let is_null = list.is_null_at(idx);
286286

287287
if is_null {
@@ -319,7 +319,7 @@ pub fn append_list_element<T: ArrayBuilder>(
319319
.values()
320320
.as_any_mut()
321321
.downcast_mut::<StructBuilder>()
322-
.unwrap();
322+
.expect("StructBuilder");
323323
let is_null = list.is_null_at(idx);
324324

325325
let nested_row = if is_null {

‎core/src/execution/shuffle/row.rs‎

Lines changed: 23 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ use arrow_array::{
3939
types::Int32Type,
4040
Array, ArrayRef, RecordBatch, RecordBatchOptions,
4141
};
42-
use arrow_schema::{DataType, Field, Schema, TimeUnit};
42+
use arrow_schema::{ArrowError, DataType, Field, Schema, TimeUnit};
4343
use jni::sys::{jint, jlong};
4444
use std::{
4545
fs::OpenOptions,
@@ -275,7 +275,10 @@ impl SparkUnsafeRow {
275275

276276
macro_rules! downcast_builder {
277277
($builder_type:ty, $builder:expr) => {
278-
$builder.into_box_any().downcast::<$builder_type>().unwrap()
278+
$builder
279+
.into_box_any()
280+
.downcast::<$builder_type>()
281+
.expect(stringify!($builder_type))
279282
};
280283
}
281284

@@ -284,7 +287,7 @@ macro_rules! downcast_builder_ref {
284287
$builder
285288
.as_any_mut()
286289
.downcast_mut::<$builder_type>()
287-
.unwrap()
290+
.expect(stringify!($builder_type))
288291
};
289292
}
290293

@@ -348,8 +351,7 @@ pub(crate) fn append_field(
348351
$field,
349352
field_builder,
350353
&row.get_map(idx),
351-
)
352-
.unwrap();
354+
)?;
353355
}
354356
}
355357
}};
@@ -378,8 +380,7 @@ pub(crate) fn append_field(
378380
$element_dt,
379381
field_builder,
380382
&row.get_array(idx),
381-
)
382-
.unwrap()
383+
)?
383384
}
384385
}
385386
}};
@@ -1057,7 +1058,7 @@ pub(crate) fn append_columns(
10571058
let element_builder = builder
10581059
.as_any_mut()
10591060
.downcast_mut::<$builder_type>()
1060-
.unwrap();
1061+
.expect(stringify!($builder_type));
10611062
let mut row = SparkUnsafeRow::new(schema);
10621063

10631064
for i in row_start..row_end {
@@ -1084,7 +1085,7 @@ pub(crate) fn append_columns(
10841085
let list_builder = builder
10851086
.as_any_mut()
10861087
.downcast_mut::<ListBuilder<$builder_type>>()
1087-
.unwrap();
1088+
.expect(stringify!($builder_type));
10881089
let mut row = SparkUnsafeRow::new(schema);
10891090

10901091
for i in row_start..row_end {
@@ -1103,8 +1104,7 @@ pub(crate) fn append_columns(
11031104
$element_dt,
11041105
list_builder,
11051106
&row.get_array(column_idx),
1106-
)
1107-
.unwrap()
1107+
)?
11081108
}
11091109
}
11101110
}};
@@ -1116,7 +1116,11 @@ pub(crate) fn append_columns(
11161116
let map_builder = builder
11171117
.as_any_mut()
11181118
.downcast_mut::<MapBuilder<$key_builder_type, $value_builder_type>>()
1119-
.unwrap();
1119+
.expect(&format!(
1120+
"MapBuilder<{},{}>",
1121+
stringify!($key_builder_type),
1122+
stringify!($value_builder_type)
1123+
));
11201124
let mut row = SparkUnsafeRow::new(schema);
11211125

11221126
for i in row_start..row_end {
@@ -1135,8 +1139,7 @@ pub(crate) fn append_columns(
11351139
$field,
11361140
map_builder,
11371141
&row.get_map(column_idx),
1138-
)
1139-
.unwrap()
1142+
)?
11401143
}
11411144
}
11421145
}};
@@ -1148,7 +1151,7 @@ pub(crate) fn append_columns(
11481151
let struct_builder = builder
11491152
.as_any_mut()
11501153
.downcast_mut::<StructBuilder>()
1151-
.unwrap();
1154+
.expect("StructBuilder");
11521155
let mut row = SparkUnsafeRow::new(schema);
11531156

11541157
for i in row_start..row_end {
@@ -3347,7 +3350,7 @@ pub fn process_sorted_row_partition(
33473350
.zip(schema.iter())
33483351
.map(|(builder, datatype)| builder_to_array(builder, datatype, prefer_dictionary_ratio))
33493352
.collect();
3350-
let batch = make_batch(array_refs?, n);
3353+
let batch = make_batch(array_refs?, n)?;
33513354

33523355
let mut frozen: Vec<u8> = vec![];
33533356
let mut cursor = Cursor::new(&mut frozen);
@@ -3382,7 +3385,7 @@ fn builder_to_array(
33823385
let builder = builder
33833386
.as_any_mut()
33843387
.downcast_mut::<StringDictionaryBuilder<Int32Type>>()
3385-
.unwrap();
3388+
.expect("StringDictionaryBuilder<Int32Type>");
33863389

33873390
let dict_array = builder.finish();
33883391
let num_keys = dict_array.keys().len();
@@ -3401,7 +3404,7 @@ fn builder_to_array(
34013404
let builder = builder
34023405
.as_any_mut()
34033406
.downcast_mut::<BinaryDictionaryBuilder<Int32Type>>()
3404-
.unwrap();
3407+
.expect("BinaryDictionaryBuilder<Int32Type>");
34053408

34063409
let dict_array = builder.finish();
34073410
let num_keys = dict_array.keys().len();
@@ -3420,7 +3423,7 @@ fn builder_to_array(
34203423
}
34213424
}
34223425

3423-
fn make_batch(arrays: Vec<ArrayRef>, row_count: usize) -> RecordBatch {
3426+
fn make_batch(arrays: Vec<ArrayRef>, row_count: usize) -> Result<RecordBatch, ArrowError> {
34243427
let mut dict_id = 0;
34253428
let fields = arrays
34263429
.iter()
@@ -3442,5 +3445,5 @@ fn make_batch(arrays: Vec<ArrayRef>, row_count: usize) -> RecordBatch {
34423445
.collect::<Vec<_>>();
34433446
let schema = Arc::new(Schema::new(fields));
34443447
let options = RecordBatchOptions::new().with_row_count(Option::from(row_count));
3445-
RecordBatch::try_new_with_options(schema, arrays, &options).unwrap()
3448+
RecordBatch::try_new_with_options(schema, arrays, &options)
34463449
}

0 commit comments

Comments
 (0)