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

Jump to bottom

Add SessionContext::read_batches #9197

Merged
merged 7 commits into from
Feb 14, 2024
Merged

Add SessionContext::read_batches #9197

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
d3e162b
feat: issue #9157 adding record_batches for Vec<BatchRecord>
Lordworms Feb 12, 2024
7f0f734
fix bugs
Lordworms Feb 12, 2024
769dbfd
optimize code and tests
Lordworms Feb 12, 2024
c18b8b9
optimize test
Lordworms Feb 12, 2024
3d34459
optimize tests
Lordworms Feb 12, 2024
e13670e
abandon useless schema
Lordworms Feb 12, 2024
455a5d9
collect into a single batches
Lordworms Feb 14, 2024
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
21 changes: 20 additions & 1 deletion datafusion/core/src/execution/context/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -934,7 +934,26 @@ impl SessionContext {
.build()?,
))
}

/// Create a [`DataFrame`] for reading a [`Vec[`RecordBatch`]`]
pub fn read_batches(
&self,
batches: impl IntoIterator<Item = RecordBatch>,
) -> Result<DataFrame> {
// check schema uniqueness
let mut batches = batches.into_iter().peekable();
let schema: SchemaRef = batches.peek().unwrap().schema().clone();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there any way we can avoid an unwrap here?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Perhaps something like

Suggested change
let schema: SchemaRef = batches.peek().unwrap().schema().clone();
let schema: SchemaRef = if let Some(schema) = batches.peek() {
schema.clone()
} else {
// use empty schema)
Arc::new(Schema::new())
};

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Got it!

let provider =
MemTable::try_new(schema, batches.map(|batch| vec![batch]).collect())?;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, is there any difference to mapping each batch into its own vec (which seemingly represents a different partition?) vs just collecting the batches into a single vec and passing in a single vec?

e.g. something like

Suggested change
let provider =
MemTable::try_new(schema, batches.map(|batch| vec![batch]).collect())?;
let provider =
MemTable::try_new(schema, vec![batches.collect()])?;

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is a difference (as you say the code in this PR makes its own partition). I think you are right that a single partition might be better (and DataFusion will repartition the plan into multiple partitions) if necessary

Is this something you can do @Lordworms ? Otherwise we can merge this PR as is and make it as a follow on change.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sure, so what we need to do is to flatten all the batches into a vec? I wonder in what situation the datafusion would do the repartition? Is there any doc for this?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TLDR is that it is added by https://docs.rs/datafusion/latest/datafusion/physical_optimizer/enforce_distribution/struct.EnforceDistribution.html -- the rules about when it happens are non trivial but basically are "when it helps"

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, I would read about it and commit the changes soon.

Ok(DataFrame::new(
self.state(),
LogicalPlanBuilder::scan(
UNNAMED_TABLE,
provider_as_source(Arc::new(provider)),
None,
)?
.build()?,
))
}
/// Registers a [`ListingTable`] that can assemble multiple files
/// from locations in an [`ObjectStore`] instance into a single
/// table.
Expand Down
84 changes: 84 additions & 0 deletions datafusion/core/tests/dataframe/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ use arrow::{
},
record_batch::RecordBatch,
};
use arrow_array::Float32Array;
use arrow_schema::ArrowError;
use std::sync::Arc;

Expand Down Expand Up @@ -1431,6 +1432,89 @@ async fn unnest_analyze_metrics() -> Result<()> {

Ok(())
}
#[tokio::test]
async fn test_read_batches() -> Result<()> {
let config = SessionConfig::new();
let runtime = Arc::new(RuntimeEnv::default());
let state = SessionState::new_with_config_rt(config, runtime);
let ctx = SessionContext::new_with_state(state);

let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("number", DataType::Float32, false),
]));

let batches = vec![
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])),
Arc::new(Float32Array::from(vec![1.12, 3.40, 2.33, 9.10, 6.66])),
],
)
.unwrap(),
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![3, 4, 5])),
Arc::new(Float32Array::from(vec![1.11, 2.22, 3.33])),
],
)
.unwrap(),
RecordBatch::try_new(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor: I suggest making the test shorter by feeding in 2 batches rather than 5

schema.clone(),
vec![
Arc::new(Int32Array::from(vec![0, 1, 2])),
Arc::new(Float32Array::from(vec![4.44, 5.02, 6.03])),
],
)
.unwrap(),
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![0, 1, 3])),
Arc::new(Float32Array::from(vec![6.01, 2.02, 3.03])),
],
)
.unwrap(),
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![3, 1, 2])),
Arc::new(Float32Array::from(vec![1000.01, 2.02, 3.03])),
],
)
.unwrap(),
];
let df = ctx.read_batches(batches).unwrap();
df.clone().show().await.unwrap();
let result = df.collect().await?;
let expected = vec![
"+----+---------+",
"| id | number |",
"+----+---------+",
"| 1 | 1.12 |",
"| 2 | 3.4 |",
"| 3 | 2.33 |",
"| 4 | 9.1 |",
"| 5 | 6.66 |",
"| 3 | 1.11 |",
"| 4 | 2.22 |",
"| 5 | 3.33 |",
"| 0 | 4.44 |",
"| 1 | 5.02 |",
"| 2 | 6.03 |",
"| 0 | 6.01 |",
"| 1 | 2.02 |",
"| 3 | 3.03 |",
"| 3 | 1000.01 |",
"| 1 | 2.02 |",
"| 2 | 3.03 |",
"+----+---------+",
];
assert_batches_sorted_eq!(expected, &result);
Ok(())
}

#[tokio::test]
async fn consecutive_projection_same_schema() -> Result<()> {
Expand Down
Loading