Implement filter creation

[maximedion2]

As of now, the filter push down logic is implemented in the zarr stream, however creating a filter from a physical expression (e.g. WHERE a > 2 AND b <5) is not yet implemented. Basically we need to implement this method, https://github.com/datafusion-contrib/arrow-zarr/blob/zarrs_refactor/src/zarr_store_opener/filter.rs#L58.

Now I remember datafusion had some helpers for that, but they seemed to be a bit too complicated for the use case here, and for my first pass at this I had copied some code over (it was not public in datafusion) and re-implemented a simplified version in this repo. But that was a while ago, I need to check in the more recent datafusion version, and I need to reconsider using the helpers datafusion provides out of the box.

One this is implemented, we need to add test cases everywhere, should be pretty easy, we can use the same test data I've been creating on-the-fly in test cases, and add a lat/lon filter.

[kylebarron]

implement this method, zarrs_refactor/src/zarr_store_opener/filter.rs#L58.

Do you have a permalink for that? It looks like the line of code got messed up after later commits.

Now I remember datafusion had some helpers for that

Do you know what these were? I thought datafusion only had format-specific implementations of these and not totally generic helpers? I thought the TableProvider implementation just gives you a list of Expr that you can downcast on?

[maximedion2]

implement this method, zarrs_refactor/src/zarr_store_opener/filter.rs#L58.

Do you have a permalink for that? It looks like the line of code got messed up after later commits.

Now I remember datafusion had some helpers for that

Do you know what these were? I thought datafusion only had format-specific implementations of these and not totally generic helpers? I thought the TableProvider implementation just gives you a list of Expr that you can downcast on?

Oh yeah, this should point to the right line

arrow-zarr/src/zarr_store_opener/filter.rs

Line 57 in 961f805

pub fn new(_physical_expr: Arc<dyn PhysicalExpr>) -> Self {

, sorry about that.

Yes so the execute method on the TableProvider trait takes in a &[Expr] for the filters, you're on your own after that. However, last time I looked at how it was done for parquet, and iirc I copied the helpers to this repo, simplified them a bit and used that. The code was from a few datafusion versions ago, but here's what I had done,

arrow-zarr/src/datafusion/helpers.rs

Lines 48 to 269 in 5bcdea1

	pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool {
	let mut is_applicable = true;
	expr.apply(&mut |expr| match expr {
	Expr::Column(datafusion_common::Column { ref name, .. }) => {
	is_applicable &= col_names.contains(name);
	if is_applicable {
	Ok(VisitRecursion::Skip)
	} else {
	Ok(VisitRecursion::Stop)
	}
	}
	Expr::Literal(_)
	| Expr::Alias(_)
	| Expr::OuterReferenceColumn(_, _)
	| Expr::ScalarVariable(_, _)
	| Expr::Not(_)
	| Expr::IsNotNull(_)
	| Expr::IsNull(_)
	| Expr::IsTrue(_)
	| Expr::IsFalse(_)
	| Expr::IsUnknown(_)
	| Expr::IsNotTrue(_)
	| Expr::IsNotFalse(_)
	| Expr::IsNotUnknown(_)
	| Expr::Negative(_)
	| Expr::Cast { .. }
	| Expr::TryCast { .. }
	| Expr::BinaryExpr { .. }
	| Expr::Between { .. }
	| Expr::Like { .. }
	| Expr::SimilarTo { .. }
	| Expr::InList { .. }
	| Expr::Exists { .. }
	| Expr::InSubquery(_)
	| Expr::ScalarSubquery(_)
	| Expr::GetIndexedField { .. }
	| Expr::GroupingSet(_)
	| Expr::Case { .. } => Ok(VisitRecursion::Continue),
	Expr::ScalarFunction(scalar_function) => match &scalar_function.func_def {
	ScalarFunctionDefinition::BuiltIn(fun) => match fun.volatility() {
	Volatility::Immutable => Ok(VisitRecursion::Continue),
	Volatility::Stable | Volatility::Volatile => {
	is_applicable = false;
	Ok(VisitRecursion::Stop)
	}
	},
	ScalarFunctionDefinition::UDF(fun) => match fun.signature().volatility {
	Volatility::Immutable => Ok(VisitRecursion::Continue),
	Volatility::Stable | Volatility::Volatile => {
	is_applicable = false;
	Ok(VisitRecursion::Stop)
	}
	},
	ScalarFunctionDefinition::Name(_) => {
	internal_err!("Function `Expr` with name should be resolved.")
	}
	},
	Expr::AggregateFunction { .. }
	| Expr::Sort { .. }
	| Expr::WindowFunction { .. }
	| Expr::Wildcard { .. }
	| Expr::Unnest { .. }
	| Expr::Placeholder(_) => {
	is_applicable = false;
	Ok(VisitRecursion::Stop)
	}
	})
	.unwrap();
	is_applicable
	}
	// Below is all the logic necessary (I think) to convert a PhysicalExpr into a ZarrChunkFilter.
	// The logic is mostly copied from datafusion, and is simplified here for the zarr use case.
	pub struct ZarrFilterCandidate {
	expr: Arc<dyn PhysicalExpr>,
	projection: Vec<usize>,
	}
	struct ZarrFilterCandidateBuilder<'a> {
	expr: Arc<dyn PhysicalExpr>,
	file_schema: &'a Schema,
	required_column_indices: BTreeSet<usize>,
	projected_columns: bool,
	}
	impl<'a> ZarrFilterCandidateBuilder<'a> {
	pub fn new(expr: Arc<dyn PhysicalExpr>, file_schema: &'a Schema) -> Self {
	Self {
	expr,
	file_schema,
	required_column_indices: BTreeSet::default(),
	projected_columns: false,
	}
	}
	pub fn build(mut self) -> DataFusionResult<Option<ZarrFilterCandidate>> {
	let expr = self.expr.clone().rewrite(&mut self)?;
	// if we are dealing with a projected column, which here means it's
	// a partitioned column, we don't produce a filter for it.
	if self.projected_columns {
	return Ok(None);
	}
	Ok(Some(ZarrFilterCandidate {
	expr,
	projection: self.required_column_indices.into_iter().collect(),
	}))
	}
	}
	impl TreeNodeRewriter for ZarrFilterCandidateBuilder<'_> {
	type N = Arc<dyn PhysicalExpr>;
	fn pre_visit(&mut self, node: &Arc<dyn PhysicalExpr>) -> DataFusionResult<RewriteRecursion> {
	if let Some(column) = node.as_any().downcast_ref::<Column>() {
	if let Ok(idx) = self.file_schema.index_of(column.name()) {
	self.required_column_indices.insert(idx);
	} else {
	// set the flag is we detect that the column is not in the file schema. for the
	// zarr implementation, this would mean that the column is actually a partitioned
	// column, and we shouldn't be pushing down a filter for it.
	// TODO handle cases where a filter contains a column that doesn't exist (not even
	// as a partition).
	self.projected_columns = true;
	return Ok(RewriteRecursion::Stop);
	}
	}
	Ok(RewriteRecursion::Continue)
	}
	fn mutate(&mut self, expr: Arc<dyn PhysicalExpr>) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
	Ok(expr)
	}
	}
	#[derive(Clone)]
	pub struct ZarrDatafusionArrowPredicate {
	physical_expr: Arc<dyn PhysicalExpr>,
	projection_mask: ZarrProjection,
	projection: Vec<String>,
	}
	impl ZarrDatafusionArrowPredicate {
	pub fn new(candidate: ZarrFilterCandidate, schema: &Schema) -> DataFusionResult<Self> {
	let cols: Vec<_> = candidate
	.projection
	.iter()
	.map(|idx| schema.field(*idx).name().to_string())
	.collect();
	let schema = Arc::new(schema.project(&candidate.projection)?);
	let physical_expr = reassign_predicate_columns(candidate.expr, &schema, true)?;
	Ok(Self {
	physical_expr,
	projection_mask: ZarrProjection::keep(cols.clone()),
	projection: cols,
	})
	}
	}
	impl ZarrArrowPredicate for ZarrDatafusionArrowPredicate {
	fn projection(&self) -> &ZarrProjection {
	&self.projection_mask
	}
	fn evaluate(&mut self, batch: &RecordBatch) -> Result<BooleanArray, ArrowError> {
	let index_projection = self
	.projection
	.iter()
	.map(|col| batch.schema().index_of(col))
	.collect::<Result<Vec<_>, _>>()?;
	let batch = batch.project(&index_projection[..])?;
	match self
	.physical_expr
	.evaluate(&batch)
	.and_then(|v| v.into_array(batch.num_rows()))
	{
	Ok(array) => {
	let bool_arr = as_boolean_array(&array)?.clone();
	Ok(bool_arr)
	}
	Err(e) => Err(ArrowError::ComputeError(format!(
	"Error evaluating filter predicate: {e:?}"
	))),
	}
	}
	}
	pub(crate) fn build_row_filter(
	expr: &Arc<dyn PhysicalExpr>,
	file_schema: &Schema,
	) -> DataFusionResult<Option<ZarrChunkFilter>> {
	let predicates = split_conjunction(expr);
	let candidates: Vec<ZarrFilterCandidate> = predicates
	.into_iter()
	.flat_map(|expr| {
	ZarrFilterCandidateBuilder::new(expr.clone(), file_schema)
	.build()
	.unwrap_or_default()
	})
	.collect();
	if candidates.is_empty() {
	Ok(None)
	} else {
	let mut filters: Vec<Box<dyn ZarrArrowPredicate>> = vec![];
	for candidate in candidates {
	let filter = ZarrDatafusionArrowPredicate::new(candidate, file_schema)?;
	filters.push(Box::new(filter));
	}
	let chunk_filter = ZarrChunkFilter::new(filters);
	Ok(Some(chunk_filter))
	}
	}

. It's not pretty but that was working.

Now that I think about it more, my comment was misleading, I didn't copy them over to simplify them, I copied them over because I had to, since all those functions and structs were not public, and then I decided oh I might as well simplify that a bit because I don't need all that logic.