Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
139 changes: 108 additions & 31 deletions vortex-datafusion/src/convert/exprs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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 {
Expand All @@ -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 {
Expand Down Expand Up @@ -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 {
Expand All @@ -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)]
Expand All @@ -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;
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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)]
Expand Down Expand Up @@ -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<()> {
Expand Down
42 changes: 33 additions & 9 deletions vortex-datafusion/src/persistent/format.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
/// });
Expand All @@ -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.
///
Expand Down Expand Up @@ -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()
/// }));
///
Expand Down Expand Up @@ -263,6 +271,7 @@ impl VortexFormatFactory {
///
/// let factory = VortexFormatFactory::new().with_options(VortexTableOptions {
/// projection_pushdown: true,
/// predicate_pushdown: true,
/// ..Default::default()
/// });
/// # let _ = factory;
Expand Down Expand Up @@ -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 _
}
}

Expand Down Expand Up @@ -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()
Expand All @@ -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(())
}
}
Loading
Loading