You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@fluss.apache.org by ab...@apache.org on 2026/10/02 10:48:51 UTC
(fluss) branch main updated: [rust] Send table stats requests to bucket leaders (#4529)
This is an automated email from the ASF dual-hosted git repository.
fresh-borzoni pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/fluss.git
The following commit(s) were added to refs/heads/main by this push:
new b4fe91b4b2 [rust] Send table stats requests to bucket leaders (#4529)
b4fe91b4b2 is described below
commit b4fe91b4b272faa6561cb04adb9189d81e6aa47b
Author: Anton Borisov <14...@users.noreply.github.com>
AuthorDate: Fri Oct 2 11:48:46 2026 +0100
[rust] Send table stats requests to bucket leaders (#4529)
---
fluss-rust/crates/fluss/src/client/admin.rs | 174 +++++++++++++++++++--
fluss-rust/crates/fluss/tests/integration/admin.rs | 137 +++++++++++++++-
2 files changed, 295 insertions(+), 16 deletions(-)
diff --git a/fluss-rust/crates/fluss/src/client/admin.rs b/fluss-rust/crates/fluss/src/client/admin.rs
index 6e7ab70f4d..c1de1812c3 100644
--- a/fluss-rust/crates/fluss/src/client/admin.rs
+++ b/fluss-rust/crates/fluss/src/client/admin.rs
@@ -16,7 +16,7 @@
// under the License.
use crate::client::metadata::Metadata;
-use crate::cluster::ServerNode;
+use crate::cluster::{Cluster, ServerNode};
use crate::metadata::{
AclFilter, AclInfo, AcquireKvSnapshotLeaseResult, ActiveKvSnapshots, AlterConfig,
AlterTableChanges, BucketStatsRequest, ClusterHealth, CreateAclResult, DatabaseDescriptor,
@@ -46,6 +46,7 @@ use crate::rpc::{RpcClient, ServerConnection};
use crate::error::{Error, Result};
use crate::proto::GetTableInfoResponse;
use crate::{BucketId, PartitionId, SnapshotId, TableId};
+use futures::future::try_join_all;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use tokio::task::JoinHandle;
@@ -562,22 +563,73 @@ impl FlussAdmin {
}
/// Get table statistics for buckets. Pass empty `target_columns` to request stats for all columns.
+ ///
+ /// The table must be in the metadata cache, as it is after [`Self::get_table_info`].
pub async fn get_table_stats(
&self,
table_id: TableId,
buckets_req: Vec<BucketStatsRequest>,
target_columns: Vec<i32>,
) -> Result<TableStats> {
- let response = self
- .admin_gateway()
- .await?
- .request(GetTableStatsRequest::new(
- table_id,
- buckets_req,
- target_columns,
- ))
+ if buckets_req.is_empty() {
+ return Ok(TableStats {
+ buckets: Vec::new(),
+ });
+ }
+ let table_path = self
+ .metadata
+ .get_cluster()
+ .get_table_path_by_id(table_id)
+ .cloned()
+ .ok_or_else(|| Error::IllegalArgument {
+ message: format!("Table id {table_id} is not in the metadata cache."),
+ })?;
+ // Nothing else refreshes a cached partition's leaders, so refresh them with the table.
+ let partition_ids: Vec<PartitionId> = buckets_req
+ .iter()
+ .filter_map(|bucket| bucket.partition_id)
+ .collect::<HashSet<_>>()
+ .into_iter()
+ .collect();
+ self.metadata
+ .update_tables_metadata(
+ &HashSet::from([&table_path]),
+ &HashSet::new(),
+ partition_ids,
+ )
.await?;
- Ok(TableStats::from_pb(&response))
+
+ let cluster = self.metadata.get_cluster();
+ let mut requests = Vec::new();
+ for (leader_id, buckets) in group_stats_requests_by_leader(&cluster, table_id, buckets_req)?
+ {
+ let tablet_server = cluster
+ .get_tablet_server(leader_id)
+ .cloned()
+ .ok_or_else(|| {
+ Error::leader_not_available(format!(
+ "Tablet server {leader_id} is not found in metadata cache."
+ ))
+ })?;
+ let request = GetTableStatsRequest::new(table_id, buckets, target_columns.clone());
+ requests.push((tablet_server, request));
+ }
+ let responses = try_join_all(requests.into_iter().map(
+ |(tablet_server, request)| async move {
+ self.rpc_client
+ .get_connection(&tablet_server)
+ .await?
+ .request(request)
+ .await
+ },
+ ))
+ .await?;
+ Ok(TableStats {
+ buckets: responses
+ .iter()
+ .flat_map(|response| TableStats::from_pb(response).buckets)
+ .collect(),
+ })
}
/// Get the latest KV snapshots for a table (optionally scoped to one partition).
@@ -887,10 +939,108 @@ impl FlussAdmin {
}
}
+fn group_stats_requests_by_leader(
+ cluster: &Cluster,
+ table_id: TableId,
+ buckets: Vec<BucketStatsRequest>,
+) -> Result<HashMap<i32, Vec<BucketStatsRequest>>> {
+ let mut buckets_by_leader: HashMap<i32, Vec<BucketStatsRequest>> = HashMap::new();
+ for bucket in buckets {
+ let table_bucket =
+ TableBucket::new_with_partition(table_id, bucket.partition_id, bucket.bucket_id);
+ let leader = cluster.leader_for(&table_bucket).ok_or_else(|| {
+ Error::leader_not_available(format!(
+ "No leader found for table bucket: {table_bucket}."
+ ))
+ })?;
+ buckets_by_leader
+ .entry(leader.id())
+ .or_default()
+ .push(bucket);
+ }
+ Ok(buckets_by_leader)
+}
+
#[cfg(test)]
mod tests {
- use super::parse_table_descriptor;
- use crate::error::Error;
+ use super::{group_stats_requests_by_leader, parse_table_descriptor};
+ use crate::cluster::{BucketLocation, Cluster, ServerNode, ServerType};
+ use crate::error::{Error, FlussError};
+ use crate::metadata::{BucketStatsRequest, PhysicalTablePath, TableBucket, TablePath};
+ use crate::test_utils::build_table_info;
+ use std::collections::HashMap;
+ use std::sync::Arc;
+
+ const TABLE_ID: i64 = 1;
+
+ fn cluster(leaders: [Option<i32>; 4]) -> Cluster {
+ let table_path = TablePath::new("db", "tbl");
+ let physical_path = Arc::new(PhysicalTablePath::of(Arc::new(table_path.clone())));
+ let servers: HashMap<i32, ServerNode> = (1..=2)
+ .map(|id| {
+ let node =
+ ServerNode::new(id, "127.0.0.1".to_string(), 9000, ServerType::TabletServer);
+ (id, node)
+ })
+ .collect();
+ let locations: Vec<BucketLocation> = leaders
+ .iter()
+ .enumerate()
+ .map(|(bucket, leader)| {
+ BucketLocation::new(
+ TableBucket::new(TABLE_ID, bucket as i32),
+ leader.map(|id| servers[&id].clone()),
+ Arc::clone(&physical_path),
+ )
+ })
+ .collect();
+ Cluster::new(
+ None,
+ servers,
+ HashMap::from([(physical_path, locations.clone())]),
+ locations
+ .into_iter()
+ .map(|location| (location.table_bucket.clone(), location))
+ .collect(),
+ HashMap::from([(table_path.clone(), TABLE_ID)]),
+ HashMap::from([(
+ table_path.clone(),
+ build_table_info(table_path, TABLE_ID, 4),
+ )]),
+ HashMap::new(),
+ )
+ }
+
+ fn all_buckets() -> Vec<BucketStatsRequest> {
+ (0..4)
+ .map(|bucket| BucketStatsRequest::new(None, bucket))
+ .collect()
+ }
+
+ #[test]
+ fn table_stats_requests_are_grouped_by_bucket_leader() {
+ let cluster = cluster([Some(1), Some(2), Some(1), Some(2)]);
+ let grouped = group_stats_requests_by_leader(&cluster, TABLE_ID, all_buckets()).unwrap();
+ let bucket_ids = |leader: i32| -> Vec<i32> {
+ grouped[&leader]
+ .iter()
+ .map(|bucket| bucket.bucket_id)
+ .collect()
+ };
+ assert_eq!(grouped.len(), 2);
+ assert_eq!(bucket_ids(1), vec![0, 2]);
+ assert_eq!(bucket_ids(2), vec![1, 3]);
+ }
+
+ #[test]
+ fn a_bucket_without_a_leader_fails_the_table_stats_request() {
+ let cluster = cluster([Some(1), None, Some(1), Some(1)]);
+ let error = group_stats_requests_by_leader(&cluster, TABLE_ID, all_buckets()).unwrap_err();
+ assert_eq!(
+ error.api_error(),
+ Some(FlussError::LeaderNotAvailableException)
+ );
+ }
#[test]
fn test_parse_table_descriptor_rejects_malformed_json() {
diff --git a/fluss-rust/crates/fluss/tests/integration/admin.rs b/fluss-rust/crates/fluss/tests/integration/admin.rs
index 64218abdf3..d93d7ea87d 100644
--- a/fluss-rust/crates/fluss/tests/integration/admin.rs
+++ b/fluss-rust/crates/fluss/tests/integration/admin.rs
@@ -17,14 +17,18 @@
#[cfg(test)]
mod admin_test {
- use crate::integration::utils::get_shared_cluster;
- use fluss::client::FlussConnection;
+ use crate::integration::utils::{
+ create_partitions, create_table, get_shared_cluster, wait_for_partitions_ready,
+ wait_for_table_buckets_ready,
+ };
+ use fluss::client::{FlussAdmin, FlussConnection};
use fluss::config::Config;
use fluss::error::FlussError;
use fluss::metadata::{
- DataTypes, DatabaseDescriptorBuilder, KvFormat, LogFormat, PartitionSpec, Schema,
- TableDescriptor, TablePath,
+ BucketStatsRequest, DataTypes, DatabaseDescriptorBuilder, KvFormat, LogFormat,
+ PartitionSpec, Schema, TableDescriptor, TablePath,
};
+ use fluss::row::GenericRow;
use std::collections::HashMap;
#[tokio::test]
@@ -684,4 +688,129 @@ mod admin_test {
admin.drop_table(&table_path, true).await.unwrap();
admin.drop_database(db_name, true, true).await.unwrap();
}
+
+ async fn row_count(admin: &FlussAdmin, table_id: i64, buckets: Vec<BucketStatsRequest>) -> i64 {
+ let stats = admin
+ .get_table_stats(table_id, buckets, vec![])
+ .await
+ .expect("table stats");
+ stats
+ .buckets
+ .iter()
+ .map(|bucket| {
+ assert_eq!(bucket.error, None, "{bucket:?}");
+ bucket.row_count.expect("row count")
+ })
+ .sum()
+ }
+
+ fn id_row(id: i32, name: &str) -> GenericRow<'static> {
+ let mut row = GenericRow::new(2);
+ row.set_field(0, id);
+ row.set_field(1, name.to_string());
+ row
+ }
+
+ #[tokio::test]
+ async fn test_get_table_stats() {
+ let cluster = get_shared_cluster();
+ let connection = cluster.get_fluss_connection().await;
+ let admin = connection.get_admin().expect("should get admin");
+ let buckets = |partition_id| {
+ (0..3)
+ .map(|bucket| BucketStatsRequest::new(partition_id, bucket))
+ .collect::<Vec<_>>()
+ };
+
+ let log_path = TablePath::new("fluss", "test_table_stats_log");
+ let log_descriptor = TableDescriptor::builder()
+ .schema(
+ Schema::builder()
+ .column("id", DataTypes::int())
+ .column("name", DataTypes::string())
+ .build()
+ .unwrap(),
+ )
+ .distributed_by(Some(3), vec!["id".to_string()])
+ .build()
+ .unwrap();
+ create_table(&admin, &log_path, &log_descriptor).await;
+ wait_for_table_buckets_ready(&admin, &log_path, &[0, 1, 2]).await;
+ let log_table = connection.get_table(&log_path).await.unwrap();
+ let writer = log_table.new_append().unwrap().create_writer().unwrap();
+ for id in 0..7 {
+ writer.append(&id_row(id, "v")).unwrap();
+ }
+ writer.flush().await.unwrap();
+ let table_id = log_table.get_table_info().table_id;
+ assert_eq!(row_count(&admin, table_id, buckets(None)).await, 7);
+
+ let pk_path = TablePath::new("fluss", "test_table_stats_pk");
+ let pk_descriptor = TableDescriptor::builder()
+ .schema(
+ Schema::builder()
+ .column("id", DataTypes::int())
+ .column("name", DataTypes::string())
+ .primary_key(vec!["id"])
+ .unwrap()
+ .build()
+ .unwrap(),
+ )
+ .distributed_by(Some(3), vec![])
+ .build()
+ .unwrap();
+ create_table(&admin, &pk_path, &pk_descriptor).await;
+ wait_for_table_buckets_ready(&admin, &pk_path, &[0, 1, 2]).await;
+ let pk_table = connection.get_table(&pk_path).await.unwrap();
+ let writer = pk_table.new_upsert().unwrap().create_writer().unwrap();
+ for id in (0..10).chain(0..3) {
+ writer.upsert(&id_row(id, "v")).unwrap();
+ }
+ writer.flush().await.unwrap();
+ let table_id = pk_table.get_table_info().table_id;
+ assert_eq!(row_count(&admin, table_id, buckets(None)).await, 10);
+
+ let partitioned_path = TablePath::new("fluss", "test_table_stats_partitioned");
+ let partitioned_descriptor = TableDescriptor::builder()
+ .schema(
+ Schema::builder()
+ .column("id", DataTypes::int())
+ .column("name", DataTypes::string())
+ .primary_key(vec!["id", "name"])
+ .unwrap()
+ .build()
+ .unwrap(),
+ )
+ .partitioned_by(vec!["name"])
+ .distributed_by(Some(3), vec![])
+ .build()
+ .unwrap();
+ create_table(&admin, &partitioned_path, &partitioned_descriptor).await;
+ create_partitions(&admin, &partitioned_path, "name", &["a", "b"]).await;
+ wait_for_partitions_ready(&admin, &partitioned_path, &["a", "b"]).await;
+ let partitioned_table = connection.get_table(&partitioned_path).await.unwrap();
+ let writer = partitioned_table
+ .new_upsert()
+ .unwrap()
+ .create_writer()
+ .unwrap();
+ for id in 0..5 {
+ writer.upsert(&id_row(id, "a")).unwrap();
+ writer.upsert(&id_row(id, "b")).unwrap();
+ }
+ writer.flush().await.unwrap();
+ let table_id = partitioned_table.get_table_info().table_id;
+ let partition_buckets = admin
+ .list_partition_infos(&partitioned_path)
+ .await
+ .unwrap()
+ .iter()
+ .flat_map(|partition| buckets(Some(partition.get_partition_id())))
+ .collect();
+ assert_eq!(row_count(&admin, table_id, partition_buckets).await, 10);
+
+ for path in [&log_path, &pk_path, &partitioned_path] {
+ admin.drop_table(path, true).await.unwrap();
+ }
+ }
}