andygrove commented on code in PR #2262:
URL:
https://github.com/apache/datafusion-ballista/pull/2262#discussion_r3744107892
##########
ballista/core/src/execution_plans/mod.rs:
##########
@@ -47,6 +48,7 @@ pub use
ordered_range_repartition::OrderedRangeRepartitionExec;
pub use partitioned_bounded_window_agg::PartitionedBoundedWindowAggExec;
pub use per_partition_filter::{PerPartitionFilterExec,
range_partition_predicates};
pub use plan_algebra::{preserves_distribution, preserves_partitioning};
+pub use range_filter::RangeFilterExec;
Review Comment:
Small composability nit. `RangeBound` is `pub` inside `range_filter`, but
`range_filter` itself is a private module and only `RangeFilterExec` is
re-exported here, so downstream crates can't name the type. It still compiles
for them because a type alias is transparent and they can spell out
`Vec<(Option<ScalarValue>, Option<ScalarValue>)>`, but rustdoc renders an
unlinkable name in the `try_new_resolved` and `raw_bounds` signatures, which is
not a great first impression for anyone building on the operator.
One line fixes it:
```rust
pub use range_filter::{RangeBound, RangeFilterExec};
```
##########
ballista/core/src/execution_plans/range_filter.rs:
##########
@@ -0,0 +1,1095 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Filter inputs with per-input-partition half-open bounds.
+//!
+//! `execute(k)` applies the predicate
+//!
+//! ```text
+//! raw_bounds[k].0 - halo_lo <= routing_expr < raw_bounds[k].1 + halo_hi
+//! ```
+//!
+//! `None` on either bound means unbounded on that side (virtual ±∞). Zero halo
+//! (`halo_lo == halo_hi == 0`) recovers the exact range-repartition trim used
+//! above `ShuffleReaderExec` for hash-agg correctness; non-zero halo widens
+//! each partition's read range to include a boundary "context" band
+//! (`WindowFrame` PRECEDING/FOLLOWING for bounded RANGE frames
+//!
+//! # Separation of concerns
+//!
+//! RFE is a pure per-partition filter: the scheduler decides which
+//! (input-partition → half-open cut range) mapping applies (see
+//! `resolve_range_filter_bounds` in the AQE adapter) and hands the
+//! **unwidened** ranges here. Halos live on RFE — the parallel-window
+//! rewrite rule plants them at plan time — and RFE widens the incoming
+//! ranges by its own halos at [`RangeFilterExec::resolve_bounds`] time. This
means the
+//! scheduler stays halo-blind at the RFE boundary.
+//!
+//! # Late-binding bounds
+//!
+//! `raw_bounds` is `Arc<Mutex<Option<Vec<...>>>>` — the ParallelWindow rewrite
+//! rule plants a `RangeFilterExec` at plan time, well before the runtime cuts
+//! are known. The scheduler calls [`RangeFilterExec::resolve_bounds`] after
+//! stage 0's `RuntimeStatsExec` reports have been merged. `execute` refuses
+//! while bounds are unresolved; serialization refuses too — over-the-wire
+//! plans always ship with bounds bound.
+//!
+//! # Type generality
+//!
+//! `ScalarValue` at the API + serde surface. The internal fast path is
+//! Float64-only today (matches URRE/ORRE T-Digest); widening to other
+//! numeric primitives is a KLL-migration follow-up that
+//! will land without breaking callers.
+
+use std::fmt::{self, Debug, Formatter};
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use datafusion::arrow::array::{Array, RecordBatch};
+use datafusion::arrow::compute::filter_record_batch;
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::common::cast::{as_boolean_array, as_float64_array};
+use datafusion::common::{Result, Statistics, internal_err};
+use datafusion::execution::TaskContext;
+use datafusion::logical_expr::Operator;
+use datafusion::physical_expr::expressions::{BinaryExpr, Literal};
+use datafusion::physical_expr::{Distribution, OrderingRequirements,
PhysicalExpr};
+use datafusion::physical_plan::execution_plan::CardinalityEffect;
+use datafusion::physical_plan::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet,
+};
+use datafusion::physical_plan::stream::{
+ EmptyRecordBatchStream, RecordBatchStreamAdapter,
+};
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
PlanProperties,
+ RecordBatchStream, SendableRecordBatchStream,
+};
+use datafusion::scalar::ScalarValue;
+use futures::{Stream, StreamExt, ready};
+use parking_lot::Mutex;
+
+/// Half-open `[lo, hi)` bound for one input partition. `None` on either side
+/// means unbounded (virtual ±∞).
+pub type RangeBound = (Option<ScalarValue>, Option<ScalarValue>);
+
+/// Bounds after halo widening. Float64-only internally today — see the
+/// "Type generality" section in the module doc.
+type WidenedBound = (Option<f64>, Option<f64>);
+
+/// Both raw and widened bounds. `raw` is preserved for serialization; the
+/// executor consumes `widened`.
+struct BoundsState {
+ raw: Vec<RangeBound>,
+ widened: Vec<WidenedBound>,
+}
+
+/// Filter over an ordered input with a per-input-partition half-open range
+/// predicate, widened by the operator's halo. Range logic (cuts →
per-partition
+/// half-open ranges → task-slice) lives scheduler-side; RFE is the runtime
+/// filter that applies the resolved bounds.
+pub struct RangeFilterExec {
+ input: Arc<dyn ExecutionPlan>,
+ routing_expr: Arc<dyn PhysicalExpr>,
+ /// Lower halo — subtracted from each partition's `lo` at widen time.
+ halo_lo: ScalarValue,
+ /// Upper halo — added from each partition's `hi` at widen time.
+ halo_hi: ScalarValue,
+ /// Late-bound: `None` until [`RangeFilterExec::resolve_bounds`];
`execute` and serde
+ /// refuse while unresolved.
+ bounds: Arc<Mutex<Option<BoundsState>>>,
+ /// True when `input.output_ordering()` leads with `routing_expr` in
+ /// ascending order. Enables the min/max fast path + binary-search slice
+ /// in [`RangeFilterStream`] — with a sorted input, per-batch first/last
+ /// values bound the entire batch, so most batches never touch
+ /// `filter_record_batch`.
+ sorted_on_key: bool,
+ properties: Arc<PlanProperties>,
+ metrics: ExecutionPlanMetricsSet,
+}
+
+impl RangeFilterExec {
+ /// Construct with bounds pending (rule path — the ParallelWindow rewrite
+ /// plants the operator at plan time; the scheduler resolves bounds after
+ /// stage 0's stats reports merge).
+ ///
+ /// # Arguments
+ ///
+ /// * `input` - upstream operator; its partition count fixes the eventual
+ /// `raw_bounds.len()`.
+ /// * `routing_expr` - numeric physical expression each row is bucketed by.
+ /// * `halo_lo`, `halo_hi` - non-negative widening amounts applied by
+ /// [`RangeFilterExec::resolve_bounds`]. Both must be finite Float64
today.
+ pub fn try_new_pending(
+ input: Arc<dyn ExecutionPlan>,
+ routing_expr: Arc<dyn PhysicalExpr>,
+ halo_lo: ScalarValue,
+ halo_hi: ScalarValue,
+ ) -> Result<Self> {
+ Self::try_new_inner(input, routing_expr, halo_lo, halo_hi, None)
+ }
+
+ /// Construct with bounds already known. Used by wire decode and by
+ /// task-restriction (task builder slices raw bounds parallel to the input
+ /// restriction, then hands them here as a fresh operator).
+ ///
+ /// # Arguments
+ ///
+ /// * `input`, `routing_expr`, `halo_lo`, `halo_hi` - same as
+ /// [`Self::try_new_pending`].
+ /// * `raw_bounds` - one half-open cut range per input partition. Widening
+ /// by halos happens internally; caller passes unwidened.
+ pub fn try_new_resolved(
+ input: Arc<dyn ExecutionPlan>,
+ routing_expr: Arc<dyn PhysicalExpr>,
+ halo_lo: ScalarValue,
+ halo_hi: ScalarValue,
+ raw_bounds: Vec<RangeBound>,
+ ) -> Result<Self> {
+ Self::try_new_inner(input, routing_expr, halo_lo, halo_hi,
Some(raw_bounds))
+ }
+
+ fn try_new_inner(
+ input: Arc<dyn ExecutionPlan>,
+ routing_expr: Arc<dyn PhysicalExpr>,
+ halo_lo: ScalarValue,
+ halo_hi: ScalarValue,
+ raw_bounds: Option<Vec<RangeBound>>,
+ ) -> Result<Self> {
+ let schema = input.schema();
+ let expr_type = routing_expr.data_type(&schema)?;
+ if !expr_type.is_numeric() {
+ return internal_err!(
+ "RangeFilterExec: routing_expr must be numeric, got
{expr_type}"
+ );
+ }
+ let halo_lo_f64 = as_f64(&halo_lo)?;
+ let halo_hi_f64 = as_f64(&halo_hi)?;
+ if !halo_lo_f64.is_finite() || halo_lo_f64 < 0.0 {
+ return internal_err!(
+ "RangeFilterExec: halo_lo must be finite and non-negative, got
{halo_lo_f64}"
+ );
+ }
+ if !halo_hi_f64.is_finite() || halo_hi_f64 < 0.0 {
+ return internal_err!(
+ "RangeFilterExec: halo_hi must be finite and non-negative, got
{halo_hi_f64}"
+ );
+ }
+ let bounds_state = raw_bounds
+ .map(|raw| build_bounds_state(&input, raw, halo_lo_f64,
halo_hi_f64))
+ .transpose()?;
+ let properties = Arc::new(PlanProperties::new(
+ input.equivalence_properties().clone(),
+ input.output_partitioning().clone(),
+ input.pipeline_behavior(),
+ input.boundedness(),
+ ));
+ let sorted_on_key = input
+ .output_ordering()
+ .map(|ord| {
+ let first = ord.first();
+ first.expr.as_ref() == routing_expr.as_ref() &&
!first.options.descending
+ })
+ .unwrap_or(false);
+ Ok(Self {
+ input,
+ routing_expr,
+ halo_lo,
+ halo_hi,
+ bounds: Arc::new(Mutex::new(bounds_state)),
+ sorted_on_key,
+ properties,
+ metrics: ExecutionPlanMetricsSet::new(),
+ })
+ }
+
+ /// Idempotent overwrite. Called by the scheduler once stage-0 sketches
+ /// merge into cuts and the adapter has projected those cuts onto per-input
+ /// partition half-open ranges. Widens by RFE's halos before caching.
+ pub fn resolve_bounds(&self, raw_bounds: Vec<RangeBound>) -> Result<()> {
+ let halo_lo = as_f64(&self.halo_lo)?;
+ let halo_hi = as_f64(&self.halo_hi)?;
+ let state = build_bounds_state(&self.input, raw_bounds, halo_lo,
halo_hi)?;
+ self.bounds.lock().replace(state);
+ Ok(())
+ }
+
+ /// Snapshot the unwidened bounds. `None` before
[`RangeFilterExec::resolve_bounds`];
+ /// `Some` after. Callers that need the widened form should either call
+ /// [`Self::widened_bounds`] or expand `raw_bounds[k] ± (halo_lo, halo_hi)`
+ /// themselves.
+ pub fn raw_bounds(&self) -> Option<Vec<RangeBound>> {
+ self.bounds.lock().as_ref().map(|s| s.raw.clone())
+ }
+
+ /// Snapshot the halo-widened bounds — what the runtime filter actually
+ /// applies. Convenience for callers that already have the halos and just
+ /// want the resolved form.
+ pub fn widened_bounds(&self) -> Option<Vec<WidenedBound>> {
Review Comment:
Same idea as the `RangeBound` note over in `mod.rs`. `widened_bounds` is
public but `WidenedBound` on line 93 is a private alias, so the rendered
signature reads as a type nobody outside the crate can refer to. Either make
the alias `pub` and export it alongside `RangeBound`, or just spell the tuple
out in the return type here.
Not a blocker either way. Easier to sort out now while the operator has no
callers than after it ships.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
