jayshrivastava commented on code in PR #24501:
URL: https://github.com/apache/datafusion/pull/24501#discussion_r3853294425
##########
datafusion/physical-expr/src/partitioning.rs:
##########
@@ -1644,6 +1651,174 @@ mod tests {
Ok(())
}
+ #[test]
+ fn check_monotonic_transform_rejects_order_reversing_negation() ->
Result<()> {
+ let int_fixture = PartitioningTestFixture::int64(&["x"])?;
+ let neg: Arc<dyn PhysicalExpr> =
Arc::new(NegativeExpr::new(int_fixture.col(0)));
+ assert!(
+ !int_fixture
+ .eq_properties
+ .check_monotonic_transform(&neg, &int_fixture.col(0)),
+ "-x reverses ASC to DESC"
+ );
+ assert!(
+ !int_fixture
+ .eq_properties
+ .check_monotonic_transform(&int_fixture.col(0),
&int_fixture.col(0)),
+ "identity is not a transform"
+ );
+
+ let ts_fixture = PartitioningTestFixture::new(vec![(
+ "timestamp",
+ DataType::Timestamp(TimeUnit::Nanosecond, None),
+ )])?;
+ let date_bin = date_bin_of(ts_fixture.col(0), 60_000_000_000);
+ let date_trunc = date_trunc_of(ts_fixture.col(0), "hour");
+ assert!(
+ ts_fixture
+ .eq_properties
+ .check_monotonic_transform(&date_bin, &ts_fixture.col(0)),
+ "date_bin preserves ASC"
+ );
+ assert!(
+ ts_fixture
+ .eq_properties
+ .check_monotonic_transform(&date_trunc, &ts_fixture.col(0)),
+ "date_trunc preserves ASC"
+ );
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_range_partitioning_project_keeps_range_key() -> Result<()> {
+ let fixture = PartitioningTestFixture::new(vec![(
+ "a",
+ DataType::Timestamp(TimeUnit::Nanosecond, None),
+ )])?;
+ let hour_ns = 1_704_070_800_000_000_000i64;
+ let date_bin = date_bin_of(fixture.col(0), 60_000_000_000);
+ let aligned = fixture.range_partitioning([0],
vec![ts_ns_split(hour_ns)]);
+
+ // SELECT a, date_bin(a) — range key is kept, so Range stays on a.
+ let a_target: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+ let bin_target: Arc<dyn PhysicalExpr> =
Arc::new(Column::new("date_bin", 1));
+ let keep_a_mapping = ProjectionMapping::from_iter([
+ (
+ fixture.col(0),
+ ProjectionTargets::from(vec![(Arc::clone(&a_target), 0)]),
+ ),
+ (
+ Arc::clone(&date_bin),
+ ProjectionTargets::from(vec![(Arc::clone(&bin_target), 1)]),
+ ),
+ ]);
+ let projected = aligned.project(&keep_a_mapping,
&fixture.eq_properties);
+ assert_eq!(
+ projected.to_string(),
+ "Range([a@0 ASC], [(1704070800000000000)], 2)"
+ );
+
+ // SELECT a AS b, date_bin(a) AS bucket — alias of a preserves Range.
+ let b_target: Arc<dyn PhysicalExpr> = Arc::new(Column::new("b", 0));
+ let bucket_target: Arc<dyn PhysicalExpr> =
Arc::new(Column::new("bucket", 1));
+ let alias_mapping = ProjectionMapping::from_iter([
+ (
+ fixture.col(0),
+ ProjectionTargets::from(vec![(Arc::clone(&b_target), 0)]),
+ ),
+ (
+ Arc::clone(&date_bin),
+ ProjectionTargets::from(vec![(Arc::clone(&bucket_target), 1)]),
+ ),
+ ]);
+ let projected = aligned.project(&alias_mapping,
&fixture.eq_properties);
+ assert_eq!(
+ projected.to_string(),
+ "Range([b@0 ASC], [(1704070800000000000)], 2)"
+ );
+
+ Ok(())
+ }
+
+ #[test]
+ fn
test_range_partitioning_project_skips_non_monotonic_and_straddling_bins()
+ -> Result<()> {
+ let fixture = PartitioningTestFixture::new(vec![(
+ "a",
+ DataType::Timestamp(TimeUnit::Nanosecond, None),
+ )])?;
+ let hour_ns = 1_704_070_800_000_000_000i64;
+ let aligned = fixture.range_partitioning([0],
vec![ts_ns_split(hour_ns)]);
+ let date_bin_60s = date_bin_of(fixture.col(0), 60_000_000_000);
+ // 70s does not divide the hour split, so bins straddle the file
groups.
+ let date_bin_70s = date_bin_of(fixture.col(0), 70_000_000_000);
+ let neg: Arc<dyn PhysicalExpr> =
Arc::new(NegativeExpr::new(fixture.col(0)));
+
+ // SELECT -a, date_bin(60s, a) — skip the order-reversing source.
+ let neg_target: Arc<dyn PhysicalExpr> = Arc::new(Column::new("neg",
0));
+ let bin_target: Arc<dyn PhysicalExpr> = Arc::new(Column::new("bin",
1));
+ let mapping = ProjectionMapping::from_iter([
+ (
+ Arc::clone(&neg),
+ ProjectionTargets::from(vec![(Arc::clone(&neg_target), 0)]),
+ ),
+ (
+ Arc::clone(&date_bin_60s),
+ ProjectionTargets::from(vec![(Arc::clone(&bin_target), 1)]),
+ ),
+ ]);
+ let projected = aligned.project(&mapping, &fixture.eq_properties);
+ assert_eq!(
+ projected.to_string(),
+ "Range([bin@1 ASC], [(1704070800000000000)], 2)"
+ );
+
+ // SELECT -a, date_bin(60s, a) AS bin1, date_bin(60s, a) AS bin2
+ let bin1: Arc<dyn PhysicalExpr> = Arc::new(Column::new("bin1", 1));
+ let bin2: Arc<dyn PhysicalExpr> = Arc::new(Column::new("bin2", 2));
+ let mapping = ProjectionMapping::from_iter([
+ (
+ Arc::clone(&neg),
+ ProjectionTargets::from(vec![(Arc::clone(&neg_target), 0)]),
+ ),
+ (
+ Arc::clone(&date_bin_60s),
+ ProjectionTargets::from(vec![
+ (Arc::clone(&bin1), 1),
+ (Arc::clone(&bin2), 2),
+ ]),
+ ),
+ ]);
+ let projected = aligned.project(&mapping, &fixture.eq_properties);
+ assert_eq!(
+ projected.to_string(),
+ "Range([bin1@1 ASC], [(1704070800000000000)], 2)"
+ );
+
+ // SELECT date_bin(70s, a), date_bin(60s, a) — first bin straddles,
+ // second is aligned, so Range is preserved through the 60s bin.
Review Comment:
Nice.
##########
datafusion/physical-expr/src/partitioning.rs:
##########
@@ -577,13 +582,15 @@ impl Partitioning {
allow_subset,
);
if satisfaction == PartitioningSatisfaction::NotSatisfied
- && allow_subset
&& range_monotonic_fn_satisfies_keys(
range,
required_exprs,
eq_properties,
)
{
+ // A monotonic transform of the range key (for example
+ // `date_bin(timestamp)`) is not Hash-subset logic, so
+ // this is independent of `allow_subset`.
Review Comment:
I think this is too simple. There's 4 explicit cases we care about:
Assume we are range partitioned on `a`
1. `select baz, a from foo group by a, baz`
- This is `PartitioningSatisfaction::Subset` because `a` is a subset of the
grouping key `a, baz`.
2. `select baz, a from foo group by a`
- This is `PartitioningSatisfaction::Exact` because `a` is the entire of the
grouping key `a`.
3. `select baz, time_bin(a, 60) as bin from foo group by bin, baz`
- This is `Subset`
4. `select baz, time_bin(a, 60) as bin from foo group by bin`
- This is `Exact`
I assume 1 and 2 above happen on this line `let satisfaction =
Self::key_satisfaction(...)` but your code needs to distinguish cases 3 and 4
here. Right now, you always return `PartitioningSatisfaction::Subset` but in
some cases, we should return `PartitioningSatisfaction::Exact`.
--
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]
