Propegate output_ordering to scan (vortex-data#4909)

Fix a bug where datafusion scans the provide `output_ordering` assume
the stream itself is provided in-order, and we explicitly don't
guarantee the stream we return from files is in-order.

Signed-off-by: Adam Gutglick <[email protected]>

Author: AdamGS

vortex-datafusion/src/convert/exprs.rs

@@ -372,12 +372,14 @@ mod tests {
 
         let result = ExprRef::try_from_df(&like_expr).unwrap();
 
-        assert_snapshot!(result.display_tree().to_string(), @r#"
-        Like
-        ├── child: GetItem(text_col)
-        │   └── Root
-        └── pattern: Literal(value: "test%", dtype: utf8)
-        "#);
+        insta::allow_duplicates! {
+            assert_snapshot!(result.display_tree().to_string(), @r#"
+            Like
+            ├── child: GetItem(text_col)
+            │   └── Root
+            └── pattern: Literal(value: "test%", dtype: utf8)
+            "#);
+        }
     }
 
     #[rstest]

vortex-datafusion/src/persistent/mod.rs

@@ -44,6 +44,7 @@ mod tests {
     use datafusion::prelude::SessionContext;
     use datafusion_datasource::file_format::format_as_file_type;
     use datafusion_expr::LogicalPlanBuilder;
+    use datafusion_physical_plan::display::DisplayableExecutionPlan;
     use insta::assert_snapshot;
     use rstest::rstest;
     use tempfile::{TempDir, tempdir};
@@ -188,4 +189,80 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn create_table_ordered_by() -> anyhow::Result<()> {
+        let dir = TempDir::new().unwrap();
+
+        let factory: VortexFormatFactory = VortexFormatFactory::new();
+        let mut session_state_builder = SessionStateBuilder::new().with_default_features();
+        register_vortex_format_factory(factory, &mut session_state_builder);
+        let session = SessionContext::new_with_state(session_state_builder.build());
+
+        // Vortex
+        session
+            .sql(&format!(
+                "CREATE EXTERNAL TABLE my_tbl_vx \
+                (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \
+                STORED AS vortex  \
+                WITH ORDER (c1 ASC)
+                LOCATION '{}/vx/'",
+                dir.path().to_str().unwrap()
+            ))
+            .await?;
+
+        session
+            .sql("INSERT INTO my_tbl_vx VALUES ('air', 5), ('balloon', 42)")
+            .await?
+            .collect()
+            .await?;
+
+        session
+            .sql("INSERT INTO my_tbl_vx VALUES ('zebra', 5)")
+            .await?
+            .collect()
+            .await?;
+
+        session
+            .sql("INSERT INTO my_tbl_vx VALUES ('texas', 2000), ('alabama', 2000)")
+            .await?
+            .collect()
+            .await?;
+
+        let df = session
+            .sql("SELECT * FROM my_tbl_vx ORDER BY c1 ASC limit 3")
+            .await?;
+        let (state, plan) = df.clone().into_parts();
+        let physical_plan = state.create_physical_plan(&plan).await?;
+
+        insta::assert_snapshot!(DisplayableExecutionPlan::new(physical_plan.as_ref())
+                .tree_render().to_string(), @r"
+        ┌───────────────────────────┐
+        │  SortPreservingMergeExec  │
+        │    --------------------   │
+        │  c1 ASC NULLS LASTlimit:  │
+        │             3             │
+        └─────────────┬─────────────┘
+        ┌─────────────┴─────────────┐
+        │       DataSourceExec      │
+        │    --------------------   │
+        │          files: 3         │
+        │       format: vortex      │
+        └───────────────────────────┘
+        ");
+
+        let r = df.collect().await?;
+
+        insta::assert_snapshot!(pretty_format_batches(&r)?.to_string(), @r"
+        +---------+------+
+        | c1      | c2   |
+        +---------+------+
+        | air     | 5    |
+        | alabama | 2000 |
+        | balloon | 42   |
+        +---------+------+
+        ");
+
+        Ok(())
+    }
 }

vortex-datafusion/src/persistent/opener.rs

@@ -55,6 +55,8 @@ pub(crate) struct VortexOpener {
     pub limit: Option<usize>,
     pub metrics: VortexMetrics,
     pub layout_readers: Arc<DashMap<Path, Weak<dyn LayoutReader>>>,
+    /// Whether the query has output ordering specified
+    pub has_output_ordering: bool,
 }
 
 impl FileOpener for VortexOpener {
@@ -71,6 +73,7 @@ impl FileOpener for VortexOpener {
         let limit = self.limit;
         let metrics = self.metrics.clone();
         let layout_reader = self.layout_readers.clone();
+        let has_output_ordering = self.has_output_ordering;
 
         let projected_schema = match projection.as_ref() {
             None => logical_schema.clone(),
@@ -224,7 +227,7 @@ impl FileOpener for VortexOpener {
                 .with_metrics(metrics)
                 .with_projection(projection_expr)
                 .with_some_filter(filter)
-                .with_ordered(false)
+                .with_ordered(has_output_ordering)
                 .map(|chunk| chunk.to_struct().into_record_batch())
                 .into_stream()
                 .map_err(|e| {
@@ -432,6 +435,7 @@ mod tests {
             limit: None,
             metrics: Default::default(),
             layout_readers: Default::default(),
+            has_output_ordering: false,
         };
 
         // filter matches partition value
@@ -512,6 +516,7 @@ mod tests {
             limit: None,
             metrics: Default::default(),
             layout_readers: Default::default(),
+            has_output_ordering: false,
         };
 
         let filter = col("a").lt(lit(100_i32));

vortex-datafusion/src/persistent/source.rs

@@ -148,6 +148,7 @@ impl FileSource for VortexSource {
             limit: base_config.limit,
             metrics: partition_metrics,
             layout_readers: self.layout_readers.clone(),
+            has_output_ordering: !base_config.output_ordering.is_empty(),
         };
 
         Arc::new(opener)