You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by tu...@apache.org on 2022/05/18 18:34:28 UTC
[arrow-datafusion] branch master updated: Fix avro tests (#2570) (#2571)
This is an automated email from the ASF dual-hosted git repository.
tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/master by this push:
new a08d26eef Fix avro tests (#2570) (#2571)
a08d26eef is described below
commit a08d26eef39bcb2adac527e5c260d31f473fca79
Author: Raphael Taylor-Davies <17...@users.noreply.github.com>
AuthorDate: Wed May 18 19:34:24 2022 +0100
Fix avro tests (#2570) (#2571)
---
.github/workflows/rust.yml | 2 +-
.../core/src/avro_to_arrow/arrow_array_reader.rs | 1 +
datafusion/core/src/datasource/file_format/avro.rs | 5 ++-
.../core/src/physical_plan/file_format/avro.rs | 39 ++++++++++++++--------
datafusion/core/tests/sql/avro.rs | 2 +-
5 files changed, 31 insertions(+), 18 deletions(-)
diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index 4a4e32f90..ce1bc32f8 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -115,7 +115,7 @@ jobs:
run: |
export ARROW_TEST_DATA=$(pwd)/testing/data
export PARQUET_TEST_DATA=$(pwd)/parquet-testing/data
- cargo test
+ cargo test --features avro
# test datafusion examples
cd datafusion-examples
cargo run --example csv_sql
diff --git a/datafusion/core/src/avro_to_arrow/arrow_array_reader.rs b/datafusion/core/src/avro_to_arrow/arrow_array_reader.rs
index 9d5552954..dfaa98555 100644
--- a/datafusion/core/src/avro_to_arrow/arrow_array_reader.rs
+++ b/datafusion/core/src/avro_to_arrow/arrow_array_reader.rs
@@ -861,6 +861,7 @@ fn flatten_string_values(values: &[&Value]) -> Vec<Option<String>> {
values
.iter()
.flat_map(|row| {
+ let row = maybe_resolve_union(row);
if let Value::Array(values) = row {
values
.iter()
diff --git a/datafusion/core/src/datasource/file_format/avro.rs b/datafusion/core/src/datasource/file_format/avro.rs
index 954c28a77..882066206 100644
--- a/datafusion/core/src/datasource/file_format/avro.rs
+++ b/datafusion/core/src/datasource/file_format/avro.rs
@@ -87,7 +87,6 @@ mod tests {
use super::*;
use crate::datasource::listing::local_unpartitioned_file;
- use crate::execution::runtime_env::{RuntimeConfig, RuntimeEnv};
use crate::prelude::{SessionConfig, SessionContext};
use arrow::array::{
BinaryArray, BooleanArray, Float32Array, Float64Array, Int32Array,
@@ -99,10 +98,10 @@ mod tests {
async fn read_small_batches() -> Result<()> {
let config = SessionConfig::new().with_batch_size(2);
let ctx = SessionContext::with_config(config);
- let task_ctx = session_ctx.task_ctx();
+ let task_ctx = ctx.task_ctx();
let projection = None;
let exec = get_exec("alltypes_plain.avro", &projection, None).await?;
- let stream = exec.execute(0, task_ctx).await?;
+ let stream = exec.execute(0, task_ctx)?;
let tt_batches = stream
.map(|batch| {
diff --git a/datafusion/core/src/physical_plan/file_format/avro.rs b/datafusion/core/src/physical_plan/file_format/avro.rs
index eed0161a2..a25ad60e1 100644
--- a/datafusion/core/src/physical_plan/file_format/avro.rs
+++ b/datafusion/core/src/physical_plan/file_format/avro.rs
@@ -18,7 +18,7 @@
//! Execution plan for reading line-delimited Avro files
#[cfg(feature = "avro")]
use crate::avro_to_arrow;
-use crate::error::{DataFusionError, Result};
+use crate::error::Result;
use crate::physical_plan::expressions::PhysicalSortExpr;
use crate::physical_plan::{
DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
@@ -98,7 +98,7 @@ impl ExecutionPlan for AvroExec {
_partition: usize,
_context: Arc<TaskContext>,
) -> Result<SendableRecordBatchStream> {
- Err(DataFusionError::NotImplemented(
+ Err(crate::error::DataFusionError::NotImplemented(
"Cannot execute avro plan without avro feature enabled".to_string(),
))
}
@@ -165,17 +165,17 @@ impl ExecutionPlan for AvroExec {
#[cfg(test)]
#[cfg(feature = "avro")]
mod tests {
- use crate::datasource::object_store::local::{
- local_object_reader_stream, LocalFileSystem,
- };
use crate::datasource::{
file_format::{avro::AvroFormat, FileFormat},
listing::local_unpartitioned_file,
};
+ use crate::prelude::SessionContext;
use crate::scalar::ScalarValue;
use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_data_access::object_store::local::{
+ local_object_reader_stream, LocalFileSystem,
+ };
use futures::StreamExt;
- use sqlparser::ast::ObjectType::Schema;
use super::*;
@@ -196,7 +196,11 @@ mod tests {
});
assert_eq!(avro_exec.output_partitioning().partition_count(), 1);
- let mut results = avro_exec.execute(0).await.expect("plan execution failed");
+ let ctx = SessionContext::new();
+ let mut results = avro_exec
+ .execute(0, ctx.task_ctx())
+ .expect("plan execution failed");
+
let batch = results
.next()
.await
@@ -237,27 +241,32 @@ mod tests {
let testdata = crate::test_util::arrow_test_data();
let filename = format!("{}/avro/alltypes_plain.avro", testdata);
let actual_schema = AvroFormat {}
- .infer_schema(local_object_reader_stream(vec![filename]))
+ .infer_schema(local_object_reader_stream(vec![filename.clone()]))
.await?;
let mut fields = actual_schema.fields().clone();
fields.push(Field::new("missing_col", DataType::Int32, true));
let file_schema = Arc::new(Schema::new(fields));
+ // Include the missing column in the projection
+ let projection = Some(vec![0, 1, 2, actual_schema.fields().len()]);
let avro_exec = AvroExec::new(FileScanConfig {
object_store: Arc::new(LocalFileSystem {}),
file_groups: vec![vec![local_unpartitioned_file(filename.clone())]],
file_schema,
statistics: Statistics::default(),
- // Include the missing column in the projection
- projection: Some(vec![0, 1, 2, file_schema.fields().len()]),
+ projection,
limit: None,
table_partition_cols: vec![],
});
assert_eq!(avro_exec.output_partitioning().partition_count(), 1);
- let mut results = avro_exec.execute(0).await.expect("plan execution failed");
+ let ctx = SessionContext::new();
+ let mut results = avro_exec
+ .execute(0, ctx.task_ctx())
+ .expect("plan execution failed");
+
let batch = results
.next()
.await
@@ -310,14 +319,18 @@ mod tests {
projection: Some(vec![0, 1, file_schema.fields().len(), 2]),
object_store: Arc::new(LocalFileSystem {}),
file_groups: vec![vec![partitioned_file]],
- file_schema: file_schema,
+ file_schema,
statistics: Statistics::default(),
limit: None,
table_partition_cols: vec!["date".to_owned()],
});
assert_eq!(avro_exec.output_partitioning().partition_count(), 1);
- let mut results = avro_exec.execute(0).await.expect("plan execution failed");
+ let ctx = SessionContext::new();
+ let mut results = avro_exec
+ .execute(0, ctx.task_ctx())
+ .expect("plan execution failed");
+
let batch = results
.next()
.await
diff --git a/datafusion/core/tests/sql/avro.rs b/datafusion/core/tests/sql/avro.rs
index f5c25dbad..0102cf3b7 100644
--- a/datafusion/core/tests/sql/avro.rs
+++ b/datafusion/core/tests/sql/avro.rs
@@ -124,7 +124,7 @@ async fn avro_single_nan_schema() {
let plan = ctx.create_logical_plan(sql).unwrap();
let plan = ctx.optimize(&plan).unwrap();
let plan = ctx.create_physical_plan(&plan).await.unwrap();
- let runtime = ctx.state.lock().runtime_env.clone();
+ let runtime = ctx.task_ctx();
let results = collect(plan, runtime).await.unwrap();
for batch in results {
assert_eq!(1, batch.num_rows());