Skip to content

Commit

Permalink
fix: avoid send empty record batch to client (apache#796)
Browse files Browse the repository at this point in the history
  • Loading branch information
ShiKaiWi authored Mar 31, 2023
1 parent e01e891 commit 3d97c8e
Show file tree
Hide file tree
Showing 3 changed files with 17 additions and 2 deletions.
12 changes: 11 additions & 1 deletion common_types/src/record_batch.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// Copyright 2022 CeresDB Project Authors. Licensed under Apache-2.0.
// Copyright 2022-2023 CeresDB Project Authors. Licensed under Apache-2.0.

//! Record batch

Expand Down Expand Up @@ -234,23 +234,33 @@ impl RecordBatch {
&self.schema
}

#[inline]
pub fn is_empty(&self) -> bool {
self.num_rows() == 0
}

// REQUIRE: index is valid
#[inline]
pub fn column(&self, index: usize) -> &ColumnBlock {
&self.data.column_blocks[index]
}

#[inline]
pub fn num_columns(&self) -> usize {
self.schema.num_columns()
}

#[inline]
pub fn num_rows(&self) -> usize {
self.data.num_rows()
}

#[inline]
pub fn as_arrow_record_batch(&self) -> &ArrowRecordBatch {
&self.data.arrow_record_batch
}

#[inline]
pub fn into_arrow_record_batch(self) -> ArrowRecordBatch {
self.data.arrow_record_batch
}
Expand Down
3 changes: 2 additions & 1 deletion components/arrow_ext/src/ipc.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// Copyright 2022 CeresDB Project Authors. Licensed under Apache-2.0.
// Copyright 2022-2023 CeresDB Project Authors. Licensed under Apache-2.0.

//! Utilities for `RecordBatch` serialization using Arrow IPC

Expand Down Expand Up @@ -109,6 +109,7 @@ impl RecordBatchesEncoder {
let stream_writer = if let Some(v) = &mut self.stream_writer {
v
} else {
// TODO: pre-allocate the buffer.
let buffer: Vec<u8> = Vec::new();
let stream_writer =
StreamWriter::try_new(buffer, &batch.schema()).context(ArrowError)?;
Expand Down
4 changes: 4 additions & 0 deletions server/src/proxy/grpc/sql_query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -453,6 +453,10 @@ impl QueryResponseWriter {
}

pub fn write(&mut self, batch: &RecordBatch) -> Result<()> {
if batch.is_empty() {
return Ok(());
}

self.encoder
.write(batch.as_arrow_record_batch())
.box_err()
Expand Down

0 comments on commit 3d97c8e

Please sign in to comment.