feat: logical Node for find files

crates/core/src/delta_datafusion/find_files/logical.rs

@@ -0,0 +1,47 @@
+use std::collections::HashSet;
+
+use datafusion_common::DFSchemaRef;
+use datafusion_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore};
+
+#[derive(Debug, Hash, Eq, PartialEq, Clone)]
+pub struct FindFilesNode {
+    pub id: String,
+    pub input: LogicalPlan,
+    pub predicates: Vec<Expr>,
+    pub files: Vec<String>,
+    pub schema: DFSchemaRef,
+}
+
+impl UserDefinedLogicalNodeCore for FindFilesNode {
+    fn name(&self) -> &str {
+        "FindFiles"
+    }
+
+    fn inputs(&self) -> Vec<&LogicalPlan> {
+        vec![]
+    }
+
+    fn schema(&self) -> &DFSchemaRef {
+        &self.schema
+    }
+
+    fn expressions(&self) -> Vec<Expr> {
+        vec![]
+    }
+
+    fn prevent_predicate_push_down_columns(&self) -> HashSet<String> {
+        HashSet::new()
+    }
+
+    fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(
+            f,
+            "FindFiles id={}, predicate={:?}, files={:?}",
+            &self.id, self.predicates, self.files
+        )
+    }
+
+    fn from_template(&self, _exprs: &[Expr], _inputs: &[LogicalPlan]) -> Self {
+        self.clone()
+    }
+}

crates/core/src/delta_datafusion/find_files/mod.rs

@@ -0,0 +1,160 @@
+use std::sync::Arc;
+
+use arrow_array::{RecordBatch, StringArray};
+use arrow_array::cast::AsArray;
+use arrow_schema::{DataType, Field, Schema};
+use arrow_schema::SchemaBuilder;
+use async_trait::async_trait;
+use datafusion::execution::context::{QueryPlanner, SessionState};
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner};
+use datafusion::prelude::{ParquetReadOptions, SessionContext};
+use datafusion_common::Result;
+use datafusion_expr::{Expr, LogicalPlan, UserDefinedLogicalNode};
+
+use crate::delta_datafusion::find_files::logical::FindFilesNode;
+use crate::delta_datafusion::find_files::physical::FindFilesExec;
+use crate::delta_datafusion::PATH_COLUMN;
+
+pub mod logical;
+pub mod physical;
+
+#[inline]
+fn only_file_path_schema() -> Arc<Schema> {
+    let mut builder = SchemaBuilder::new();
+    builder.push(Field::new(PATH_COLUMN, DataType::Utf8, false));
+    Arc::new(builder.finish())
+}
+
+struct FindFilesPlannerExtension {}
+
+struct FindFilesPlanner {}
+
+#[async_trait]
+impl ExtensionPlanner for FindFilesPlannerExtension {
+    async fn plan_extension(
+        &self,
+        _planner: &dyn PhysicalPlanner,
+        node: &dyn UserDefinedLogicalNode,
+        _logical_inputs: &[&LogicalPlan],
+        _physical_inputs: &[Arc<dyn ExecutionPlan>],
+        _session_state: &SessionState,
+    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+        if let Some(node) = node.as_any().downcast_ref::<FindFilesNode>() {
+            dbg!(&node.files, &node.predicates);
+            let schema = Arc::new(Schema::from(node.schema.as_ref()));
+
+            return Ok(Some(Arc::new(FindFilesExec::new(
+                node.files.clone(),
+                node.predicates[0].clone(),
+                schema,
+            )?)));
+        }
+        Ok(None)
+    }
+}
+
+#[async_trait]
+impl QueryPlanner for FindFilesPlanner {
+    async fn create_physical_plan(
+        &self,
+        logical_plan: &LogicalPlan,
+        session_state: &SessionState,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let planner = Arc::new(Box::new(DefaultPhysicalPlanner::with_extension_planners(
+            vec![Arc::new(FindFilesPlannerExtension {})],
+        )));
+        planner
+            .create_physical_plan(logical_plan, session_state)
+            .await
+    }
+}
+
+async fn scan_memory_table_batch(batch: RecordBatch, predicate: Expr) -> Result<RecordBatch> {
+    let ctx = SessionContext::new();
+    let mut batches = vec![];
+
+    if let Some(column) = batch.column_by_name(PATH_COLUMN) {
+        let mut column_iter = column.as_string::<i32>().into_iter();
+        while let Some(Some(row)) = column_iter.next() {
+            let df = ctx
+                .read_parquet(row, ParquetReadOptions::default())
+                .await?
+                .filter(predicate.to_owned())?;
+            if df.count().await? > 0 {
+                batches.push(row);
+            }
+        }
+    }
+    let str_array = Arc::new(StringArray::from(batches));
+    RecordBatch::try_new(only_file_path_schema(), vec![str_array]).map_err(Into::into)
+}
+
+#[cfg(test)]
+pub mod tests {
+    use std::sync::Arc;
+
+    use arrow_cast::pretty::print_batches;
+    use arrow_schema::{DataType, Field, Fields, Schema, SchemaBuilder};
+    use datafusion::prelude::{DataFrame, SessionContext};
+    use datafusion_common::ToDFSchema;
+    use datafusion_expr::{col, Extension, lit, LogicalPlan, LogicalPlanBuilder};
+
+    use crate::{DeltaOps, DeltaResult};
+    use crate::delta_datafusion::find_files::FindFilesPlanner;
+    use crate::delta_datafusion::find_files::logical::FindFilesNode;
+    use crate::delta_datafusion::PATH_COLUMN;
+    use crate::operations::collect_sendable_stream;
+    use crate::writer::test_utils::{create_bare_table, get_record_batch};
+
+    #[inline]
+    fn find_files_schema(fields: &Fields) -> Arc<Schema> {
+        let mut builder = SchemaBuilder::from(fields);
+        builder.reverse();
+        builder.push(Field::new(PATH_COLUMN, DataType::Utf8, false));
+        builder.reverse();
+        Arc::new(builder.finish())
+    }
+
+    async fn make_table() -> DeltaOps {
+        let batch = get_record_batch(None, false);
+        let write = DeltaOps(create_bare_table())
+            .write(vec![batch.clone()])
+            .await
+            .unwrap();
+        DeltaOps(write)
+    }
+
+    #[tokio::test]
+    pub async fn test_find_files() -> DeltaResult<()> {
+        let ctx = SessionContext::new();
+        let state = ctx
+            .state()
+            .with_query_planner(Arc::new(FindFilesPlanner {}));
+        let table = make_table().await;
+        let files = table.0.get_file_uris()?.collect::<Vec<String>>();
+        let plan = LogicalPlanBuilder::empty(false).build()?;
+
+        let schema = find_files_schema(table.0.snapshot()?.arrow_schema()?.fields()).to_dfschema_ref()?;
+        let find_files_node = LogicalPlan::Extension(Extension {
+            node: Arc::new(FindFilesNode {
+                id: "my_cool_id".to_string(),
+                input: plan,
+                predicates: vec![col("id").eq(lit("A"))],
+                files,
+                schema,
+            }),
+        });
+        let df = DataFrame::new(state.clone(), find_files_node);
+        let p = state
+            .clone()
+            .create_physical_plan(df.logical_plan())
+            .await
+            .unwrap();
+
+        let e = p.execute(0, state.task_ctx())?;
+        let s = collect_sendable_stream(e).await.unwrap();
+        print_batches(&s)?;
+        Ok(())
+    }
+}

crates/core/src/delta_datafusion/find_files/physical.rs

@@ -0,0 +1,124 @@
+use std::any::Any;
+use std::fmt::{Debug, Formatter};
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use arrow_array::{RecordBatch, StringArray};
+use arrow_schema::SchemaRef;
+use datafusion::error::Result;
+use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext};
+use datafusion::physical_plan::memory::MemoryStream;
+use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan};
+use datafusion_expr::Expr;
+use datafusion_physical_expr::{Partitioning, PhysicalSortExpr};
+use futures::stream::BoxStream;
+use futures::{Stream, StreamExt, TryStreamExt};
+
+use crate::delta_datafusion::find_files::{only_file_path_schema, scan_memory_table_batch};
+
+pub struct FindFilesExec {
+    schema: SchemaRef,
+    files: Vec<String>,
+    predicate: Expr,
+}
+
+impl FindFilesExec {
+    pub fn new(files: Vec<String>, predicate: Expr, schema: SchemaRef) -> Result<Self> {
+        Ok(Self {
+            schema,
+            files,
+            predicate,
+        })
+    }
+}
+
+struct FindFilesStream<'a> {
+    mem_stream: BoxStream<'a, Result<RecordBatch>>,
+}
+
+impl<'a> FindFilesStream<'a> {
+    pub fn new(mem_stream: BoxStream<'a, Result<RecordBatch>>) -> Result<Self> {
+        Ok(Self { mem_stream })
+    }
+}
+
+impl<'a> RecordBatchStream for FindFilesStream<'a> {
+    fn schema(&self) -> SchemaRef {
+        only_file_path_schema()
+    }
+}
+
+impl<'a> Stream for FindFilesStream<'a> {
+    type Item = Result<RecordBatch>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.as_mut().mem_stream.poll_next_unpin(cx)
+    }
+}
+
+impl Debug for FindFilesExec {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "FindFilesExec[schema={:?}, files={:?}]",
+            self.schema, self.files
+        )
+    }
+}
+
+impl DisplayAs for FindFilesExec {
+    fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
+        write!(
+            f,
+            "FindFilesExec[schema={:?}, files={:?}]",
+            self.schema, self.files
+        )
+    }
+}
+
+impl ExecutionPlan for FindFilesExec {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn schema(&self) -> SchemaRef {
+        self.schema.clone()
+    }
+
+    fn output_partitioning(&self) -> Partitioning {
+        Partitioning::UnknownPartitioning(0)
+    }
+
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
+    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
+        vec![]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        _children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        Ok(self)
+    }
+
+    fn execute(
+        &self,
+        _partition: usize,
+        _context: Arc<TaskContext>,
+    ) -> Result<SendableRecordBatchStream> {
+        let array = Arc::new(StringArray::from(self.files.clone()));
+        let record_batch = RecordBatch::try_new(only_file_path_schema(), vec![array])?;
+        let predicate = self.predicate.clone();
+        let mem_stream =
+            MemoryStream::try_new(vec![record_batch.clone()], only_file_path_schema(), None)?
+                .and_then(move |batch| scan_memory_table_batch(batch, predicate.clone()))
+                .boxed();
+
+        Ok(Box::pin(FindFilesStream::new(mem_stream)?))
+    }
+}

crates/core/src/delta_datafusion/mod.rs

@@ -87,6 +87,8 @@ pub mod expr;
 pub mod logical;
 pub mod physical;
 
+mod find_files;
+
 impl From<DeltaTableError> for DataFusionError {
     fn from(err: DeltaTableError) -> Self {
         match err {
@@ -1145,6 +1147,7 @@ impl TreeNodeVisitor for FindFilesExprProperties {
     }
 }
 
+#[derive(Debug, Hash, Eq, PartialEq)]
 /// Representing the result of the [find_files] function.
 pub struct FindFiles {
     /// A list of `Add` objects that match the given predicate
@@ -1198,7 +1201,7 @@ fn join_batches_with_add_actions(
     Ok(files)
 }
 
-/// Determine which files contain a record that statisfies the predicate
+/// Determine which files contain a record that satisfies the predicate
 pub(crate) async fn find_files_scan<'a>(
     snapshot: &DeltaTableState,
     log_store: LogStoreRef,

crates/core/src/table/state_arrow.rs

@@ -91,7 +91,7 @@ impl DeltaTableState {
                     .fields
                     .iter()
                     .map(|field| Cow::Owned(field.name().clone()))
-                    .zip(partition_cols_batch.columns().iter().map(Arc::clone)),
+                    .zip(partition_cols_batch.columns().iter().cloned()),
             )
         }
 
@@ -103,7 +103,7 @@ impl DeltaTableState {
                     .fields
                     .iter()
                     .map(|field| Cow::Owned(field.name().clone()))
-                    .zip(stats.columns().iter().map(Arc::clone)),
+                    .zip(stats.columns().iter().cloned()),
             );
         }
         if files.iter().any(|add| add.deletion_vector.is_some()) {
@@ -114,7 +114,7 @@ impl DeltaTableState {
                     .fields
                     .iter()
                     .map(|field| Cow::Owned(field.name().clone()))
-                    .zip(delvs.columns().iter().map(Arc::clone)),
+                    .zip(delvs.columns().iter().cloned()),
             );
         }
         if files.iter().any(|add| {
@@ -129,7 +129,7 @@ impl DeltaTableState {
                     .fields
                     .iter()
                     .map(|field| Cow::Owned(field.name().clone()))
-                    .zip(tags.columns().iter().map(Arc::clone)),
+                    .zip(tags.columns().iter().cloned()),
             );
         }