feat: support table sample

Cargo.toml

@@ -164,6 +164,7 @@ pbjson-types = "0.7"
 insta = { version = "1.43.1", features = ["glob", "filters"] }
 prost = "0.13.1"
 rand = "0.9"
+rand_distr = "0.5"
 recursive = "0.1.1"
 regex = "1.8"
 rstest = "0.25.0"

datafusion/core/Cargo.toml

@@ -157,7 +157,7 @@ env_logger = { workspace = true }
 insta = { workspace = true }
 paste = "^1.0"
 rand = { workspace = true, features = ["small_rng"] }
-rand_distr = "0.5"
+rand_distr = { workspace = true }
 regex = { workspace = true }
 rstest = { workspace = true }
 serde_json = { workspace = true }

datafusion/core/src/dataframe/mod.rs

@@ -2262,6 +2262,50 @@ impl DataFrame {
         let df = ctx.read_batch(batch)?;
         Ok(df)
     }
+
+    /// Sample a fraction of the rows from the DataFrame.
+    ///
+    /// # Arguments
+    /// * `fraction` - The fraction of rows to sample.
+    /// * `with_replacement` - Whether to sample with replacement.
+    /// * `seed` - The seed for the random number generator.
+    ///
+    /// # Example
+    /// ```
+    /// use datafusion::prelude::*;
+    /// # use datafusion::error::Result;
+    /// # #[tokio::main]
+    /// # async fn main() -> Result<()> {
+    /// let df = dataframe!(
+    ///    "id" => [1, 2, 3],
+    ///    "name" => ["foo", "bar", "baz"]
+    ///  )?;
+    /// let df = df.sample(0.5, Some(true), Some(42))?;
+    /// df.show().await?;
+    /// # Ok(())
+    /// # }
+    /// // +----+------+
+    /// // | id | name |
+    /// // +----+------+
+    /// // | 3  | baz  |
+    /// // +----+------+
+    /// ```
+    ///
+    pub fn sample(
+        self,
+        fraction: f64,
+        with_replacement: Option<bool>,
+        seed: Option<u64>,
+    ) -> Result<Self> {
+        let plan = LogicalPlanBuilder::from(self.plan)
+            .sample(fraction, with_replacement, seed)?
+            .build()?;
+        Ok(DataFrame {
+            session_state: self.session_state,
+            plan,
+            projection_requires_validation: self.projection_requires_validation,
+        })
+    }
 }
 
 /// Macro for creating DataFrame.

datafusion/core/src/physical_planner.rs

@@ -79,7 +79,7 @@ use datafusion_expr::expr_rewriter::unnormalize_cols;
 use datafusion_expr::logical_plan::builder::wrap_projection_for_join_if_necessary;
 use datafusion_expr::{
     Analyze, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension, FetchType,
-    Filter, JoinType, RecursiveQuery, SkipType, StringifiedPlan, WindowFrame,
+    Filter, JoinType, RecursiveQuery, Sample, SkipType, StringifiedPlan, WindowFrame,
     WindowFrameBound, WriteOp,
 };
 use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
@@ -93,6 +93,7 @@ use datafusion_physical_plan::execution_plan::InvariantLevel;
 use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
 use datafusion_physical_plan::recursive_query::RecursiveQueryExec;
 use datafusion_physical_plan::unnest::ListUnnest;
+use datafusion_physical_plan::SampleExec;
 use sqlparser::ast::NullTreatment;
 
 use async_trait::async_trait;
@@ -902,6 +903,23 @@ impl DefaultPhysicalPlanner {
 
                 Arc::new(GlobalLimitExec::new(input, skip, fetch))
             }
+            LogicalPlan::Sample(Sample {
+                lower_bound,
+                upper_bound,
+                seed,
+                with_replacement,
+                ..
+            }) => {
+                let input = children.one()?;
+                let sample = SampleExec::try_new(
+                    input,
+                    *lower_bound,
+                    *upper_bound,
+                    *with_replacement,
+                    *seed,
+                )?;
+                Arc::new(sample)
+            }
             LogicalPlan::Unnest(Unnest {
                 list_type_columns,
                 struct_type_columns,

datafusion/core/tests/dataframe/mod.rs

@@ -6137,3 +6137,73 @@ async fn test_dataframe_macro() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn test_dataframe_sample() -> Result<()> {
+    let df = dataframe!(
+        "a" => (1..40).collect::<Vec<_>>(),
+    )?;
+
+    // Test sampling 20% of rows with replacement
+    let df_sampled = df.clone().sample(0.2, Some(true), Some(42))?;
+    assert_batches_eq!(
+        #[rustfmt::skip]
+        &[
+            "+----+", 
+            "| a  |", 
+            "+----+", 
+            "| 8  |", 
+            "| 10 |", 
+            "| 19 |", 
+            "| 29 |",
+            "| 29 |", 
+            "| 36 |", 
+            "+----+",
+        ],
+        &df_sampled.collect().await?
+    );
+
+    // Test sampling 20% of rows without replacement
+    let df_sampled = df.clone().sample(0.2, Some(false), Some(42))?;
+    assert_batches_eq!(
+        #[rustfmt::skip]
+        &[
+            "+----+", 
+            "| a  |", 
+            "+----+", 
+            "| 5  |", 
+            "| 9  |", 
+            "| 10 |", 
+            "| 14 |",
+            "| 17 |", 
+            "| 19 |", 
+            "| 24 |", 
+            "| 39 |", 
+            "+----+",
+        ],
+        &df_sampled.collect().await?
+    );
+
+    // Test sampling with None parameters (should use defaults)
+    let df_sampled_default = df.clone().sample(0.2, None, Some(42))?;
+    assert_batches_eq!(
+        #[rustfmt::skip]
+        &[
+            "+----+",
+            "| a  |", 
+            "+----+", 
+            "| 5  |", 
+            "| 9  |", 
+            "| 10 |", 
+            "| 14 |",
+            "| 17 |", 
+            "| 19 |", 
+            "| 24 |", 
+            "| 39 |", 
+            "+----+",
+        ],
+        &df_sampled_default.collect().await?
+    );
+
+    Ok(())
+}

datafusion/expr/Cargo.toml

@@ -52,6 +52,7 @@ datafusion-functions-window-common = { workspace = true }
 datafusion-physical-expr-common = { workspace = true }
 indexmap = { workspace = true }
 paste = "^1.0"
+rand = { workspace = true }
 recursive = { workspace = true, optional = true }
 serde_json = { workspace = true }
 sqlparser = { workspace = true }

datafusion/expr/src/logical_plan/builder.rs

@@ -30,6 +30,7 @@ use crate::expr_rewriter::{
     normalize_col_with_schemas_and_ambiguity_check, normalize_cols, normalize_sorts,
     rewrite_sort_cols_by_aggs,
 };
+use crate::logical_plan::plan::Sample;
 use crate::logical_plan::{
     Aggregate, Analyze, Distinct, DistinctOn, EmptyRelation, Explain, Filter, Join,
     JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Prepare,
@@ -1474,6 +1475,21 @@ impl LogicalPlanBuilder {
         unnest_with_options(Arc::unwrap_or_clone(self.plan), columns, options)
             .map(Self::new)
     }
+
+    pub fn sample(
+        self,
+        fraction: f64,
+        with_replacement: Option<bool>,
+        seed: Option<u64>,
+    ) -> Result<Self> {
+        Ok(Self::new(LogicalPlan::Sample(Sample {
+            input: self.plan,
+            lower_bound: 0.0,
+            upper_bound: fraction,
+            with_replacement: with_replacement.unwrap_or(false),
+            seed: seed.unwrap_or_else(rand::random),
+        })))
+    }
 }
 
 impl From<LogicalPlan> for LogicalPlanBuilder {

datafusion/expr/src/logical_plan/display.rs

@@ -20,6 +20,7 @@
 use std::collections::HashMap;
 use std::fmt;
 
+use crate::logical_plan::plan::Sample;
 use crate::{
     expr_vec_fmt, Aggregate, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr,
     Filter, Join, Limit, LogicalPlan, Partitioning, Projection, RecursiveQuery,
@@ -650,6 +651,21 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> {
                     "StructColumn": expr_vec_fmt!(struct_type_columns),
                 })
             }
+            LogicalPlan::Sample(Sample {
+                input: _,
+                lower_bound,
+                upper_bound,
+                with_replacement,
+                seed,
+            }) => {
+                json!({
+                    "Node Type": "Sample",
+                    "Lower Bound": lower_bound,
+                    "Upper Bound": upper_bound,
+                    "With Replacement": with_replacement,
+                    "Seed": seed,
+                })
+            }
         }
     }
 }

datafusion/expr/src/logical_plan/mod.rs

@@ -40,8 +40,8 @@ pub use plan::{
     projection_schema, Aggregate, Analyze, ColumnUnnestList, DescribeTable, Distinct,
     DistinctOn, EmptyRelation, Explain, ExplainFormat, Extension, FetchType, Filter,
     Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType,
-    Projection, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, Subquery,
-    SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window,
+    Projection, RecursiveQuery, Repartition, Sample, SkipType, Sort, StringifiedPlan,
+    Subquery, SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window,
 };
 pub use statement::{
     Deallocate, Execute, Prepare, SetVariable, Statement, TransactionAccessMode,