You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/11 01:58:13 UTC
(datafusion) 01/01: fix(parquet): remap sorting columns for partitioned writes (#24211)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24211-d5bd10d8042772539630e7a00275c7b66cac95da
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit a9b61abf04b59ff08e7254b598b07f4fa058c5d1
Author: xudong.w <wx...@gmail.com>
AuthorDate: Tue Aug 11 09:57:49 2026 +0800
fix(parquet): remap sorting columns for partitioned writes (#24211)
## Which issue does this PR close?
- Closes #24210.
## Rationale for this change
When Hive partition columns are not kept in Parquet files,
`FileSinkConfig::output_schema()` still describes the sink input while
`get_writer_schema(&conf)` removes those partition columns.
`ParquetFormat::create_writer_physical_plan` currently converts the
input ordering directly to Parquet `sorting_columns`, so its column
indices can reference removed columns or the wrong positions in the
written schema. Reading that metadata can then panic in
`SchemaDescriptor::column` with an out-of-bounds index.
The execution ordering and the Parquet footer ordering have different
schema domains: `DataSinkExec` must retain the original input ordering,
while `ParquetSink` metadata must use indices from the actual writer
schema.
## What changes are included in this PR?
- Derive Parquet `sorting_columns` from both the input schema and
`get_writer_schema(&conf)`.
- Omit ordering keys that are removed from the written file, such as
Hive partition columns.
- Remap retained ordering keys to their writer-schema indices while
preserving sort direction and null ordering.
- Keep the original `order_requirements` unchanged for `DataSinkExec`.
- Avoid writing an empty `sorting_columns` list when every ordering key
is removed.
- Add unit and end-to-end regression coverage for partitioned Parquet
writes.
## Are these changes tested?
Yes. The new tests cover the writer-schema remapping directly and verify
the resulting Parquet footer after a partitioned write.
## Are there any user-facing changes?
There are no public API changes. Newly written partitioned Parquet files
now contain valid `sorting_columns` metadata based on their physical
file schema. Existing files with invalid metadata are unchanged and must
be rewritten separately.
---
datafusion/core/tests/parquet/ordering.rs | 80 +++++++++++++++++++
datafusion/datasource-parquet/src/file_format.rs | 10 ++-
datafusion/datasource-parquet/src/metadata.rs | 97 +++++++++++++++++++++---
3 files changed, 177 insertions(+), 10 deletions(-)
diff --git a/datafusion/core/tests/parquet/ordering.rs b/datafusion/core/tests/parquet/ordering.rs
index faecb4ca6a..1bdad7f593 100644
--- a/datafusion/core/tests/parquet/ordering.rs
+++ b/datafusion/core/tests/parquet/ordering.rs
@@ -101,3 +101,83 @@ async fn test_create_table_with_order_writes_sorting_columns() -> Result<()> {
Ok(())
}
+
+/// Test that partition columns are removed and remaining column indices are
+/// remapped when writing sorting_columns to Parquet metadata.
+#[tokio::test]
+async fn test_partitioned_table_remaps_sorting_columns() -> Result<()> {
+ use parquet::file::reader::FileReader;
+ use parquet::file::serialized_reader::SerializedFileReader;
+ use std::fs::File;
+
+ let ctx = SessionContext::new();
+ let tmp_dir = tempdir()?;
+ let table_path = tmp_dir.path().join("sorted_partitioned_table");
+ std::fs::create_dir_all(&table_path)?;
+
+ let create_table_sql = format!(
+ "CREATE EXTERNAL TABLE sorted_partitioned_data (a INT, b VARCHAR, part VARCHAR) \
+ STORED AS PARQUET \
+ LOCATION '{}' \
+ PARTITIONED BY (part) \
+ WITH ORDER (part ASC NULLS FIRST, a ASC NULLS FIRST, b DESC NULLS LAST)",
+ table_path.display()
+ );
+ ctx.sql(&create_table_sql).await?;
+
+ ctx.sql(
+ "INSERT INTO sorted_partitioned_data VALUES \
+ (2, 'c', 'x'), (1, 'a', 'x'), (1, 'b', 'x')",
+ )
+ .await?
+ .collect()
+ .await?;
+
+ let parquet_file = find_parquet_file(&table_path)?
+ .expect("expected a Parquet file in the partition directory");
+
+ let file = File::open(parquet_file)?;
+ let reader = SerializedFileReader::new(file)?;
+ let metadata = reader.metadata();
+ let parquet_schema = metadata.file_metadata().schema_descr();
+ assert_eq!(parquet_schema.num_columns(), 2);
+ assert_eq!(parquet_schema.column(0).name(), "a");
+ assert_eq!(parquet_schema.column(1).name(), "b");
+
+ let sorting = metadata
+ .row_group(0)
+ .sorting_columns()
+ .expect("expected sorting_columns in row group metadata");
+ assert_eq!(sorting.len(), 2);
+
+ assert_eq!(sorting[0].column_idx, 0);
+ assert!(!sorting[0].descending);
+ assert!(sorting[0].nulls_first);
+
+ assert_eq!(sorting[1].column_idx, 1);
+ assert!(sorting[1].descending);
+ assert!(!sorting[1].nulls_first);
+
+ Ok(())
+}
+
+fn find_parquet_file(
+ path: &std::path::Path,
+) -> std::io::Result<Option<std::path::PathBuf>> {
+ for entry in std::fs::read_dir(path)? {
+ let entry = entry?;
+ let path = entry.path();
+ if path.is_dir() {
+ if let Some(path) = find_parquet_file(&path)? {
+ return Ok(Some(path));
+ }
+ } else if path
+ .extension()
+ .is_some_and(|extension| extension == "parquet")
+ {
+ return Ok(Some(path));
+ }
+ }
+
+ Ok(None)
+}
diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs
index 8d19bedeb7..fb67c85e44 100644
--- a/datafusion/datasource-parquet/src/file_format.rs
+++ b/datafusion/datasource-parquet/src/file_format.rs
@@ -50,6 +50,7 @@ use datafusion_common::{
use datafusion_datasource::file::FileSource;
use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
use datafusion_datasource::sink::DataSinkExec;
+use datafusion_datasource::write::get_writer_schema;
use datafusion_expr::dml::InsertOp;
use datafusion_physical_expr_common::sort_expr::{LexOrdering, LexRequirement};
use datafusion_physical_plan::ExecutionPlan;
@@ -534,11 +535,18 @@ impl FileFormat for ParquetFormat {
// Convert ordering requirements to Parquet SortingColumns for file metadata
let sorting_columns = if let Some(ref requirements) = order_requirements {
let ordering: LexOrdering = requirements.clone().into();
+ let writer_schema = get_writer_schema(&conf);
// In cases like `COPY (... ORDER BY ...) TO ...` the ORDER BY clause
// may not be compatible with Parquet sorting columns (e.g. ordering on `random()`).
// So if we cannot create a Parquet sorting column from the ordering requirement,
// we skip setting sorting columns on the Parquet sink.
- lex_ordering_to_sorting_columns(&ordering).ok()
+ lex_ordering_to_sorting_columns(
+ &ordering,
+ conf.output_schema(),
+ &writer_schema,
+ )
+ .ok()
+ .filter(|columns| !columns.is_empty())
} else {
None
};
diff --git a/datafusion/datasource-parquet/src/metadata.rs b/datafusion/datasource-parquet/src/metadata.rs
index 3294ee00f1..38db1e22ee 100644
--- a/datafusion/datasource-parquet/src/metadata.rs
+++ b/datafusion/datasource-parquet/src/metadata.rs
@@ -28,6 +28,7 @@ use datafusion_common::encryption::FileDecryptionProperties;
use datafusion_common::stats::Precision;
use datafusion_common::{
ColumnStatistics, DataFusionError, HashMap, Result, ScalarValue, Statistics,
+ internal_datafusion_err,
};
use datafusion_execution::cache::cache_manager::{
CachedFileMetadataEntry, FileMetadata, FileMetadataCache,
@@ -952,10 +953,15 @@ impl FileMetadata for CachedParquetMetaData {
/// Convert a [`PhysicalSortExpr`] to a Parquet [`SortingColumn`].
///
-/// Returns `Err` if the expression is not a simple column reference.
+/// Returns `Ok(None)` if the referenced column is not in `writer_schema`, such as a
+/// hive partition column that is removed before writing the Parquet file. Returns
+/// `Err` if the expression is not a simple column reference or references a column
+/// outside `input_schema`.
pub(crate) fn sort_expr_to_sorting_column(
sort_expr: &PhysicalSortExpr,
-) -> Result<SortingColumn> {
+ input_schema: &Schema,
+ writer_schema: &Schema,
+) -> Result<Option<SortingColumn>> {
let column = sort_expr.expr.downcast_ref::<Column>().ok_or_else(|| {
DataFusionError::Plan(format!(
"Parquet sorting_columns only supports simple column references, \
@@ -964,27 +970,49 @@ pub(crate) fn sort_expr_to_sorting_column(
))
})?;
- let column_idx: i32 = column.index().try_into().map_err(|_| {
+ let input_field = input_schema.fields().get(column.index()).ok_or_else(|| {
+ internal_datafusion_err!(
+ "Parquet sorting column '{}' references index {} but the input schema has {} columns",
+ column.name(),
+ column.index(),
+ input_schema.fields().len()
+ )
+ })?;
+ let Some((writer_index, _)) = writer_schema.column_with_name(input_field.name())
+ else {
+ return Ok(None);
+ };
+
+ let column_idx: i32 = writer_index.try_into().map_err(|_| {
DataFusionError::Plan(format!(
- "Column index {} is too large to be represented as i32",
- column.index()
+ "Column index {writer_index} is too large to be represented as i32"
))
})?;
- Ok(SortingColumn {
+ Ok(Some(SortingColumn {
column_idx,
descending: sort_expr.options.descending,
nulls_first: sort_expr.options.nulls_first,
- })
+ }))
}
/// Convert a [`LexOrdering`] to `Vec<SortingColumn>` for Parquet.
///
-/// Returns `Err` if any expression is not a simple column reference.
+/// Columns that are not present in `writer_schema` are omitted from the resulting
+/// metadata. Returns `Err` if any expression is not a simple column reference or
+/// references a column outside `input_schema`.
pub(crate) fn lex_ordering_to_sorting_columns(
ordering: &LexOrdering,
+ input_schema: &Schema,
+ writer_schema: &Schema,
) -> Result<Vec<SortingColumn>> {
- ordering.iter().map(sort_expr_to_sorting_column).collect()
+ ordering
+ .iter()
+ .filter_map(|sort_expr| {
+ sort_expr_to_sorting_column(sort_expr, input_schema, writer_schema)
+ .transpose()
+ })
+ .collect()
}
/// Extracts ordering information from Parquet metadata.
@@ -1058,6 +1086,57 @@ fn sorting_columns_to_physical_exprs(
mod tests {
use super::*;
use arrow::array::Int32Array;
+ use arrow::compute::SortOptions;
+ use arrow::datatypes::Field;
+
+ #[test]
+ fn test_lex_ordering_to_sorting_columns_uses_writer_schema() -> Result<()> {
+ let input_schema = Schema::new(vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("part", DataType::Utf8, true),
+ Field::new("b", DataType::Utf8, true),
+ ]);
+ let writer_schema = Schema::new(vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Utf8, true),
+ ]);
+ let ordering = LexOrdering::new(vec![
+ PhysicalSortExpr::new(
+ Arc::new(Column::new("part", 1)),
+ SortOptions::default(),
+ ),
+ PhysicalSortExpr::new(Arc::new(Column::new("a", 0)), SortOptions::default()),
+ PhysicalSortExpr::new(
+ Arc::new(Column::new("b", 2)),
+ SortOptions {
+ descending: true,
+ nulls_first: false,
+ },
+ ),
+ ])
+ .unwrap();
+
+ let sorting_columns =
+ lex_ordering_to_sorting_columns(&ordering, &input_schema, &writer_schema)?;
+
+ assert_eq!(
+ sorting_columns,
+ vec![
+ SortingColumn {
+ column_idx: 0,
+ descending: false,
+ nulls_first: true,
+ },
+ SortingColumn {
+ column_idx: 1,
+ descending: true,
+ nulls_first: false,
+ },
+ ]
+ );
+
+ Ok(())
+ }
#[test]
fn test_has_any_exact_match() {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org