Skip to content

Commit

Permalink
Showing 3 changed files with 17 additions and 8 deletions.
10 changes: 7 additions & 3 deletions datafusion/common/src/utils.rs
Original file line number Diff line number Diff line change
@@ -25,7 +25,7 @@ use arrow::compute;
use arrow::compute::{partition, SortColumn, SortOptions};
use arrow::datatypes::{Field, SchemaRef, UInt32Type};
use arrow::record_batch::RecordBatch;
use arrow_array::{Array, LargeListArray, ListArray};
use arrow_array::{Array, LargeListArray, ListArray, RecordBatchOptions};
use arrow_schema::DataType;
use sqlparser::ast::Ident;
use sqlparser::dialect::GenericDialect;
@@ -90,8 +90,12 @@ pub fn get_record_batch_at_indices(
indices: &PrimitiveArray<UInt32Type>,
) -> Result<RecordBatch> {
let new_columns = get_arrayref_at_indices(record_batch.columns(), indices)?;
RecordBatch::try_new(record_batch.schema(), new_columns)
.map_err(DataFusionError::ArrowError)
RecordBatch::try_new_with_options(
record_batch.schema(),
new_columns,
&RecordBatchOptions::new().with_row_count(Some(indices.len())),
)
.map_err(DataFusionError::ArrowError)
}

/// This function compares two tuples depending on the given sort options.
Original file line number Diff line number Diff line change
@@ -40,7 +40,7 @@ use crate::{
};

use arrow::{
array::{Array, ArrayRef, UInt32Builder},
array::{Array, ArrayRef, RecordBatchOptions, UInt32Builder},
compute::{concat, concat_batches, sort_to_indices},
datatypes::{Schema, SchemaBuilder, SchemaRef},
record_batch::RecordBatch,
@@ -1026,8 +1026,11 @@ impl BoundedWindowAggStream {
.iter()
.map(|elem| elem.slice(n_out, n_to_keep))
.collect::<Vec<_>>();
self.input_buffer =
RecordBatch::try_new(self.input_buffer.schema(), batch_to_keep)?;
self.input_buffer = RecordBatch::try_new_with_options(
self.input_buffer.schema(),
batch_to_keep,
&RecordBatchOptions::new().with_row_count(Some(n_to_keep)),
)?;
Ok(())
}

6 changes: 4 additions & 2 deletions datafusion/sqllogictest/test_files/window.slt
Original file line number Diff line number Diff line change
@@ -3771,10 +3771,12 @@ select a,
1 1
2 1

# TODO: this works in Postgres which returns [1, 1].
query error DataFusion error: Arrow error: Invalid argument error: must either specify a row count or at least one column
query I
select rank() over (RANGE between UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) rnk
from (select 1 a union select 2 a) q;
----
1
1

query II
select a,

0 comments on commit ecb7c7d

Please sign in to comment.