fix: bug in sending data more than 4MB

Eshanatnight · Eshanatnight · commit 82f2aa9ffe18 · 2024-05-07T18:02:26.000+05:30
Need to update the `push_logs_unchecked` function as event processing
has changed
Need to clean up the `do_get` function for airplane
diff --git a/server/src/event/writer.rs b/server/src/event/writer.rs
@@ -67,7 +67,6 @@ impl Writer {
         Ok(())
     }
 
-
     fn push_mem(&mut self, schema_key: &str, rb: RecordBatch) -> Result<(), StreamWriterError> {
         self.mem.push(schema_key, rb);
         Ok(())
@@ -95,8 +94,8 @@ impl WriterTable {
                         stream_name,
                         schema_key,
                         record,
-                    parsed_timestamp,
-                )?;
+                        parsed_timestamp,
+                    )?;
                 } else {
                     stream_writer
                         .lock()
@@ -115,15 +114,11 @@ impl WriterTable {
                             stream_name,
                             schema_key,
                             record,
-                        parsed_timestamp,
-                    )?;
+                            parsed_timestamp,
+                        )?;
                     } else {
                         writer.lock().unwrap().push_mem(stream_name, record)?;
                     }
-                } else if CONFIG.parseable.mode != Mode::Query {
-                    let mut writer = Writer::default();
-                    writer.push(stream_name, schema_key, record, parsed_timestamp)?;
-                    map.insert(stream_name.to_owned(), Mutex::new(writer));
                 } else {
                     let mut writer = Writer::default();
                     writer.push(stream_name, schema_key, record, parsed_timestamp)?;
diff --git a/server/src/handlers/airplane.rs b/server/src/handlers/airplane.rs
@@ -1,10 +1,10 @@
 use arrow_array::RecordBatch;
+use arrow_flight::encode::FlightDataEncoderBuilder;
 use arrow_flight::flight_service_server::FlightServiceServer;
 use arrow_flight::PollInfo;
 use arrow_schema::{ArrowError, Schema};
 use arrow_select::concat::concat_batches;
 use chrono::Utc;
-use crossterm::event;
 use datafusion::common::tree_node::TreeNode;
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -28,14 +28,14 @@ use crate::handlers::livetail::cross_origin_config;
 use crate::handlers::http::query::{authorize_and_set_filter_tags, into_query};
 use crate::query::{TableScanVisitor, QUERY_SESSION};
 use crate::storage::object_storage::commit_schema_to_storage;
-use crate::utils::arrow::flight::{get_query_from_ticket, run_do_get_rpc};
+use crate::utils::arrow::flight::{add_temporary_events, get_query_from_ticket, run_do_get_rpc};
 use arrow_flight::{
     flight_service_server::FlightService, Action, ActionType, Criteria, Empty, FlightData,
     FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse, PutResult, SchemaAsIpc,
     SchemaResult, Ticket,
 };
-use arrow_ipc::writer::{DictionaryTracker, IpcDataGenerator, IpcWriteOptions};
-use futures::stream::BoxStream;
+use arrow_ipc::writer::IpcWriteOptions;
+use futures::{stream, TryStreamExt};
 use tonic::{Request, Response, Status, Streaming};
 
 use crate::handlers::livetail::extract_session_key;
@@ -50,13 +50,13 @@ pub struct AirServiceImpl {}
 
 #[tonic::async_trait]
 impl FlightService for AirServiceImpl {
-    type HandshakeStream = BoxStream<'static, Result<HandshakeResponse, Status>>;
-    type ListFlightsStream = BoxStream<'static, Result<FlightInfo, Status>>;
-    type DoGetStream = BoxStream<'static, Result<FlightData, Status>>;
-    type DoPutStream = BoxStream<'static, Result<PutResult, Status>>;
-    type DoActionStream = BoxStream<'static, Result<arrow_flight::Result, Status>>;
-    type ListActionsStream = BoxStream<'static, Result<ActionType, Status>>;
-    type DoExchangeStream = BoxStream<'static, Result<FlightData, Status>>;
+    type HandshakeStream = stream::BoxStream<'static, Result<HandshakeResponse, Status>>;
+    type ListFlightsStream = stream::BoxStream<'static, Result<FlightInfo, Status>>;
+    type DoGetStream = stream::BoxStream<'static, Result<FlightData, Status>>;
+    type DoPutStream = stream::BoxStream<'static, Result<PutResult, Status>>;
+    type DoActionStream = stream::BoxStream<'static, Result<arrow_flight::Result, Status>>;
+    type ListActionsStream = stream::BoxStream<'static, Result<ActionType, Status>>;
+    type DoExchangeStream = stream::BoxStream<'static, Result<FlightData, Status>>;
 
     async fn handshake(
         &self,
@@ -175,24 +175,26 @@ impl FlightService for AirServiceImpl {
                 let mut batches = run_do_get_rpc(im, sql.clone()).await?;
                 minute_result.append(&mut batches);
             }
-            let mr = minute_result.iter().map(|rb| rb).collect::<Vec<_>>();
-            let schema = STREAM_INFO
-                .schema(&stream_name)
-                .map_err(|err| Status::failed_precondition(format!("Metadata Error: {}", err)))?;
-            let rb = concat_batches(&schema, mr)
-                .map_err(|err| Status::failed_precondition(format!("ArrowError: {}", err)))?;
-
-            let event = push_logs_unchecked(rb, &stream_name)
-                .await
-                .map_err(|err| Status::internal(err.to_string()))?;
-            let mut events = vec![];
-            for batch in minute_result {
-                events.push(
-                    push_logs_unchecked(batch, &stream_name)
-                        .await
-                        .map_err(|err| Status::internal(err.to_string()))?,
-                );
-            }
+            let mr = minute_result.iter().collect::<Vec<_>>();
+            let event = add_temporary_events(&stream_name, mr).await?;
+
+            // let schema = STREAM_INFO
+            //     .schema(&stream_name)
+            //     .map_err(|err| Status::failed_precondition(format!("Metadata Error: {}", err)))?;
+            // let rb = concat_batches(&schema, mr)
+            //     .map_err(|err| Status::failed_precondition(format!("ArrowError: {}", err)))?;
+            //
+            // let event = push_logs_unchecked(rb, &stream_name)
+            //     .await
+            //     .map_err(|err| Status::internal(err.to_string()))?;
+            // let mut events = vec![];
+            // for batch in minute_result {
+            //     events.push(
+            //         push_logs_unchecked(batch, &stream_name)
+            //             .await
+            //             .map_err(|err| Status::internal(err.to_string()))?,
+            //     );
+            // }
             Some(event)
         } else {
             None
@@ -216,22 +218,30 @@ impl FlightService for AirServiceImpl {
             .map(|batch| batch.schema())
             .map(|s| s.as_ref().clone())
             .collect::<Vec<_>>();
-
         let schema = Schema::try_merge(schemas).map_err(|err| Status::internal(err.to_string()))?;
-        let options = datafusion::arrow::ipc::writer::IpcWriteOptions::default();
-        let schema_flight_data = SchemaAsIpc::new(&schema, &options);
-
-        let mut flights = vec![FlightData::from(schema_flight_data)];
-        let encoder = IpcDataGenerator::default();
-        let mut tracker = DictionaryTracker::new(false);
-        for batch in &results {
-            let (flight_dictionaries, flight_batch) = encoder
-                .encoded_batch(batch, &mut tracker, &options)
-                .map_err(|e| Status::internal(e.to_string()))?;
-            flights.extend(flight_dictionaries.into_iter().map(Into::into));
-            flights.push(flight_batch.into());
-        }
-        let output = futures::stream::iter(flights.into_iter().map(Ok));
+        let input_stream = futures::stream::iter(results.into_iter().map(Ok));
+
+        let flight_data_stream = FlightDataEncoderBuilder::new()
+            .with_max_flight_data_size(usize::MAX)
+            .with_schema(schema.into())
+            .build(input_stream);
+
+        let flight_data_stream = flight_data_stream.map_err(|err| Status::unknown(err.to_string()));
+
+        // let options = datafusion::arrow::ipc::writer::IpcWriteOptions::default();
+        // let schema_flight_data = SchemaAsIpc::new(&schema, &options);
+        //
+        // let mut flights = vec![FlightData::from(schema_flight_data)];
+        // let encoder = IpcDataGenerator::default();
+        // let mut tracker = DictionaryTracker::new(false);
+        // for batch in &results {
+        //     let (flight_dictionaries, flight_batch) = encoder
+        //         .encoded_batch(batch, &mut tracker, &options)
+        //         .map_err(|e| Status::internal(e.to_string()))?;
+        //     flights.extend(flight_dictionaries.into_iter().map(Into::into));
+        //     flights.push(flight_batch.into());
+        // }
+        // let output = futures::stream::iter(flights.into_iter().map(Ok));
         if let Some(events) = events {
             events.clear(&stream_name);
             // for event in events {
@@ -244,7 +254,9 @@ impl FlightService for AirServiceImpl {
             .with_label_values(&[&format!("flight-query-{}", stream_name)])
             .observe(time);
 
-        Ok(Response::new(Box::pin(output) as Self::DoGetStream))
+        Ok(Response::new(
+            Box::pin(flight_data_stream) as Self::DoGetStream
+        ))
     }
 
     async fn do_put(
diff --git a/server/src/handlers/http/ingest.rs b/server/src/handlers/http/ingest.rs
@@ -102,13 +102,23 @@ pub async fn push_logs_unchecked(
     batches: RecordBatch,
     stream_name: &str,
 ) -> Result<event::Event, PostError> {
-    todo!("timepartition fix");
+    todo!("need to fix");
+
+    // let glob_storage = CONFIG.storage().get_object_store();
+    // let object_store_format = glob_storage
+    //     .get_object_store_format(&stream_name)
+    //     .await
+    //     .map_err(|_| PostError::StreamNotFound(stream_name.to_string()))?;
+
+    // let time_partition = object_store_format.time_partition;
+    // let time_partition_limit = object_store_format.time_partition_limit;
 
     // let event = event::Event {
     //     rb: batches,
     //     stream_name: stream_name.to_string(),
     //     origin_format: "json",
     //     origin_size: 0,
+    //     time_partition,
     //     is_first_event: true,
     // };
     // event.process_unchecked()?;
diff --git a/server/src/utils/arrow/flight.rs b/server/src/utils/arrow/flight.rs
@@ -1,10 +1,14 @@
+use crate::event::Event;
+use crate::handlers::http::ingest::push_logs_unchecked;
 use crate::handlers::http::query::Query as QueryJson;
+use crate::metadata::STREAM_INFO;
 use crate::{
     handlers::http::modal::IngestorMetadata,
     option::{Mode, CONFIG},
 };
 use arrow_array::RecordBatch;
 use arrow_flight::Ticket;
+use arrow_select::concat::concat_batches;
 use futures::TryStreamExt;
 use serde_json::Value as JsonValue;
 use tonic::{Request, Status};
@@ -66,3 +70,34 @@ pub async fn run_do_get_rpc(im: IngestorMetadata, sql: String) -> Result<Vec<Rec
 
     Ok(response.try_collect().await?)
 }
+
+/// all the records from the ingesters are concatinated into one event and pushed to memory
+pub async fn add_temporary_events(
+    stream_name: &str,
+    minute_result: Vec<&RecordBatch>,
+) -> Result<
+    //Vec<Event>
+    Event,
+    Status,
+> {
+    let schema = STREAM_INFO
+        .schema(stream_name)
+        .map_err(|err| Status::failed_precondition(format!("Metadata Error: {}", err)))?;
+    let rb = concat_batches(&schema, minute_result)
+        .map_err(|err| Status::failed_precondition(format!("ArrowError: {}", err)))?;
+
+    let event = push_logs_unchecked(rb, stream_name)
+        .await
+        .map_err(|err| Status::internal(err.to_string()))?;
+    // let mut events = vec![];
+    // for batch in minute_result {
+    //     events.push(
+    //         push_logs_unchecked(batch, &stream_name)
+    //             .await
+    //             .map_err(|err| Status::internal(err.to_string()))?,
+    //     );
+    // }
+    //
+    // Ok(events)
+    Ok(event)
+}
