Skip to content

Commit f027e5f

Browse files
jimexistjimexist
authored
add from_slice trait to ease arrow2 migration (#1588)
* add from_slice trait to ease arrow2 migration * update more
1 parent 6f7b2d2 commit f027e5f

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

42 files changed

+375
-274
lines changed

‎datafusion-cli/src/print_format.rs

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,7 @@ mod tests {
120120
use super::*;
121121
use arrow::array::Int32Array;
122122
use arrow::datatypes::{DataType, Field, Schema};
123+
use datafusion::from_slice::FromSlice;
123124
use std::sync::Arc;
124125

125126
#[test]
@@ -168,9 +169,9 @@ mod tests {
168169
let batch = RecordBatch::try_new(
169170
schema,
170171
vec![
171-
Arc::new(Int32Array::from(vec![1, 2, 3])),
172-
Arc::new(Int32Array::from(vec![4, 5, 6])),
173-
Arc::new(Int32Array::from(vec![7, 8, 9])),
172+
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
173+
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
174+
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
174175
],
175176
)
176177
.unwrap();
@@ -198,9 +199,9 @@ mod tests {
198199
let batch = RecordBatch::try_new(
199200
schema,
200201
vec![
201-
Arc::new(Int32Array::from(vec![1, 2, 3])),
202-
Arc::new(Int32Array::from(vec![4, 5, 6])),
203-
Arc::new(Int32Array::from(vec![7, 8, 9])),
202+
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
203+
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
204+
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
204205
],
205206
)
206207
.unwrap();

‎datafusion-examples/examples/dataframe_in_memory.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,9 @@ use std::sync::Arc;
2020
use datafusion::arrow::array::{Int32Array, StringArray};
2121
use datafusion::arrow::datatypes::{DataType, Field, Schema};
2222
use datafusion::arrow::record_batch::RecordBatch;
23-
2423
use datafusion::datasource::MemTable;
2524
use datafusion::error::Result;
25+
use datafusion::from_slice::FromSlice;
2626
use datafusion::prelude::*;
2727

2828
/// This example demonstrates how to use the DataFrame API against in-memory data.
@@ -39,7 +39,7 @@ async fn main() -> Result<()> {
3939
schema.clone(),
4040
vec![
4141
Arc::new(StringArray::from(vec!["a", "b", "c", "d"])),
42-
Arc::new(Int32Array::from(vec![1, 10, 10, 100])),
42+
Arc::new(Int32Array::from_slice(&[1, 10, 10, 100])),
4343
],
4444
)?;
4545

‎datafusion-examples/examples/simple_udaf.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ use datafusion::arrow::{
2222
record_batch::RecordBatch,
2323
};
2424

25+
use datafusion::from_slice::FromSlice;
2526
use datafusion::physical_plan::functions::Volatility;
2627
use datafusion::{error::Result, logical_plan::create_udaf, physical_plan::Accumulator};
2728
use datafusion::{prelude::*, scalar::ScalarValue};
@@ -37,11 +38,11 @@ fn create_context() -> Result<ExecutionContext> {
3738
// define data in two partitions
3839
let batch1 = RecordBatch::try_new(
3940
schema.clone(),
40-
vec![Arc::new(Float32Array::from(vec![2.0, 4.0, 8.0]))],
41+
vec![Arc::new(Float32Array::from_slice(&[2.0, 4.0, 8.0]))],
4142
)?;
4243
let batch2 = RecordBatch::try_new(
4344
schema.clone(),
44-
vec![Arc::new(Float32Array::from(vec![64.0]))],
45+
vec![Arc::new(Float32Array::from_slice(&[64.0]))],
4546
)?;
4647

4748
// declare a new context. In spark API, this corresponds to a new spark SQLsession

‎datafusion-examples/examples/simple_udf.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ use datafusion::{
2424
physical_plan::functions::Volatility,
2525
};
2626

27+
use datafusion::from_slice::FromSlice;
2728
use datafusion::prelude::*;
2829
use datafusion::{error::Result, physical_plan::functions::make_scalar_function};
2930
use std::sync::Arc;
@@ -42,8 +43,8 @@ fn create_context() -> Result<ExecutionContext> {
4243
let batch = RecordBatch::try_new(
4344
schema.clone(),
4445
vec![
45-
Arc::new(Float32Array::from(vec![2.1, 3.1, 4.1, 5.1])),
46-
Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0])),
46+
Arc::new(Float32Array::from_slice(&[2.1, 3.1, 4.1, 5.1])),
47+
Arc::new(Float64Array::from_slice(&[1.0, 2.0, 3.0, 4.0])),
4748
],
4849
)?;
4950

‎datafusion/benches/data_utils/mod.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ use arrow::{
2727
};
2828
use datafusion::datasource::MemTable;
2929
use datafusion::error::Result;
30+
use datafusion::from_slice::FromSlice;
3031
use rand::rngs::StdRng;
3132
use rand::seq::SliceRandom;
3233
use rand::{Rng, SeedableRng};
@@ -128,7 +129,7 @@ fn create_record_batch(
128129
schema,
129130
vec![
130131
Arc::new(StringArray::from(keys)),
131-
Arc::new(Float32Array::from(vec![i as f32; batch_size])),
132+
Arc::new(Float32Array::from_slice(&vec![i as f32; batch_size])),
132133
Arc::new(Float64Array::from(values)),
133134
Arc::new(UInt64Array::from(integer_values_wide)),
134135
Arc::new(UInt64Array::from(integer_values_narrow)),

‎datafusion/benches/filter_query_sql.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ use arrow::{
2121
record_batch::RecordBatch,
2222
};
2323
use criterion::{criterion_group, criterion_main, Criterion};
24+
use datafusion::from_slice::FromSlice;
2425
use datafusion::prelude::ExecutionContext;
2526
use datafusion::{datasource::MemTable, error::Result};
2627
use futures::executor::block_on;
@@ -48,8 +49,8 @@ fn create_context(array_len: usize, batch_size: usize) -> Result<ExecutionContex
4849
RecordBatch::try_new(
4950
schema.clone(),
5051
vec![
51-
Arc::new(Float32Array::from(vec![i as f32; batch_size])),
52-
Arc::new(Float64Array::from(vec![i as f64; batch_size])),
52+
Arc::new(Float32Array::from_slice(&vec![i as f32; batch_size])),
53+
Arc::new(Float64Array::from_slice(&vec![i as f64; batch_size])),
5354
],
5455
)
5556
.unwrap()

‎datafusion/benches/math_query_sql.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,10 @@ use arrow::{
3131
datatypes::{DataType, Field, Schema},
3232
record_batch::RecordBatch,
3333
};
34-
use datafusion::error::Result;
35-
3634
use datafusion::datasource::MemTable;
35+
use datafusion::error::Result;
3736
use datafusion::execution::context::ExecutionContext;
37+
use datafusion::from_slice::FromSlice;
3838

3939
fn query(ctx: Arc<Mutex<ExecutionContext>>, sql: &str) {
4040
let rt = Runtime::new().unwrap();
@@ -60,8 +60,8 @@ fn create_context(
6060
RecordBatch::try_new(
6161
schema.clone(),
6262
vec![
63-
Arc::new(Float32Array::from(vec![i as f32; batch_size])),
64-
Arc::new(Float64Array::from(vec![i as f64; batch_size])),
63+
Arc::new(Float32Array::from_slice(&vec![i as f32; batch_size])),
64+
Arc::new(Float64Array::from_slice(&vec![i as f64; batch_size])),
6565
],
6666
)
6767
.unwrap()

‎datafusion/src/datasource/memory.rs

Lines changed: 21 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,7 @@ impl TableProvider for MemTable {
146146
#[cfg(test)]
147147
mod tests {
148148
use super::*;
149+
use crate::from_slice::FromSlice;
149150
use arrow::array::Int32Array;
150151
use arrow::datatypes::{DataType, Field, Schema};
151152
use futures::StreamExt;
@@ -164,9 +165,9 @@ mod tests {
164165
let batch = RecordBatch::try_new(
165166
schema.clone(),
166167
vec![
167-
Arc::new(Int32Array::from(vec![1, 2, 3])),
168-
Arc::new(Int32Array::from(vec![4, 5, 6])),
169-
Arc::new(Int32Array::from(vec![7, 8, 9])),
168+
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
169+
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
170+
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
170171
Arc::new(Int32Array::from(vec![None, None, Some(9)])),
171172
],
172173
)?;
@@ -197,9 +198,9 @@ mod tests {
197198
let batch = RecordBatch::try_new(
198199
schema.clone(),
199200
vec![
200-
Arc::new(Int32Array::from(vec![1, 2, 3])),
201-
Arc::new(Int32Array::from(vec![4, 5, 6])),
202-
Arc::new(Int32Array::from(vec![7, 8, 9])),
201+
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
202+
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
203+
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
203204
],
204205
)?;
205206

@@ -225,9 +226,9 @@ mod tests {
225226
let batch = RecordBatch::try_new(
226227
schema.clone(),
227228
vec![
228-
Arc::new(Int32Array::from(vec![1, 2, 3])),
229-
Arc::new(Int32Array::from(vec![4, 5, 6])),
230-
Arc::new(Int32Array::from(vec![7, 8, 9])),
229+
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
230+
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
231+
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
231232
],
232233
)?;
233234

@@ -262,9 +263,9 @@ mod tests {
262263
let batch = RecordBatch::try_new(
263264
schema1,
264265
vec![
265-
Arc::new(Int32Array::from(vec![1, 2, 3])),
266-
Arc::new(Int32Array::from(vec![4, 5, 6])),
267-
Arc::new(Int32Array::from(vec![7, 8, 9])),
266+
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
267+
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
268+
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
268269
],
269270
)?;
270271

@@ -295,8 +296,8 @@ mod tests {
295296
let batch = RecordBatch::try_new(
296297
schema1,
297298
vec![
298-
Arc::new(Int32Array::from(vec![1, 2, 3])),
299-
Arc::new(Int32Array::from(vec![7, 5, 9])),
299+
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
300+
Arc::new(Int32Array::from_slice(&[7, 5, 9])),
300301
],
301302
)?;
302303

@@ -339,18 +340,18 @@ mod tests {
339340
let batch1 = RecordBatch::try_new(
340341
Arc::new(schema1),
341342
vec![
342-
Arc::new(Int32Array::from(vec![1, 2, 3])),
343-
Arc::new(Int32Array::from(vec![4, 5, 6])),
344-
Arc::new(Int32Array::from(vec![7, 8, 9])),
343+
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
344+
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
345+
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
345346
],
346347
)?;
347348

348349
let batch2 = RecordBatch::try_new(
349350
Arc::new(schema2),
350351
vec![
351-
Arc::new(Int32Array::from(vec![1, 2, 3])),
352-
Arc::new(Int32Array::from(vec![4, 5, 6])),
353-
Arc::new(Int32Array::from(vec![7, 8, 9])),
352+
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
353+
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
354+
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
354355
],
355356
)?;
356357

‎datafusion/src/execution/context.rs

Lines changed: 23 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1214,6 +1214,7 @@ impl FunctionRegistry for ExecutionContextState {
12141214
#[cfg(test)]
12151215
mod tests {
12161216
use super::*;
1217+
use crate::from_slice::FromSlice;
12171218
use crate::logical_plan::plan::Projection;
12181219
use crate::logical_plan::TableScan;
12191220
use crate::logical_plan::{binary_expr, lit, Operator};
@@ -1514,9 +1515,9 @@ mod tests {
15141515
let partitions = vec![vec![RecordBatch::try_new(
15151516
schema.clone(),
15161517
vec![
1517-
Arc::new(Int32Array::from(vec![1, 10, 10, 100])),
1518-
Arc::new(Int32Array::from(vec![2, 12, 12, 120])),
1519-
Arc::new(Int32Array::from(vec![3, 12, 12, 120])),
1518+
Arc::new(Int32Array::from_slice(&[1, 10, 10, 100])),
1519+
Arc::new(Int32Array::from_slice(&[2, 12, 12, 120])),
1520+
Arc::new(Int32Array::from_slice(&[3, 12, 12, 120])),
15201521
],
15211522
)?]];
15221523

@@ -2928,43 +2929,43 @@ mod tests {
29282929
let type_values = vec![
29292930
(
29302931
DataType::Int8,
2931-
Arc::new(Int8Array::from(vec![1])) as ArrayRef,
2932+
Arc::new(Int8Array::from_slice(&[1])) as ArrayRef,
29322933
),
29332934
(
29342935
DataType::Int16,
2935-
Arc::new(Int16Array::from(vec![1])) as ArrayRef,
2936+
Arc::new(Int16Array::from_slice(&[1])) as ArrayRef,
29362937
),
29372938
(
29382939
DataType::Int32,
2939-
Arc::new(Int32Array::from(vec![1])) as ArrayRef,
2940+
Arc::new(Int32Array::from_slice(&[1])) as ArrayRef,
29402941
),
29412942
(
29422943
DataType::Int64,
2943-
Arc::new(Int64Array::from(vec![1])) as ArrayRef,
2944+
Arc::new(Int64Array::from_slice(&[1])) as ArrayRef,
29442945
),
29452946
(
29462947
DataType::UInt8,
2947-
Arc::new(UInt8Array::from(vec![1])) as ArrayRef,
2948+
Arc::new(UInt8Array::from_slice(&[1])) as ArrayRef,
29482949
),
29492950
(
29502951
DataType::UInt16,
2951-
Arc::new(UInt16Array::from(vec![1])) as ArrayRef,
2952+
Arc::new(UInt16Array::from_slice(&[1])) as ArrayRef,
29522953
),
29532954
(
29542955
DataType::UInt32,
2955-
Arc::new(UInt32Array::from(vec![1])) as ArrayRef,
2956+
Arc::new(UInt32Array::from_slice(&[1])) as ArrayRef,
29562957
),
29572958
(
29582959
DataType::UInt64,
2959-
Arc::new(UInt64Array::from(vec![1])) as ArrayRef,
2960+
Arc::new(UInt64Array::from_slice(&[1])) as ArrayRef,
29602961
),
29612962
(
29622963
DataType::Float32,
2963-
Arc::new(Float32Array::from(vec![1.0_f32])) as ArrayRef,
2964+
Arc::new(Float32Array::from_slice(&[1.0_f32])) as ArrayRef,
29642965
),
29652966
(
29662967
DataType::Float64,
2967-
Arc::new(Float64Array::from(vec![1.0_f64])) as ArrayRef,
2968+
Arc::new(Float64Array::from_slice(&[1.0_f64])) as ArrayRef,
29682969
),
29692970
];
29702971

@@ -3278,8 +3279,8 @@ mod tests {
32783279
let batch = RecordBatch::try_new(
32793280
Arc::new(schema.clone()),
32803281
vec![
3281-
Arc::new(Int32Array::from(vec![1, 10, 10, 100])),
3282-
Arc::new(Int32Array::from(vec![2, 12, 12, 120])),
3282+
Arc::new(Int32Array::from_slice(&[1, 10, 10, 100])),
3283+
Arc::new(Int32Array::from_slice(&[2, 12, 12, 120])),
32833284
],
32843285
)?;
32853286

@@ -3379,11 +3380,11 @@ mod tests {
33793380

33803381
let batch1 = RecordBatch::try_new(
33813382
Arc::new(schema.clone()),
3382-
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
3383+
vec![Arc::new(Int32Array::from_slice(&[1, 2, 3]))],
33833384
)?;
33843385
let batch2 = RecordBatch::try_new(
33853386
Arc::new(schema.clone()),
3386-
vec![Arc::new(Int32Array::from(vec![4, 5]))],
3387+
vec![Arc::new(Int32Array::from_slice(&[4, 5]))],
33873388
)?;
33883389

33893390
let mut ctx = ExecutionContext::new();
@@ -3416,11 +3417,11 @@ mod tests {
34163417

34173418
let batch1 = RecordBatch::try_new(
34183419
Arc::new(schema.clone()),
3419-
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
3420+
vec![Arc::new(Int32Array::from_slice(&[1, 2, 3]))],
34203421
)?;
34213422
let batch2 = RecordBatch::try_new(
34223423
Arc::new(schema.clone()),
3423-
vec![Arc::new(Int32Array::from(vec![4, 5]))],
3424+
vec![Arc::new(Int32Array::from_slice(&[4, 5]))],
34243425
)?;
34253426

34263427
let mut ctx = ExecutionContext::new();
@@ -3880,8 +3881,8 @@ mod tests {
38803881
let batch = RecordBatch::try_new(
38813882
Arc::new(schema.clone()),
38823883
vec![
3883-
Arc::new(Int32Array::from(vec![1])),
3884-
Arc::new(Float64Array::from(vec![1.0])),
3884+
Arc::new(Int32Array::from_slice(&[1])),
3885+
Arc::new(Float64Array::from_slice(&[1.0])),
38853886
Arc::new(StringArray::from(vec![Some("foo")])),
38863887
Arc::new(LargeStringArray::from(vec![Some("bar")])),
38873888
Arc::new(BinaryArray::from(vec![b"foo" as &[u8]])),
@@ -4148,7 +4149,7 @@ mod tests {
41484149
.unwrap();
41494150

41504151
// create mock record batch
4151-
let ids = Arc::new(Int32Array::from(vec![i as i32]));
4152+
let ids = Arc::new(Int32Array::from_slice(&[i as i32]));
41524153
let names = Arc::new(StringArray::from(vec!["test"]));
41534154
let rec_batch =
41544155
RecordBatch::try_new(schema.clone(), vec![ids, names]).unwrap();

0 commit comments

Comments
 (0)