Implement partition_statistics API for NestedLoopJoinExec (#19468)

## Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->

- Part of #15873

## Rationale for this change

After the introduction of the `partition_statistics` API there are some
operators which don't support this new API, the goal is to update the
operators with new API.

<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->

## What changes are included in this PR?

- Updated `NestedLoopJoinExec` with the new `partition_statistics`

<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->

## Are these changes tested?

- Added relevant tests 

<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code

If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->

## Are there any user-facing changes?

<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->

<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->

Author: kumarUjjawal

datafusion/core/tests/physical_optimizer/partition_statistics.rs

@@ -46,7 +46,7 @@ mod test {
     use datafusion_physical_plan::common::compute_record_batch_statistics;
     use datafusion_physical_plan::empty::EmptyExec;
     use datafusion_physical_plan::filter::FilterExec;
-    use datafusion_physical_plan::joins::CrossJoinExec;
+    use datafusion_physical_plan::joins::{CrossJoinExec, NestedLoopJoinExec};
     use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
     use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
     use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
@@ -638,6 +638,81 @@ mod test {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_statistic_by_partition_of_nested_loop_join() -> Result<()> {
+        use datafusion_expr::JoinType;
+
+        let left_scan = create_scan_exec_with_statistics(None, Some(2)).await;
+        let left_scan_coalesced: Arc<dyn ExecutionPlan> =
+            Arc::new(CoalescePartitionsExec::new(left_scan));
+
+        let right_scan = create_scan_exec_with_statistics(None, Some(2)).await;
+
+        let nested_loop_join: Arc<dyn ExecutionPlan> =
+            Arc::new(NestedLoopJoinExec::try_new(
+                left_scan_coalesced,
+                right_scan,
+                None,
+                &JoinType::RightSemi,
+                None,
+            )?);
+
+        // Test partition_statistics(None) - returns overall statistics
+        // For RightSemi join, output columns come from right side only
+        let full_statistics = nested_loop_join.partition_statistics(None)?;
+        // With empty join columns, estimate_join_statistics returns Inexact row count
+        // based on the outer side (right side for RightSemi)
+        let mut expected_full_statistics = create_partition_statistics(
+            4,
+            32,
+            1,
+            4,
+            Some((DATE_2025_03_01, DATE_2025_03_04)),
+        );
+        expected_full_statistics.num_rows = Precision::Inexact(4);
+        expected_full_statistics.total_byte_size = Precision::Absent;
+        assert_eq!(full_statistics, expected_full_statistics);
+
+        // Test partition_statistics(Some(idx)) - returns partition-specific statistics
+        // Partition 1: ids [3,4], dates [2025-03-01, 2025-03-02]
+        let mut expected_statistic_partition_1 = create_partition_statistics(
+            2,
+            16,
+            3,
+            4,
+            Some((DATE_2025_03_01, DATE_2025_03_02)),
+        );
+        expected_statistic_partition_1.num_rows = Precision::Inexact(2);
+        expected_statistic_partition_1.total_byte_size = Precision::Absent;
+
+        // Partition 2: ids [1,2], dates [2025-03-03, 2025-03-04]
+        let mut expected_statistic_partition_2 = create_partition_statistics(
+            2,
+            16,
+            1,
+            2,
+            Some((DATE_2025_03_03, DATE_2025_03_04)),
+        );
+        expected_statistic_partition_2.num_rows = Precision::Inexact(2);
+        expected_statistic_partition_2.total_byte_size = Precision::Absent;
+
+        let statistics = (0..nested_loop_join.output_partitioning().partition_count())
+            .map(|idx| nested_loop_join.partition_statistics(Some(idx)))
+            .collect::<Result<Vec<_>>>()?;
+        assert_eq!(statistics.len(), 2);
+        assert_eq!(statistics[0], expected_statistic_partition_1);
+        assert_eq!(statistics[1], expected_statistic_partition_2);
+
+        // Check the statistics_by_partition with real results
+        let expected_stats = vec![
+            ExpectedStatistics::NonEmpty(3, 4, 2),
+            ExpectedStatistics::NonEmpty(1, 2, 2),
+        ];
+        validate_statistics_with_data(nested_loop_join, expected_stats, 0).await?;
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_statistic_by_partition_of_coalesce_batches() -> Result<()> {
         let scan = create_scan_exec_with_statistics(None, Some(2)).await;

datafusion/physical-plan/src/joins/nested_loop_join.rs

@@ -550,17 +550,34 @@ impl ExecutionPlan for NestedLoopJoinExec {
     }
 
     fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
-        if partition.is_some() {
-            return Ok(Statistics::new_unknown(&self.schema()));
-        }
+        // NestedLoopJoinExec is designed for joins without equijoin keys in the
+        // ON clause (e.g., `t1 JOIN t2 ON (t1.v1 + t2.v1) % 2 = 0`). Any join
+        // predicates are stored in `self.filter`, but `estimate_join_statistics`
+        // currently doesn't support selectivity estimation for such arbitrary
+        // filter expressions. We pass an empty join column list, which means
+        // the cardinality estimation cannot use column statistics and returns
+        // unknown row counts.
         let join_columns = Vec::new();
-        estimate_join_statistics(
-            self.left.partition_statistics(None)?,
-            self.right.partition_statistics(None)?,
+
+        // Left side is always a single partition (Distribution::SinglePartition),
+        // so we always request overall stats with `None`. Right side can have
+        // multiple partitions, so we forward the partition parameter to get
+        // partition-specific statistics when requested.
+        let left_stats = self.left.partition_statistics(None)?;
+        let right_stats = match partition {
+            Some(partition) => self.right.partition_statistics(Some(partition))?,
+            None => self.right.partition_statistics(None)?,
+        };
+
+        let stats = estimate_join_statistics(
+            left_stats,
+            right_stats,
             &join_columns,
             &self.join_type,
-            &self.schema(),
-        )
+            &self.join_schema,
+        )?;
+
+        Ok(stats.project(self.projection.as_ref()))
     }
 
     /// Tries to push `projection` down through `nested_loop_join`. If possible, performs the

datafusion/physical-plan/src/joins/utils.rs

@@ -411,7 +411,35 @@ struct PartialJoinStatistics {
     pub column_statistics: Vec<ColumnStatistics>,
 }
 
-/// Estimate the statistics for the given join's output.
+/// Estimates the output statistics for a join operation based on input statistics.
+///
+/// # Statistics Propagation
+///
+/// This function estimates join output statistics using the following approach:
+/// - **Row count estimation**: Uses the `on` parameter (equijoin keys) to estimate
+///   output cardinality via [`estimate_join_cardinality`]. The estimation is based on
+///   column-level statistics (distinct counts, min/max values) of the join keys.
+/// - **Column statistics**: Combines column statistics from both inputs. For join types
+///   that preserve all columns (Inner, Left, Right, Full), statistics from both sides
+///   are concatenated. For semi/anti joins, only the relevant side's statistics are kept.
+/// - **Byte size**: Always returns `Precision::Absent` as join output size is difficult
+///   to estimate without knowing the actual data.
+///
+/// # The `on` Parameter
+///
+/// The `on` parameter represents equijoin keys (e.g., `t1.id = t2.id`). When `on` is
+/// empty (as in NestedLoopJoinExec which handles non-equijoin predicates), the
+/// cardinality estimation cannot compute selectivity from join keys, and this function
+/// returns unknown statistics (`num_rows: Precision::Absent`).
+///
+/// # Limitations
+///
+/// - Does not account for selectivity of arbitrary join filter expressions
+///   (e.g., `(t1.v1 + t2.v1) % 2 = 0`). Such filters, common in NestedLoopJoinExec,
+///   are not factored into the cardinality estimation.
+/// - Column statistics for the output are simply combined from inputs without
+///   adjusting for join selectivity (acknowledged in the code as needing
+///   "filter selectivity analysis").
 pub(crate) fn estimate_join_statistics(
     left_stats: Statistics,
     right_stats: Statistics,