Add range partitioning sqllogictest fixture

datafusion/physical-expr/src/partitioning.rs

@@ -171,6 +171,10 @@ impl Display for Partitioning {
 /// Values equal to split point `i` belong to partition `i + 1`, so interior
 /// partitions are lower-inclusive and upper-exclusive.
 ///
+/// If a source declares range partitioning, it is responsible for placing each
+/// row in the partition described by the split points, DataFusion will not validate this is
+/// upheld.
+///
 /// For a single range key:
 ///
 /// ```text

datafusion/sqllogictest/Cargo.toml

@@ -47,6 +47,7 @@ bytes = { workspace = true, optional = true }
 chrono = { workspace = true, optional = true }
 clap = { version = "4.5.60", features = ["derive", "env"] }
 datafusion = { workspace = true, default-features = true, features = ["avro"] }
+datafusion-datasource = { workspace = true }
 datafusion-spark = { workspace = true, features = ["core"] }
 datafusion-substrait = { workspace = true, default-features = true, optional = true }
 futures = { workspace = true }

datafusion/sqllogictest/src/test_context.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use std::collections::HashMap;
+use std::fmt;
 use std::fs::File;
 use std::io::Write;
 use std::path::Path;
@@ -28,6 +29,7 @@ use arrow::array::{
     TimestampNanosecondArray, UInt32Array, UnionArray,
 };
 use arrow::buffer::ScalarBuffer;
+use arrow::compute::SortOptions;
 use arrow::datatypes::{
     DataType, Field, FieldRef, Fields, Schema, SchemaRef, TimeUnit, UInt32Type,
     UnionFields,
@@ -36,20 +38,31 @@ use arrow::record_batch::RecordBatch;
 use datafusion::catalog::{
     CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider, SchemaProvider, Session,
 };
-use datafusion::common::{DataFusionError, Result, not_impl_err};
+use datafusion::common::{DataFusionError, Result, ScalarValue, not_impl_err};
+use datafusion::datasource::source::{DataSource, DataSourceExec};
+use datafusion::execution::context::TaskContext;
 use datafusion::functions::math::abs;
 use datafusion::logical_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl};
 use datafusion::logical_expr::planner::TypePlanner;
 use datafusion::logical_expr::{
     ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
     Volatility, create_udf,
 };
-use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_expr::expressions::col as physical_col;
+use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr};
+use datafusion::physical_plan::execution_plan::SchedulingType;
+use datafusion::physical_plan::projection::ProjectionExprs;
+use datafusion::physical_plan::{
+    DisplayFormatType, ExecutionPlan, Partitioning, RangePartitioning,
+    SendableRecordBatchStream, SplitPoint, Statistics, project_schema,
+};
 use datafusion::prelude::*;
 use datafusion::{
     datasource::{MemTable, TableProvider, TableType},
     prelude::{CsvReadOptions, SessionContext},
 };
+use datafusion_datasource::memory::MemorySourceConfig;
 use datafusion_spark::SessionStateBuilderSpark;
 
 use crate::is_spark_path;
@@ -167,6 +180,10 @@ impl TestContext {
                 info!("Registering table with many types");
                 register_table_with_many_types(test_ctx.session_ctx()).await;
             }
+            "range_partitioning.slt" => {
+                info!("Registering range partitioned table");
+                register_range_partitioned_table(test_ctx.session_ctx());
+            }
             "metadata.slt" | "arrow_field.slt" => {
                 info!("Registering metadata table tables");
                 register_metadata_tables(test_ctx.session_ctx()).await;
@@ -286,6 +303,212 @@ fn register_strict_schema_provider(ctx: &SessionContext) {
     );
 }
 
+// ==============================================================================
+// Range Partitioned Table (sqllogictest-only)
+// ==============================================================================
+
+#[derive(Debug)]
+struct RangePartitionedTable {
+    schema: SchemaRef,
+    partitions: Vec<Vec<RecordBatch>>,
+    range_column_index: usize,
+    split_points: Vec<SplitPoint>,
+}
+
+#[async_trait]
+impl TableProvider for RangePartitionedTable {
+    fn schema(&self) -> SchemaRef {
+        Arc::clone(&self.schema)
+    }
+
+    fn table_type(&self) -> TableType {
+        TableType::Base
+    }
+
+    async fn scan(
+        &self,
+        state: &dyn Session,
+        projection: Option<&Vec<usize>>,
+        _filters: &[Expr],
+        _limit: Option<usize>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let projected_schema = project_schema(&self.schema, projection)?;
+        let mut source = MemorySourceConfig::try_new(
+            &self.partitions,
+            Arc::clone(&self.schema),
+            projection.cloned(),
+        )?;
+        source = source.with_show_sizes(state.config_options().explain.show_sizes);
+
+        let output_partitioning =
+            self.output_partitioning(projection, &projected_schema)?;
+        let source = RangePartitionedSource {
+            inner: source,
+            output_partitioning,
+        };
+
+        Ok(DataSourceExec::from_data_source(source))
+    }
+}
+
+impl RangePartitionedTable {
+    fn output_partitioning(
+        &self,
+        projection: Option<&Vec<usize>>,
+        projected_schema: &SchemaRef,
+    ) -> Result<Partitioning> {
+        let Some(projected_range_index) =
+            projected_index(self.range_column_index, projection)
+        else {
+            return Ok(Partitioning::UnknownPartitioning(self.partitions.len()));
+        };
+
+        let range_column = projected_schema.field(projected_range_index).name();
+        let ordering = LexOrdering::new(vec![PhysicalSortExpr::new(
+            physical_col(range_column, projected_schema)?,
+            SortOptions::default(),
+        )])
+        .expect("range ordering should not be empty");
+
+        Ok(Partitioning::Range(RangePartitioning::try_new(
+            ordering,
+            self.split_points.clone(),
+        )?))
+    }
+}
+
+fn projected_index(
+    column_index: usize,
+    projection: Option<&Vec<usize>>,
+) -> Option<usize> {
+    projection
+        .map(|projection| projection.iter().position(|idx| *idx == column_index))
+        .unwrap_or(Some(column_index))
+}
+
+#[derive(Clone, Debug)]
+struct RangePartitionedSource {
+    inner: MemorySourceConfig,
+    output_partitioning: Partitioning,
+}
+
+impl DataSource for RangePartitionedSource {
+    fn open(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> Result<SendableRecordBatchStream> {
+        self.inner.open(partition, context)
+    }
+
+    fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
+        self.inner.fmt_as(t, f)?;
+        match t {
+            DisplayFormatType::Default | DisplayFormatType::Verbose => {
+                write!(f, ", output_partitioning={}", self.output_partitioning)
+            }
+            DisplayFormatType::TreeRender => Ok(()),
+        }
+    }
+
+    fn output_partitioning(&self) -> Partitioning {
+        self.output_partitioning.clone()
+    }
+
+    fn eq_properties(&self) -> EquivalenceProperties {
+        self.inner.eq_properties()
+    }
+
+    fn scheduling_type(&self) -> SchedulingType {
+        self.inner.scheduling_type()
+    }
+
+    fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
+        self.inner.partition_statistics(partition)
+    }
+
+    fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn DataSource>> {
+        Some(Arc::new(Self {
+            inner: self.inner.clone().with_limit(limit),
+            output_partitioning: self.output_partitioning.clone(),
+        }))
+    }
+
+    fn fetch(&self) -> Option<usize> {
+        self.inner.fetch()
+    }
+
+    fn try_swapping_with_projection(
+        &self,
+        _projection: &ProjectionExprs,
+    ) -> Result<Option<Arc<dyn DataSource>>> {
+        // Range partitioning metadata is projection-sensitive. This fixture
+        // computes it in TableProvider::scan, so do not rewrite later
+        // ProjectionExec nodes into the source.
+        Ok(None)
+    }
+}
+
+fn register_range_partitioned_table(ctx: &SessionContext) {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("range_key", DataType::Int32, false),
+        Field::new("non_range_key", DataType::Int32, false),
+        Field::new("value", DataType::Int32, false),
+    ]));
+    let partitions = vec![
+        vec![range_partition_batch(&schema, &[1, 5], &[1, 2], &[10, 50])],
+        vec![range_partition_batch(
+            &schema,
+            &[10, 15],
+            &[1, 2],
+            &[100, 150],
+        )],
+        vec![range_partition_batch(
+            &schema,
+            &[20, 25],
+            &[1, 2],
+            &[200, 250],
+        )],
+        vec![range_partition_batch(
+            &schema,
+            &[30, 35],
+            &[1, 2],
+            &[300, 350],
+        )],
+    ];
+    let split_points = vec![
+        SplitPoint::new(vec![ScalarValue::Int32(Some(10))]),
+        SplitPoint::new(vec![ScalarValue::Int32(Some(20))]),
+        SplitPoint::new(vec![ScalarValue::Int32(Some(30))]),
+    ];
+    let table = RangePartitionedTable {
+        schema,
+        partitions,
+        range_column_index: 0,
+        split_points,
+    };
+
+    ctx.register_table("range_partitioned", Arc::new(table))
+        .expect("range partitioned table registration should succeed");
+}
+
+fn range_partition_batch(
+    schema: &SchemaRef,
+    range_key: &[i32],
+    non_range_key: &[i32],
+    value: &[i32],
+) -> RecordBatch {
+    RecordBatch::try_new(
+        Arc::clone(schema),
+        vec![
+            Arc::new(Int32Array::from(range_key.to_vec())),
+            Arc::new(Int32Array::from(non_range_key.to_vec())),
+            Arc::new(Int32Array::from(value.to_vec())),
+        ],
+    )
+    .expect("range partition batch should be valid")
+}
+
 #[cfg(feature = "avro")]
 pub async fn register_avro_tables(ctx: &mut TestContext) {
     use datafusion::prelude::AvroReadOptions;