Misc upgrades to VortexSource (#8501)

## Summary

1. Add support for pushing down `byte_length()`
2. Better Display for `VortexSink`
3. Support disabling predicate pushdown

---------

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

vortex-datafusion/src/convert/exprs.rs

@@ -11,6 +11,7 @@ use datafusion_common::tree_node::TreeNode;
 use datafusion_common::tree_node::TreeNodeRecursion;
 use datafusion_expr::Operator as DFOperator;
 use datafusion_functions::core::getfield::GetFieldFunc;
+use datafusion_functions::string::octet_length::OctetLengthFunc;
 use datafusion_physical_expr::PhysicalExpr;
 use datafusion_physical_expr::ScalarFunctionExpr;
 use datafusion_physical_expr::projection::ProjectionExpr;
@@ -24,6 +25,7 @@ use vortex::dtype::Nullability;
 use vortex::dtype::arrow::FromArrowType;
 use vortex::expr::Expression;
 use vortex::expr::and_collect;
+use vortex::expr::byte_length;
 use vortex::expr::cast;
 use vortex::expr::get_item;
 use vortex::expr::is_not_null;
@@ -111,8 +113,28 @@ pub trait ExpressionConvertor: Send + Sync {
 pub struct DefaultExpressionConvertor {}
 
 impl DefaultExpressionConvertor {
+    /// Attempts to convert DataFusion's `octet_length` function to Vortex `byte_length`.
+    fn try_convert_octet_length(&self, scalar_fn: &ScalarFunctionExpr) -> DFResult<Expression> {
+        let [input] = scalar_fn.args() else {
+            return Err(exec_datafusion_err!(
+                "octet_length requires exactly one argument"
+            ));
+        };
+
+        let input = self.convert(input.as_ref())?;
+        let return_dtype =
+            DType::from_arrow((scalar_fn.return_type(), scalar_fn.nullable().into()));
+        Ok(cast(byte_length(input), return_dtype))
+    }
+
     /// Attempts to convert a DataFusion ScalarFunctionExpr to a Vortex expression.
     fn try_convert_scalar_function(&self, scalar_fn: &ScalarFunctionExpr) -> DFResult<Expression> {
+        if let Some(octet_length_fn) =
+            ScalarFunctionExpr::try_downcast_func::<OctetLengthFunc>(scalar_fn)
+        {
+            return self.try_convert_octet_length(octet_length_fn);
+        }
+
         if let Some(get_field_fn) = ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(scalar_fn)
         {
             // DataFusion's GetFieldFunc flattens nested field access into a single call
@@ -289,7 +311,7 @@ impl ExpressionConvertor for DefaultExpressionConvertor {
             let r = projection_expr.expr.apply(|node| {
                 // We only pull column children of scalar functions that we can't push into the scan.
                 if let Some(scalar_fn_expr) = node.downcast_ref::<ScalarFunctionExpr>()
-                    && !can_scalar_fn_be_pushed_down(scalar_fn_expr)
+                    && !can_scalar_fn_be_pushed_down(scalar_fn_expr, input_schema)
                 {
                     scan_projection.extend(
                         collect_columns(node)
@@ -305,8 +327,8 @@ impl ExpressionConvertor for DefaultExpressionConvertor {
                 // Vortex expects a perfect match so we don't push it down.
                 if let Some(binary_expr) = node.downcast_ref::<df_expr::BinaryExpr>()
                     && binary_expr.op().is_numerical_operators()
-                    && (is_decimal(&binary_expr.left().data_type(input_schema)?)
-                        && is_decimal(&binary_expr.right().data_type(input_schema)?))
+                    && binary_expr.left().data_type(input_schema)?.is_decimal()
+                    && binary_expr.right().data_type(input_schema)?.is_decimal()
                 {
                     scan_projection.extend(
                         collect_columns(node)
@@ -430,7 +452,7 @@ fn can_be_pushed_down_impl(expr: &Arc<dyn PhysicalExpr>, schema: &Schema) -> boo
                 .iter()
                 .all(|e| can_be_pushed_down_impl(e, schema))
     } else if let Some(scalar_fn) = expr.downcast_ref::<ScalarFunctionExpr>() {
-        can_scalar_fn_be_pushed_down(scalar_fn)
+        can_scalar_fn_be_pushed_down(scalar_fn, schema)
     } else if let Some(case_expr) = expr.downcast_ref::<df_expr::CaseExpr>() {
         can_case_be_pushed_down(case_expr, schema)
     } else {
@@ -454,9 +476,10 @@ fn is_convertible_expr(expr: &Arc<dyn PhysicalExpr>) -> bool {
         || expr.downcast_ref::<df_expr::IsNullExpr>().is_some()
         || expr.downcast_ref::<df_expr::IsNotNullExpr>().is_some()
         || expr.downcast_ref::<df_expr::InListExpr>().is_some()
-        || expr
-            .downcast_ref::<ScalarFunctionExpr>()
-            .is_some_and(|sf| ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(sf).is_some())
+        || expr.downcast_ref::<ScalarFunctionExpr>().is_some_and(|sf| {
+            ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(sf).is_some()
+                || ScalarFunctionExpr::try_downcast_func::<OctetLengthFunc>(sf).is_some()
+        })
 }
 
 fn can_binary_be_pushed_down(binary: &df_expr::BinaryExpr, schema: &Schema) -> bool {
@@ -502,20 +525,11 @@ fn supported_data_types(dt: &DataType) -> bool {
 
     let is_supported = dt.is_null()
         || dt.is_numeric()
+        || dt.is_binary()
+        || dt.is_string()
         || matches!(
             dt,
-            Boolean
-                | Utf8
-                | LargeUtf8
-                | Utf8View
-                | Binary
-                | LargeBinary
-                | BinaryView
-                | Date32
-                | Date64
-                | Timestamp(_, _)
-                | Time32(_)
-                | Time64(_)
+            Boolean | Date32 | Date64 | Timestamp(_, _) | Time32(_) | Time64(_)
         );
 
     if !is_supported {
@@ -526,20 +540,30 @@ fn supported_data_types(dt: &DataType) -> bool {
 }
 
 /// Checks if a scalar function can be pushed down.
-/// Currently only GetFieldFunc is supported.
-fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr) -> bool {
-    ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(scalar_fn).is_some()
+/// Currently GetFieldFunc and OctetLengthFunc are supported.
+fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool {
+    if ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(scalar_fn).is_some() {
+        return true;
+    }
+
+    ScalarFunctionExpr::try_downcast_func::<OctetLengthFunc>(scalar_fn)
+        .is_some_and(|octet_length| can_octet_length_be_pushed_down(octet_length, schema))
 }
 
-// TODO(adam): Replace with `DataType::is_decimal` once its released.
-fn is_decimal(dt: &DataType) -> bool {
-    matches!(
-        dt,
-        DataType::Decimal32(_, _)
-            | DataType::Decimal64(_, _)
-            | DataType::Decimal128(_, _)
-            | DataType::Decimal256(_, _)
-    )
+fn can_octet_length_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool {
+    let [input] = scalar_fn.args() else {
+        return false;
+    };
+
+    input.data_type(schema).as_ref().is_ok_and(|data_type| {
+        let dt = if let DataType::Dictionary(_, value_type) = data_type {
+            value_type.as_ref()
+        } else {
+            data_type
+        };
+
+        dt.is_binary() || dt.is_string()
+    }) && can_be_pushed_down_impl(input, schema)
 }
 
 #[cfg(test)]
@@ -553,7 +577,9 @@ mod tests {
     use datafusion::arrow::array::AsArray;
     use datafusion::arrow::datatypes::Int32Type;
     use datafusion_common::ScalarValue;
+    use datafusion_common::config::ConfigOptions;
     use datafusion_expr::Operator as DFOperator;
+    use datafusion_expr::ScalarUDF;
     use datafusion_physical_expr::PhysicalExpr;
     use datafusion_physical_plan::expressions as df_expr;
     use insta::assert_snapshot;
@@ -582,6 +608,18 @@ mod tests {
         ])
     }
 
+    fn octet_length_expr(input: Arc<dyn PhysicalExpr>, schema: &Schema) -> Arc<dyn PhysicalExpr> {
+        Arc::new(
+            ScalarFunctionExpr::try_new(
+                Arc::new(ScalarUDF::from(OctetLengthFunc::new())),
+                vec![input],
+                schema,
+                Arc::new(ConfigOptions::new()),
+            )
+            .unwrap(),
+        )
+    }
+
     #[test]
     fn test_make_vortex_predicate_empty() {
         let expr_convertor = DefaultExpressionConvertor::default();
@@ -711,6 +749,23 @@ mod tests {
         );
     }
 
+    #[rstest]
+    fn test_expr_from_df_octet_length(test_schema: Schema) {
+        let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc<dyn PhysicalExpr>;
+        let octet_length = octet_length_expr(expr, &test_schema);
+
+        let result = DefaultExpressionConvertor::default()
+            .convert(octet_length.as_ref())
+            .unwrap();
+
+        assert_snapshot!(result.display_tree().to_string(), @r"
+        vortex.cast(i32?)
+        └── input: vortex.byte_length()
+            └── input: vortex.get_item(name)
+                └── input: vortex.root()
+        ");
+    }
+
     #[rstest]
     // Supported types
     #[case::null(DataType::Null, true)]
@@ -865,6 +920,28 @@ mod tests {
         assert!(!can_be_pushed_down_impl(&like_expr, &test_schema));
     }
 
+    #[rstest]
+    fn test_can_be_pushed_down_octet_length_supported(test_schema: Schema) {
+        let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc<dyn PhysicalExpr>;
+        let octet_length = octet_length_expr(expr, &test_schema);
+
+        assert!(can_be_pushed_down_impl(&octet_length, &test_schema));
+    }
+
+    #[rstest]
+    fn test_can_be_pushed_down_octet_length_unsupported_operand(test_schema: Schema) {
+        let expr = Arc::new(df_expr::Column::new("unsupported_list", 5)) as Arc<dyn PhysicalExpr>;
+        let octet_length = Arc::new(ScalarFunctionExpr::new(
+            "octet_length",
+            Arc::new(ScalarUDF::from(OctetLengthFunc::new())),
+            vec![expr],
+            Arc::new(Field::new("octet_length", DataType::Int32, true)),
+            Arc::new(ConfigOptions::new()),
+        )) as Arc<dyn PhysicalExpr>;
+
+        assert!(!can_be_pushed_down_impl(&octet_length, &test_schema));
+    }
+
     // https://github.com/vortex-data/vortex/issues/6211
     #[tokio::test]
     async fn test_cast_int_to_string() -> anyhow::Result<()> {

vortex-datafusion/src/persistent/format.rs

@@ -146,6 +146,7 @@ config_namespace! {
     ///
     /// let factory = VortexFormatFactory::new().with_options(VortexTableOptions {
     ///     projection_pushdown: true,
+    ///     predicate_pushdown: true,
     ///     scan_concurrency: Some(8),
     ///     ..Default::default()
     /// });
@@ -165,6 +166,12 @@ config_namespace! {
         /// the scan. When disabled, Vortex reads only the referenced columns and
         /// all expressions are evaluated after the scan.
         pub projection_pushdown: bool, default = false
+        /// Whether to enable predicate pushdown into the underlying Vortex scan.
+        ///
+        /// When enabled, supported filters are evaluated during the scan. When
+        /// disabled, DataFusion evaluates filters after the scan, while
+        /// `VortexSource` can still use the full predicate for file pruning.
+        pub predicate_pushdown: bool, default = true
         /// The intra-partition scan concurrency, controlling the number of row splits to process
         /// concurrently per-thread within each file.
         ///
@@ -198,6 +205,7 @@ impl Eq for VortexTableOptions {}
 ///
 /// let factory = Arc::new(VortexFormatFactory::new().with_options(VortexTableOptions {
 ///     projection_pushdown: true,
+///     predicate_pushdown: true,
 ///     ..Default::default()
 /// }));
 ///
@@ -263,6 +271,7 @@ impl VortexFormatFactory {
     ///
     /// let factory = VortexFormatFactory::new().with_options(VortexTableOptions {
     ///     projection_pushdown: true,
+    ///     predicate_pushdown: true,
     ///     ..Default::default()
     /// });
     /// # let _ = factory;
@@ -617,14 +626,9 @@ impl FileFormat for VortexFormat {
     }
 
     fn file_source(&self, table_schema: TableSchema) -> Arc<dyn FileSource> {
-        let mut source = VortexSource::new(table_schema, self.session.clone())
-            .with_projection_pushdown(self.opts.projection_pushdown);
-
-        if let Some(scan_concurrency) = self.opts.scan_concurrency {
-            source = source.with_scan_concurrency(scan_concurrency);
-        }
-
-        Arc::new(source) as _
+        Arc::new(
+            VortexSource::new(table_schema, self.session.clone()).with_options(self.opts.clone()),
+        ) as _
     }
 }
 
@@ -682,7 +686,7 @@ mod tests {
                 (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \
                 STORED AS vortex \
                 LOCATION 'table/' \
-                OPTIONS( footer_initial_read_size_bytes '12345', scan_concurrency '3' );",
+                OPTIONS( footer_initial_read_size_bytes '12345', predicate_pushdown 'false', scan_concurrency '3' );",
             )
             .await?
             .collect()
@@ -699,4 +703,24 @@ mod tests {
         let format = VortexFormat::new_with_options(VortexSession::default(), opts);
         assert_eq!(format.options().footer_initial_read_size_bytes, 12345);
     }
+
+    #[test]
+    fn format_plumbs_source_options() -> anyhow::Result<()> {
+        let opts = VortexTableOptions {
+            projection_pushdown: true,
+            predicate_pushdown: false,
+            scan_concurrency: Some(3),
+            ..Default::default()
+        };
+        let format = VortexFormat::new_with_options(VortexSession::default(), opts.clone());
+        let table_schema = TableSchema::from_file_schema(Arc::new(Schema::empty()));
+
+        let source = format.file_source(table_schema);
+        let source = source
+            .downcast_ref::<VortexSource>()
+            .ok_or_else(|| anyhow::anyhow!("expected VortexSource"))?;
+
+        assert_eq!(source.options(), &opts);
+        Ok(())
+    }
 }