feat: Add example for implementing `SAMPLE` using extension APIs #17633

theirix · 2025-09-17T21:14:31Z

Which issue does this PR close?

Closes Support data source sampling with TABLESAMPLE #13563.

Rationale for this change

The rationale is explained in #13563 in detail with known syntax examples.

This is the third design for the table sample support.

My first design was an addition of an explicit rewrite function baked in into a select logical plan - Support data source sampling with TABLESAMPLE #16325
Second design introduced dedicated flexible logical and physical plans, but tied to datafusion core - feat: support table sample #16505
This third design abstracts the second design out of datafusion core into extensions.

What changes are included in this PR?

All changes are bundled to an example file since it is a PoC of extensibility - as discussed with @alamb in #13563 (comment) .

If the idea is viable, the code could be modularised, which is not possible in a datafusion-examples crate.

It adds several components:

a custom TableSamplePlanNode trait for a sampling logical plan
a query planner TableSampleQueryPlanner (trait QueryPlanner)
an execution plan SampleExec - mostly adapted from the second design, all kudos and thanks to @chenkovsky !
an extension planner (trait ExtensionPlanner) to build a physical plan
tests
an example runner

The setup, as seen in main, is a bit cumbersome, but it works. Building a SQL extension with access to AST and without introducing a whole new statement (as large projects like arroyo, greptime or cube introduce new syntax) is complicated.
For full modularity I would propose a few extension points to SqlToRel. It could help to avoid manual parsing / logical plan / physical plan / execution conversions and would keep the concise client syntax.

Are these changes tested?

A set of unit tests
An example with asserts

Are there any user-facing changes?

No

alamb · 2025-09-17T21:44:28Z

This looks amazing -- I will try and review it soon

alamb

Thank you @theirix -- this is an amazing PR that shows how to "support your own SQL" dialect. I am sorry for the delay in reviewing

I recommend we call this example something more general like "custom_sql_dialect" or something (we can rename it as a follow on PR)

I also had several suggestions on how we can improve / simplify this example / improve the underlying APIs as well, though they can also be done as follow ons I think, if you prefer.

What I suggest is:

Implement some of the simpler changes (move main up, add comments, simplify Eq, cleanup imports, etc)
Merge this PR
Work on improvements to API and comments as follow on PRs (you don't have to do it)

Again, this is really really cool. I have dreamed about this type of example for years 💤

alamb · 2025-09-22T14:34:07Z

datafusion-examples/examples/table_sample.rs

+    seed: u64,
+}
+
+impl Hash for TableSamplePlanNode {


I think you can derive these implementations automatically, as in

#[derive(Debug, Clone, Hash, Eq, PartialEq, PartialOrd)] struct TableSamplePlanNode { ...

Rather than provide the implementations directly

Update: I tried that and found it was due to the f64

error[E0277]: the trait bound `f64: std::hash::Hash` is not satisfied --> datafusion-examples/examples/table_sample.rs:117:5 | 113 | #[derive(Debug, Clone, Hash, Eq, PartialEq, PartialOrd)] | ---- in this derive macro expansion ... 117 | lower_bound: f64, | ^^^^^^^^^^^^^^^^ the trait `std::hash::Hash` is not implemented for `f64` |

Here is one proposal

Use newtype to get hash/eq for f64 theirix/datafusion#1

alamb · 2025-09-22T14:34:18Z

datafusion-examples/examples/table_sample.rs

+// specific language governing permissions and limitations
+// under the License.
+
+#![allow(unused_imports)]


Is this needed? I think the code will get much simpler if it is removed

Agree, I forgot to remove all these experimental includes

alamb · 2025-09-22T14:37:12Z

datafusion-examples/examples/table_sample.rs

+    expr: &ast::Expr,
+) -> Option<T> {
+    match expr {
+        ast::Expr::BinaryOp { left, op, right } => {


most of the rest of DataFusion would use an Expr here and reuse the existing simplification implementation (e.g.

datafusion/datafusion-examples/examples/expr_api.rs

Line 162 in a8d5016

let expr = col("ts").eq(to_timestamp(vec![lit("2020-09-08T12:00:00+00:00")]));

)

Tried to rework this functionality using a more generic expression simplified. Since this change is quite large, I would prefer to submit it as a different follow-up PR.

alamb · 2025-09-22T14:40:02Z

datafusion-examples/examples/table_sample.rs

+
+// Context provider for tests
+
+struct MockContextProvider<'a> {


What do you think about making SessionContextProvider public so it can be more easily reused? Having to define a very similar one here is a fairly large usability papercut in my opinion

https://github.com/apache/datafusion/blob/0181c7906b47f9b17fbde92bf30b2d939a80230c/datafusion/core/src/execution/session_state.rs#L1667-L1666

That's what I initially thought, but I decided not to touch the core. All in-crate tests were able to reuse SessionContextProvider. It would be reasonable and much more ergonomic to provide this functionality to the example and extensions as well.

I changed the visibility of this struct and added a public constructor while keeping the internals private. Please let me know if I should do it in a different PR.

alamb · 2025-09-22T14:54:32Z

datafusion-examples/examples/table_sample.rs

+    )
+    .await?;
+
+    let table_source = provider_as_source(ctx.table_provider("alltypes_plain").await?);


if we made SessionContextProvider public this would be much simpler i think

alamb · 2025-09-22T14:54:40Z

datafusion-examples/examples/table_sample.rs

+    }
+}
+
+#[tokio::main]


I suggest we move this main to the top of the file so people can see what this example does / how it can be used, before having to wade through all the implementation details

datafusion-examples/examples/table_sample.rs

alamb · 2025-09-22T14:56:18Z

datafusion-examples/examples/table_sample.rs

+        match poll {
+            Poll::Ready(Some(Ok(batch))) => {
+                let start = baseline_metrics.elapsed_compute().clone();
+                let result = self.sampler.sample(&batch);
+                let _timer = start.timer();
+                Poll::Ready(Some(result))
+            }
+            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
+            Poll::Ready(None) => Poll::Ready(None),
+            Poll::Pending => Poll::Pending,
+        }


You can use the ready! macro here to make this more concise it you want

Something like

match ready!(poll) { Some(Ok(..)) ... }

It's really useful, thank you. Switched to it and recorded additional metrics

alamb · 2025-09-22T14:57:18Z

datafusion-examples/examples/table_sample.rs

+use rand::{Rng, SeedableRng};
+use rand_distr::{Distribution, Poisson};
+
+/// This example demonstrates the table sample support.


I think this undersells somewhat what this example does

How about something more like this:

Suggested change

/// This example demonstrates the table sample support.

/// This example demonstrates how to extend DataFusion's SQL parser to recognize

/// other syntax.

///

/// This example shows how to extend the DataFusion SQL planner to support the

/// `TABLESAMPLE` clause in SQL queries and then use a custom user defined node

/// to implement the sampling logic in the physical plan.

Use newtype to get hash/eq for f64

Co-authored-by: Andrew Lamb <[email protected]>

Signed-off-by: theirix <[email protected]>

theirix · 2025-09-23T22:06:25Z

Thank you @theirix -- this is an amazing PR that shows how to "support your own SQL" dialect. I am sorry for the delay in reviewing

I recommend we call this example something more general like "custom_sql_dialect" or something (we can rename it as a follow on PR)

I also had several suggestions on how we can improve / simplify this example / improve the underlying APIs as well, though they can also be done as follow ons I think, if you prefer.

What I suggest is:
1. Implement some of the simpler changes (move `main` up, add comments, simplify `Eq`, cleanup imports, etc)

2. Merge this PR

3. Work on improvements to API and comments as follow on PRs (you don't have to do it)
Again, this is really really cool. I have dreamed about this type of example for years 💤

Thank you! I am glad it helps showcase Datafusion's extensibility. It wasn't immediately clear to me how to approach it, so it could be helpful for developers as a reference or a starting point.

I have attempted to address most of the review suggestions, including a small core API change. I would prefer to move with smaller iterations with follow-ups.

vegarsti · 2025-09-27T06:17:39Z

datafusion-examples/examples/table_sample.rs

+use futures::{ready, TryStreamExt};
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+use rand_distr::{Distribution, Poisson};


Drive by comment: Would these imports be tidier if they were in the same block? There's several sections of datafusion, several sections of arrow.

Thank you for reviewing! I cleaned up comments and reordered imports - now it'd make much more sense!

vegarsti

Just noting some stuff from reading through this out of interest. Sorry if you're already on top of these changes!

vegarsti · 2025-09-27T06:19:04Z

datafusion-examples/examples/table_sample.rs

+    // Classical way
+    // let sql_to_rel = SqlToRel::new(&context_provider);
+    // let logical_plan = sql_to_rel.sql_statement_to_plan(statement.clone())?;


Classical way to do what?

vegarsti · 2025-09-27T06:19:12Z

datafusion-examples/examples/table_sample.rs

+    // Execute via standard sql call - doesn't work
+    // let df = ctx.sql(sql).await?;
+    // let batches = df.collect().await?;


vegarsti · 2025-09-27T06:19:57Z

datafusion-examples/examples/table_sample.rs

+    info!("Batches: {:?}", &batches);
+
+    let result_string = pretty_format_batches(&batches)
+        // pretty_format_batches_with_schema(table_source.schema(), &batches)


vegarsti · 2025-09-27T06:23:06Z

datafusion-examples/examples/table_sample.rs

+    Ok(())
+}
+
+/// Hashable and comparible f64 for sampling bounds


Suggested change

/// Hashable and comparible f64 for sampling bounds

/// Hashable and comparable f64 for sampling bounds

alamb · 2025-09-29T10:49:55Z

I am hoping to help push this PR along this week. Thank you for your patience, as always, @theirix

geoffreyclaude · 2025-09-30T14:21:48Z

Hi @theirix!
@alamb linked me to your PR as I'm looking to do something very similar (extend the DataFusion SQL to support MatchRecognize in an extension.)
From what I understand of your code, it only works if the TABLESAMPLE is at the root of the query, so any TABLESAMPLE inside a CTE or sides of a JOIN won't get transformed, is that correct?

For the next steps, I have a work in progress implementation of a similar idea, but which allows injecting any number of custom relation planners, and supports nested custom relations as well. I've added your TABLESAMPLE example in it and it works great (and a mini-MATCH_RECOGNIZE for @alamb too 😄)

If you want to give it a look, the draft PR is #17843 .

theirix · 2025-09-30T16:07:16Z

Hello @geoffreyclaude ! I reckon it was only tested for the root sample clause so far. Let me check it soon

theirix added 3 commits September 17, 2025 19:58

Add ctor, rand_distr dependencies

1676c0a

Implement SQL table sample as an extension

f6b73bd

Format

edceba3

theirix marked this pull request as ready for review September 17, 2025 21:48

Use newtype to get hash/eq for f64

2a6f0d6

alamb mentioned this pull request Sep 22, 2025

Use newtype to get hash/eq for f64 theirix/datafusion#1

Merged

alamb reviewed Sep 22, 2025

View reviewed changes

alamb mentioned this pull request Sep 23, 2025

Blog on Extending SQL to create own SQL Dialects apache/datafusion-site#97

Open

theirix and others added 8 commits September 23, 2025 19:25

Merge pull request #1 from alamb/alamb/bound

4b09619

Use newtype to get hash/eq for f64

Apply docs suggestions from code review

66a68ec

Co-authored-by: Andrew Lamb <[email protected]>

Optimise imports

ef4752d

Signed-off-by: theirix <[email protected]>

Move main

c98add8

Signed-off-by: theirix <[email protected]>

Fixup newtype casts

6e6e4a2

Signed-off-by: theirix <[email protected]>

Use poll helper and record metrics

6482b52

Signed-off-by: theirix <[email protected]>

Make SessionContextProvider public

eefeaf6

Reuse SessionContextProvider in table sample code

4a8d8d9

github-actions bot added the core Core DataFusion crate label Sep 23, 2025

Reformat SessionContextProvider

8b13b0e

Fix clippy lints

e63de51

vegarsti reviewed Sep 27, 2025

View reviewed changes

theirix added 4 commits September 27, 2025 09:36

Rework comments

591c6a0

Streamline imports

003ccb6

Better separation of ast vs logical expressions

71a1452

Fix typo

500bab9

alamb mentioned this pull request Sep 29, 2025

Examples of extending SQL syntax #17824

Open

alamb changed the title ~~feat: data source sampling via extension~~ feat: Add example for implementing SAMPLE using extension APIs Sep 29, 2025


		// Context provider for tests

		struct MockContextProvider<'a> {

-/// This example demonstrates the table sample support.
+/// This example demonstrates how to extend DataFusion's SQL parser to recognize
+/// other syntax.
+///
+/// This example shows how to extend the DataFusion SQL planner to support the
+/// `TABLESAMPLE` clause in SQL queries and then use a custom user defined node
+/// to implement the sampling logic in the physical plan.

	/// Hashable and comparible f64 for sampling bounds
	/// Hashable and comparable f64 for sampling bounds

feat: Add example for implementing SAMPLE using extension APIs #17633

Are you sure you want to change the base?

feat: Add example for implementing SAMPLE using extension APIs #17633

Uh oh!

Conversation

theirix commented Sep 17, 2025

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

alamb commented Sep 17, 2025

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

theirix commented Sep 23, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vegarsti left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb commented Sep 29, 2025

Uh oh!

geoffreyclaude commented Sep 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

theirix commented Sep 30, 2025

Uh oh!

Uh oh!

feat: Add example for implementing `SAMPLE` using extension APIs #17633

feat: Add example for implementing `SAMPLE` using extension APIs #17633

geoffreyclaude commented Sep 30, 2025 •

edited

Loading