diff --git a/crates/integrations/datafusion/src/format_partition_analyze.rs b/crates/integrations/datafusion/src/format_partition_analyze.rs new file mode 100644 index 000000000..71a8c6104 --- /dev/null +++ b/crates/integrations/datafusion/src/format_partition_analyze.rs @@ -0,0 +1,196 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! ANALYZE TABLE for Format Tables with catalog-managed partitions. + +use std::collections::HashSet; + +use datafusion::error::{DataFusionError, Result as DFResult}; +use datafusion::prelude::DataFrame; +use datafusion::sql::sqlparser::ast::{Analyze, Expr as SqlExpr}; +use paimon::table::FormatTablePartitionStatsCollector; + +use crate::error::to_datafusion_error; +use crate::format_partition_ddl::{ + ensure_catalog_managed_format_table, has_custom_location, parse_format_partition_spec, +}; +use crate::sql_context::{ + normalize_schema_identifier, ok_result, partition_assignment, SQLContext, +}; + +/// `ANALYZE TABLE t [PARTITION (...)] COMPUTE STATISTICS [NOSCAN]` on a Format Table with +/// catalog-managed partitions. Mirrors Java `PaimonAnalyzeFormatTablePartitionsCommand`. +pub(crate) async fn execute_analyze( + ctx: &SQLContext, + analyze: &Analyze, + enable_ident_normalization: bool, +) -> DFResult { + let Some(table_name) = &analyze.table_name else { + return Err(DataFusionError::Plan( + "ANALYZE requires a table name".to_string(), + )); + }; + if analyze.for_columns || !analyze.columns.is_empty() { + return Err(DataFusionError::NotImplemented( + "ANALYZE TABLE ... FOR COLUMNS is not supported: a Format Table has nowhere to \ + keep column statistics" + .to_string(), + )); + } + if analyze.cache_metadata { + return Err(DataFusionError::NotImplemented( + "ANALYZE TABLE ... CACHE METADATA is not supported".to_string(), + )); + } + if !analyze.compute_statistics { + return Err(DataFusionError::Plan( + "ANALYZE TABLE requires COMPUTE STATISTICS".to_string(), + )); + } + SQLContext::ensure_partition_command_target(table_name, "ANALYZE TABLE")?; + let (catalog, _catalog_name, identifier) = ctx.resolve_catalog_and_table(table_name)?; + let table = catalog + .get_table(&identifier) + .await + .map_err(to_datafusion_error)?; + ensure_catalog_managed_format_table(&table, "ANALYZE TABLE")?; + let prefix = analyze_partition_prefix( + analyze.partitions.as_deref().unwrap_or_default(), + &table, + enable_ident_normalization, + )?; + + let selected = catalog + .list_partitions(&identifier) + .await + .map_err(to_datafusion_error)? + .into_iter() + .filter(|partition| { + prefix + .iter() + .all(|(key, value)| partition.spec.get(key) == Some(value)) + }) + .collect::>(); + if selected.is_empty() && !prefix.is_empty() { + return Err(DataFusionError::Plan(format!( + "Partition {prefix:?} does not exist in table {}", + identifier.full_name() + ))); + } + let custom_located = selected + .iter() + .filter(|partition| has_custom_location(partition)) + .map(|partition| &partition.spec) + .collect::>(); + if !custom_located.is_empty() { + return Err(DataFusionError::NotImplemented(format!( + "ANALYZE TABLE cannot measure partitions with a custom location in Format Table \ + {}: {custom_located:?}", + identifier.full_name() + ))); + } + if selected.is_empty() { + return ok_result(ctx.ctx()); + } + + let specs = selected + .into_iter() + .map(|partition| partition.spec) + .collect::>(); + let statistics = FormatTablePartitionStatsCollector::new( + &table, + !analyze.noscan, + format_table_statistics_parallelism(ctx), + ) + .collect(&specs) + .await + .map_err(to_datafusion_error)?; + catalog + .create_partitions_with_statistics(&identifier, specs, true, Some(statistics), true) + .await + .map_err(to_datafusion_error)?; + ok_result(ctx.ctx()) +} + +/// `format-table.statistics.parallelism` from the session (`SET 'paimon.'`), default 8. +/// A value below one is read as one. +fn format_table_statistics_parallelism(ctx: &SQLContext) -> usize { + const KEY: &str = "format-table.statistics.parallelism"; + ctx.dynamic_options() + .read() + .unwrap() + .get(KEY) + .and_then(|value| value.trim().parse::().ok()) + .map(|value| value.max(1) as usize) + .unwrap_or(8) +} + +/// The values an `ANALYZE ... PARTITION (...)` clause fixes, in partition-key order; valued +/// columns must be a leading run of the keys, so `PARTITION (hour = '00')` is rejected. +fn analyze_partition_prefix( + expressions: &[SqlExpr], + table: &paimon::Table, + enable_ident_normalization: bool, +) -> DFResult> { + let partition_keys = table.schema().partition_keys(); + let mut named = HashSet::with_capacity(expressions.len()); + let mut assignments = Vec::with_capacity(expressions.len()); + for expression in expressions { + let column = match expression { + SqlExpr::Identifier(identifier) => { + normalize_schema_identifier(identifier, enable_ident_normalization) + } + other => { + let (column, _) = partition_assignment(other, enable_ident_normalization)?; + assignments.push(other.clone()); + column + } + }; + if !partition_keys.contains(&column) { + return Err(DataFusionError::Plan(format!( + "Column '{column}' is not a partition column" + ))); + } + if !named.insert(column.clone()) { + return Err(DataFusionError::Plan(format!( + "Duplicate partition column '{column}'" + ))); + } + } + let spec = parse_format_partition_spec( + &assignments, + table, + false, + Some("ANALYZE TABLE"), + enable_ident_normalization, + )?; + let leading = partition_keys + .iter() + .take_while(|key| spec.contains_key(key.as_str())) + .count(); + if leading != spec.len() { + return Err(DataFusionError::Plan(format!( + "ANALYZE TABLE {} PARTITION must give values for a leading run of its partition \ + columns {partition_keys:?}", + table.identifier().full_name() + ))); + } + Ok(partition_keys[..leading] + .iter() + .map(|key| (key.clone(), spec[key].clone())) + .collect()) +} diff --git a/crates/integrations/datafusion/src/format_partition_ddl.rs b/crates/integrations/datafusion/src/format_partition_ddl.rs index b611c7798..5c08a359e 100644 --- a/crates/integrations/datafusion/src/format_partition_ddl.rs +++ b/crates/integrations/datafusion/src/format_partition_ddl.rs @@ -360,7 +360,7 @@ pub(crate) async fn drop_catalog_managed_partitions( /// Whether the catalog registered a partition at a location of its own rather than under the /// table directory. -fn has_custom_location(partition: &paimon::spec::Partition) -> bool { +pub(crate) fn has_custom_location(partition: &paimon::spec::Partition) -> bool { partition .options .as_ref() @@ -386,9 +386,9 @@ pub(crate) fn ensure_catalog_managed_format_table( Ok(()) } -/// `mutating_operation` names the statement when it changes partitions (ADD or DROP PARTITION), -/// which refuses a blank string for a string partition column. -fn parse_format_partition_spec( +/// `mutating_operation` names the statement when it changes partitions or their statistics (ADD +/// or DROP PARTITION, ANALYZE TABLE), which refuses a blank string for a string partition column. +pub(crate) fn parse_format_partition_spec( exprs: &[SqlExpr], table: &paimon::Table, require_complete: bool, diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index dca1d4097..ad1e8c32d 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -43,6 +43,7 @@ mod catalog; mod delete; mod error; mod filter_pushdown; +mod format_partition_analyze; mod format_partition_ddl; mod format_partition_repair; #[cfg(feature = "fulltext")] diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 09cf9764a..41a425ebc 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -436,7 +436,6 @@ impl SQLContext { } } - #[cfg(test)] pub(crate) fn dynamic_options(&self) -> &DynamicOptions { &self.dynamic_options } @@ -611,6 +610,14 @@ impl SQLContext { .await } Statement::Msck(msck) => crate::format_partition_repair::execute_msck(self, msck).await, + Statement::Analyze(analyze) => { + crate::format_partition_analyze::execute_analyze( + self, + analyze, + enable_ident_normalization, + ) + .await + } Statement::CreateView(create_view) => { if create_view.temporary { // Temporary views are always handled by us (Paimon catalog temp storage) diff --git a/crates/integrations/datafusion/tests/rest_format_partition_sql.rs b/crates/integrations/datafusion/tests/rest_format_partition_sql.rs index 9a3073f7e..e07b1f4ba 100644 --- a/crates/integrations/datafusion/tests/rest_format_partition_sql.rs +++ b/crates/integrations/datafusion/tests/rest_format_partition_sql.rs @@ -21,13 +21,17 @@ mod common; mod mock_server; use std::collections::HashMap; +use std::path::Path; use std::sync::Arc; +use arrow_array::{Int64Array, RecordBatch}; +use arrow_schema::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use paimon::api::ConfigResponse; use paimon::catalog::RESTCatalog; use paimon::spec::{BigIntType, BooleanType, DataType, DateType, IntType, Schema, VarCharType}; use paimon::{CatalogOptions, Options}; use paimon_datafusion::SQLContext; +use parquet::arrow::ArrowWriter; use tempfile::TempDir; use mock_server::{start_mock_server, RESTServer}; @@ -130,6 +134,66 @@ fn spec(values: &[(&str, &str)]) -> HashMap { .collect() } +const UNKNOWN: i64 = paimon::spec::Partition::UNKNOWN; + +/// The partitions the catalog holds, by partition name with keys in name order. +fn partition_statistics(server: &RESTServer) -> HashMap { + server + .table_partitions(DATABASE, TABLE) + .into_iter() + .map(|partition| { + let mut entries = partition.spec.iter().collect::>(); + entries.sort(); + let name = entries + .into_iter() + .map(|(key, value)| format!("{key}={value}")) + .collect::>() + .join("/"); + (name, partition) + }) + .collect() +} + +fn counts(partition: &paimon::spec::Partition) -> (i64, i64) { + (partition.record_count, partition.file_count) +} + +fn write_ids(directory: &Path, ids: &[i64]) { + write_ids_file(&directory.join("part-0.parquet"), ids); +} + +fn write_ids_file(path: &Path, ids: &[i64]) { + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "id", + ArrowDataType::Int64, + true, + )])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(ids.to_vec()))], + ) + .unwrap(); + let file = std::fs::File::create(path).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); +} + +async fn ids(context: &SQLContext, sql: &str) -> Vec { + let mut ids = Vec::new(); + for batch in context.sql(sql).await.unwrap().collect().await.unwrap() { + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + ids.extend(values.iter().flatten()); + } + ids.sort_unstable(); + ids +} + #[cfg(not(windows))] #[tokio::test] async fn test_partition_commands_update_rest_metadata_and_directories() { @@ -593,6 +657,280 @@ async fn test_msck_repair_keeps_a_partition_at_a_custom_location() { ); } +#[cfg(not(windows))] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_analyze_measures_registered_partitions_and_replaces_their_statistics() { + let (temp_dir, server, context) = dt_hh_table(&[("a", "00"), ("a", "01"), ("b", "00")]).await; + write_ids_file(&temp_dir.path().join("dt=a/hh=00/part-0.parquet"), &[1, 2]); + write_ids_file(&temp_dir.path().join("dt=a/hh=00/part-1.parquet"), &[3]); + write_ids_file(&temp_dir.path().join("dt=a/hh=01/part-0.parquet"), &[4]); + + // NOSCAN measures what a listing gives and leaves the row counts as they were. + common::exec( + &context, + &format!("ANALYZE TABLE {TABLE_NAME} COMPUTE STATISTICS NOSCAN"), + ) + .await; + let measured = partition_statistics(&server); + assert_eq!(counts(&measured["dt=a/hh=00"]), (UNKNOWN, 2)); + assert_eq!(counts(&measured["dt=a/hh=01"]), (UNKNOWN, 1)); + assert_eq!(counts(&measured["dt=b/hh=00"]), (UNKNOWN, 0)); + assert!(measured["dt=a/hh=00"].file_size_in_bytes > 0); + assert!(measured["dt=a/hh=00"].last_file_creation_time > 0); + assert_eq!(measured["dt=b/hh=00"].file_size_in_bytes, 0); + assert_eq!(measured["dt=b/hh=00"].last_file_creation_time, UNKNOWN); + + // A full ANALYZE reads every footer, and an empty partition holds exactly no rows. + common::exec( + &context, + &format!("ANALYZE TABLE {TABLE_NAME} COMPUTE STATISTICS"), + ) + .await; + let measured = partition_statistics(&server); + assert_eq!(counts(&measured["dt=a/hh=00"]), (3, 2)); + assert_eq!(counts(&measured["dt=a/hh=01"]), (1, 1)); + assert_eq!(counts(&measured["dt=b/hh=00"]), (0, 0)); + + // A later NOSCAN keeps the known row counts, and measuring again replaces rather than adds. + common::exec( + &context, + &format!("ANALYZE TABLE {TABLE_NAME} COMPUTE STATISTICS NOSCAN"), + ) + .await; + let remeasured = partition_statistics(&server); + assert_eq!(counts(&remeasured["dt=a/hh=00"]), (3, 2)); + assert_eq!(counts(&remeasured["dt=a/hh=01"]), (1, 1)); + assert_eq!( + remeasured["dt=a/hh=00"].file_size_in_bytes, + measured["dt=a/hh=00"].file_size_in_bytes + ); + + let calls = server.create_partitions_calls(); + let (_, _, request) = calls.last().unwrap(); + assert!(request.ignore_if_exists); + assert_eq!(request.replace_statistics, Some(true)); + assert_eq!(server.table_partition_specs(DATABASE, TABLE).len(), 3); +} + +#[cfg(not(windows))] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_analyze_partition_clause_selects_a_leading_run_of_partition_values() { + let (temp_dir, server, context) = dt_hh_table(&[("a", "00"), ("a", "01"), ("b", "00")]).await; + for directory in ["dt=a/hh=00", "dt=a/hh=01", "dt=b/hh=00"] { + write_ids(&temp_dir.path().join(directory), &[1]); + } + let file_counts = |server: &RESTServer| { + let measured = partition_statistics(server); + ["dt=a/hh=00", "dt=a/hh=01", "dt=b/hh=00"].map(|name| measured[name].file_count) + }; + + common::exec( + &context, + &format!( + "ANALYZE TABLE {TABLE_NAME} PARTITION (dt = 'a', hh = '00') COMPUTE STATISTICS NOSCAN" + ), + ) + .await; + assert_eq!(file_counts(&server), [1, UNKNOWN, UNKNOWN]); + + // A column named without a value means every value of it. + common::exec( + &context, + &format!("ANALYZE TABLE {TABLE_NAME} PARTITION (dt = 'a', hh) COMPUTE STATISTICS NOSCAN"), + ) + .await; + assert_eq!(file_counts(&server), [1, 1, UNKNOWN]); + common::exec( + &context, + &format!("ANALYZE TABLE {TABLE_NAME} PARTITION (dt, hh) COMPUTE STATISTICS NOSCAN"), + ) + .await; + assert_eq!(file_counts(&server), [1, 1, 1]); + + for (clause, message) in [ + ("PARTITION (hh = '00')", "leading run"), + ("PARTITION (id = 1)", "not a partition column"), + ("PARTITION (dt = 'zzz')", "does not exist"), + ] { + common::assert_sql_error( + &context, + &format!("ANALYZE TABLE {TABLE_NAME} {clause} COMPUTE STATISTICS NOSCAN"), + message, + ) + .await; + } +} + +#[cfg(not(windows))] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_analyze_reads_a_partition_value_as_its_column_type() { + let temp_dir = tempfile::tempdir().unwrap(); + let schema = format_table_schema(&[("p", DataType::Int(IntType::new()))]); + let (server, context) = setup_rest_table(&temp_dir, schema).await; + common::exec( + &context, + &format!("ALTER TABLE {TABLE_NAME} ADD PARTITION (p = 1)"), + ) + .await; + write_ids(&temp_dir.path().join("p=1"), &[1]); + + common::exec( + &context, + &format!("ANALYZE TABLE {TABLE_NAME} PARTITION (p = '01') COMPUTE STATISTICS NOSCAN"), + ) + .await; + + assert_eq!(partition_statistics(&server)["p=1"].file_count, 1); +} + +#[cfg(not(windows))] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_analyze_and_scan_count_only_the_files_a_reader_returns() { + let temp_dir = tempfile::tempdir().unwrap(); + let (server, context) = + setup_rest_table(&temp_dir, format_table_schema(&[("dt", varchar())])).await; + common::exec( + &context, + &format!("ALTER TABLE {TABLE_NAME} ADD PARTITION (dt = 'a')"), + ) + .await; + let partition = temp_dir.path().join("dt=a"); + write_ids_file(&partition.join("part-0.parquet"), &[1]); + // What committers and tools leave beside the data: staging trees, markers, hidden files. + write_ids_file(&partition.join("_temporary/0/part-9.parquet"), &[9]); + write_ids_file(&partition.join("__magic_job-1/tasks/part-8.parquet"), &[8]); + write_ids_file(&partition.join(".part-7.parquet"), &[7]); + std::fs::write(partition.join("_SUCCESS"), b"").unwrap(); + std::fs::write(partition.join("notes.txt"), b"not data").unwrap(); + + common::exec( + &context, + &format!("ANALYZE TABLE {TABLE_NAME} COMPUTE STATISTICS"), + ) + .await; + + assert_eq!(counts(&partition_statistics(&server)["dt=a"]), (1, 1)); + assert_eq!( + ids( + &context, + &format!("SELECT id FROM {TABLE_NAME} WHERE dt = 'a'") + ) + .await, + vec![1] + ); +} + +#[cfg(not(windows))] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_analyze_leaves_a_row_count_unknown_rather_than_short() { + let temp_dir = tempfile::tempdir().unwrap(); + let (server, context) = + setup_rest_table(&temp_dir, format_table_schema(&[("dt", varchar())])).await; + common::exec( + &context, + &format!("ALTER TABLE {TABLE_NAME} ADD PARTITION (dt = 'a')"), + ) + .await; + write_ids_file(&temp_dir.path().join("dt=a/part-0.parquet"), &[1, 2]); + std::fs::write( + temp_dir.path().join("dt=a/part-1.parquet"), + b"not a parquet footer", + ) + .unwrap(); + + common::exec( + &context, + &format!("ANALYZE TABLE {TABLE_NAME} COMPUTE STATISTICS"), + ) + .await; + + // A sum missing one file, reported as exact, would be worse than no number. + assert_eq!(counts(&partition_statistics(&server)["dt=a"]), (UNKNOWN, 2)); +} + +#[cfg(not(windows))] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_analyze_refuses_what_it_cannot_measure() { + let temp_dir = tempfile::tempdir().unwrap(); + let (server, context) = + setup_rest_table(&temp_dir, format_table_schema(&[("dt", varchar())])).await; + for dt in ["a", "b"] { + common::exec( + &context, + &format!("ALTER TABLE {TABLE_NAME} ADD PARTITION (dt = '{dt}')"), + ) + .await; + write_ids(&temp_dir.path().join(format!("dt={dt}")), &[1]); + } + + common::assert_sql_error( + &context, + &format!("ANALYZE TABLE {TABLE_NAME} COMPUTE STATISTICS FOR COLUMNS id"), + "FOR COLUMNS", + ) + .await; + + // A blank string names the default partition, which the statement would measure instead. + common::assert_sql_error( + &context, + &format!("ANALYZE TABLE {TABLE_NAME} PARTITION (dt = '') COMPUTE STATISTICS"), + "empty or whitespace-only string for partition column 'dt'", + ) + .await; + + server.set_table_partition_options( + DATABASE, + TABLE, + &spec(&[("dt", "b")]), + HashMap::from([("path".to_string(), "file:///elsewhere/b".to_string())]), + ); + common::assert_sql_error( + &context, + &format!("ANALYZE TABLE {TABLE_NAME} COMPUTE STATISTICS NOSCAN"), + "custom location", + ) + .await; + assert!(partition_statistics(&server) + .values() + .all(|partition| partition.file_count == UNKNOWN)); + + // A non-positive parallelism is read as one rather than failing the statement. + common::exec( + &context, + "SET \"paimon.format-table.statistics.parallelism\" = '0'", + ) + .await; + common::exec( + &context, + &format!("ANALYZE TABLE {TABLE_NAME} PARTITION (dt = 'a') COMPUTE STATISTICS"), + ) + .await; + assert_eq!(counts(&partition_statistics(&server)["dt=a"]), (1, 1)); + + // Without catalog-managed partitions there is no catalog to write the numbers to. + let plain = Schema::builder() + .column("dt", varchar()) + .column("id", DataType::BigInt(BigIntType::new())) + .partition_keys(["dt"]) + .option("type", "format-table") + .option("file.format", "parquet") + .build() + .unwrap(); + server.add_table_with_schema( + DATABASE, + "plain", + plain, + &format!("file://{}/plain", temp_dir.path().display()), + ); + server.set_table_external(DATABASE, "plain", false); + common::assert_sql_error( + &context, + "ANALYZE TABLE paimon.default.plain COMPUTE STATISTICS", + "catalog-managed", + ) + .await; +} + /// `SQLContext::sql` futures have to stay `Send` for callers that box or spawn them; this stops /// compiling when a stream over borrowed items anywhere below a statement takes that away. #[allow(dead_code)] diff --git a/crates/paimon/src/arrow/format/mod.rs b/crates/paimon/src/arrow/format/mod.rs index 7a5abbd89..f67ae5651 100644 --- a/crates/paimon/src/arrow/format/mod.rs +++ b/crates/paimon/src/arrow/format/mod.rs @@ -30,7 +30,7 @@ pub(crate) use parquet::ParquetFormatWriter; use super::ParquetReadBudget; use super::RowFilterFactory; -use crate::io::{FileRead, OutputFile}; +use crate::io::{FileIO, FileRead, OutputFile}; use crate::spec::stats::BinaryTableStats; use crate::spec::{DataField, Predicate}; use crate::table::{ArrowRecordBatchStream, RowRange}; @@ -158,6 +158,31 @@ impl FormatWriteResult { } } +/// Rows in a data file of the given format, read from its footer alone, or `None` when the +/// format keeps no row count there and every row would have to be decoded to count them. +pub(crate) async fn read_file_row_count( + file_io: &FileIO, + format: &str, + path: &str, + file_size: u64, +) -> crate::Result> { + match format.to_ascii_lowercase().as_str() { + "parquet" => { + let reader = file_io.new_input(path)?.reader().await?; + parquet::read_row_count(Box::new(reader), file_size) + .await + .map(Some) + } + "orc" => { + let reader = file_io.new_input(path)?.reader().await?; + orc::read_row_count(Box::new(reader), file_size) + .await + .map(Some) + } + _ => Ok(None), + } +} + /// Create a format reader based on the file extension. #[cfg(test)] pub(crate) fn create_format_reader( diff --git a/crates/paimon/src/arrow/format/orc.rs b/crates/paimon/src/arrow/format/orc.rs index fd2e38114..116b62aaf 100644 --- a/crates/paimon/src/arrow/format/orc.rs +++ b/crates/paimon/src/arrow/format/orc.rs @@ -368,6 +368,29 @@ fn build_range_row_selection( ) } +/// Rows in an ORC file, read from its footer alone. +pub(crate) async fn read_row_count( + reader: Box, + file_size: u64, +) -> crate::Result { + let builder = ArrowReaderBuilder::try_new_async(OrcFileReader::new(file_size, reader)) + .await + .map_err(|error| Error::UnexpectedError { + message: format!("Failed to open ORC file: {error}"), + source: Some(Box::new(error)), + })?; + let rows = builder + .file_metadata() + .stripe_metadatas() + .iter() + .map(|stripe| stripe.number_of_rows()) + .sum::(); + i64::try_from(rows).map_err(|_| Error::DataInvalid { + message: format!("ORC file holds {rows} rows, more than a row count can carry"), + source: None, + }) +} + // --------------------------------------------------------------------------- // OrcFileReader — adapts paimon FileRead to orc-rust AsyncChunkReader // --------------------------------------------------------------------------- diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index 7aab8a074..583b76285 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -1964,6 +1964,23 @@ const METADATA_SIZE_HINT: usize = 512 * 1024; /// avoid excessive small IO requests whose per-request overhead dominates. const IO_BLOCK_SIZE: u64 = 4 * 1024 * 1024; +/// Rows in a Parquet file, read from its footer alone. +pub(crate) async fn read_row_count( + reader: Box, + file_size: u64, +) -> crate::Result { + let reader = ArrowFileReader::new(file_size, Arc::from(reader)); + let metadata = ParquetMetaDataReader::new() + .with_prefetch_hint(Some(METADATA_SIZE_HINT)) + .load_and_finish(reader, file_size) + .await + .map_err(|error| Error::UnexpectedError { + message: format!("Failed to read the Parquet footer: {error}"), + source: Some(Box::new(error)), + })?; + Ok(metadata.file_metadata().num_rows()) +} + impl ArrowFileReader { fn new(file_size: u64, r: Arc) -> Self { Self { file_size, r } diff --git a/crates/paimon/src/table/format_partition.rs b/crates/paimon/src/table/format_partition.rs index 0ee14316a..0f27e9a0a 100644 --- a/crates/paimon/src/table/format_partition.rs +++ b/crates/paimon/src/table/format_partition.rs @@ -325,7 +325,7 @@ fn format_partition_date(epoch_days: i32) -> Option { .map(|date| date.format("%Y-%m-%d").to_string()) } -fn is_storage_not_found(error: &crate::Error) -> bool { +pub(crate) fn is_storage_not_found(error: &crate::Error) -> bool { matches!( error, crate::Error::IoUnexpected { source, .. } diff --git a/crates/paimon/src/table/format_partition_stats.rs b/crates/paimon/src/table/format_partition_stats.rs new file mode 100644 index 000000000..6fcc86db6 --- /dev/null +++ b/crates/paimon/src/table/format_partition_stats.rs @@ -0,0 +1,186 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Measures the partitions of a Format Table with catalog-managed partitions. + +use std::collections::HashMap; + +use futures::{StreamExt, TryStreamExt}; + +use super::format_partition::FormatTablePartitionPaths; +use super::format_table_scan::{list_format_table_data_files, supported_format_table_extension}; +use super::Table; +use crate::arrow::format::read_file_row_count; +use crate::io::FileStatus; +use crate::spec::{CoreOptions, Partition, PartitionStatistics}; + +/// Measures whole partitions of a Format Table through the listing a scan uses. +/// Mirrors Java `FormatTablePartitionStatsCollector`. +#[derive(Debug)] +pub struct FormatTablePartitionStatsCollector<'a> { + table: &'a Table, + with_record_count: bool, + parallelism: usize, +} + +impl<'a> FormatTablePartitionStatsCollector<'a> { + /// Measure `table`, reading file footers for row counts only when `with_record_count` is set. + /// `parallelism` bounds listings and footer reads together; a value below one is read as one. + pub fn new(table: &'a Table, with_record_count: bool, parallelism: usize) -> Self { + Self { + table, + with_record_count, + parallelism: parallelism.max(1), + } + } + + /// Measure the given complete partition specs. The result is aligned with `partitions` one for + /// one, so it can be sent to the catalog together with the same specs. + pub async fn collect( + &self, + partitions: &[HashMap], + ) -> crate::Result> { + if partitions.is_empty() { + return Ok(Vec::new()); + } + if !self.table.has_catalog_managed_partitions() { + return Err(crate::Error::Unsupported { + message: format!( + "Format Table {} does not have catalog-managed partitions, so its partitions \ + cannot be measured", + self.table.identifier().full_name() + ), + }); + } + let options = CoreOptions::new(self.table.schema().options()); + let format_extension = supported_format_table_extension(&options.file_format())?; + let partition_paths = FormatTablePartitionPaths::new( + self.table.schema().partition_keys().iter().cloned(), + options.format_table_partition_only_value_in_path(), + ); + let table_path = options + .path() + .unwrap_or_else(|| self.table.location()) + .trim_end_matches('/'); + let directories = partitions + .iter() + .map(|spec| { + partition_paths + .relative_path(spec) + .map(|relative_path| format!("{table_path}/{relative_path}")) + }) + .collect::>>()?; + + let file_io = self.table.file_io(); + // Each future owns what it lists. A stream over borrowed items would leave the future of + // any SQL statement that measures partitions without a provable `Send`. + let listings: Vec> = futures::stream::iter(directories) + .map(|directory| async move { + // Each directory is a complete partition, so no partition level lies below it. + list_format_table_data_files(file_io, &directory, 0, format_extension).await + }) + .buffered(self.parallelism) + .try_collect() + .await?; + + let record_counts = if self.with_record_count { + self.count_rows(&options.file_format(), &listings).await + } else { + vec![Partition::UNKNOWN; listings.len()] + }; + + Ok(partitions + .iter() + .zip(&listings) + .zip(record_counts) + .map(|((spec, files), record_count)| statistics(spec, files, record_count)) + .collect()) + } + + /// The rows each listed partition holds. Every file of every partition goes through one + /// bounded stream, so a partition with many files is counted with all of it. + async fn count_rows(&self, file_format: &str, listings: &[Vec]) -> Vec { + let file_io = self.table.file_io(); + let files = listings + .iter() + .enumerate() + .flat_map(|(index, files)| { + files + .iter() + .map(move |file| (index, file.path.clone(), file.size)) + }) + .collect::>(); + let counts: Vec<(usize, Option)> = futures::stream::iter(files) + .map(|(index, path, size)| async move { + let count = match read_file_row_count(file_io, file_format, &path, size).await { + Ok(count) => count, + Err(error) => { + log::warn!( + "Failed to read the row count of {path} in table {}; the row count \ + of its partition stays unknown: {error}", + self.table.identifier().full_name() + ); + None + } + }; + (index, count) + }) + .buffered(self.parallelism) + .collect() + .await; + + // A partition with no files holds exactly zero rows; one file whose count is unknown + // makes the whole partition unknown rather than short. + let mut record_counts = vec![Some(0i64); listings.len()]; + for (index, count) in counts { + record_counts[index] = match (record_counts[index], count) { + (Some(total), Some(count)) => total.checked_add(count), + _ => None, + }; + } + record_counts + .into_iter() + .map(|count| count.unwrap_or(Partition::UNKNOWN)) + .collect() + } +} + +/// What the listed files of a partition add up to. +fn statistics( + spec: &HashMap, + files: &[FileStatus], + record_count: i64, +) -> PartitionStatistics { + let file_size_in_bytes = files + .iter() + .map(|file| i64::try_from(file.size).unwrap_or(i64::MAX)) + .fold(0i64, i64::saturating_add); + let last_file_creation_time = files + .iter() + .filter_map(|file| file.last_modified) + .map(|modified| modified.timestamp_millis()) + .max() + .unwrap_or(Partition::UNKNOWN); + PartitionStatistics { + spec: spec.clone(), + record_count, + file_size_in_bytes, + file_count: files.len() as i64, + last_file_creation_time, + total_buckets: Partition::UNKNOWN_TOTAL_BUCKETS, + } +} diff --git a/crates/paimon/src/table/format_table_scan.rs b/crates/paimon/src/table/format_table_scan.rs index de2c02e0a..b002e88a8 100644 --- a/crates/paimon/src/table/format_table_scan.rs +++ b/crates/paimon/src/table/format_table_scan.rs @@ -20,7 +20,8 @@ use std::collections::{HashMap, HashSet}; use super::format_partition::{ - format_partition_value, parse_format_partition_value, FormatTablePartitionPaths, + format_partition_value, is_storage_not_found, parse_format_partition_value, + FormatTablePartitionPaths, }; use super::{Plan, RESTEnv, ScanTrace, Table}; use crate::api::RestError; @@ -107,29 +108,22 @@ impl<'a> FormatTableScan<'a> { let partition_levels_below_root = partition_fields .len() .saturating_sub(root_segments.len().saturating_sub(table_depth)); - let statuses = self - .list_status_recursive_if_exists(&scan_root.path) - .await?; + let files = list_format_table_data_files( + self.table.file_io(), + &scan_root.path, + partition_levels_below_root, + format_extension, + ) + .await?; let mut splits = Vec::new(); - for status in statuses { - if is_hidden_below_partitions( - &root_segments, - partition_levels_below_root, - &status.path, - ) { - continue; - } - if let Some(split) = self - .status_to_split( - status, - table_path, - format_extension, - schema_id, - partition_fields, - scan_root.partition.clone(), - ) - .await? - { + for status in files { + if let Some(split) = self.status_to_split( + status, + table_path, + schema_id, + partition_fields, + scan_root.partition.clone(), + )? { splits.push(split); } } @@ -363,27 +357,11 @@ impl<'a> FormatTableScan<'a> { } } - async fn list_status_recursive_if_exists( - &self, - path: &str, - ) -> crate::Result> { - match self.table.file_io().list_status_recursive(path).await { - Ok(statuses) => Ok(statuses), - Err(err) => { - if !self.table.file_io().exists(path).await.unwrap_or(true) { - Ok(Vec::new()) - } else { - Err(err) - } - } - } - } - - async fn status_to_split( + /// The split reading one file that [`list_format_table_data_files`] returned. + fn status_to_split( &self, status: crate::io::FileStatus, table_path: &str, - format_extension: &str, schema_id: i64, partition_fields: &[DataField], known_partition: BinaryRow, @@ -393,17 +371,6 @@ impl<'a> FormatTableScan<'a> { }; let parent = parent.to_string(); let file_name = file_name.to_string(); - if !is_format_table_data_file_name(&file_name) { - return Ok(None); - } - if !file_name.to_ascii_lowercase().ends_with(format_extension) { - return Ok(None); - } - let status = if status.size == 0 { - self.table.file_io().get_status(&status.path).await? - } else { - status - }; let file_size = i64::try_from(status.size).map_err(|_| crate::Error::DataInvalid { message: format!( "Format table file '{}' is too large to fit in i64 metadata", @@ -483,6 +450,59 @@ fn is_format_table_data_file_name(file_name: &str) -> bool { !file_name.is_empty() && !file_name.starts_with('.') && !file_name.starts_with('_') } +/// The non-hidden files with the format's extension that a Format Table scan reads below `root`. +/// Only a root the store reports as not found holds no files; other listing failures are returned. +pub(crate) async fn list_format_table_data_files( + file_io: &crate::io::FileIO, + root: &str, + partition_levels_below_root: usize, + format_extension: &str, +) -> crate::Result> { + let statuses = match file_io.list_status_recursive_stream(root, None).await { + Ok(listing) => collect_listing(listing).await?, + Err(error) if is_storage_not_found(&error) => Vec::new(), + Err(error) => return Err(error), + }; + let root_segments = path_segments(root); + let mut files = Vec::with_capacity(statuses.len()); + for status in statuses { + if is_hidden_below_partitions(&root_segments, partition_levels_below_root, &status.path) { + continue; + } + let is_data_file = split_parent_and_file(&status.path).is_some_and(|(_, file_name)| { + is_format_table_data_file_name(file_name) + && file_name.to_ascii_lowercase().ends_with(format_extension) + }); + if !is_data_file { + continue; + } + let status = if status.size == 0 { + file_io.get_status(&status.path).await? + } else { + status + }; + files.push(status); + } + Ok(files) +} + +/// Every listed status, or none when the store reports the root not found before listing anything. +async fn collect_listing( + mut listing: impl futures::Stream> + Unpin, +) -> crate::Result> { + let mut statuses = Vec::new(); + while let Some(status) = listing.next().await { + match status { + Ok(status) => statuses.push(status), + Err(error) if statuses.is_empty() && is_storage_not_found(&error) => { + return Ok(statuses) + } + Err(error) => return Err(error), + } + } + Ok(statuses) +} + /// Whether a listed file is, or lies inside, an entry whose name starts with `.` or `_` below /// the partition directories, such as a committer staging tree (`_temporary`, `__magic_*`) /// whose files may never be committed. @@ -843,7 +863,7 @@ fn supported_format_table_formats() -> Vec<&'static str> { ] } -fn supported_format_table_extension(format: &str) -> crate::Result<&'static str> { +pub(crate) fn supported_format_table_extension(format: &str) -> crate::Result<&'static str> { match format.to_ascii_lowercase().as_str() { "parquet" => Ok(".parquet"), "orc" => Ok(".orc"), @@ -1093,6 +1113,82 @@ mod tests { ); } + #[tokio::test] + async fn test_data_file_listing_returns_what_a_partition_scan_reads() { + let table = format_table("memory:/data_file_listing", &["dt"], &[]); + write_files( + &table, + &[ + "dt=a/part-0.parquet", + "dt=a/_temporary/0/part-1.parquet", + "dt=a/__magic_job_1/tasks/part-2.parquet", + "dt=a/.part-3.parquet", + "dt=a/_SUCCESS", + "dt=a/notes.txt", + ], + ) + .await; + let file_names = |files: Vec| { + files + .iter() + .filter_map(|file| { + split_parent_and_file(&file.path).map(|(_, name)| name.to_string()) + }) + .collect::>() + }; + + let partition = format!("{}/dt=a", table.location()); + let listed = list_format_table_data_files(table.file_io(), &partition, 0, ".parquet") + .await + .unwrap(); + assert_eq!(file_names(listed), vec!["part-0.parquet"]); + assert_eq!( + planned_files(&table, Some(partition_set(&table, &[&[Some("a")]]))).await, + vec!["dt=a/part-0.parquet"] + ); + + // A partition whose directory is gone holds no files rather than failing the listing. + let missing = format!("{}/dt=b", table.location()); + assert!( + list_format_table_data_files(table.file_io(), &missing, 0, ".parquet") + .await + .unwrap() + .is_empty() + ); + } + + #[tokio::test] + async fn test_listing_failure_is_not_read_as_an_empty_directory() { + let file = crate::io::FileStatus { + size: 1, + is_dir: false, + path: "memory:/t/dt=a/part-0.parquet".to_string(), + last_modified: None, + }; + let failure = |kind: opendal::ErrorKind| crate::Error::IoUnexpected { + message: "list partition directory".to_string(), + source: Box::new(opendal::Error::new(kind, "injected")), + }; + // A failure after a listed file, or of any kind but not found, is not an empty list. + for listing in [ + vec![ + Ok(file.clone()), + Err(failure(opendal::ErrorKind::Unexpected)), + ], + vec![Err(failure(opendal::ErrorKind::Unexpected))], + vec![Ok(file.clone()), Err(failure(opendal::ErrorKind::NotFound))], + ] { + assert!(collect_listing(futures::stream::iter(listing)) + .await + .is_err()); + } + let missing = vec![Err(failure(opendal::ErrorKind::NotFound))]; + assert!(collect_listing(futures::stream::iter(missing)) + .await + .unwrap() + .is_empty()); + } + #[tokio::test] async fn test_concurrent_listing_keeps_the_plan_order() { let partitions = ["a", "b", "c", "d", "e", "f", "g", "h"]; diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index bd98ba8ee..00235b5c8 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -41,6 +41,7 @@ mod data_file_reader; mod data_file_writer; mod dedicated_format_file_writer; mod format_partition; +mod format_partition_stats; mod format_read_builder; mod format_table_read; mod format_table_scan; @@ -121,6 +122,7 @@ pub use data_evolution_writer::{DataEvolutionDeleteWriter, DataEvolutionWriter}; pub use format_partition::{ format_partition_value, parse_format_partition_value, FormatTablePartitionPaths, }; +pub use format_partition_stats::FormatTablePartitionStatsCollector; #[cfg(feature = "fulltext")] pub use full_text_search_builder::FullTextSearchBuilder; use futures::stream::BoxStream; diff --git a/crates/paimon/tests/mock_server.rs b/crates/paimon/tests/mock_server.rs index 523541b82..b5b3bf56b 100644 --- a/crates/paimon/tests/mock_server.rs +++ b/crates/paimon/tests/mock_server.rs @@ -1455,6 +1455,17 @@ impl RESTServer { .unwrap_or_default() } + /// Return the partitions registered for a table, statistics included, in registration order. + pub fn table_partitions(&self, database: &str, table: &str) -> Vec { + self.inner + .lock() + .unwrap() + .partitions + .get(&format!("{database}.{table}")) + .cloned() + .unwrap_or_default() + } + /// Set whether a stored table is external. pub fn set_table_external(&self, database: &str, table: &str, is_external: bool) { let key = format!("{database}.{table}"); diff --git a/docs/src/sql.md b/docs/src/sql.md index 75ac8f691..8dc81360f 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -40,7 +40,7 @@ Mosaic support is always available and currently read-only. SQL queries can read SQL support has two layers: - DataFusion provides the parser, query planner, optimizer, execution engine, expressions, scalar functions, aggregate functions, and window functions. SQL statements that `SQLContext` does not intercept are delegated to DataFusion. This includes the DataFusion SQL surface for `SELECT` queries, CTEs (including recursive CTEs), subqueries, joins including `LATERAL` joins, SQL lambda functions, grouping, `HAVING`, window clauses, `QUALIFY`, set operations, `ORDER BY`, `LIMIT`/`OFFSET`, `EXPLAIN`, information-schema commands such as `SHOW TABLES`, `DESCRIBE`, `COPY`, and ordinary `INSERT`. -- Paimon-specific table management and row-level writes are implemented by `SQLContext`. This includes Paimon `CREATE TABLE`, `ALTER TABLE`, `DROP TABLE`, `CREATE TEMPORARY TABLE`, `CREATE TEMPORARY VIEW`, REST Catalog persistent `CREATE VIEW`, `DROP VIEW`, and `CREATE FUNCTION`, `DROP TEMPORARY TABLE` / `VIEW`, `INSERT OVERWRITE ... PARTITION`, `UPDATE`, `DELETE`, `MERGE INTO`, `TRUNCATE TABLE`, `ALTER TABLE ... ADD PARTITION`, `ALTER TABLE ... DROP PARTITION`, `SHOW PARTITIONS`, `MSCK REPAIR TABLE`, `CALL sys.*`, Paimon time travel, and `SET` / `RESET 'paimon.*'`. +- Paimon-specific table management and row-level writes are implemented by `SQLContext`. This includes Paimon `CREATE TABLE`, `ALTER TABLE`, `DROP TABLE`, `CREATE TEMPORARY TABLE`, `CREATE TEMPORARY VIEW`, REST Catalog persistent `CREATE VIEW`, `DROP VIEW`, and `CREATE FUNCTION`, `DROP TEMPORARY TABLE` / `VIEW`, `INSERT OVERWRITE ... PARTITION`, `UPDATE`, `DELETE`, `MERGE INTO`, `TRUNCATE TABLE`, `ALTER TABLE ... ADD PARTITION`, `ALTER TABLE ... DROP PARTITION`, `SHOW PARTITIONS`, `MSCK REPAIR TABLE`, `ANALYZE TABLE`, `CALL sys.*`, Paimon time travel, and `SET` / `RESET 'paimon.*'`. Not every DataFusion DDL/DML statement maps to a Paimon table operation. For Paimon catalogs, `CREATE EXTERNAL TABLE`, `LOCATION`, `CREATE MATERIALIZED VIEW`, and persistent `CREATE TABLE AS SELECT` are rejected or not implemented. Persistent `CREATE FUNCTION` is supported only for the REST Catalog SQL scalar form documented below. DataFusion `COPY` can export query results to files; it does not create or commit Paimon table files. @@ -985,6 +985,41 @@ truncated view of the table into a `DROP` diff. There is no dry-run and no scope argument — repair always covers the whole table. A partition at a custom location is never unregistered by repair. +### ANALYZE TABLE + +Measure what the registered partitions hold and report it to the catalog: + +```sql +ANALYZE TABLE paimon.my_db.events COMPUTE STATISTICS NOSCAN; -- files, size, last file time +ANALYZE TABLE paimon.my_db.events COMPUTE STATISTICS; -- also row counts +ANALYZE TABLE paimon.my_db.events PARTITION (dt = '2024-01-01') COMPUTE STATISTICS; +``` + +Each partition is measured through the listing a scan uses, so it counts exactly the files +a query reads and leaves staging entries such as `_temporary` out. `NOSCAN` stops at the +listing: it reports the file count, the total size and the latest file modification time. +Without `NOSCAN` the row count is also read from each file's footer, which Parquet and ORC +keep; for other formats it stays unknown. A footer that cannot be read leaves the row count +of its partition unknown rather than short, and a partition without files holds exactly +zero rows. A field the statement does not measure, such as the row count under `NOSCAN`, +is reported as unknown (`-1`). + +The measurement replaces the statistics the catalog holds for each partition; it never +adds or removes a partition. `PARTITION (...)` must give values for a leading run of the +partition keys and selects every registered partition under them: on a `(dt, region)` +table, `PARTITION (dt = '2024-01-01')` measures every region of that date, while +`PARTITION (region = 'us')` is rejected. It is an error when no registered partition +matches. A selected partition at a custom location fails the statement, and +`FOR COLUMNS` is not supported. + +A listing failure fails the statement before anything is reported. +`format-table.statistics.parallelism` (default 8) bounds the storage requests in flight, +listings and footer reads alike. Set it for the session: + +```sql +SET 'paimon.format-table.statistics.parallelism' = '16'; +``` + ## Procedures Use `CALL` to invoke built-in procedures. All procedures are under the `sys` namespace.