feat: Execute Python queries (#575)

This removes the ability to inspect the prepared data from the Python
FFI. Instead, the Python FFI now has an `execute` method allowing you to
inspect the results of an expression, including the bare table.

This changes the execute methods a bit to have a simpler API to
interface with.

This relies on pushing "in-memory" data into the `DataContext`, which is
very brittle and hacky. For example, it currently only supports adding
data once, and requires that the schemas line up.

crates/sparrow-compiler/src/data_context.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 
 use anyhow::Context;
 use arrow::datatypes::{DataType, SchemaRef};
+use arrow::record_batch::RecordBatch;
 use sparrow_api::kaskada::v1alpha::slice_plan::Slice;
 use sparrow_api::kaskada::v1alpha::{compute_table, ComputeTable, PreparedFile, TableConfig};
 use sparrow_core::context_code;
@@ -17,7 +18,7 @@ use crate::AstDfgRef;
 ///
 /// Specifically, this holds the information about the tables
 /// available to the compilation.
-#[derive(Default, Debug)]
+#[derive(Default, Debug, Clone)]
 pub struct DataContext {
     /// Information about the groupings in the context.
     group_info: Vec<GroupInfo>,
@@ -60,6 +61,10 @@ impl DataContext {
         self.table_info.get(&id)
     }
 
+    pub fn table_info_mut(&mut self, id: TableId) -> Option<&mut TableInfo> {
+        self.table_info.get_mut(&id)
+    }
+
     pub fn tables_for_grouping(&self, id: GroupId) -> impl Iterator<Item = &TableInfo> {
         self.table_info
             .iter()
@@ -296,7 +301,7 @@ impl DataContext {
 }
 
 /// Information about groups.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct GroupInfo {
     name: String,
     key_type: DataType,
@@ -314,6 +319,8 @@ pub struct TableInfo {
     /// Each file set corresponds to the files for the table with a specific
     /// slice configuration.
     file_sets: Vec<compute_table::FileSet>,
+    /// An in-memory record batch for the contents of the table.
+    pub in_memory: Option<RecordBatch>,
 }
 
 impl TableInfo {
@@ -335,6 +342,7 @@ impl TableInfo {
             schema,
             config,
             file_sets,
+            in_memory: None,
         })
     }
 

crates/sparrow-compiler/src/dfg.rs

@@ -364,7 +364,7 @@ impl Dfg {
     }
 
     /// Runs simplifications on the graph.
-    pub(crate) fn run_simplifications(&mut self, options: &CompilerOptions) -> anyhow::Result<()> {
+    pub fn run_simplifications(&mut self, options: &CompilerOptions) -> anyhow::Result<()> {
         let span = info_span!("Running simplifications");
         let _enter = span.enter();
 
@@ -375,7 +375,7 @@ impl Dfg {
     }
 
     /// Extract the simplest representation of the node `id` from the graph.
-    pub(crate) fn extract_simplest(&self, id: Id) -> DfgExpr {
+    pub fn extract_simplest(&self, id: Id) -> DfgExpr {
         let span = info_span!("Extracting simplest DFG");
         let _enter = span.enter();
 
@@ -447,7 +447,7 @@ impl Dfg {
     /// Remove nodes that aren't needed for the `output` from the graph.
     ///
     /// Returns the new ID of the `output`.
-    pub(crate) fn prune(&mut self, output: Id) -> anyhow::Result<Id> {
+    pub fn prune(&mut self, output: Id) -> anyhow::Result<Id> {
         // The implementation is somewhat painful -- we extract a `RecExpr`, and then
         // recreate the EGraph. This has the desired property -- only referenced nodes
         // are extracted. But, it may cause the IDs to change.

crates/sparrow-compiler/src/lib.rs

@@ -41,7 +41,7 @@ mod frontend;
 mod functions;
 mod nearest_matches;
 mod options;
-mod plan;
+pub mod plan;
 mod time_domain;
 mod types;
 

crates/sparrow-compiler/src/plan.rs

@@ -24,7 +24,7 @@ const DBG_PRINT_PLAN: bool = false;
 /// TODO: The `DataContext` is used to get the table name from an ID, which is
 /// only necessary to create the `slice_plan` because it uses a name instead of
 /// an ID.
-pub(super) fn extract_plan_proto(
+pub fn extract_plan_proto(
     data_context: &DataContext,
     expr: DfgExpr,
     per_entity_behavior: PerEntityBehavior,

crates/sparrow-main/tests/e2e/fixture/query_fixture.rs

@@ -270,7 +270,6 @@ impl QueryFixture {
             plan: Some(plan),
             destination: Some(output_to),
             tables: data.tables(),
-
             ..self.execute_request.clone()
         };