add schema method

probably-nothing-labs · Aug 12, 2024 · c6ef592 · c6ef592
1 parent e1234dd
commit c6ef592
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/crates/core/src/datastream.rs b/crates/core/src/datastream.rs
@@ -1,7 +1,7 @@
 use futures::StreamExt;
 use std::{sync::Arc, time::Duration};
 
-use datafusion::common::{DataFusionError, Result};
+use datafusion::common::{DFSchema, DataFusionError, Result};
 pub use datafusion::dataframe::DataFrame;
 use datafusion::dataframe::DataFrameWriteOptions;
 use datafusion::execution::SendableRecordBatchStream;
@@ -19,7 +19,8 @@ pub struct DataStream {
 }
 
 impl DataStream {
-    pub fn schema(&self) -> Result<Self> {
+    pub fn schema(&self) -> &DFSchema {
+        self.df.schema()
     }
 
     pub fn filter(&self, predicate: Expr) -> Result<Self> {