diff --git a/python/Cargo.lock b/python/Cargo.lock
index 7b73a36aa6a..555308d942e 100644
--- a/python/Cargo.lock
+++ b/python/Cargo.lock
@@ -4557,6 +4557,7 @@ dependencies = [
  "roaring",
  "serde",
  "serde_json",
+ "smallvec",
  "snafu",
  "tantivy",
  "tempfile",
diff --git a/rust/lance-datafusion/src/exec.rs b/rust/lance-datafusion/src/exec.rs
index 5a06ebf1de5..b1600b8930a 100644
--- a/rust/lance-datafusion/src/exec.rs
+++ b/rust/lance-datafusion/src/exec.rs
@@ -732,7 +732,7 @@ pub trait SessionContextExt {
     ) -> datafusion::common::Result<DataFrame>;
 }
 
-struct OneShotPartitionStream {
+pub struct OneShotPartitionStream {
     data: Arc<Mutex<Option<SendableRecordBatchStream>>>,
     schema: Arc<ArrowSchema>,
 }
@@ -748,7 +748,7 @@ impl std::fmt::Debug for OneShotPartitionStream {
 }
 
 impl OneShotPartitionStream {
-    fn new(data: SendableRecordBatchStream) -> Self {
+    pub fn new(data: SendableRecordBatchStream) -> Self {
         let schema = data.schema();
         Self {
             data: Arc::new(Mutex::new(Some(data))),
diff --git a/rust/lance/src/datafusion/dataframe.rs b/rust/lance/src/datafusion/dataframe.rs
index edb4ea05a68..76c4ac99e97 100644
--- a/rust/lance/src/datafusion/dataframe.rs
+++ b/rust/lance/src/datafusion/dataframe.rs
@@ -182,13 +182,13 @@ pub trait SessionContextExt {
     ) -> datafusion::common::Result<DataFrame>;
 }
 
-struct OneShotPartitionStream {
+pub struct OneShotPartitionStream {
     data: Arc<Mutex<Option<SendableRecordBatchStream>>>,
     schema: Arc<Schema>,
 }
 
 impl OneShotPartitionStream {
-    fn new(data: SendableRecordBatchStream) -> Self {
+    pub fn new(data: SendableRecordBatchStream) -> Self {
         let schema = data.schema();
         Self {
             data: Arc::new(Mutex::new(Some(data))),