Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add from_slice trait to ease arrow2 migration #1588

Merged
merged 2 commits into from
Jan 17, 2022
Merged

add from_slice trait to ease arrow2 migration #1588

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cbfa8ec
add from_slice trait to ease arrow2 migration
jimexist Jan 16, 2022
32eeb94
update more
jimexist Jan 16, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 7 additions & 6 deletions datafusion-cli/src/print_format.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,7 @@ mod tests {
use super::*;
use arrow::array::Int32Array;
use arrow::datatypes::{DataType, Field, Schema};
use datafusion::from_slice::FromSlice;
use std::sync::Arc;

#[test]
Expand Down Expand Up @@ -168,9 +169,9 @@ mod tests {
let batch = RecordBatch::try_new(
schema,
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(Int32Array::from(vec![4, 5, 6])),
Arc::new(Int32Array::from(vec![7, 8, 9])),
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
],
)
.unwrap();
Expand Down Expand Up @@ -198,9 +199,9 @@ mod tests {
let batch = RecordBatch::try_new(
schema,
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(Int32Array::from(vec![4, 5, 6])),
Arc::new(Int32Array::from(vec![7, 8, 9])),
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
],
)
.unwrap();
Expand Down
4 changes: 2 additions & 2 deletions datafusion-examples/examples/dataframe_in_memory.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,9 @@ use std::sync::Arc;
use datafusion::arrow::array::{Int32Array, StringArray};
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::arrow::record_batch::RecordBatch;

use datafusion::datasource::MemTable;
use datafusion::error::Result;
use datafusion::from_slice::FromSlice;
use datafusion::prelude::*;

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

Expand Down
5 changes: 3 additions & 2 deletions datafusion-examples/examples/simple_udaf.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ use datafusion::arrow::{
record_batch::RecordBatch,
};

use datafusion::from_slice::FromSlice;
use datafusion::physical_plan::functions::Volatility;
use datafusion::{error::Result, logical_plan::create_udaf, physical_plan::Accumulator};
use datafusion::{prelude::*, scalar::ScalarValue};
Expand All @@ -37,11 +38,11 @@ fn create_context() -> Result<ExecutionContext> {
// define data in two partitions
let batch1 = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Float32Array::from(vec![2.0, 4.0, 8.0]))],
vec![Arc::new(Float32Array::from_slice(&[2.0, 4.0, 8.0]))],
)?;
let batch2 = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Float32Array::from(vec![64.0]))],
vec![Arc::new(Float32Array::from_slice(&[64.0]))],
)?;

// declare a new context. In spark API, this corresponds to a new spark SQLsession
Expand Down
5 changes: 3 additions & 2 deletions datafusion-examples/examples/simple_udf.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ use datafusion::{
physical_plan::functions::Volatility,
};

use datafusion::from_slice::FromSlice;
use datafusion::prelude::*;
use datafusion::{error::Result, physical_plan::functions::make_scalar_function};
use std::sync::Arc;
Expand All @@ -42,8 +43,8 @@ fn create_context() -> Result<ExecutionContext> {
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Float32Array::from(vec![2.1, 3.1, 4.1, 5.1])),
Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0])),
Arc::new(Float32Array::from_slice(&[2.1, 3.1, 4.1, 5.1])),
Arc::new(Float64Array::from_slice(&[1.0, 2.0, 3.0, 4.0])),
],
)?;

Expand Down
3 changes: 2 additions & 1 deletion datafusion/benches/data_utils/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ use arrow::{
};
use datafusion::datasource::MemTable;
use datafusion::error::Result;
use datafusion::from_slice::FromSlice;
use rand::rngs::StdRng;
use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng};
Expand Down Expand Up @@ -128,7 +129,7 @@ fn create_record_batch(
schema,
vec![
Arc::new(StringArray::from(keys)),
Arc::new(Float32Array::from(vec![i as f32; batch_size])),
Arc::new(Float32Array::from_slice(&vec![i as f32; batch_size])),
Arc::new(Float64Array::from(values)),
Arc::new(UInt64Array::from(integer_values_wide)),
Arc::new(UInt64Array::from(integer_values_narrow)),
Expand Down
5 changes: 3 additions & 2 deletions datafusion/benches/filter_query_sql.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ use arrow::{
record_batch::RecordBatch,
};
use criterion::{criterion_group, criterion_main, Criterion};
use datafusion::from_slice::FromSlice;
use datafusion::prelude::ExecutionContext;
use datafusion::{datasource::MemTable, error::Result};
use futures::executor::block_on;
Expand Down Expand Up @@ -48,8 +49,8 @@ fn create_context(array_len: usize, batch_size: usize) -> Result<ExecutionContex
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Float32Array::from(vec![i as f32; batch_size])),
Arc::new(Float64Array::from(vec![i as f64; batch_size])),
Arc::new(Float32Array::from_slice(&vec![i as f32; batch_size])),
Arc::new(Float64Array::from_slice(&vec![i as f64; batch_size])),
],
)
.unwrap()
Expand Down
8 changes: 4 additions & 4 deletions datafusion/benches/math_query_sql.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,10 @@ use arrow::{
datatypes::{DataType, Field, Schema},
record_batch::RecordBatch,
};
use datafusion::error::Result;

use datafusion::datasource::MemTable;
use datafusion::error::Result;
use datafusion::execution::context::ExecutionContext;
use datafusion::from_slice::FromSlice;

fn query(ctx: Arc<Mutex<ExecutionContext>>, sql: &str) {
let rt = Runtime::new().unwrap();
Expand All @@ -60,8 +60,8 @@ fn create_context(
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Float32Array::from(vec![i as f32; batch_size])),
Arc::new(Float64Array::from(vec![i as f64; batch_size])),
Arc::new(Float32Array::from_slice(&vec![i as f32; batch_size])),
Arc::new(Float64Array::from_slice(&vec![i as f64; batch_size])),
],
)
.unwrap()
Expand Down
41 changes: 21 additions & 20 deletions datafusion/src/datasource/memory.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,7 @@ impl TableProvider for MemTable {
#[cfg(test)]
mod tests {
use super::*;
use crate::from_slice::FromSlice;
use arrow::array::Int32Array;
use arrow::datatypes::{DataType, Field, Schema};
use futures::StreamExt;
Expand All @@ -164,9 +165,9 @@ mod tests {
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(Int32Array::from(vec![4, 5, 6])),
Arc::new(Int32Array::from(vec![7, 8, 9])),
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
Arc::new(Int32Array::from(vec![None, None, Some(9)])),
],
)?;
Expand Down Expand Up @@ -197,9 +198,9 @@ mod tests {
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(Int32Array::from(vec![4, 5, 6])),
Arc::new(Int32Array::from(vec![7, 8, 9])),
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
],
)?;

Expand All @@ -225,9 +226,9 @@ mod tests {
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(Int32Array::from(vec![4, 5, 6])),
Arc::new(Int32Array::from(vec![7, 8, 9])),
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
],
)?;

Expand Down Expand Up @@ -262,9 +263,9 @@ mod tests {
let batch = RecordBatch::try_new(
schema1,
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(Int32Array::from(vec![4, 5, 6])),
Arc::new(Int32Array::from(vec![7, 8, 9])),
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
],
)?;

Expand Down Expand Up @@ -295,8 +296,8 @@ mod tests {
let batch = RecordBatch::try_new(
schema1,
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(Int32Array::from(vec![7, 5, 9])),
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
Arc::new(Int32Array::from_slice(&[7, 5, 9])),
],
)?;

Expand Down Expand Up @@ -339,18 +340,18 @@ mod tests {
let batch1 = RecordBatch::try_new(
Arc::new(schema1),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(Int32Array::from(vec![4, 5, 6])),
Arc::new(Int32Array::from(vec![7, 8, 9])),
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
],
)?;

let batch2 = RecordBatch::try_new(
Arc::new(schema2),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(Int32Array::from(vec![4, 5, 6])),
Arc::new(Int32Array::from(vec![7, 8, 9])),
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
],
)?;

Expand Down
45 changes: 23 additions & 22 deletions datafusion/src/execution/context.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1214,6 +1214,7 @@ impl FunctionRegistry for ExecutionContextState {
#[cfg(test)]
mod tests {
use super::*;
use crate::from_slice::FromSlice;
use crate::logical_plan::plan::Projection;
use crate::logical_plan::TableScan;
use crate::logical_plan::{binary_expr, lit, Operator};
Expand Down Expand Up @@ -1514,9 +1515,9 @@ mod tests {
let partitions = vec![vec![RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 10, 10, 100])),
Arc::new(Int32Array::from(vec![2, 12, 12, 120])),
Arc::new(Int32Array::from(vec![3, 12, 12, 120])),
Arc::new(Int32Array::from_slice(&[1, 10, 10, 100])),
Arc::new(Int32Array::from_slice(&[2, 12, 12, 120])),
Arc::new(Int32Array::from_slice(&[3, 12, 12, 120])),
],
)?]];

Expand Down Expand Up @@ -2928,43 +2929,43 @@ mod tests {
let type_values = vec![
(
DataType::Int8,
Arc::new(Int8Array::from(vec![1])) as ArrayRef,
Arc::new(Int8Array::from_slice(&[1])) as ArrayRef,
),
(
DataType::Int16,
Arc::new(Int16Array::from(vec![1])) as ArrayRef,
Arc::new(Int16Array::from_slice(&[1])) as ArrayRef,
),
(
DataType::Int32,
Arc::new(Int32Array::from(vec![1])) as ArrayRef,
Arc::new(Int32Array::from_slice(&[1])) as ArrayRef,
),
(
DataType::Int64,
Arc::new(Int64Array::from(vec![1])) as ArrayRef,
Arc::new(Int64Array::from_slice(&[1])) as ArrayRef,
),
(
DataType::UInt8,
Arc::new(UInt8Array::from(vec![1])) as ArrayRef,
Arc::new(UInt8Array::from_slice(&[1])) as ArrayRef,
),
(
DataType::UInt16,
Arc::new(UInt16Array::from(vec![1])) as ArrayRef,
Arc::new(UInt16Array::from_slice(&[1])) as ArrayRef,
),
(
DataType::UInt32,
Arc::new(UInt32Array::from(vec![1])) as ArrayRef,
Arc::new(UInt32Array::from_slice(&[1])) as ArrayRef,
),
(
DataType::UInt64,
Arc::new(UInt64Array::from(vec![1])) as ArrayRef,
Arc::new(UInt64Array::from_slice(&[1])) as ArrayRef,
),
(
DataType::Float32,
Arc::new(Float32Array::from(vec![1.0_f32])) as ArrayRef,
Arc::new(Float32Array::from_slice(&[1.0_f32])) as ArrayRef,
),
(
DataType::Float64,
Arc::new(Float64Array::from(vec![1.0_f64])) as ArrayRef,
Arc::new(Float64Array::from_slice(&[1.0_f64])) as ArrayRef,
),
];

Expand Down Expand Up @@ -3278,8 +3279,8 @@ mod tests {
let batch = RecordBatch::try_new(
Arc::new(schema.clone()),
vec![
Arc::new(Int32Array::from(vec![1, 10, 10, 100])),
Arc::new(Int32Array::from(vec![2, 12, 12, 120])),
Arc::new(Int32Array::from_slice(&[1, 10, 10, 100])),
Arc::new(Int32Array::from_slice(&[2, 12, 12, 120])),
],
)?;

Expand Down Expand Up @@ -3379,11 +3380,11 @@ mod tests {

let batch1 = RecordBatch::try_new(
Arc::new(schema.clone()),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
vec![Arc::new(Int32Array::from_slice(&[1, 2, 3]))],
)?;
let batch2 = RecordBatch::try_new(
Arc::new(schema.clone()),
vec![Arc::new(Int32Array::from(vec![4, 5]))],
vec![Arc::new(Int32Array::from_slice(&[4, 5]))],
)?;

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

let batch1 = RecordBatch::try_new(
Arc::new(schema.clone()),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
vec![Arc::new(Int32Array::from_slice(&[1, 2, 3]))],
)?;
let batch2 = RecordBatch::try_new(
Arc::new(schema.clone()),
vec![Arc::new(Int32Array::from(vec![4, 5]))],
vec![Arc::new(Int32Array::from_slice(&[4, 5]))],
)?;

let mut ctx = ExecutionContext::new();
Expand Down Expand Up @@ -3880,8 +3881,8 @@ mod tests {
let batch = RecordBatch::try_new(
Arc::new(schema.clone()),
vec![
Arc::new(Int32Array::from(vec![1])),
Arc::new(Float64Array::from(vec![1.0])),
Arc::new(Int32Array::from_slice(&[1])),
Arc::new(Float64Array::from_slice(&[1.0])),
Arc::new(StringArray::from(vec![Some("foo")])),
Arc::new(LargeStringArray::from(vec![Some("bar")])),
Arc::new(BinaryArray::from(vec![b"foo" as &[u8]])),
Expand Down Expand Up @@ -4148,7 +4149,7 @@ mod tests {
.unwrap();

// create mock record batch
let ids = Arc::new(Int32Array::from(vec![i as i32]));
let ids = Arc::new(Int32Array::from_slice(&[i as i32]));
let names = Arc::new(StringArray::from(vec!["test"]));
let rec_batch =
RecordBatch::try_new(schema.clone(), vec![ids, names]).unwrap();
Expand Down
Loading