vortex_datafusion/persistent/
source.rs1use std::fmt::Formatter;
5use std::ops::Range;
6use std::sync::Arc;
7use std::sync::Weak;
8
9use datafusion_common::Result as DFResult;
10use datafusion_common::config::ConfigOptions;
11use datafusion_datasource::TableSchema;
12use datafusion_datasource::file::FileSource;
13use datafusion_datasource::file_scan_config::FileScanConfig;
14use datafusion_datasource::file_stream::FileOpener;
15use datafusion_execution::cache::cache_manager::FileMetadataCache;
16use datafusion_physical_expr::EquivalenceProperties;
17use datafusion_physical_expr::PhysicalExprRef;
18use datafusion_physical_expr::PhysicalSortExpr;
19use datafusion_physical_expr::conjunction;
20use datafusion_physical_expr::projection::ProjectionExprs;
21use datafusion_physical_expr_adapter::DefaultPhysicalExprAdapterFactory;
22use datafusion_physical_expr_common::physical_expr::fmt_sql;
23use datafusion_physical_plan::DisplayFormatType;
24use datafusion_physical_plan::PhysicalExpr;
25use datafusion_physical_plan::SortOrderPushdownResult;
26use datafusion_physical_plan::filter_pushdown::FilterPushdownPropagation;
27use datafusion_physical_plan::filter_pushdown::PushedDown;
28use datafusion_physical_plan::filter_pushdown::PushedDownPredicate;
29use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
30use object_store::ObjectStore;
31use object_store::path::Path;
32use vortex::error::VortexExpect;
33use vortex::file::VORTEX_FILE_EXTENSION;
34use vortex::layout::LayoutReader;
35use vortex::metrics::DefaultMetricsRegistry;
36use vortex::metrics::MetricsRegistry;
37use vortex::session::VortexSession;
38use vortex_utils::aliases::dash_map::DashMap;
39
40use super::opener::VortexOpener;
41use crate::VortexTableOptions;
42use crate::convert::exprs::DefaultExpressionConvertor;
43use crate::convert::exprs::ExpressionConvertor;
44use crate::persistent::reader::DefaultVortexReaderFactory;
45use crate::persistent::reader::VortexReaderFactory;
46
47#[derive(Clone)]
183pub struct VortexSource {
184 pub(crate) session: VortexSession,
185 pub(crate) table_schema: TableSchema,
186 pub(crate) projection: ProjectionExprs,
187 pub(crate) full_predicate: Option<PhysicalExprRef>,
190 pub(crate) vortex_predicate: Option<PhysicalExprRef>,
193 pub(crate) batch_size: Option<usize>,
194 _unused_df_metrics: ExecutionPlanMetricsSet,
195 layout_readers: Arc<DashMap<Path, Weak<dyn LayoutReader>>>,
199 natural_split_ranges: Arc<DashMap<Path, Arc<[Range<u64>]>>>,
201 expression_convertor: Arc<dyn ExpressionConvertor>,
202 pub(crate) vortex_reader_factory: Option<Arc<dyn VortexReaderFactory>>,
203 pub(crate) ordered: bool,
204 vx_metrics_registry: Arc<dyn MetricsRegistry>,
205 file_metadata_cache: Option<Arc<dyn FileMetadataCache>>,
206 options: VortexTableOptions,
208}
209
210impl VortexSource {
211 pub fn new(table_schema: TableSchema, session: VortexSession) -> Self {
220 let full_schema = table_schema.table_schema();
221 let indices = (0..full_schema.fields().len()).collect::<Vec<_>>();
222 let projection = ProjectionExprs::from_indices(&indices, full_schema);
223 let expression_convertor = Arc::new(DefaultExpressionConvertor::new(session.clone()));
224
225 Self {
226 session,
227 table_schema,
228 projection,
229 full_predicate: None,
230 vortex_predicate: None,
231 batch_size: None,
232 _unused_df_metrics: Default::default(),
233 layout_readers: Arc::new(DashMap::default()),
234 natural_split_ranges: Arc::new(DashMap::default()),
235 expression_convertor,
236 vortex_reader_factory: None,
237 vx_metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
238 file_metadata_cache: None,
239 ordered: false,
240 options: VortexTableOptions::default(),
241 }
242 }
243
244 pub fn with_projection_pushdown(mut self, enabled: bool) -> Self {
250 self.options.projection_pushdown = enabled;
251 self
252 }
253
254 pub fn with_predicate_pushdown(mut self, enabled: bool) -> Self {
259 self.options.predicate_pushdown = enabled;
260 self
261 }
262
263 pub fn with_expression_convertor(
269 mut self,
270 expr_convertor: Arc<dyn ExpressionConvertor>,
271 ) -> Self {
272 self.expression_convertor = expr_convertor;
273 self
274 }
275
276 pub fn with_vortex_reader_factory(
283 mut self,
284 vortex_reader_factory: Arc<dyn VortexReaderFactory>,
285 ) -> Self {
286 self.vortex_reader_factory = Some(vortex_reader_factory);
287 self
288 }
289
290 pub fn metrics_registry(&self) -> &Arc<dyn MetricsRegistry> {
296 &self.vx_metrics_registry
297 }
298
299 pub fn with_file_metadata_cache(
301 mut self,
302 file_metadata_cache: Arc<dyn FileMetadataCache>,
303 ) -> Self {
304 self.file_metadata_cache = Some(file_metadata_cache);
305 self
306 }
307
308 pub fn with_scan_concurrency(mut self, scan_concurrency: usize) -> Self {
312 self.options.scan_concurrency = Some(scan_concurrency);
313 self
314 }
315
316 pub fn options(&self) -> &VortexTableOptions {
318 &self.options
319 }
320
321 pub fn with_options(mut self, opts: VortexTableOptions) -> Self {
323 self.options = opts;
324 self
325 }
326
327 pub fn predicate(&self) -> Option<&Arc<dyn PhysicalExpr>> {
329 self.vortex_predicate.as_ref()
330 }
331
332 fn create_vortex_opener(
333 &self,
334 object_store: Arc<dyn ObjectStore>,
335 base_config: &FileScanConfig,
336 partition: usize,
337 ) -> DFResult<VortexOpener> {
338 let batch_size = self
339 .batch_size
340 .vortex_expect("batch_size must be supplied to VortexSource");
341
342 let expr_adapter_factory = base_config
343 .expr_adapter_factory
344 .clone()
345 .unwrap_or_else(|| Arc::new(DefaultPhysicalExprAdapterFactory));
346
347 let vortex_reader_factory = self
348 .vortex_reader_factory
349 .clone()
350 .unwrap_or_else(|| Arc::new(DefaultVortexReaderFactory::new(object_store)));
351
352 let opener = VortexOpener {
353 partition,
354 session: self.session.clone(),
355 vortex_reader_factory,
356 projection: self.projection.clone(),
357 filter: self.vortex_predicate.clone(),
358 file_pruning_predicate: self.full_predicate.clone(),
359 expr_adapter_factory,
360 table_schema: self.table_schema.clone(),
361 batch_size,
362 limit: base_config.limit.map(|l| l as u64),
363 metrics_registry: Arc::clone(&self.vx_metrics_registry),
364 layout_readers: Arc::clone(&self.layout_readers),
365 natural_split_ranges: Arc::clone(&self.natural_split_ranges),
366 has_output_ordering: !base_config.output_ordering.is_empty() || self.ordered,
367 expression_convertor: Arc::clone(&self.expression_convertor),
368 file_metadata_cache: self.file_metadata_cache.clone(),
369 projection_pushdown: self.options.projection_pushdown,
370 scan_concurrency: self.options.scan_concurrency,
371 };
372
373 Ok(opener)
374 }
375}
376
377impl FileSource for VortexSource {
378 fn create_file_opener(
379 &self,
380 object_store: Arc<dyn ObjectStore>,
381 base_config: &FileScanConfig,
382 partition: usize,
383 ) -> DFResult<Arc<dyn FileOpener>> {
384 Ok(Arc::new(self.create_vortex_opener(
385 object_store,
386 base_config,
387 partition,
388 )?))
389 }
390
391 fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> {
392 let mut source = self.clone();
393 source.batch_size = Some(batch_size);
394 Arc::new(source)
395 }
396
397 fn filter(&self) -> Option<Arc<dyn PhysicalExpr>> {
398 self.vortex_predicate.clone()
399 }
400
401 fn metrics(&self) -> &ExecutionPlanMetricsSet {
402 &self._unused_df_metrics
403 }
404
405 fn file_type(&self) -> &str {
406 VORTEX_FILE_EXTENSION
407 }
408
409 fn try_pushdown_sort(
410 &self,
411 order: &[PhysicalSortExpr],
412 eq_properties: &EquivalenceProperties,
413 ) -> DFResult<SortOrderPushdownResult<Arc<dyn FileSource>>> {
414 if order.is_empty() {
415 return Ok(SortOrderPushdownResult::Unsupported);
416 }
417
418 if eq_properties.ordering_satisfy(order.iter().cloned())? {
419 let mut this = self.clone();
420 this.ordered = true;
421
422 return Ok(SortOrderPushdownResult::Exact {
423 inner: Arc::new(this) as Arc<dyn FileSource>,
424 });
425 }
426
427 Ok(SortOrderPushdownResult::Unsupported)
428 }
429
430 fn fmt_extra(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
431 match t {
432 DisplayFormatType::Default | DisplayFormatType::Verbose => {
433 if let Some(predicate) = &self.vortex_predicate {
434 write!(f, ", predicate: {predicate}")?;
435 }
436 }
437 DisplayFormatType::TreeRender => {
439 if let Some(predicate) = &self.vortex_predicate {
440 writeln!(f, "predicate={}", fmt_sql(predicate.as_ref()))?;
441 };
442 }
443 }
444 Ok(())
445 }
446
447 fn supports_repartitioning(&self) -> bool {
448 true
449 }
450
451 fn try_pushdown_filters(
452 &self,
453 filters: Vec<Arc<dyn PhysicalExpr>>,
454 _config: &ConfigOptions,
455 ) -> DFResult<FilterPushdownPropagation<Arc<dyn FileSource>>> {
456 if filters.is_empty() {
457 return Ok(FilterPushdownPropagation::with_parent_pushdown_result(
458 vec![],
459 ));
460 }
461
462 let mut source = self.clone();
463
464 source.full_predicate = match source.full_predicate {
467 Some(predicate) => Some(conjunction(
468 std::iter::once(predicate).chain(filters.clone()),
469 )),
470 None => Some(conjunction(filters.clone())),
471 };
472
473 if !source.options.predicate_pushdown {
474 return Ok(FilterPushdownPropagation::with_parent_pushdown_result(vec![
475 PushedDown::No;
476 filters.len()
477 ])
478 .with_updated_node(Arc::new(source) as _));
479 }
480
481 let supported_filters = filters
482 .into_iter()
483 .map(|expr| {
484 if self
485 .expression_convertor
486 .can_be_pushed_down(&expr, self.table_schema.file_schema())
487 {
488 PushedDownPredicate::supported(expr)
489 } else {
490 PushedDownPredicate::unsupported(expr)
491 }
492 })
493 .collect::<Vec<_>>();
494
495 if supported_filters
496 .iter()
497 .all(|p| matches!(p.discriminant, PushedDown::No))
498 {
499 return Ok(FilterPushdownPropagation::with_parent_pushdown_result(
500 vec![PushedDown::No; supported_filters.len()],
501 )
502 .with_updated_node(Arc::new(source) as _));
503 }
504
505 let supported = supported_filters
506 .iter()
507 .filter_map(|p| match p.discriminant {
508 PushedDown::Yes => Some(&p.predicate),
509 PushedDown::No => None,
510 })
511 .cloned();
512
513 let predicate = match source.vortex_predicate {
514 Some(predicate) => conjunction(std::iter::once(predicate).chain(supported)),
515 None => conjunction(supported),
516 };
517
518 tracing::debug!(%predicate, "Saving predicate");
519
520 source.vortex_predicate = Some(predicate);
521
522 Ok(FilterPushdownPropagation::with_parent_pushdown_result(
523 supported_filters.iter().map(|f| f.discriminant).collect(),
524 )
525 .with_updated_node(Arc::new(source) as _))
526 }
527
528 fn try_pushdown_projection(
529 &self,
530 projection: &ProjectionExprs,
531 ) -> DFResult<Option<Arc<dyn FileSource>>> {
532 let mut source = self.clone();
533 source.projection = self.projection.try_merge(projection)?;
534 Ok(Some(Arc::new(source)))
535 }
536
537 fn projection(&self) -> Option<&ProjectionExprs> {
538 Some(&self.projection)
539 }
540
541 fn table_schema(&self) -> &TableSchema {
542 &self.table_schema
543 }
544}
545
546#[cfg(test)]
547mod tests {
548 use arrow_schema::DataType;
549 use arrow_schema::Field;
550 use arrow_schema::Schema;
551 use datafusion_common::ScalarValue;
552 use datafusion_common::config::ConfigOptions;
553 use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
554 use datafusion_execution::object_store::ObjectStoreUrl;
555 use datafusion_expr::Operator;
556 use datafusion_expr::ScalarUDF;
557 use datafusion_functions::string::octet_length::OctetLengthFunc;
558 use datafusion_physical_expr::ScalarFunctionExpr;
559 use datafusion_physical_expr::expressions as df_expr;
560 use datafusion_physical_expr::expressions::Column;
561 use object_store::memory::InMemory;
562 use vortex::VortexSessionDefault;
563
564 use super::*;
565 use crate::convert::exprs::ProcessedProjection;
566
567 struct TrackingExpressionConvertor {
568 inner: DefaultExpressionConvertor,
569 }
570
571 impl ExpressionConvertor for TrackingExpressionConvertor {
572 fn can_be_pushed_down(&self, expr: &PhysicalExprRef, schema: &Schema) -> bool {
573 self.inner.can_be_pushed_down(expr, schema)
574 }
575
576 fn convert(&self, expr: &dyn PhysicalExpr) -> DFResult<vortex::expr::Expression> {
577 self.inner.convert(expr)
578 }
579
580 fn split_projection(
581 &self,
582 source_projection: ProjectionExprs,
583 input_schema: &Schema,
584 output_schema: &Schema,
585 ) -> DFResult<ProcessedProjection> {
586 self.inner
587 .split_projection(source_projection, input_schema, output_schema)
588 }
589
590 fn no_pushdown_projection(
591 &self,
592 source_projection: ProjectionExprs,
593 input_schema: &Schema,
594 ) -> DFResult<ProcessedProjection> {
595 self.inner
596 .no_pushdown_projection(source_projection, input_schema)
597 }
598 }
599
600 fn sort_column(name: &str, index: usize) -> PhysicalSortExpr {
601 let expr: PhysicalExprRef = Arc::new(Column::new(name, index));
602 PhysicalSortExpr::new_default(expr)
603 }
604
605 fn sort_test_schema() -> Arc<Schema> {
606 Arc::new(Schema::new(vec![
607 Field::new("a", DataType::Int32, false),
608 Field::new("b", DataType::Int32, false),
609 ]))
610 }
611
612 fn sort_test_source(schema: Arc<Schema>) -> VortexSource {
613 VortexSource::new(
614 TableSchema::from_file_schema(schema),
615 VortexSession::default(),
616 )
617 }
618
619 fn octet_length_filter(schema: &Schema) -> PhysicalExprRef {
620 let name = Arc::new(Column::new("name", 0)) as PhysicalExprRef;
621 let octet_length = Arc::new(
622 ScalarFunctionExpr::try_new(
623 Arc::new(ScalarUDF::from(OctetLengthFunc::new())),
624 vec![name],
625 schema,
626 Arc::new(ConfigOptions::new()),
627 )
628 .unwrap(),
629 ) as PhysicalExprRef;
630 let one = Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1)))) as PhysicalExprRef;
631
632 Arc::new(df_expr::BinaryExpr::new(octet_length, Operator::Gt, one)) as PhysicalExprRef
633 }
634
635 fn assert_ordered_source(inner: Arc<dyn FileSource>) -> anyhow::Result<()> {
636 let source = inner
637 .downcast_ref::<VortexSource>()
638 .ok_or_else(|| anyhow::anyhow!("expected VortexSource"))?;
639
640 assert!(source.ordered);
641 Ok(())
642 }
643
644 #[test]
645 fn try_pushdown_sort_returns_exact_when_ordering_is_satisfied() -> anyhow::Result<()> {
646 let schema = sort_test_schema();
647 let source = sort_test_source(Arc::clone(&schema));
648 let order = vec![sort_column("a", 0), sort_column("b", 1)];
649 let eq_properties = EquivalenceProperties::new_with_orderings(schema, [order.clone()]);
650
651 let result = source.try_pushdown_sort(&order, &eq_properties)?;
652
653 match result {
654 SortOrderPushdownResult::Exact { inner } => assert_ordered_source(inner)?,
655 SortOrderPushdownResult::Inexact { .. } | SortOrderPushdownResult::Unsupported => {
656 anyhow::bail!("expected exact sort pushdown")
657 }
658 }
659 assert!(!source.ordered);
660 Ok(())
661 }
662
663 #[test]
664 fn create_vortex_opener_preserves_expression_convertor() -> anyhow::Result<()> {
665 let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
666 let expression_convertor = Arc::new(TrackingExpressionConvertor {
667 inner: DefaultExpressionConvertor::default(),
668 }) as Arc<dyn ExpressionConvertor>;
669
670 let mut source = VortexSource::new(
671 TableSchema::from_file_schema(file_schema),
672 VortexSession::default(),
673 )
674 .with_expression_convertor(Arc::clone(&expression_convertor));
675 source.batch_size = Some(100);
676
677 let config = FileScanConfigBuilder::new(
678 ObjectStoreUrl::local_filesystem(),
679 Arc::new(source.clone()),
680 )
681 .build();
682
683 let opener = source.create_vortex_opener(
684 Arc::new(InMemory::new()) as Arc<dyn ObjectStore>,
685 &config,
686 0,
687 )?;
688
689 assert!(Arc::ptr_eq(
690 &opener.expression_convertor,
691 &expression_convertor
692 ));
693 Ok(())
694 }
695
696 #[test]
697 fn try_pushdown_filters_accepts_octet_length() -> anyhow::Result<()> {
698 let schema = Arc::new(Schema::new(vec![Field::new("name", DataType::Utf8, false)]));
699 let source = sort_test_source(Arc::clone(&schema));
700 let filter = octet_length_filter(&schema);
701
702 let result = source.try_pushdown_filters(vec![filter], &ConfigOptions::new())?;
703
704 assert!(matches!(result.filters.as_slice(), [PushedDown::Yes]));
705 let updated_source = result
706 .updated_node
707 .ok_or_else(|| anyhow::anyhow!("expected updated VortexSource"))?
708 .downcast_ref::<VortexSource>()
709 .ok_or_else(|| anyhow::anyhow!("expected VortexSource"))?
710 .clone();
711 assert!(updated_source.vortex_predicate.is_some());
712 Ok(())
713 }
714
715 #[test]
716 fn try_pushdown_filters_respects_disabled_predicate_pushdown() -> anyhow::Result<()> {
717 let schema = Arc::new(Schema::new(vec![Field::new("name", DataType::Utf8, false)]));
718 let source = sort_test_source(Arc::clone(&schema)).with_predicate_pushdown(false);
719 let filter = octet_length_filter(&schema);
720
721 let result = source.try_pushdown_filters(vec![filter], &ConfigOptions::new())?;
722
723 assert!(matches!(result.filters.as_slice(), [PushedDown::No]));
724 let updated_source = result
725 .updated_node
726 .ok_or_else(|| anyhow::anyhow!("expected updated VortexSource"))?
727 .downcast_ref::<VortexSource>()
728 .ok_or_else(|| anyhow::anyhow!("expected VortexSource"))?
729 .clone();
730 assert!(updated_source.full_predicate.is_some());
731 assert!(updated_source.vortex_predicate.is_none());
732 Ok(())
733 }
734}