[PERF] Enable Predicates in Parquet Reader

daft/daft.pyi

@@ -570,6 +570,7 @@ def read_parquet(
     start_offset: int | None = None,
     num_rows: int | None = None,
     row_groups: list[int] | None = None,
+    predicate: PyExpr | None = None,
     io_config: IOConfig | None = None,
     multithreaded_io: bool | None = None,
     coerce_int96_timestamp_unit: PyTimeUnit | None = None,
@@ -580,6 +581,7 @@ def read_parquet_bulk(
     start_offset: int | None = None,
     num_rows: int | None = None,
     row_groups: list[list[int] | None] | None = None,
+    predicate: PyExpr | None = None,
     io_config: IOConfig | None = None,
     num_parallel_tasks: int | None = 128,
     multithreaded_io: bool | None = None,
@@ -963,6 +965,7 @@ class PyMicroPartition:
         start_offset: int | None = None,
         num_rows: int | None = None,
         row_groups: list[int] | None = None,
+        predicate: PyExpr | None = None,
         io_config: IOConfig | None = None,
         multithreaded_io: bool | None = None,
         coerce_int96_timestamp_unit: PyTimeUnit = PyTimeUnit.nanoseconds(),
@@ -975,6 +978,7 @@ class PyMicroPartition:
         start_offset: int | None = None,
         num_rows: int | None = None,
         row_groups: list[list[int] | None] | None = None,
+        predicate: PyExpr | None = None,
         io_config: IOConfig | None = None,
         num_parallel_tasks: int | None = None,
         multithreaded_io: bool | None = None,

daft/table/micropartition.py

@@ -313,6 +313,7 @@ def read_parquet(
         start_offset: int | None = None,
         num_rows: int | None = None,
         row_groups: list[int] | None = None,
+        predicate: Expression | None = None,
         io_config: IOConfig | None = None,
         multithreaded_io: bool | None = None,
         coerce_int96_timestamp_unit: TimeUnit = TimeUnit.ns(),
@@ -324,6 +325,7 @@ def read_parquet(
                 start_offset,
                 num_rows,
                 row_groups,
+                predicate._expr if predicate is not None else None,
                 io_config,
                 multithreaded_io,
                 coerce_int96_timestamp_unit._timeunit,
@@ -338,6 +340,7 @@ def read_parquet_bulk(
         start_offset: int | None = None,
         num_rows: int | None = None,
         row_groups_per_path: list[list[int] | None] | None = None,
+        predicate: Expression | None = None,
         io_config: IOConfig | None = None,
         num_parallel_tasks: int | None = 128,
         multithreaded_io: bool | None = None,
@@ -350,6 +353,7 @@ def read_parquet_bulk(
                 start_offset,
                 num_rows,
                 row_groups_per_path,
+                predicate._expr if predicate is not None else None,
                 io_config,
                 num_parallel_tasks,
                 multithreaded_io,

daft/table/table.py

@@ -392,6 +392,7 @@ def read_parquet(
         start_offset: int | None = None,
         num_rows: int | None = None,
         row_groups: list[int] | None = None,
+        predicate: Expression | None = None,
         io_config: IOConfig | None = None,
         multithreaded_io: bool | None = None,
         coerce_int96_timestamp_unit: TimeUnit = TimeUnit.ns(),
@@ -403,6 +404,7 @@ def read_parquet(
                 start_offset=start_offset,
                 num_rows=num_rows,
                 row_groups=row_groups,
+                predicate=predicate._expr if predicate is not None else None,
                 io_config=io_config,
                 multithreaded_io=multithreaded_io,
                 coerce_int96_timestamp_unit=coerce_int96_timestamp_unit._timeunit,
@@ -417,6 +419,7 @@ def read_parquet_bulk(
         start_offset: int | None = None,
         num_rows: int | None = None,
         row_groups_per_path: list[list[int] | None] | None = None,
+        predicate: Expression | None = None,
         io_config: IOConfig | None = None,
         num_parallel_tasks: int | None = 128,
         multithreaded_io: bool | None = None,
@@ -428,6 +431,7 @@ def read_parquet_bulk(
             start_offset=start_offset,
             num_rows=num_rows,
             row_groups=row_groups_per_path,
+            predicate=predicate._expr if predicate is not None else None,
             io_config=io_config,
             num_parallel_tasks=num_parallel_tasks,
             multithreaded_io=multithreaded_io,

src/daft-micropartition/src/micropartition.rs

@@ -8,6 +8,7 @@ use common_error::DaftResult;
 use daft_core::schema::{Schema, SchemaRef};
 
 use daft_csv::{CsvConvertOptions, CsvParseOptions, CsvReadOptions};
+use daft_dsl::ExprRef;
 use daft_json::{JsonConvertOptions, JsonParseOptions, JsonReadOptions};
 use daft_parquet::read::{
     read_parquet_bulk, read_parquet_metadata_bulk, ParquetSchemaInferenceOptions,
@@ -120,6 +121,7 @@ fn materialize_scan_task(
                         None,
                         scan_task.pushdowns.limit,
                         row_groups,
+                        scan_task.pushdowns.filters.clone(),
                         io_client.clone(),
                         io_stats,
                         8,
@@ -410,6 +412,7 @@ impl MicroPartition {
                     None,
                     scan_task.pushdowns.limit,
                     row_groups,
+                    scan_task.pushdowns.filters.clone(),
                     cfg.io_config
                         .clone()
                         .map(|c| Arc::new(c.clone()))
@@ -707,6 +710,7 @@ pub(crate) fn read_parquet_into_micropartition(
     start_offset: Option<usize>,
     num_rows: Option<usize>,
     row_groups: Option<Vec<Option<Vec<i64>>>>,
+    predicate: Option<ExprRef>,
     io_config: Arc<IOConfig>,
     io_stats: Option<IOStatsRef>,
     num_parallel_tasks: usize,
@@ -720,6 +724,46 @@ pub(crate) fn read_parquet_into_micropartition(
     // Run the required I/O to retrieve all the Parquet FileMetaData
     let runtime_handle = daft_io::get_runtime(multithreaded_io)?;
     let io_client = daft_io::get_io_client(multithreaded_io, io_config.clone())?;
+
+    if let Some(predicate) = predicate {
+        // We have a predicate, so we will perform eager read only reading what row groups we need.
+        let all_tables = read_parquet_bulk(
+            uris,
+            columns,
+            None,
+            num_rows,
+            row_groups,
+            Some(predicate.clone()),
+            io_client,
+            io_stats,
+            num_parallel_tasks,
+            runtime_handle,
+            schema_infer_options,
+        )?;
+
+        let unioned_schema = all_tables
+            .iter()
+            .map(|t| t.schema.clone())
+            .try_reduce(|l, r| DaftResult::Ok(l.union(&r)?.into()))?;
+        let full_daft_schema = unioned_schema.expect("we need at least 1 schema");
+        // Hack to avoid to owned schema
+        let full_daft_schema = Schema {
+            fields: full_daft_schema.fields.clone(),
+        };
+        let pruned_daft_schema = prune_fields_from_schema(full_daft_schema, columns)?;
+
+        let all_tables = all_tables
+            .into_iter()
+            .map(|t| t.cast_to_schema(&pruned_daft_schema))
+            .collect::<DaftResult<Vec<_>>>()?;
+        // TODO: we can pass in stats here to optimize downstream workloads such as join
+        return Ok(MicroPartition::new_loaded(
+            Arc::new(pruned_daft_schema),
+            all_tables.into(),
+            None,
+        ));
+    }
+
     let meta_io_client = io_client.clone();
     let meta_io_stats = io_stats.clone();
     let metadata = runtime_handle.block_on(async move {
@@ -847,6 +891,7 @@ pub(crate) fn read_parquet_into_micropartition(
             start_offset,
             num_rows,
             row_groups,
+            None,
             io_client,
             io_stats,
             num_parallel_tasks,

src/daft-micropartition/src/ops/filter.rs

@@ -8,15 +8,15 @@ use crate::{micropartition::MicroPartition, DaftCoreComputeSnafu};
 use daft_stats::TruthValue;
 
 impl MicroPartition {
-    pub fn filter(&self, predicate: &[Expr]) -> DaftResult<Self> {
+    pub fn filter<E: AsRef<Expr>>(&self, predicate: &[E]) -> DaftResult<Self> {
         let io_stats = IOStatsContext::new("MicroPartition::filter");
         if predicate.is_empty() {
             return Ok(Self::empty(Some(self.schema.clone())));
         }
         if let Some(statistics) = &self.statistics {
             let folded_expr = predicate
                 .iter()
-                .cloned()
+                .map(|e| e.as_ref().clone())
                 .reduce(|a, b| a.and(&b))
                 .expect("should have at least 1 expr");
             let eval_result = statistics.eval_expression(&folded_expr)?;

src/daft-micropartition/src/python.rs

@@ -433,6 +433,7 @@ impl PyMicroPartition {
         start_offset: Option<usize>,
         num_rows: Option<usize>,
         row_groups: Option<Vec<i64>>,
+        predicate: Option<PyExpr>,
         io_config: Option<IOConfig>,
         multithreaded_io: Option<bool>,
         coerce_int96_timestamp_unit: Option<PyTimeUnit>,
@@ -451,6 +452,7 @@ impl PyMicroPartition {
                 start_offset,
                 num_rows,
                 row_groups.map(|rg| vec![Some(rg)]),
+                predicate.map(|e| e.expr.into()),
                 io_config,
                 Some(io_stats),
                 1,
@@ -470,6 +472,7 @@ impl PyMicroPartition {
         start_offset: Option<usize>,
         num_rows: Option<usize>,
         row_groups: Option<Vec<Option<Vec<i64>>>>,
+        predicate: Option<PyExpr>,
         io_config: Option<IOConfig>,
         num_parallel_tasks: Option<i64>,
         multithreaded_io: Option<bool>,
@@ -489,6 +492,7 @@ impl PyMicroPartition {
                 start_offset,
                 num_rows,
                 row_groups,
+                predicate.map(|e| e.expr.into()),
                 io_config,
                 Some(io_stats),
                 num_parallel_tasks.unwrap_or(128) as usize,

src/daft-parquet/Cargo.toml

@@ -5,6 +5,7 @@ async-stream = {workspace = true}
 bytes = {workspace = true}
 common-error = {path = "../common/error", default-features = false}
 daft-core = {path = "../daft-core", default-features = false}
+daft-dsl = {path = "../daft-dsl", default-features = false}
 daft-io = {path = "../daft-io", default-features = false}
 daft-stats = {path = "../daft-stats", default-features = false}
 daft-table = {path = "../daft-table", default-features = false}
@@ -24,7 +25,7 @@ tokio-util = {workspace = true}
 
 [features]
 default = ["python"]
-python = ["dep:pyo3", "dep:pyo3-log", "common-error/python", "daft-core/python", "daft-io/python", "daft-table/python", "daft-stats/python"]
+python = ["dep:pyo3", "dep:pyo3-log", "common-error/python", "daft-core/python", "daft-io/python", "daft-table/python", "daft-stats/python", "daft-dsl/python"]
 
 [package]
 edition = {workspace = true}