Skip to content

add more logs to flight decoder/encoder for debugging #66

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

Draft
fsdvh wants to merge 5 commits into
base: v53
Choose a base branch
from
Draft

add more logs to flight decoder/encoder for debugging #66

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
02b1164
add more logs to flight decoder/encoder
fsdvh May 8, 2025
60ef285
update
fsdvh May 8, 2025
be6eac6
poll count
fsdvh May 8, 2025
70925c9
missing poll count
fsdvh May 8, 2025
e979b1c
measure poll time
fsdvh May 8, 2025
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
10 changes: 5 additions & 5 deletions arrow-flight/examples/flight_sql_server.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -458,7 +458,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
let batch = builder.build();
let stream = FlightDataEncoderBuilder::new()
.with_schema(schema)
.build(futures::stream::once(async { batch }))
.build(futures::stream::once(async { batch }), "")
.map_err(Status::from);
Ok(Response::new(Box::pin(stream)))
}
Expand All @@ -485,7 +485,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
let batch = builder.build();
let stream = FlightDataEncoderBuilder::new()
.with_schema(schema)
.build(futures::stream::once(async { batch }))
.build(futures::stream::once(async { batch }), "")
.map_err(Status::from);
Ok(Response::new(Box::pin(stream)))
}
Expand Down Expand Up @@ -525,7 +525,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
let batch = builder.build();
let stream = FlightDataEncoderBuilder::new()
.with_schema(schema)
.build(futures::stream::once(async { batch }))
.build(futures::stream::once(async { batch }), "")
.map_err(Status::from);
Ok(Response::new(Box::pin(stream)))
}
Expand All @@ -548,7 +548,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
let batch = builder.build();
let stream = FlightDataEncoderBuilder::new()
.with_schema(schema)
.build(futures::stream::once(async { batch }))
.build(futures::stream::once(async { batch }), "")
.map_err(Status::from);
Ok(Response::new(Box::pin(stream)))
}
Expand Down Expand Up @@ -602,7 +602,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
let batch = builder.build();
let stream = FlightDataEncoderBuilder::new()
.with_schema(schema)
.build(futures::stream::once(async { batch }))
.build(futures::stream::once(async { batch }), "")
.map_err(Status::from);
Ok(Response::new(Box::pin(stream)))
}
Expand Down
6 changes: 5 additions & 1 deletion arrow-flight/src/client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,7 @@ impl FlightClient {

Ok(FlightRecordBatchStream::new_from_flight_data(
response_stream.map_err(FlightError::Tonic),
"",
)
.with_headers(md)
.with_trailers(trailers))
Expand Down Expand Up @@ -429,7 +430,10 @@ impl FlightClient {
let error_stream = FallibleTonicResponseStream::new(receiver, response_stream);

// combine the response from the server and any error from the client
Ok(FlightRecordBatchStream::new_from_flight_data(error_stream))
Ok(FlightRecordBatchStream::new_from_flight_data(
error_stream,
"",
))
}

/// Make a `ListFlights` call to the server with the provided
Expand Down
42 changes: 37 additions & 5 deletions arrow-flight/src/decode.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ use arrow_buffer::Buffer;
use arrow_schema::{Schema, SchemaRef};
use bytes::Bytes;
use futures::{ready, stream::BoxStream, Stream, StreamExt};
use std::{collections::HashMap, fmt::Debug, pin::Pin, sync::Arc, task::Poll};
use std::{collections::HashMap, fmt::Debug, pin::Pin, sync::Arc, task::Poll, time::Instant};
use tonic::metadata::MetadataMap;

use crate::error::{FlightError, Result};
Expand Down Expand Up @@ -101,12 +101,12 @@ impl FlightRecordBatchStream {
}

/// Create a new [`FlightRecordBatchStream`] from a stream of [`FlightData`]
pub fn new_from_flight_data<S>(inner: S) -> Self
pub fn new_from_flight_data<S>(inner: S, reader_id: &str) -> Self
where
S: Stream<Item = Result<FlightData>> + Send + 'static,
{
Self {
inner: FlightDataDecoder::new(inner),
inner: FlightDataDecoder::new(inner, reader_id),
headers: MetadataMap::default(),
trailers: None,
}
Expand Down Expand Up @@ -234,6 +234,10 @@ pub struct FlightDataDecoder {
state: Option<FlightStreamState>,
/// Seen the end of the inner stream?
done: bool,

reader_id: String,

poll_count: usize,
}

impl Debug for FlightDataDecoder {
Expand All @@ -248,14 +252,16 @@ impl Debug for FlightDataDecoder {

impl FlightDataDecoder {
/// Create a new wrapper around the stream of [`FlightData`]
pub fn new<S>(response: S) -> Self
pub fn new<S>(response: S, reader_id: &str) -> Self
where
S: Stream<Item = Result<FlightData>> + Send + 'static,
{
Self {
state: None,
response: response.boxed(),
done: false,
reader_id: reader_id.to_string(),
poll_count: 0,
}
}

Expand Down Expand Up @@ -354,20 +360,46 @@ impl futures::Stream for FlightDataDecoder {
cx: &mut std::task::Context<'_>,
) -> Poll<Option<Self::Item>> {
if self.done {
println!(
"flight data decoder - stream done, {}/{}",
self.reader_id, self.poll_count
);
return Poll::Ready(None);
}
self.poll_count += 1;
loop {
println!(
"flight data decoder polling next, {}/{}",
self.reader_id, self.poll_count
);
let now = Instant::now();
let res = ready!(self.response.poll_next_unpin(cx));
println!(
"flight data decoder polled next, {}/{} - took {:?}",
self.reader_id,
self.poll_count,
now.elapsed()
);

return Poll::Ready(match res {
None => {
self.done = true;
println!(
"flight data decoder inner is exhausted, {}/{}",
self.reader_id, self.poll_count
);
None // inner is exhausted
}
Some(data) => Some(match data {
Err(e) => Err(e),
Ok(data) => match self.extract_message(data) {
Ok(Some(extracted)) => Ok(extracted),
Ok(Some(extracted)) => {
println!(
"flight data decoder message extracted, {}/{}",
self.reader_id, self.poll_count
);
Ok(extracted)
}
Ok(None) => continue, // Need next input message
Err(e) => Err(e),
},
Expand Down
73 changes: 57 additions & 16 deletions arrow-flight/src/encode.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -238,7 +238,7 @@ impl FlightDataEncoderBuilder {
/// of [`FlightData`], consuming self.
///
/// See example on [`Self`] and [`FlightDataEncoder`] for more details
pub fn build<S>(self, input: S) -> FlightDataEncoder
pub fn build<S>(self, input: S, reader_id: &str) -> FlightDataEncoder
where
S: Stream<Item = Result<RecordBatch>> + Send + 'static,
{
Expand All @@ -259,6 +259,7 @@ impl FlightDataEncoderBuilder {
app_metadata,
descriptor,
dictionary_handling,
reader_id,
)
}
}
Expand Down Expand Up @@ -287,9 +288,14 @@ pub struct FlightDataEncoder {
/// Deterimines how `DictionaryArray`s are encoded for transport.
/// See [`DictionaryHandling`] for more information.
dictionary_handling: DictionaryHandling,

reader_id: String,

poll_count: usize,
}

impl FlightDataEncoder {
#[allow(clippy::too_many_arguments)]
fn new(
inner: BoxStream<'static, Result<RecordBatch>>,
schema: Option<SchemaRef>,
Expand All @@ -298,6 +304,7 @@ impl FlightDataEncoder {
app_metadata: Bytes,
descriptor: Option<FlightDescriptor>,
dictionary_handling: DictionaryHandling,
reader_id: &str,
) -> Self {
let mut encoder = Self {
inner,
Expand All @@ -312,6 +319,8 @@ impl FlightDataEncoder {
done: false,
descriptor,
dictionary_handling,
reader_id: reader_id.to_string(),
poll_count: 0,
};

// If schema is known up front, enqueue it immediately
Expand Down Expand Up @@ -398,24 +407,49 @@ impl Stream for FlightDataEncoder {
cx: &mut std::task::Context<'_>,
) -> Poll<Option<Self::Item>> {
loop {
self.poll_count += 1;

println!(
"flight data encoder polling next, {}/{}",
self.reader_id, self.poll_count
);
if self.done && self.queue.is_empty() {
println!(
"flight data encoder stream done, no more data to send, {}/{}",
self.reader_id, self.poll_count
);
return Poll::Ready(None);
}

// Any messages queued to send?
if let Some(data) = self.queue.pop_front() {
println!(
"flight data encoder sending queued message, {}/{}",
self.reader_id, self.poll_count
);
return Poll::Ready(Some(Ok(data)));
}

// Get next batch
let now = std::time::Instant::now();
let batch = ready!(self.inner.poll_next_unpin(cx));
println!(
"flight data encoder polled next, {}/{} - took {:?}",
self.reader_id,
self.poll_count,
now.elapsed()
);

match batch {
None => {
// inner is done
self.done = true;
// queue must also be empty so we are done
assert!(self.queue.is_empty());
println!(
"flight data encoder stream done, no more data to send, {}/{}",
self.reader_id, self.poll_count
);
return Poll::Ready(None);
}
Some(Err(e)) => {
Expand All @@ -425,6 +459,10 @@ impl Stream for FlightDataEncoder {
return Poll::Ready(Some(Err(e)));
}
Some(Ok(batch)) => {
println!(
"flight data encoder got batch, {}/{}",
self.reader_id, self.poll_count
);
// had data, encode into the queue
if let Err(e) = self.encode_batch(batch) {
self.done = true;
Expand Down Expand Up @@ -790,8 +828,8 @@ mod tests {

let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2)]);

let encoder = FlightDataEncoderBuilder::default().build(stream);
let mut decoder = FlightDataDecoder::new(encoder);
let encoder = FlightDataEncoderBuilder::default().build(stream, "");
let mut decoder = FlightDataDecoder::new(encoder, "");
let expected_schema = Schema::new(vec![Field::new("dict", DataType::Utf8, false)]);
let expected_schema = Arc::new(expected_schema);
let mut expected_arrays = vec![
Expand Down Expand Up @@ -851,7 +889,7 @@ mod tests {

let encoder = FlightDataEncoderBuilder::default()
.with_schema(schema)
.build(stream);
.build(stream, "");
let expected_schema =
Arc::new(Schema::new(vec![Field::new("dict", DataType::Utf8, false)]));
assert_eq!(Some(expected_schema), encoder.known_schema())
Expand All @@ -876,7 +914,7 @@ mod tests {
let encoder = FlightDataEncoderBuilder::default()
.with_dictionary_handling(DictionaryHandling::Resend)
.with_schema(schema.clone())
.build(stream);
.build(stream, "");
assert_eq!(Some(schema), encoder.known_schema())
}

Expand Down Expand Up @@ -929,9 +967,9 @@ mod tests {

let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2)]);

let encoder = FlightDataEncoderBuilder::default().build(stream);
let encoder = FlightDataEncoderBuilder::default().build(stream, "");

let mut decoder = FlightDataDecoder::new(encoder);
let mut decoder = FlightDataDecoder::new(encoder, "");
let expected_schema = Schema::new(vec![Field::new_list(
"dict_list",
Field::new("item", DataType::Utf8, true),
Expand Down Expand Up @@ -1031,9 +1069,9 @@ mod tests {

let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2)]);

let encoder = FlightDataEncoderBuilder::default().build(stream);
let encoder = FlightDataEncoderBuilder::default().build(stream, "");

let mut decoder = FlightDataDecoder::new(encoder);
let mut decoder = FlightDataDecoder::new(encoder, "");
let expected_schema = Schema::new(vec![Field::new_struct(
"struct",
vec![Field::new_list(
Expand Down Expand Up @@ -1212,9 +1250,9 @@ mod tests {

let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2), Ok(batch3)]);

let encoder = FlightDataEncoderBuilder::default().build(stream);
let encoder = FlightDataEncoderBuilder::default().build(stream, "");

let mut decoder = FlightDataDecoder::new(encoder);
let mut decoder = FlightDataDecoder::new(encoder, "");

let hydrated_struct_fields = vec![Field::new_list(
"dict_list",
Expand Down Expand Up @@ -1427,9 +1465,9 @@ mod tests {

let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2)]);

let encoder = FlightDataEncoderBuilder::default().build(stream);
let encoder = FlightDataEncoderBuilder::default().build(stream, "");

let mut decoder = FlightDataDecoder::new(encoder);
let mut decoder = FlightDataDecoder::new(encoder, "");
let expected_schema = Schema::new(vec![Field::new_map(
"dict_map",
"entries",
Expand Down Expand Up @@ -1540,11 +1578,14 @@ mod tests {
let encoder = FlightDataEncoderBuilder::default()
.with_options(IpcWriteOptions::default().with_preserve_dict_id(false))
.with_dictionary_handling(DictionaryHandling::Resend)
.build(futures::stream::iter(batches.clone().into_iter().map(Ok)));
.build(
futures::stream::iter(batches.clone().into_iter().map(Ok)),
"",
);

let mut expected_batches = batches.drain(..);

let mut decoder = FlightDataDecoder::new(encoder);
let mut decoder = FlightDataDecoder::new(encoder, "");
while let Some(decoded) = decoder.next().await {
let decoded = decoded.unwrap();
match decoded.payload {
Expand Down Expand Up @@ -1841,7 +1882,7 @@ mod tests {
.with_max_flight_data_size(max_flight_data_size)
// use 8-byte alignment - default alignment is 64 which produces bigger ipc data
.with_options(IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap())
.build(futures::stream::iter([Ok(batch.clone())]));
.build(futures::stream::iter([Ok(batch.clone())]), "");

let mut i = 0;
while let Some(data) = stream.next().await.transpose().unwrap() {
Expand Down
Loading