rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-12 15:02:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

merge: adapt state merge tests to DataFusion 55

Browse Source 
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-09-16 13:08:44 +08:00
parent 1a59df7f1d 3a2a084d0d
commit ab3071ef23
352 changed files with 11747 additions and 3053 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/integration.yml
+3 -4
View File
@@ -781,10 +781,9 @@ jobs:
kubectl -n kafka-cluster apply \
-f .github/actions/setup-greptimedb-cluster/kafka-wal-helper.yaml
kubectl -n kafka-cluster wait \
--for=condition=Ready \
pod -l app=kafka-wal-helper \
--timeout=120s
kubectl rollout status deployment/kafka-wal-helper \
--timeout=120s \
-n kafka-cluster
- name: Print etcd info
shell: bash
run: kubectl get all --show-labels -n etcd-cluster
Cargo.lock
Generated
+459 -414
View File
File diff suppressed because it is too large Load Diff
Cargo.toml
+47 -44
View File
@@ -100,13 +100,13 @@ rust.unexpected_cfgs = { level = "warn", check-cfg = ['cfg(tokio_unstable)'] }
# See for more detaiils: https://github.com/rust-lang/cargo/issues/11329
ahash = { version = "0.8", features = ["compile-time-rng"] }
aquamarine = "0.6"
arrow = { version = "58.3", features = ["prettyprint"] }
arrow-array = { version = "58.3", default-features = false, features = ["chrono-tz"] }
arrow-buffer = "58.3"
arrow-cast = "58.3"
arrow-flight = "58.3"
arrow-ipc = { version = "58.3", default-features = false, features = ["lz4", "zstd"] }
arrow-schema = { version = "58.3", features = ["serde"] }
arrow = { version = "=59.2.0", features = ["prettyprint"] }
arrow-array = { version = "=59.2.0", default-features = false, features = ["chrono-tz"] }
arrow-buffer = "=59.2.0"
arrow-cast = "=59.2.0"
arrow-flight = "=59.2.0"
arrow-ipc = { version = "=59.2.0", default-features = false, features = ["lz4", "zstd"] }
arrow-schema = { version = "=59.2.0", features = ["serde"] }
async-stream = "0.3"
async-trait = "0.1"
# Remember to update axum-extra, axum-macros when updating axum
@@ -128,22 +128,22 @@ const_format = "0.2"
criterion = "0.7"
crossbeam-utils = "0.8"
dashmap = "6.1"
datafusion = "=53.1.0"
datafusion-common = "=53.1.0"
datafusion-datasource = "=53.1.0"
datafusion-expr = "=53.1.0"
datafusion-expr-common = "=53.1.0"
datafusion-functions = "=53.1.0"
datafusion-functions-aggregate-common = "=53.1.0"
datafusion-functions-window-common = "=53.1.0"
datafusion-optimizer = "=53.1.0"
datafusion-orc = { git = "https://github.com/datafusion-contrib/datafusion-orc.git", rev = "6c07fa282dc8d62db2aa4ded06ab55485efc811a" }
datafusion-pg-catalog = "0.17.3"
datafusion-physical-expr = "=53.1.0"
datafusion-physical-plan = "=53.1.0"
datafusion-proto = "=53.1.0"
datafusion-sql = "=53.1.0"
datafusion-substrait = "=53.1.0"
datafusion = "=55.0.0"
datafusion-common = "=55.0.0"
datafusion-datasource = "=55.0.0"
datafusion-expr = "=55.0.0"
datafusion-expr-common = "=55.0.0"
datafusion-functions = "=55.0.0"
datafusion-functions-aggregate-common = "=55.0.0"
datafusion-functions-window-common = "=55.0.0"
datafusion-optimizer = "=55.0.0"
datafusion-orc = "0.10.0"
datafusion-pg-catalog = "0.18.3"
datafusion-physical-expr = "=55.0.0"
datafusion-physical-plan = "=55.0.0"
datafusion-proto = "=55.0.0"
datafusion-sql = "=55.0.0"
datafusion-substrait = "=55.0.0"
datafusion_object_store = { package = "object_store", version = "0.13.2" }
deadpool = "0.12"
deadpool-postgres = "0.14"
@@ -195,10 +195,10 @@ otel-arrow-rust = { git = "https://github.com/GreptimeTeam/otel-arrow", rev = "5
"server",
] }
parking_lot = "0.12"
parquet = { version = "58.3", default-features = false, features = ["arrow", "async", "object_store"] }
parquet-variant = "58.3"
parquet-variant-compute = "58.3"
parquet-variant-json = "58.3"
parquet = { version = "=59.2.0", default-features = false, features = ["arrow", "async", "object_store"] }
parquet-variant = "=59.2.0"
parquet-variant-compute = "=59.2.0"
parquet-variant-json = "=59.2.0"
paste = "1.0"
pin-project = "1.0"
pretty_assertions = "1.4.0"
@@ -239,7 +239,7 @@ simd-json = "0.15"
similar-asserts = "1.6.0"
smallvec = { version = "1", features = ["serde"] }
snafu = "0.8"
sqlparser = { version = "0.61.0", default-features = false, features = ["std", "visitor", "serde"] }
sqlparser = { version = "0.62.0", default-features = false, features = ["std", "visitor", "serde"] }
sqlx = { version = "0.8", default-features = false, features = [
"any",
"macros",
@@ -353,22 +353,25 @@ git = "https://github.com/GreptimeTeam/greptime-meter.git"
rev = "5618e779cf2bb4755b499c630fba4c35e91898cb"
[patch.crates-io]
datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" }
sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "2aefa08a8d69c96eec2d6d6703598a009bba6e4c" } # on branch v0.61.x
substrait = { git = "https://github.com/GreptimeTeam/substrait-rs.git", rev = "91ec978b0649417ad3da8390e7a515baec723b1b" }
datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "b66dca3260c8e04f314c8036cdfbb08f1a564d95" }
datafusion-pg-catalog = { git = "https://github.com/GreptimeTeam/datafusion-postgres.git", rev = "3c77e6c32b8db80635a0d2f4b318a36b31170bc3" }
arrow-pg = { git = "https://github.com/GreptimeTeam/datafusion-postgres.git", rev = "3c77e6c32b8db80635a0d2f4b318a36b31170bc3" }
sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "9e9019bb1c7040ed956f654e39378dd42ab17884" }
# Temporary: use the GreptimeTeam fork of tikv-jemalloc-sys embedding
# jemalloc 5.3.1 + backport of 54f22c83 ("Initialize TSD tcache before
config/config.md
+6
View File
@@ -165,6 +165,9 @@
| `region_engine.mito.worker_request_batch_size` | Integer | `64` | Max batch size for a worker to handle requests. |
| `region_engine.mito.manifest_checkpoint_distance` | Integer | `10` | Number of meta action updated to trigger a new checkpoint for the manifest. |
| `region_engine.mito.compress_manifest` | Bool | `false` | Whether to compress manifest and checkpoint file by gzip (default false). |
| `region_engine.mito.experimental_enable_series_index` | Bool | `false` | Under development; do not enable. Whether to enable series indexes.<br/>Indexes are stored on the local filesystem under `{data_home}/series_index`. |
| `region_engine.mito.experimental_series_index_maintenance_interval` | String | `5m` | Interval between series-index maintenance runs. Zero uses the default of 5 min. |
| `region_engine.mito.experimental_series_index_bucket_width` | String | `5days` | Requested minimum series-index bucket width (default: 5 days), rounded up to<br/>an exact multiple of each region's compaction time window. |
| `region_engine.mito.max_background_flushes` | Integer | Auto | Max number of running background flush jobs (default: 1/2 of cpu cores). |
| `region_engine.mito.max_background_compactions` | Integer | Auto | Max number of running background compaction jobs (default: 1/4 of cpu cores). |
| `region_engine.mito.max_background_purges` | Integer | Auto | Max number of running background purge jobs (default: number of cpu cores). |
@@ -612,6 +615,9 @@
| `region_engine.mito.experimental_manifest_keep_removed_file_count` | Integer | `256` | Number of removed files to keep in manifest's `removed_files` field before also<br/>remove them from `removed_files`. Mostly for debugging purpose.<br/>If set to 0, it will only use `keep_removed_file_ttl` to decide when to remove files<br/>from `removed_files` field. |
| `region_engine.mito.experimental_manifest_keep_removed_file_ttl` | String | `1h` | How long to keep removed files in the `removed_files` field of manifest<br/>after they are removed from manifest.<br/>files will only be removed from `removed_files` field<br/>if both `keep_removed_file_count` and `keep_removed_file_ttl` is reached. |
| `region_engine.mito.compress_manifest` | Bool | `false` | Whether to compress manifest and checkpoint file by gzip (default false). |
| `region_engine.mito.experimental_enable_series_index` | Bool | `false` | Under development; do not enable. Whether to enable series indexes.<br/>Indexes are stored on the local filesystem under `{data_home}/series_index`. |
| `region_engine.mito.experimental_series_index_maintenance_interval` | String | `5m` | Interval between series-index maintenance runs. Zero uses the default of 5 min. |
| `region_engine.mito.experimental_series_index_bucket_width` | String | `5days` | Requested minimum series-index bucket width (default: 5 days), rounded up to<br/>an exact multiple of each region's compaction time window. |
| `region_engine.mito.max_background_flushes` | Integer | Auto | Max number of running background flush jobs (default: 1/2 of cpu cores). |
| `region_engine.mito.max_background_compactions` | Integer | Auto | Max number of running background compaction jobs (default: 1/4 of cpu cores). |
| `region_engine.mito.max_background_purges` | Integer | Auto | Max number of running background purge jobs (default: number of cpu cores). |
config/datanode.example.toml
+11
View File
@@ -474,6 +474,17 @@ experimental_manifest_keep_removed_file_ttl = "1h"
## Whether to compress manifest and checkpoint file by gzip (default false).
compress_manifest = false
## Under development; do not enable. Whether to enable series indexes.
## Indexes are stored on the local filesystem under `{data_home}/series_index`.
#+ experimental_enable_series_index = false
## Interval between series-index maintenance runs. Zero uses the default of 5 min.
#+ experimental_series_index_maintenance_interval = "5m"
## Requested minimum series-index bucket width (default: 5 days), rounded up to
## an exact multiple of each region's compaction time window.
#+ experimental_series_index_bucket_width = "5days"
## Max number of running background flush jobs (default: 1/2 of cpu cores).
## @toml2docs:none-default="Auto"
#+ max_background_flushes = 4
config/standalone.example.toml
+11
View File
@@ -635,6 +635,17 @@ manifest_checkpoint_distance = 10
## Whether to compress manifest and checkpoint file by gzip (default false).
compress_manifest = false
## Under development; do not enable. Whether to enable series indexes.
## Indexes are stored on the local filesystem under `{data_home}/series_index`.
#+ experimental_enable_series_index = false
## Interval between series-index maintenance runs. Zero uses the default of 5 min.
#+ experimental_series_index_maintenance_interval = "5m"
## Requested minimum series-index bucket width (default: 5 days), rounded up to
## an exact multiple of each region's compaction time window.
#+ experimental_series_index_bucket_width = "5days"
## Max number of running background flush jobs (default: 1/2 of cpu cores).
## @toml2docs:none-default="Auto"
#+ max_background_flushes = 4
src/api/src/helper.rs
+34 -1
View File
@@ -449,7 +449,14 @@ impl TryFrom<ConcreteDataType> for ColumnDataTypeWrapper {
}),
JsonFormat::Json2(native_type) => {
if native_type.is_null() {
None
Some(ColumnDataTypeExtension {
type_ext: Some(TypeExt::JsonNativeType(Box::new(
JsonNativeTypeExtension {
datatype: ColumnDataType::Json as i32,
datatype_extension: None,
},
))),
})
} else {
let concrete_type =
ConcreteDataType::from_arrow_type(&native_type.as_arrow_type());
@@ -1291,6 +1298,32 @@ mod tests {
assert!(values.bool_values.is_empty());
}
#[test]
fn test_json2_unknown_type_encoding() {
let datatype = ConcreteDataType::json2(JsonNativeType::Null);
let wrapper = ColumnDataTypeWrapper::try_from(datatype.clone()).unwrap();
assert_eq!(
wrapper.to_parts(),
(
ColumnDataType::Json,
Some(ColumnDataTypeExtension {
type_ext: Some(TypeExt::JsonNativeType(Box::new(JsonNativeTypeExtension {
datatype: ColumnDataType::Json as i32,
datatype_extension: None,
}))),
}),
)
);
assert_eq!(ConcreteDataType::from(wrapper), datatype);
for extension in [None, Some(ColumnDataTypeExtension::default())] {
assert_eq!(
ConcreteDataType::from(ColumnDataTypeWrapper::new(ColumnDataType::Json, extension)),
datatype
);
}
}
#[test]
fn test_concrete_datatype_from_column_datatype() {
assert_eq!(
src/catalog/src/information_extension.rs
+10 -5
View File
@@ -27,8 +27,10 @@ use common_query::request::QueryRequest;
use common_recordbatch::adapter::{AsyncRecordBatchStreamAdapter, DfRecordBatchStreamAdapter};
use common_recordbatch::util::{ChainedRecordBatchStream, LimitedRecordBatchStream};
use common_recordbatch::{DfSendableRecordBatchStream, SendableRecordBatchStream};
use datafusion::common::Result;
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::execution::TaskContext;
use datafusion::physical_expr::{EquivalenceProperties, Partitioning};
use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef;
@@ -157,10 +159,6 @@ impl ExecutionPlan for DistributedInspectExec {
"DistributedInspectExec"
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn schema(&self) -> ArrowSchemaRef {
self.arrow_schema.clone()
}
@@ -173,6 +171,13 @@ impl ExecutionPlan for DistributedInspectExec {
vec![]
}
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
) -> Result<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,
src/catalog/src/system_schema/information_schema.rs
+6
View File
@@ -135,6 +135,9 @@ lazy_static! {
GLOBAL_STATUS,
SESSION_STATUS,
PARTITIONS,
PLUGINS,
USER_PRIVILEGES,
PROCESSLIST,
];
}
@@ -237,6 +240,9 @@ impl SystemSchemaProviderInner for InformationSchemaProvider {
TABLE_PRIVILEGES => setup_memory_table!(TABLE_PRIVILEGES),
GLOBAL_STATUS => setup_memory_table!(GLOBAL_STATUS),
SESSION_STATUS => setup_memory_table!(SESSION_STATUS),
PLUGINS => setup_memory_table!(PLUGINS),
USER_PRIVILEGES => setup_memory_table!(USER_PRIVILEGES),
PROCESSLIST => setup_memory_table!(PROCESSLIST),
KEY_COLUMN_USAGE => Some(Arc::new(InformationSchemaKeyColumnUsage::new(
self.catalog_name.clone(),
self.catalog_manager.clone(),
src/catalog/src/system_schema/information_schema/information_memory_table.rs
+41
View File
@@ -383,6 +383,47 @@ pub(super) fn get_schema_columns(table_name: &str) -> (SchemaRef, Vec<VectorRef>
vec![],
),
// GreptimeDB has no pluggable components.
PLUGINS => (
string_columns(&[
"PLUGIN_NAME",
"PLUGIN_VERSION",
"PLUGIN_STATUS",
"PLUGIN_TYPE",
"PLUGIN_TYPE_VERSION",
"PLUGIN_LIBRARY",
"PLUGIN_LIBRARY_VERSION",
"PLUGIN_AUTHOR",
"PLUGIN_DESCRIPTION",
"PLUGIN_LICENSE",
"LOAD_OPTION",
]),
vec![],
),
// Privileges are not exposed through `information_schema`, same as the other
// `*_privileges` tables above.
USER_PRIVILEGES => (
string_columns(&["GRANTEE", "TABLE_CATALOG", "PRIVILEGE_TYPE", "IS_GRANTABLE"]),
vec![],
),
// Sessions are reported through `information_schema.process_list` and
// `SHOW PROCESSLIST`; this table carries the MySQL column shape only.
PROCESSLIST => (
vec![
bigint_column("ID"),
string_column("USER"),
string_column("HOST"),
string_column("DB"),
string_column("COMMAND"),
bigint_column("TIME"),
string_column("STATE"),
string_column("INFO"),
],
vec![],
),
_ => unreachable!("Unknown table in information_schema: {}", table_name),
};
src/catalog/src/system_schema/information_schema/ssts.rs
+1 -1
View File
@@ -336,7 +336,7 @@ mod tests {
};
let plan = table.scan_to_plan(request).unwrap().unwrap();
assert!(plan.as_any().is::<EmptyExec>());
assert!(plan.as_ref().is::<EmptyExec>());
assert_eq!(1, plan.schema().fields().len());
}
}
src/catalog/src/system_schema/information_schema/table_names.rs
+5
View File
@@ -55,3 +55,8 @@ pub const SSTS_INDEX_META: &str = "ssts_index_meta";
pub const TABLE_SEMANTICS: &str = "table_semantics";
pub const STATISTICS: &str = "statistics";
pub const RECYCLE_BIN: &str = "recycle_bin";
pub const PLUGINS: &str = "plugins";
pub const USER_PRIVILEGES: &str = "user_privileges";
/// MySQL's session list. GreptimeDB reports its own sessions through [`PROCESS_LIST`];
/// this table only exists so MySQL tooling finds the name and column shape it expects.
pub const PROCESSLIST: &str = "processlist";
src/catalog/src/table_source/dummy_catalog.rs
-13
View File
@@ -14,7 +14,6 @@
//! Dummy catalog for region server.
use std::any::Any;
use std::fmt;
use std::sync::Arc;
@@ -64,10 +63,6 @@ impl fmt::Debug for DummyCatalogList {
}
impl CatalogProviderList for DummyCatalogList {
fn as_any(&self) -> &dyn Any {
self
}
fn register_catalog(
&self,
_name: String,
@@ -98,10 +93,6 @@ struct DummyCatalogProvider {
}
impl CatalogProvider for DummyCatalogProvider {
fn as_any(&self) -> &dyn Any {
self
}
fn schema_names(&self) -> Vec<String> {
vec![]
}
@@ -135,10 +126,6 @@ struct DummySchemaProvider {
#[async_trait]
impl SchemaProvider for DummySchemaProvider {
fn as_any(&self) -> &dyn Any {
self
}
fn table_names(&self) -> Vec<String> {
vec![]
}
src/cmd/src/bin/query_perf_fixture/inspect_footer.rs
+73 -9
View File
@@ -14,18 +14,21 @@
use std::collections::BTreeSet;
use std::fs;
use std::ops::Range;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use clap::Args as ClapArgs;
use datafusion_object_store::path::Path as StorePath;
use datafusion_object_store::{ObjectMeta, ObjectStore};
use futures::{StreamExt, TryStreamExt};
use datafusion_object_store::{ObjectMeta, ObjectStore, ObjectStoreExt};
use futures::future::BoxFuture;
use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt};
use object_store::config::ObjectStoreConfig;
use object_store::factory::new_raw_object_store;
use object_store::services::Fs;
use parquet::arrow::async_reader::ParquetObjectReader;
use parquet::file::metadata::ParquetMetaDataReader;
use parquet::arrow::async_reader::AsyncFileReader;
use parquet::errors::{ParquetError, Result as ParquetResult};
use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
use serde::{Deserialize, Serialize};
/// Same shape as `query_regression_runner::model::DestinationConfig`. The two
@@ -100,6 +103,70 @@ struct ListedFile {
relative_path: String,
}
/// An asynchronous Parquet reader backed directly by an object store.
///
/// The file size is retained from the listing so footer reads use bounded
/// ranges without an additional stat/head request.
#[derive(Clone, Debug)]
struct ObjectStoreReader {
store: Arc<dyn ObjectStore>,
path: StorePath,
file_size: u64,
}
impl ObjectStoreReader {
fn new(store: Arc<dyn ObjectStore>, path: StorePath, file_size: u64) -> Self {
Self {
store,
path,
file_size,
}
}
}
fn to_parquet_error(error: datafusion_object_store::Error) -> ParquetError {
ParquetError::External(Box::new(error))
}
impl AsyncFileReader for ObjectStoreReader {
fn get_bytes(
&mut self,
range: Range<u64>,
) -> BoxFuture<'_, ParquetResult<prost::bytes::Bytes>> {
self.store
.get_range(&self.path, range)
.map_err(to_parquet_error)
.boxed()
}
fn get_byte_ranges(
&mut self,
ranges: Vec<Range<u64>>,
) -> BoxFuture<'_, ParquetResult<Vec<prost::bytes::Bytes>>> {
async move {
self.store
.get_ranges(&self.path, &ranges)
.await
.map_err(to_parquet_error)
}
.boxed()
}
fn get_metadata<'a>(
&'a mut self,
_options: Option<&'a parquet::arrow::arrow_reader::ArrowReaderOptions>,
) -> BoxFuture<'a, ParquetResult<Arc<ParquetMetaData>>> {
let file_size = self.file_size;
async move {
let metadata = ParquetMetaDataReader::new()
.load_and_finish(self, file_size)
.await?;
Ok(Arc::new(metadata))
}
.boxed()
}
}
pub(super) async fn run_inspect_footer(
args: InspectFooterArgs,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
@@ -217,11 +284,8 @@ async fn inspect_file(
file: &ListedFile,
column: &str,
) -> Result<FooterFileReport, Box<dyn std::error::Error + Send + Sync>> {
let mut reader =
ParquetObjectReader::new(store, file.location.clone()).with_file_size(file.size);
let metadata = ParquetMetaDataReader::new()
.load_and_finish(&mut reader, file.size)
.await?;
let mut reader = ObjectStoreReader::new(store, file.location.clone(), file.size);
let metadata = reader.get_metadata(None).await?;
let file_metadata = metadata.file_metadata();
let row_groups = metadata.row_groups();
let mut columns = Vec::new();
src/common/catalog/src/consts.rs
+6
View File
@@ -122,6 +122,12 @@ pub const INFORMATION_SCHEMA_STATISTICS_TABLE_ID: u32 = 43;
pub const INFORMATION_SCHEMA_RECYCLE_BIN_TABLE_ID: u32 = 44;
/// id for information_schema.flow_statistics
pub const INFORMATION_SCHEMA_FLOW_STATISTICS_TABLE_ID: u32 = 45;
/// id for information_schema.PLUGINS
pub const INFORMATION_SCHEMA_PLUGINS_TABLE_ID: u32 = 47;
/// id for information_schema.USER_PRIVILEGES
pub const INFORMATION_SCHEMA_USER_PRIVILEGES_TABLE_ID: u32 = 48;
/// id for information_schema.PROCESSLIST (for mysql)
pub const INFORMATION_SCHEMA_PROCESSLIST_TABLE_ID: u32 = 49;
// ----- End of information_schema tables -----
src/common/datasource/Cargo.toml
+2 -1
View File
@@ -34,7 +34,7 @@ futures.workspace = true
lazy_static.workspace = true
object-store.workspace = true
object_store_opendal.workspace = true
orc-rust = { version = "0.8", default-features = false, features = ["async"] }
orc-rust = { version = "0.9", default-features = false, features = ["async"] }
parquet.workspace = true
paste.workspace = true
regex.workspace = true
@@ -47,3 +47,4 @@ url.workspace = true
[dev-dependencies]
common-test-util.workspace = true
object-store = { workspace = true, features = ["testing"] }
src/common/datasource/src/error.rs
+12
View File
@@ -27,6 +27,15 @@ use url::ParseError;
#[snafu(visibility(pub))]
#[stack_trace_debug]
pub enum Error {
#[snafu(display("Parquet write cancelled"))]
ParquetWriteCancelled {},
#[snafu(display("Parquet writer limits must be positive"))]
InvalidParquetWriterLimits {},
#[snafu(display("Parquet writer resource limit exceeded: {reason}"))]
ParquetWriterResource { reason: String },
#[snafu(display("Unsupported compression type: {}", compression_type))]
UnsupportedCompressionType {
compression_type: String,
@@ -275,6 +284,9 @@ impl ErrorExt for Error {
fn status_code(&self) -> StatusCode {
use Error::*;
match self {
ParquetWriteCancelled {} => StatusCode::Cancelled,
InvalidParquetWriterLimits {} => StatusCode::InvalidArguments,
ParquetWriterResource { .. } => StatusCode::Suspended,
BuildBackend { .. }
| ListObjects { .. }
| ReadObject { .. }
src/common/datasource/src/file_format.rs
+6 -2
View File
@@ -35,7 +35,7 @@ use datafusion::datasource::file_format::file_compression_type::FileCompressionT
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::object_store::ObjectStoreUrl;
use datafusion::datasource::physical_plan::{
FileGroup, FileOpenFuture, FileScanConfigBuilder, FileSource, FileStream,
FileGroup, FileOpenFuture, FileScanConfigBuilder, FileSource, FileStreamBuilder,
};
use datafusion::error::{DataFusionError, Result as DataFusionResult};
use datafusion::physical_plan::SendableRecordBatchStream;
@@ -321,7 +321,11 @@ pub async fn file_to_stream(
let store = Arc::new(object_store_opendal::OpendalStore::new(store.clone()));
let file_opener = config.file_source().create_file_opener(store, &config, 0)?;
let stream = FileStream::new(&config, 0, file_opener, &ExecutionPlanMetricsSet::new())?;
let stream = FileStreamBuilder::new(&config)
.with_partition(0)
.with_file_opener(file_opener)
.with_metrics(&ExecutionPlanMetricsSet::new())
.build()?;
Ok(Box::pin(stream))
}
src/common/datasource/src/file_format/parquet.rs
+9 -52
View File
@@ -29,22 +29,17 @@ use datafusion::parquet::file::metadata::{
use datafusion::physical_plan::SendableRecordBatchStream;
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion_datasource::PartitionedFile;
use datatypes::schema::SchemaRef;
use futures::StreamExt;
use futures::future::BoxFuture;
use object_store::{FuturesAsyncReader, ObjectStore};
use parquet::arrow::AsyncArrowWriter;
use parquet::arrow::arrow_reader::ArrowReaderOptions;
use parquet::basic::{Compression, Encoding, ZstdLevel};
use parquet::file::properties::{WriterProperties, WriterPropertiesBuilder};
use parquet::schema::types::ColumnPath;
use snafu::ResultExt;
use tokio_util::compat::{Compat, FuturesAsyncReadCompatExt, FuturesAsyncWriteCompatExt};
use tokio_util::compat::{Compat, FuturesAsyncReadCompatExt};
use crate::DEFAULT_WRITE_BUFFER_SIZE;
use crate::buffered_writer::{ArrowWriterCloser, DfRecordBatchEncoder};
use crate::error::{self, Result, WriteObjectSnafu, WriteParquetSnafu};
use crate::error::{self, Result};
use crate::file_format::FileFormat;
use crate::parquet_writer::ParquetFileWriter;
use crate::share_buffer::SharedBuffer;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
@@ -212,61 +207,23 @@ impl ArrowWriterCloser for ArrowWriter<SharedBuffer> {
/// Returns number of rows written.
pub async fn stream_to_parquet(
mut stream: SendableRecordBatchStream,
schema: datatypes::schema::SchemaRef,
store: ObjectStore,
path: &str,
concurrency: usize,
) -> Result<usize> {
let write_props = column_wise_config(
WriterProperties::builder()
.set_compression(Compression::ZSTD(ZstdLevel::default()))
.set_statistics_truncate_length(None)
.set_column_index_truncate_length(None),
schema,
)
.build();
let inner_writer = store
.writer_with(path)
.concurrent(concurrency)
.chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize)
.await
.map(|w| w.into_futures_async_write().compat_write())
.context(WriteObjectSnafu { path })?;
let mut writer = AsyncArrowWriter::try_new(inner_writer, stream.schema(), Some(write_props))
.context(WriteParquetSnafu { path })?;
let mut writer =
ParquetFileWriter::open(stream.schema(), store, path, concurrency, None).await?;
let mut rows_written = 0;
while let Some(batch) = stream.next().await {
let batch = batch.context(error::ReadRecordBatchSnafu)?;
writer
.write(&batch)
.await
.context(WriteParquetSnafu { path })?;
rows_written += batch.num_rows();
let rows = batch.num_rows();
writer.write(batch, None).await?;
rows_written += rows;
}
writer.close().await.context(WriteParquetSnafu { path })?;
writer.finish(None).await?;
Ok(rows_written)
}
/// Customizes per-column properties.
fn column_wise_config(
mut props: WriterPropertiesBuilder,
schema: SchemaRef,
) -> WriterPropertiesBuilder {
// Disable dictionary for timestamp column, since for increasing timestamp column,
// the dictionary pages will be larger than data pages.
for col in schema.column_schemas() {
if col.data_type.is_timestamp() {
let path = ColumnPath::new(vec![col.name.clone()]);
props = props
.set_column_dictionary_enabled(path.clone(), false)
.set_column_encoding(path, Encoding::DELTA_BINARY_PACKED)
}
}
props
}
#[cfg(test)]
mod tests {
use common_test_util::find_workspace_path;
src/common/datasource/src/file_format/tests.rs
+10 -11
View File
@@ -19,7 +19,7 @@ use std::{assert_matches, vec};
use common_test_util::find_workspace_path;
use datafusion::assert_batches_eq;
use datafusion::datasource::physical_plan::{
CsvSource, FileScanConfig, FileSource, FileStream, JsonSource, ParquetSource,
CsvSource, FileScanConfig, FileSource, FileStreamBuilder, JsonSource, ParquetSource,
};
use datafusion::datasource::source::DataSourceExec;
use datafusion::execution::context::TaskContext;
@@ -50,16 +50,15 @@ impl Test<'_> {
.create_file_opener(store, &self.config, 0)
.unwrap();
let result = FileStream::new(
&self.config,
0,
file_opener,
&ExecutionPlanMetricsSet::new(),
)
.unwrap()
.map(|b| b.unwrap())
.collect::<Vec<_>>()
.await;
let result = FileStreamBuilder::new(&self.config)
.with_partition(0)
.with_file_opener(file_opener)
.with_metrics(&ExecutionPlanMetricsSet::new())
.build()
.unwrap()
.map(|b| b.unwrap())
.collect::<Vec<_>>()
.await;
assert_batches_eq!(self.expected, &result);
}
src/common/datasource/src/parquet_writer.rs
+446 -1
View File
@@ -12,11 +12,215 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use arrow::datatypes::{DataType, SchemaRef};
use arrow::record_batch::RecordBatch;
use bytes::Bytes;
use futures::future::BoxFuture;
use object_store::Writer;
use object_store::{ObjectStore, Writer};
use parquet::arrow::ArrowWriter;
use parquet::arrow::async_writer::AsyncFileWriter;
use parquet::basic::{Compression, Encoding, ZstdLevel};
use parquet::errors::ParquetError;
use parquet::file::properties::WriterProperties;
use parquet::schema::types::ColumnPath;
use snafu::{IntoError, ResultExt, ensure};
use tokio_util::sync::CancellationToken;
use crate::DEFAULT_WRITE_BUFFER_SIZE;
use crate::error::{self, Result};
/// Limits for one Parquet file. Flush thresholds are not hard memory caps.
#[derive(Clone, Copy, Debug)]
pub struct ParquetWriterLimits {
/// Maximum rows in a row group.
pub row_group_rows: usize,
/// Flush when encoder memory or encoded size reaches this threshold.
pub flush_threshold_bytes: usize,
/// Maximum row groups retained in the file footer.
pub max_row_groups: usize,
}
/// Encodes batches on a blocking runtime and writes one object-store file.
/// Callers must await each operation before aborting; dropping an in-flight
/// operation can leave a filesystem worker running after cleanup.
pub struct ParquetFileWriter {
encoder: Option<ArrowWriter<Vec<u8>>>,
sink: Writer,
store: ObjectStore,
path: String,
limits: Option<ParquetWriterLimits>,
}
impl ParquetFileWriter {
/// Open a file using COPY's encoding settings. Destination ownership and
/// overwrite policy belong to the caller. None preserves Parquet's defaults.
pub async fn open(
schema: SchemaRef,
store: ObjectStore,
path: &str,
concurrency: usize,
limits: Option<ParquetWriterLimits>,
) -> Result<Self> {
let mut props = WriterProperties::builder()
.set_compression(Compression::ZSTD(ZstdLevel::default()))
.set_statistics_truncate_length(None)
.set_column_index_truncate_length(None);
if let Some(limits) = limits {
ensure!(
limits.row_group_rows > 0
&& limits.flush_threshold_bytes > 0
&& limits.max_row_groups > 0,
error::InvalidParquetWriterLimitsSnafu
);
props = props
.set_max_row_group_row_count(Some(limits.row_group_rows))
.set_max_row_group_bytes(None);
}
for field in schema.fields() {
if matches!(field.data_type(), DataType::Timestamp(_, _)) {
let column = ColumnPath::new(vec![field.name().clone()]);
props = props
.set_column_dictionary_enabled(column.clone(), false)
.set_column_encoding(column, Encoding::DELTA_BINARY_PACKED);
}
}
let encoder = ArrowWriter::try_new(Vec::new(), schema, Some(props.build()))
.context(error::WriteParquetSnafu { path })?;
let sink = store
.writer_with(path)
.concurrent(concurrency)
.chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize)
.await
.context(error::WriteObjectSnafu { path })?;
Ok(Self {
encoder: Some(encoder),
sink,
store,
path: path.to_owned(),
limits,
})
}
/// Write a batch, enforcing file limits across batch and row-group boundaries.
pub async fn write(
&mut self,
batch: RecordBatch,
cancellation: Option<&CancellationToken>,
) -> Result<()> {
let mut offset = 0;
while offset < batch.num_rows() {
check_cancelled(cancellation)?;
let mut encoder = self.encoder.take().ok_or_else(|| {
error::WriteParquetSnafu { path: &self.path }
.into_error(ParquetError::General("Parquet writer is closed".into()))
})?;
let len = self.limits.map_or(batch.num_rows() - offset, |limits| {
(limits.row_group_rows - encoder.in_progress_rows()).min(batch.num_rows() - offset)
});
let slice = batch.slice(offset, len);
let limits = self.limits;
let path = self.path.clone();
let (encoder, bytes) = common_runtime::spawn_blocking_global(move || {
if let Some(limits) = limits {
ensure!(
encoder.flushed_row_groups().len() < limits.max_row_groups,
error::ParquetWriterResourceSnafu {
reason: "Parquet row-group metadata budget exceeded"
}
);
}
encoder
.write(&slice)
.context(error::WriteParquetSnafu { path: &path })?;
if limits.is_some_and(|limits| {
encoder.memory_size() >= limits.flush_threshold_bytes
|| encoder.in_progress_size() >= limits.flush_threshold_bytes
}) {
encoder
.flush()
.context(error::WriteParquetSnafu { path: &path })?;
}
// Draining preserves ArrowWriter's cumulative file offsets.
let bytes = std::mem::take(encoder.inner_mut());
Ok::<_, error::Error>((encoder, bytes))
})
.await
.context(error::JoinHandleSnafu)??;
self.encoder = Some(encoder);
check_cancelled(cancellation)?;
self.write_bytes(bytes).await?;
check_cancelled(cancellation)?;
offset += len;
}
Ok(())
}
async fn write_bytes(&mut self, bytes: Vec<u8>) -> Result<()> {
if !bytes.is_empty() {
self.sink
.write(bytes)
.await
.context(error::WriteObjectSnafu { path: &self.path })?;
}
Ok(())
}
/// Write the footer and close the file. Retains the handle for cleanup on error.
pub async fn finish(&mut self, cancellation: Option<&CancellationToken>) -> Result<()> {
check_cancelled(cancellation)?;
let mut encoder = self.encoder.take().ok_or_else(|| {
error::WriteParquetSnafu { path: &self.path }
.into_error(ParquetError::General("Parquet writer is closed".into()))
})?;
let path = self.path.clone();
let bytes = common_runtime::spawn_blocking_global(move || {
encoder
.finish()
.context(error::WriteParquetSnafu { path: &path })?;
Ok::<_, error::Error>(std::mem::take(encoder.inner_mut()))
})
.await
.context(error::JoinHandleSnafu)??;
self.write_bytes(bytes).await?;
check_cancelled(cancellation)?;
self.sink
.close()
.await
.context(error::WriteObjectSnafu { path: &self.path })?;
check_cancelled(cancellation)?;
Ok(())
}
/// Abort an exclusively owned file after all in-flight operations have completed.
/// If the backend cannot abort, deletion assumes this attempt owns the path.
pub async fn abort(mut self) -> Result<()> {
let result = self.sink.abort().await;
if result
.as_ref()
.is_err_and(|e| e.kind() == object_store::ErrorKind::Unsupported)
{
let store = self.store.clone();
let path = self.path.clone();
// Secure filesystem writers require dropping the handle before deletion.
drop(self);
store
.delete(&path)
.await
.context(error::WriteObjectSnafu { path })?;
} else {
result.context(error::WriteObjectSnafu { path: &self.path })?;
}
Ok(())
}
}
fn check_cancelled(cancellation: Option<&CancellationToken>) -> Result<()> {
ensure!(
cancellation.is_none_or(|token| !token.is_cancelled()),
error::ParquetWriteCancelledSnafu
);
Ok(())
}
/// Bridges opendal [Writer] with parquet [AsyncFileWriter].
pub struct AsyncWriter {
@@ -50,3 +254,244 @@ impl AsyncFileWriter for AsyncWriter {
})
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use arrow::array::{Int64Array, TimestampMillisecondArray};
use common_error::ext::ErrorExt;
use common_error::status_code::StatusCode;
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use super::*;
use crate::file_format::parquet::stream_to_parquet;
fn batch() -> RecordBatch {
RecordBatch::try_from_iter([
(
"value",
Arc::new(Int64Array::from(vec![Some(1), None, Some(3), Some(4)]))
as arrow::array::ArrayRef,
),
(
"ts",
Arc::new(TimestampMillisecondArray::from(vec![1, 2, 3, 4])),
),
])
.unwrap()
}
async fn read(store: &ObjectStore, path: &str) -> ParquetRecordBatchReaderBuilder<Bytes> {
ParquetRecordBatchReaderBuilder::try_new(store.read(path).await.unwrap().to_bytes())
.unwrap()
}
#[tokio::test]
async fn row_and_byte_limits_split_batches() {
let store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
let batch = batch();
for (row_group_rows, flush_threshold_bytes, split) in [(2, usize::MAX, 3), (100, 1, 2)] {
let mut writer = ParquetFileWriter::open(
batch.schema(),
store.clone(),
"groups.parquet",
1,
Some(ParquetWriterLimits {
row_group_rows,
flush_threshold_bytes,
max_row_groups: 2,
}),
)
.await
.unwrap();
writer.write(batch.slice(0, split), None).await.unwrap();
writer
.write(batch.slice(split, batch.num_rows() - split), None)
.await
.unwrap();
writer.finish(None).await.unwrap();
let reader = read(&store, "groups.parquet").await;
assert_eq!(reader.metadata().num_row_groups(), 2);
for group in reader.metadata().row_groups() {
assert_eq!(group.num_rows(), 2);
let encodings = group.column(1).encodings().collect::<Vec<_>>();
assert!(encodings.contains(&Encoding::DELTA_BINARY_PACKED));
assert!(!encodings.contains(&Encoding::RLE_DICTIONARY));
assert_eq!(
group.column(0).compression(),
Compression::ZSTD(ZstdLevel::default())
);
}
let actual = reader
.build()
.unwrap()
.collect::<std::result::Result<Vec<_>, _>>()
.unwrap();
assert_eq!(
arrow::compute::concat_batches(&batch.schema(), &actual).unwrap(),
batch
);
}
}
#[tokio::test]
async fn oversized_batch_stops_at_footer_limit_and_can_abort() {
let store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
let batch = batch();
let mut writer = ParquetFileWriter::open(
batch.schema(),
store.clone(),
"limited.parquet",
1,
Some(ParquetWriterLimits {
row_group_rows: 2,
flush_threshold_bytes: usize::MAX,
max_row_groups: 1,
}),
)
.await
.unwrap();
let err = writer.write(batch, None).await.unwrap_err();
assert!(matches!(err, error::Error::ParquetWriterResource { .. }));
assert_eq!(err.status_code(), StatusCode::Suspended);
writer.abort().await.unwrap();
assert!(!store.exists("limited.parquet").await.unwrap());
}
#[tokio::test]
async fn copy_stream_preserves_values_and_empty_schema() {
let store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
let batch = batch();
for (path, batches) in [
("copy.parquet", vec![batch.clone()]),
("empty.parquet", vec![]),
] {
let stream = RecordBatchStreamAdapter::new(
batch.schema(),
futures::stream::iter(batches.clone().into_iter().map(Ok)),
);
assert_eq!(
stream_to_parquet(Box::pin(stream), store.clone(), path, 2)
.await
.unwrap(),
batches.iter().map(RecordBatch::num_rows).sum::<usize>()
);
let reader = read(&store, path).await;
assert_eq!(reader.schema().fields(), batch.schema().fields());
let actual = reader
.build()
.unwrap()
.collect::<std::result::Result<Vec<_>, _>>()
.unwrap();
assert_eq!(actual, batches);
}
}
#[tokio::test]
async fn failed_copy_preserves_untouched_securefs_destination() {
let directory = common_test_util::temp_dir::create_temp_dir("copy_existing");
let path = directory.path().join("existing.parquet");
std::fs::write(&path, b"original bytes").unwrap();
let access = crate::object_store::LocalFileAccess::sandboxed(directory.path()).unwrap();
let store = crate::object_store::build_backend_for_write(
&format!("{}/", directory.path().display()),
&Default::default(),
&access,
)
.await
.unwrap();
let stream = RecordBatchStreamAdapter::new(
batch().schema(),
futures::stream::iter(vec![Err(datafusion::error::DataFusionError::Execution(
"injected input failure".into(),
))]),
);
let err = stream_to_parquet(Box::pin(stream), store, "existing.parquet", 1)
.await
.unwrap_err();
assert!(matches!(err, error::Error::ReadRecordBatch { .. }));
assert_eq!(std::fs::read(path).unwrap(), b"original bytes");
}
struct PausedFooterWriter {
inner: object_store::layers::mock::oio::Writer,
paused: bool,
started: Arc<tokio::sync::Notify>,
release: Arc<tokio::sync::Notify>,
closed: Arc<std::sync::atomic::AtomicBool>,
}
impl object_store::layers::mock::oio::Write for PausedFooterWriter {
async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> {
if !self.paused {
self.paused = true;
self.started.notify_one();
self.release.notified().await;
}
self.inner.write(bytes).await
}
async fn close(&mut self) -> object_store::Result<object_store::layers::mock::Metadata> {
self.closed.store(true, std::sync::atomic::Ordering::SeqCst);
self.inner.close().await
}
async fn abort(&mut self) -> object_store::Result<()> {
self.inner.abort().await
}
}
#[tokio::test]
async fn cancellation_during_footer_write_waits_then_aborts_before_close() {
use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory};
let started = Arc::new(tokio::sync::Notify::new());
let release = Arc::new(tokio::sync::Notify::new());
let closed = Arc::new(std::sync::atomic::AtomicBool::new(false));
let factory: MockWriterFactory = Arc::new({
let (started, release, closed) = (started.clone(), release.clone(), closed.clone());
move |_, _, inner| {
Box::new(PausedFooterWriter {
inner,
paused: false,
started: started.clone(),
release: release.clone(),
closed: closed.clone(),
})
}
});
let store = ObjectStore::new(object_store::services::Memory::default())
.unwrap()
.layer(
MockLayerBuilder::default()
.writer_factory(factory)
.build()
.unwrap(),
);
let mut writer =
ParquetFileWriter::open(batch().schema(), store.clone(), "cancel.parquet", 1, None)
.await
.unwrap();
// Force footer bytes through the sink before its close operation.
writer.sink = store.writer_with("cancel.parquet").chunk(1).await.unwrap();
let cancellation = CancellationToken::new();
let result = {
let finish = writer.finish(Some(&cancellation));
tokio::pin!(finish);
tokio::select! {
result = &mut finish => panic!("finished before footer write: {result:?}"),
_ = started.notified() => {},
}
cancellation.cancel();
assert!(futures::poll!(&mut finish).is_pending());
release.notify_one();
finish.await
};
assert!(matches!(
result,
Err(error::Error::ParquetWriteCancelled {})
));
assert!(!closed.load(std::sync::atomic::Ordering::SeqCst));
writer.abort().await.unwrap();
assert!(!store.exists("cancel.parquet").await.unwrap());
}
}
src/common/datasource/src/test_util.rs
+13 -9
View File
@@ -20,7 +20,7 @@ use datafusion::datasource::file_format::file_compression_type::FileCompressionT
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::object_store::ObjectStoreUrl;
use datafusion::datasource::physical_plan::{
CsvSource, FileGroup, FileScanConfig, FileScanConfigBuilder, FileSource, FileStream,
CsvSource, FileGroup, FileScanConfig, FileScanConfigBuilder, FileSource, FileStreamBuilder,
JsonOpener, JsonSource,
};
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
@@ -110,13 +110,12 @@ pub async fn setup_stream_to_json_test(origin_path: &str, threshold: impl Fn(usi
let size = store.read(origin_path).await.unwrap().len();
let config = scan_config(None, origin_path, Arc::new(JsonSource::new(schema)));
let stream = FileStream::new(
&config,
0,
Arc::new(json_opener),
&ExecutionPlanMetricsSet::new(),
)
.unwrap();
let stream = FileStreamBuilder::new(&config)
.with_partition(0)
.with_file_opener(Arc::new(json_opener))
.with_metrics(&ExecutionPlanMetricsSet::new())
.build()
.unwrap();
let (tmp_store, dir) = test_tmp_store("test_stream_to_json");
@@ -162,7 +161,12 @@ pub async fn setup_stream_to_csv_test(
0,
)
.unwrap();
let stream = FileStream::new(&config, 0, csv_opener, &ExecutionPlanMetricsSet::new()).unwrap();
let stream = FileStreamBuilder::new(&config)
.with_partition(0)
.with_file_opener(csv_opener)
.with_metrics(&ExecutionPlanMetricsSet::new())
.build()
.unwrap();
let (tmp_store, dir) = test_tmp_store("test_stream_to_csv");
src/common/function/src/aggrs/aggr_wrapper.rs
+19 -21
View File
@@ -33,15 +33,15 @@ use datafusion::functions_aggregate::count::Count;
use datafusion::functions_aggregate::min_max::{Max, Min};
use datafusion::optimizer::AnalyzerRule;
use datafusion::optimizer::analyzer::type_coercion::TypeCoercion;
use datafusion::physical_planner::create_aggregate_expr_and_maybe_filter;
use datafusion_common::{Column, ScalarValue};
use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams};
use datafusion_expr::function::StateFieldsArgs;
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion_expr::{
Accumulator, Aggregate, AggregateUDF, AggregateUDFImpl, EmitTo, Expr, ExprSchemable,
GroupsAccumulator, LogicalPlan, Signature,
};
use datafusion_physical_expr::aggregate::AggregateFunctionExpr;
use datafusion_physical_expr::aggregate::{AggregateFunctionExpr, LoweredAggregateBuilder};
use datatypes::arrow::datatypes::{DataType, Field};
use crate::aggrs::aggr_wrapper::fix_order::FixStateUdafOrderingAnalyzer;
@@ -210,12 +210,25 @@ impl StateMergeHelper {
lower_aggr_exprs.push(expr);
// then create the merge function using the physical expression of the original aggregate function
let (original_phy_expr, _filter, _ordering) = create_aggregate_expr_and_maybe_filter(
let (name, human_display) = match aggr_expr {
Expr::Alias(alias) => (alias.name.clone(), aggr_expr.human_display().to_string()),
Expr::AggregateFunction(_) => (
aggr_expr.schema_name().to_string(),
aggr_expr.human_display().to_string(),
),
_ => unreachable!("aggregate expression was validated above"),
};
let original_phy_expr = LoweredAggregateBuilder::new(
aggr_expr,
aggr.input.schema(),
aggr.input.schema().as_arrow(),
&Default::default(),
)?;
&PhysicalPlanningContext::default(),
)
.with_name(name)
.with_human_display(human_display)
.build()?
.aggregate;
let merge_func = MergeWrapper::new(
(*aggr_func.func).clone(),
@@ -371,9 +384,6 @@ impl AggregateUDFImpl for StateWrapper {
Ok(Box::new(StateGroupsAccum::new(inner, state_type)?))
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
self.name.as_str()
}
@@ -443,7 +453,7 @@ impl AggregateUDFImpl for StateWrapper {
&self,
statistics_args: &datafusion_expr::StatisticsArgs,
) -> Option<ScalarValue> {
let inner = self.inner().inner().as_any();
let inner = self.inner().inner();
// only count/min/max need special handling here, for getting result from statistics
// the result of count/min/max is also the result of count_state so can return directly
let can_use_stat = inner.is::<Count>() || inner.is::<Max>() || inner.is::<Min>();
@@ -572,11 +582,10 @@ impl GroupsAccumulator for StateGroupsAccum {
&mut self,
values: &[ArrayRef],
group_indices: &[usize],
opt_filter: Option<&BooleanArray>,
total_num_groups: usize,
) -> datafusion_common::Result<()> {
self.inner
.merge_batch(values, group_indices, opt_filter, total_num_groups)
.merge_batch(values, group_indices, total_num_groups)
}
fn evaluate(&mut self, emit_to: EmitTo) -> datafusion_common::Result<ArrayRef> {
@@ -596,10 +605,6 @@ impl GroupsAccumulator for StateGroupsAccum {
self.inner.convert_to_state(values, opt_filter)
}
fn supports_convert_to_state(&self) -> bool {
self.inner.supports_convert_to_state()
}
fn size(&self) -> usize {
self.inner.size()
}
@@ -800,10 +805,6 @@ impl AggregateUDFImpl for DeltaMergeWrapper {
}))
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
&self.name
}
@@ -973,9 +974,6 @@ impl AggregateUDFImpl for MergeWrapper {
Ok(Box::new(MergeAccum::new(inner_accum, &fields)))
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
self.name.as_str()
}
src/common/function/src/aggrs/aggr_wrapper/fix_order.rs
-1
View File
@@ -144,7 +144,6 @@ fn rewrite_expr(
let Some(old_state_wrapper) = aggregate_function
.func
.inner()
.as_any()
.downcast_ref::<StateWrapper>()
else {
return Ok(Transformed::no(Expr::AggregateFunction(aggregate_function)));
src/common/function/src/aggrs/aggr_wrapper/tests.rs
+36 -61
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
@@ -25,6 +24,7 @@ use arrow::record_batch::RecordBatch;
use arrow_schema::SchemaRef;
use common_telemetry::init_default_ut_logging;
use datafusion::catalog::{Session, TableProvider};
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::datasource::DefaultTableSource;
use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext};
use datafusion::functions_aggregate::average::avg_udaf;
@@ -39,16 +39,16 @@ use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner};
use datafusion::prelude::SessionContext;
use datafusion_common::arrow::array::AsArray;
use datafusion_common::arrow::datatypes::{Float64Type, UInt64Type};
use datafusion_common::{Column, TableReference};
use datafusion_common::{Column, Result, TableReference};
use datafusion_expr::expr::{AggregateFunction, NullTreatment};
use datafusion_expr::function::AccumulatorArgs;
use datafusion_expr::{
Aggregate, AggregateUDFImpl, ColumnarValue, Expr, LogicalPlan, ScalarFunctionArgs, SortExpr,
TableScan, TypeSignature, lit,
TableScanBuilder, TypeSignature, lit,
};
use datafusion_physical_expr::aggregate::AggregateExprBuilder;
use datafusion_physical_expr::expressions::{Column as PhysicalColumn, col, lit as physical_lit};
use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
use futures::{Stream, StreamExt as _};
use hyperloglogplus::HyperLogLog;
use pretty_assertions::assert_eq;
@@ -97,10 +97,6 @@ impl ExecutionPlan for MockInputExec {
"MockInputExec"
}
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.properties
}
@@ -109,6 +105,13 @@ impl ExecutionPlan for MockInputExec {
vec![]
}
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
) -> Result<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn with_new_children(
self: Arc<Self>,
_children: Vec<Arc<dyn ExecutionPlan>>,
@@ -203,10 +206,6 @@ impl Default for DummyTableProvider {
#[async_trait::async_trait]
impl TableProvider for DummyTableProvider {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn schema(&self) -> Arc<arrow_schema::Schema> {
self.schema.clone()
}
@@ -237,14 +236,12 @@ fn dummy_table_scan() -> LogicalPlan {
let table_provider = Arc::new(DummyTableProvider::default());
let table_source = DefaultTableSource::new(table_provider);
LogicalPlan::TableScan(
TableScan::try_new(
TableReference::bare("Number"),
Arc::new(table_source),
None,
vec![],
None,
)
.unwrap(),
TableScanBuilder::new(TableReference::bare("Number"), Arc::new(table_source))
.with_projection(None)
.with_filters(vec![])
.with_fetch(None)
.build()
.unwrap(),
)
}
@@ -252,14 +249,12 @@ fn dummy_table_scan_with_ts() -> LogicalPlan {
let table_provider = Arc::new(DummyTableProvider::with_ts(None));
let table_source = DefaultTableSource::new(table_provider);
LogicalPlan::TableScan(
TableScan::try_new(
TableReference::bare("Number"),
Arc::new(table_source),
None,
vec![],
None,
)
.unwrap(),
TableScanBuilder::new(TableReference::bare("Number"), Arc::new(table_source))
.with_projection(None)
.with_filters(vec![])
.with_fetch(None)
.build()
.unwrap(),
)
}
@@ -381,10 +376,7 @@ async fn test_sum_udaf() {
.create_physical_plan(&res.lower_state, &ctx.state())
.await
.unwrap();
let aggr_exec = phy_aggr_state_plan
.as_any()
.downcast_ref::<AggregateExec>()
.unwrap();
let aggr_exec = phy_aggr_state_plan.downcast_ref::<AggregateExec>().unwrap();
let aggr_func_expr = &aggr_exec.aggr_expr()[0];
let mut state_accum = aggr_func_expr.create_accumulator().unwrap();
@@ -414,10 +406,7 @@ async fn test_sum_udaf() {
.create_physical_plan(&res.upper_merge, &ctx.state())
.await
.unwrap();
let aggr_exec = phy_aggr_merge_plan
.as_any()
.downcast_ref::<AggregateExec>()
.unwrap();
let aggr_exec = phy_aggr_merge_plan.downcast_ref::<AggregateExec>().unwrap();
let aggr_func_expr = &aggr_exec.aggr_expr()[0];
let mut merge_accum = aggr_func_expr.create_accumulator().unwrap();
@@ -543,10 +532,7 @@ async fn test_avg_udaf() {
.create_physical_plan(&coerced_aggr_state_plan, &ctx.state())
.await
.unwrap();
let aggr_exec = phy_aggr_state_plan
.as_any()
.downcast_ref::<AggregateExec>()
.unwrap();
let aggr_exec = phy_aggr_state_plan.downcast_ref::<AggregateExec>().unwrap();
let aggr_func_expr = &aggr_exec.aggr_expr()[0];
let mut state_accum = aggr_func_expr.create_accumulator().unwrap();
@@ -582,10 +568,7 @@ async fn test_avg_udaf() {
.create_physical_plan(&res.upper_merge, &ctx.state())
.await
.unwrap();
let aggr_exec = phy_aggr_merge_plan
.as_any()
.downcast_ref::<AggregateExec>()
.unwrap();
let aggr_exec = phy_aggr_merge_plan.downcast_ref::<AggregateExec>().unwrap();
let aggr_func_expr = &aggr_exec.aggr_expr()[0];
let mut merge_accum = aggr_func_expr.create_accumulator().unwrap();
@@ -703,10 +686,7 @@ async fn test_last_value_order_by_udaf() {
.create_physical_plan(&fixed_aggr_state_plan, &ctx.state())
.await
.unwrap();
let aggr_exec = phy_aggr_state_plan
.as_any()
.downcast_ref::<AggregateExec>()
.unwrap();
let aggr_exec = phy_aggr_state_plan.downcast_ref::<AggregateExec>().unwrap();
let aggr_func_expr = &aggr_exec.aggr_expr()[0];
let merge_input_fields = vec![Arc::new(Field::new(
@@ -796,10 +776,7 @@ async fn test_last_value_order_by_udaf() {
.create_physical_plan(&res.upper_merge, &ctx.state())
.await
.unwrap();
let aggr_exec = phy_aggr_merge_plan
.as_any()
.downcast_ref::<AggregateExec>()
.unwrap();
let aggr_exec = phy_aggr_merge_plan.downcast_ref::<AggregateExec>().unwrap();
let aggr_func_expr = &aggr_exec.aggr_expr()[0];
let mut merge_accum = aggr_func_expr.create_accumulator().unwrap();
@@ -900,7 +877,7 @@ fn test_avg_state_groups_accumulator_state_merge_evaluate() {
.update_batch(&merged_values, &merged_group_indices, None, 3)
.unwrap();
merged_accum
.merge_batch(&source_state, &[1, 2, 0], None, 3)
.merge_batch(&source_state, &[1, 2, 0], 3)
.unwrap();
let result = merged_accum.evaluate(EmitTo::All).unwrap();
@@ -1324,14 +1301,12 @@ async fn test_udaf_correct_eval_result() {
);
let table_source = DefaultTableSource::new(Arc::new(table_provider));
let logical_plan = LogicalPlan::TableScan(
TableScan::try_new(
test_table_ref.clone(),
Arc::new(table_source),
None,
vec![],
None,
)
.unwrap(),
TableScanBuilder::new(test_table_ref.clone(), Arc::new(table_source))
.with_projection(None)
.with_filters(vec![])
.with_fetch(None)
.build()
.unwrap(),
);
let args = case.args;
src/common/function/src/aggrs/approximate/uddsketch.rs
-2
View File
@@ -123,7 +123,6 @@ impl UddSketchState {
fn downcast_accumulator_args(args: AccumulatorArgs) -> DfResult<(u32, f64)> {
let bucket_size = match args.exprs[0]
.as_any()
.downcast_ref::<Literal>()
.map(|lit| lit.value())
{
@@ -140,7 +139,6 @@ fn downcast_accumulator_args(args: AccumulatorArgs) -> DfResult<(u32, f64)> {
};
let error_rate = match args.exprs[1]
.as_any()
.downcast_ref::<Literal>()
.map(|lit| lit.value())
{
src/common/function/src/aggrs/count_hash.rs
+26 -13
View File
@@ -26,7 +26,7 @@ use std::sync::Arc;
use ahash::RandomState;
use datafusion_common::cast::as_list_array;
use datafusion_common::error::Result;
use datafusion_common::hash_utils::create_hashes;
use datafusion_common::hash_utils::create_hashes_with_hasher;
use datafusion_common::utils::SingleRowListArrayBuilder;
use datafusion_common::{ScalarValue, internal_err, not_impl_err};
use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
@@ -74,10 +74,6 @@ pub struct CountHash {
}
impl AggregateUDFImpl for CountHash {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
"count_hash"
}
@@ -208,7 +204,7 @@ impl GroupsAccumulator for CountHashGroupAccumulator {
let array = &values[0];
self.batch_hashes.clear();
self.batch_hashes.resize(array.len(), 0);
let hashes = create_hashes(
let hashes = create_hashes_with_hasher(
&[ArrayRef::clone(array)],
&self.random_state,
&mut self.batch_hashes,
@@ -279,7 +275,6 @@ impl GroupsAccumulator for CountHashGroupAccumulator {
&mut self,
values: &[ArrayRef],
group_indices: &[usize],
_opt_filter: Option<&BooleanArray>,
total_num_groups: usize,
) -> Result<()> {
assert_eq!(
@@ -364,10 +359,6 @@ impl GroupsAccumulator for CountHashGroupAccumulator {
Ok(vec![Arc::new(list_array)])
}
fn supports_convert_to_state(&self) -> bool {
true
}
fn size(&self) -> usize {
// Base size of the struct
let mut size = size_of::<Self>();
@@ -423,7 +414,7 @@ impl Accumulator for CountHashAccumulator {
self.batch_hashes.clear();
self.batch_hashes.resize(arr.len(), 0);
let hashes = create_hashes(
let hashes = create_hashes_with_hasher(
&[ArrayRef::clone(arr)],
&self.random_state,
&mut self.batch_hashes,
@@ -522,6 +513,28 @@ mod tests {
Ok(())
}
#[test]
fn test_count_hash_accumulator_typed_null_state_merge() -> Result<()> {
let typed_nulls = Arc::new(Int32Array::from(vec![None, None])) as ArrayRef;
let mut fresh = create_test_accumulator();
fresh.update_batch(&[typed_nulls])?;
let fresh_state = fresh.state()?;
assert_eq!(fresh.evaluate()?, ScalarValue::Int64(Some(1)));
let persisted_state = Arc::new(
SingleRowListArrayBuilder::new(Arc::new(UInt64Array::from(vec![0])) as ArrayRef)
.build_list_array(),
) as ArrayRef;
let mut restored = create_test_accumulator();
restored.merge_batch(&[persisted_state])?;
assert_eq!(restored.evaluate()?, fresh.evaluate()?);
assert_eq!(restored.state()?, fresh_state);
Ok(())
}
#[test]
fn test_count_hash_accumulator_merge() -> Result<()> {
// Accumulator 1
@@ -622,7 +635,7 @@ mod tests {
// We will merge acc1's group 0 into acc2's group 0
// and acc1's group 1 into acc2's group 2
let merge_group_indices = vec![0, 2];
acc2.merge_batch(&state1, &merge_group_indices, None, 3)?;
acc2.merge_batch(&state1, &merge_group_indices, 3)?;
let result_array = acc2.evaluate(EmitTo::All)?;
let result = result_array.as_any().downcast_ref::<Int64Array>().unwrap();
src/common/function/src/helper.rs
+27
View File
@@ -23,6 +23,33 @@ use datatypes::types::cast::cast;
use datatypes::value::ValueRef;
use snafu::{OptionExt, ResultExt};
/// Integer types accepted by geospatial function signatures.
pub(crate) const INTEGER_TYPES: &[DataType] = &[
DataType::Int8,
DataType::Int16,
DataType::Int32,
DataType::Int64,
DataType::UInt8,
DataType::UInt16,
DataType::UInt32,
DataType::UInt64,
];
/// Legacy primitive numeric signature types; Decimal values are coerced to `Float64`.
pub(crate) const NUMERICS: &[DataType] = &[
DataType::Int8,
DataType::Int16,
DataType::Int32,
DataType::Int64,
DataType::UInt8,
DataType::UInt16,
DataType::UInt32,
DataType::UInt64,
DataType::Float16,
DataType::Float32,
DataType::Float64,
];
/// Create a function signature with oneof signatures of interleaving two arguments.
pub(crate) fn one_of_sigs2(args1: Vec<DataType>, args2: Vec<DataType>) -> Signature {
let mut sigs = Vec::with_capacity(args1.len() * args2.len());
src/common/function/src/scalars/anomaly/iqr.rs
+1 -6
View File
@@ -24,7 +24,6 @@
//! When IQR = 0 (constant quartiles), returns 0.0 if value is on the fence,
//! or +inf if value is outside.
use std::any::Any;
use std::fmt::Debug;
use std::ops::Range;
use std::sync::Arc;
@@ -32,11 +31,11 @@ use std::sync::Arc;
use arrow::array::{Array, ArrayRef, Float64Array};
use arrow::datatypes::{DataType, Field, FieldRef};
use datafusion_common::{DataFusionError, Result, ScalarValue};
use datafusion_expr::type_coercion::aggregates::NUMERICS;
use datafusion_expr::{PartitionEvaluator, Signature, Volatility, WindowUDFImpl};
use datafusion_functions_window_common::field::WindowUDFFieldArgs;
use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
use crate::helper::NUMERICS;
use crate::scalars::anomaly::utils::{cast_to_f64, collect_window_values, percentile_sorted};
/// Minimum valid samples for IQR (linear-interpolated Q1 != Q3 is possible at n >= 3).
@@ -56,10 +55,6 @@ impl AnomalyScoreIqr {
}
impl WindowUDFImpl for AnomalyScoreIqr {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"anomaly_score_iqr"
}
src/common/function/src/scalars/anomaly/mad.rs
+1 -6
View File
@@ -20,7 +20,6 @@
//! When MAD = 0 (majority-constant window), returns 0.0 if value equals
//! median, or +inf otherwise.
use std::any::Any;
use std::fmt::Debug;
use std::ops::Range;
use std::sync::Arc;
@@ -28,11 +27,11 @@ use std::sync::Arc;
use arrow::array::{Array, ArrayRef, Float64Array};
use arrow::datatypes::{DataType, Field, FieldRef};
use datafusion_common::{DataFusionError, Result, ScalarValue};
use datafusion_expr::type_coercion::aggregates::NUMERICS;
use datafusion_expr::{PartitionEvaluator, Signature, Volatility, WindowUDFImpl};
use datafusion_functions_window_common::field::WindowUDFFieldArgs;
use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
use crate::helper::NUMERICS;
use crate::scalars::anomaly::utils::{
anomaly_ratio, cast_to_f64, collect_window_values, median_f64,
};
@@ -57,10 +56,6 @@ impl AnomalyScoreMad {
}
impl WindowUDFImpl for AnomalyScoreMad {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"anomaly_score_mad"
}
src/common/function/src/scalars/anomaly/zscore.rs
+1 -6
View File
@@ -19,7 +19,6 @@
//! When stddev = 0 (constant window), returns 0.0 if value equals mean,
//! or +inf otherwise.
use std::any::Any;
use std::fmt::Debug;
use std::ops::Range;
use std::sync::Arc;
@@ -27,11 +26,11 @@ use std::sync::Arc;
use arrow::array::{Array, ArrayRef, Float64Array};
use arrow::datatypes::{DataType, Field, FieldRef};
use datafusion_common::{DataFusionError, Result, ScalarValue};
use datafusion_expr::type_coercion::aggregates::NUMERICS;
use datafusion_expr::{PartitionEvaluator, Signature, Volatility, WindowUDFImpl};
use datafusion_functions_window_common::field::WindowUDFFieldArgs;
use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
use crate::helper::NUMERICS;
use crate::scalars::anomaly::utils::{anomaly_ratio, cast_to_f64, collect_window_values};
/// Minimum valid samples for zscore (stddev requires n >= 2).
@@ -51,10 +50,6 @@ impl AnomalyScoreZscore {
}
impl WindowUDFImpl for AnomalyScoreZscore {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"anomaly_score_zscore"
}
src/common/function/src/scalars/geo/geohash.rs
+3 -3
View File
@@ -22,12 +22,12 @@ use datafusion::arrow::array::{Array, AsArray, ListBuilder, StringViewBuilder};
use datafusion::arrow::datatypes::{DataType, Field, Float64Type, UInt8Type};
use datafusion::logical_expr::ColumnarValue;
use datafusion_common::DataFusionError;
use datafusion_expr::type_coercion::aggregates::INTEGERS;
use datafusion_expr::{ScalarFunctionArgs, Signature, TypeSignature, Volatility};
use geohash::Coord;
use snafu::ResultExt;
use crate::function::{Function, extract_args};
use crate::helper::INTEGER_TYPES;
use crate::scalars::geo::helpers;
fn ensure_resolution_usize(v: u8) -> datafusion_common::Result<usize> {
@@ -49,7 +49,7 @@ impl Default for GeohashFunction {
fn default() -> Self {
let mut signatures = Vec::new();
for coord_type in &[DataType::Float32, DataType::Float64] {
for resolution_type in INTEGERS {
for resolution_type in INTEGER_TYPES {
signatures.push(TypeSignature::Exact(vec![
// latitude
coord_type.clone(),
@@ -146,7 +146,7 @@ impl Default for GeohashNeighboursFunction {
fn default() -> Self {
let mut signatures = Vec::new();
for coord_type in &[DataType::Float32, DataType::Float64] {
for resolution_type in INTEGERS {
for resolution_type in INTEGER_TYPES {
signatures.push(TypeSignature::Exact(vec![
// latitude
coord_type.clone(),
src/common/function/src/scalars/geo/h3.rs
+4 -4
View File
@@ -26,7 +26,6 @@ use datafusion::arrow::compute;
use datafusion::arrow::datatypes::{Float64Type, Int64Type, UInt8Type, UInt64Type};
use datafusion::logical_expr::ColumnarValue;
use datafusion_common::{DataFusionError, ScalarValue};
use datafusion_expr::type_coercion::aggregates::INTEGERS;
use datafusion_expr::{ScalarFunctionArgs, Signature, TypeSignature, Volatility};
use datatypes::arrow::datatypes::{DataType, Field};
use derive_more::Display;
@@ -34,6 +33,7 @@ use h3o::{CellIndex, LatLng, Resolution};
use snafu::prelude::*;
use crate::function::{Function, extract_args};
use crate::helper::INTEGER_TYPES;
use crate::scalars::geo::helpers;
static CELL_TYPES: LazyLock<Vec<DataType>> =
@@ -42,11 +42,11 @@ static CELL_TYPES: LazyLock<Vec<DataType>> =
static COORDINATE_TYPES: LazyLock<Vec<DataType>> =
LazyLock::new(|| vec![DataType::Float32, DataType::Float64]);
static RESOLUTION_TYPES: &[DataType] = INTEGERS;
static RESOLUTION_TYPES: &[DataType] = INTEGER_TYPES;
static DISTANCE_TYPES: &[DataType] = INTEGERS;
static DISTANCE_TYPES: &[DataType] = INTEGER_TYPES;
static POSITION_TYPES: &[DataType] = INTEGERS;
static POSITION_TYPES: &[DataType] = INTEGER_TYPES;
/// Function that returns [h3] encoding cellid for a given geospatial coordinate.
///
src/common/function/src/scalars/geo/s2.rs
+2 -1
View File
@@ -25,6 +25,7 @@ use s2::latlng::LatLng;
use snafu::ensure;
use crate::function::{Function, extract_args};
use crate::helper::INTEGER_TYPES;
use crate::scalars::geo::helpers;
use crate::scalars::geo::helpers::ensure_and_coerce;
@@ -34,7 +35,7 @@ static CELL_TYPES: LazyLock<Vec<DataType>> =
static COORDINATE_TYPES: LazyLock<Vec<DataType>> =
LazyLock::new(|| vec![DataType::Float32, DataType::Float64]);
static LEVEL_TYPES: &[DataType] = datafusion_expr::type_coercion::aggregates::INTEGERS;
static LEVEL_TYPES: &[DataType] = INTEGER_TYPES;
/// Function that returns [s2] encoding cellid for a given geospatial coordinate.
///
src/common/function/src/scalars/json/json_get.rs
+31 -3
View File
@@ -532,7 +532,7 @@ mod tests {
use datafusion_common::arrow::datatypes::{Float64Type, Int64Type};
use datatypes::extension::json::Json2ExtensionType;
use datatypes::types::parse_string_to_jsonb;
use serde_json::json;
use serde_json::{Value, json};
use super::*;
@@ -595,6 +595,34 @@ mod tests {
})
}
fn assert_json_or_string_eq(actual: Option<&str>, expected: Option<&str>) {
let is_json_container = |value: &str| {
matches!(
serde_json::from_str::<Value>(value),
Ok(Value::Object(_) | Value::Array(_))
)
};
match (actual, expected) {
(Some(actual), Some(expected))
if is_json_container(actual) || is_json_container(expected) =>
{
let actual_value = serde_json::from_str::<Value>(actual).unwrap_or_else(|error| {
panic!("failed to parse actual JSON result {actual:?}: {error}")
});
let expected_value =
serde_json::from_str::<Value>(expected).unwrap_or_else(|error| {
panic!("failed to parse expected JSON result {expected:?}: {error}")
});
assert_eq!(
actual_value, expected_value,
"JSON result mismatch: actual {actual:?}, expected {expected:?}"
);
}
_ => assert_eq!(actual, expected),
}
}
#[test]
fn test_json_get_int() {
let json_get_int = JsonGetInt::default();
@@ -895,7 +923,7 @@ mod tests {
let result = result.as_string_view();
assert_eq!(1, result.len());
let actual = result.is_valid(0).then(|| result.value(0));
assert_eq!(actual, expect);
assert_json_or_string_eq(actual, expect);
}
}
@@ -1033,7 +1061,7 @@ mod tests {
let result = result.as_string_view();
assert_eq!(1, result.len());
let actual = result.is_valid(0).then(|| result.value(0));
assert_eq!(actual, expect);
assert_json_or_string_eq(actual, expect);
}
let json_strings = [
src/common/function/src/scalars/json/json_get_rewriter.rs
+5 -13
View File
@@ -59,10 +59,8 @@ impl FunctionRewrite for JsonGetRewriter {
// json_get(column, path, <data_type>)
// )
fn inject_type_from_cast_expr(cast: Cast) -> Result<Transformed<Expr>> {
let Cast {
expr,
mut data_type,
} = cast;
let Cast { expr, field } = cast;
let mut data_type = field.data_type().clone();
let mut json_get = match *expr {
Expr::ScalarFunction(f)
@@ -73,7 +71,7 @@ fn inject_type_from_cast_expr(cast: Cast) -> Result<Transformed<Expr>> {
expr => {
return Ok(Transformed::no(Expr::Cast(Cast {
expr: Box::new(expr),
data_type,
field,
})));
}
};
@@ -204,10 +202,7 @@ mod tests {
});
// Create a cast expression: json_get(...)::int8
let cast_expr = Expr::Cast(Cast {
expr: Box::new(json_expr),
data_type: DataType::Int8,
});
let cast_expr = Expr::Cast(Cast::new(Box::new(json_expr), DataType::Int8));
// Apply the rewriter
let result = rewriter.rewrite(cast_expr, &schema, &config).unwrap();
@@ -279,10 +274,7 @@ mod tests {
// Create an arrow cast function: cast(json_get(...), 'Int64')
// Note: ArrowCastFunc doesn't exist in this codebase, so this test uses a simple cast instead
let arrow_cast_expr = Expr::Cast(Cast {
expr: Box::new(json_get_expr),
data_type: DataType::Int64,
});
let arrow_cast_expr = Expr::Cast(Cast::new(Box::new(json_get_expr), DataType::Int64));
// Apply the rewriter
let result = rewriter.rewrite(arrow_cast_expr, &schema, &config).unwrap();
src/common/function/src/scalars/matches.rs
+7 -2
View File
@@ -24,6 +24,7 @@ use datafusion::execution::SessionStateBuilder;
use datafusion::logical_expr::{self, ColumnarValue, Expr, Volatility};
use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner};
use datafusion_common::DataFusionError;
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion_expr::{ScalarFunctionArgs, Signature};
use datatypes::arrow::array::RecordBatch;
use datatypes::arrow::datatypes::{DataType, Field};
@@ -110,8 +111,12 @@ impl MatchesFunction {
let input_schema = Self::input_schema();
let session_state = SessionStateBuilder::new().with_default_features().build();
let planner = DefaultPhysicalPlanner::default();
let physical_expr =
planner.create_physical_expr(&like_expr, &input_schema, &session_state)?;
let physical_expr = planner.create_physical_expr(
&like_expr,
&input_schema,
&session_state,
&PhysicalPlanningContext::default(),
)?;
let arrow_schema = Arc::new(input_schema.as_arrow().clone());
let input_record_batch = RecordBatch::try_new(arrow_schema, vec![data_array]).unwrap();
src/common/function/src/scalars/math/clamp.rs
+77 -2
View File
@@ -19,10 +19,10 @@ use datafusion::arrow::array::{Array, ArrayRef, AsArray, PrimitiveArray};
use datafusion::arrow::datatypes::DataType as ArrowDataType;
use datafusion::logical_expr::{ColumnarValue, Volatility};
use datafusion_common::{DataFusionError, ScalarValue, utils};
use datafusion_expr::type_coercion::aggregates::NUMERICS;
use datafusion_expr::{ScalarFunctionArgs, Signature};
use crate::function::Function;
use crate::helper::NUMERICS;
#[derive(Clone, Debug)]
pub struct ClampFunction {
@@ -338,7 +338,9 @@ mod test {
use arrow_schema::Field;
use datafusion_common::config::ConfigOptions;
use datatypes::arrow::array::{ArrayRef, Float64Array, Int64Array, UInt64Array};
use datatypes::arrow::array::{
ArrayRef, Decimal128Array, Float64Array, Int64Array, UInt64Array,
};
use datatypes::arrow_array::StringArray;
use super::*;
@@ -370,6 +372,79 @@ mod test {
impl_test_eval!(ClampMinFunction);
impl_test_eval!(ClampMaxFunction);
fn decimal_array(values: Vec<i128>) -> ColumnarValue {
ColumnarValue::Array(Arc::new(
Decimal128Array::from(values)
.with_precision_and_scale(10, 2)
.unwrap(),
))
}
fn decimal_scalar(value: i128) -> ColumnarValue {
ColumnarValue::Scalar(ScalarValue::Decimal128(Some(value), 10, 2))
}
#[allow(deprecated)]
fn evaluate_decimal(
function: &dyn Function,
args: Vec<ColumnarValue>,
) -> datafusion_common::Result<ArrayRef> {
let input_types = args
.iter()
.map(ColumnarValue::data_type)
.collect::<Vec<_>>();
let planned_types = datafusion_expr::type_coercion::functions::data_types(
function.name(),
&input_types,
function.signature(),
)?;
let args = args
.into_iter()
.zip(planned_types)
.map(|(arg, planned_type)| arg.cast_to(&planned_type, None))
.collect::<datafusion_common::Result<Vec<_>>>()?;
function
.invoke_with_args(ScalarFunctionArgs {
args,
arg_fields: vec![],
number_rows: 3,
return_field: Arc::new(Field::new("x", ArrowDataType::Float64, false)),
config_options: Arc::new(ConfigOptions::new()),
})
.and_then(|value| value.to_array(3))
}
#[test]
fn clamp_decimal_coercion_executes_as_float64() {
let test_cases = [
(
Box::new(ClampFunction::default()) as Box<dyn Function>,
vec![
decimal_array(vec![100, 300, 500]),
decimal_scalar(200),
decimal_scalar(400),
],
vec![2.0, 3.0, 4.0],
),
(
Box::new(ClampMinFunction::default()),
vec![decimal_array(vec![100, 300, 500]), decimal_scalar(200)],
vec![2.0, 3.0, 5.0],
),
(
Box::new(ClampMaxFunction::default()),
vec![decimal_array(vec![100, 300, 500]), decimal_scalar(200)],
vec![1.0, 2.0, 2.0],
),
];
for (function, args, expected) in test_cases {
let result = evaluate_decimal(function.as_ref(), args).unwrap();
let expected: ArrayRef = Arc::new(Float64Array::from(expected));
assert_eq!(expected.as_ref(), result.as_ref());
}
}
#[test]
fn clamp_i64() {
let inputs = [
src/common/function/src/scalars/math/modulo.rs
+64 -2
View File
@@ -18,10 +18,10 @@ use std::fmt::Display;
use datafusion_common::arrow::compute;
use datafusion_common::arrow::compute::kernels::numeric;
use datafusion_common::arrow::datatypes::DataType;
use datafusion_expr::type_coercion::aggregates::NUMERICS;
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility};
use crate::function::{Function, extract_args};
use crate::helper::NUMERICS;
const NAME: &str = "mod";
@@ -90,12 +90,74 @@ mod tests {
use std::sync::Arc;
use arrow_schema::Field;
use datafusion_common::ScalarValue;
use datafusion_common::arrow::array::{
AsArray, Float64Array, Int32Array, StringViewArray, UInt32Array,
AsArray, Decimal128Array, Float64Array, Int32Array, StringViewArray, UInt32Array,
};
use datafusion_common::arrow::datatypes::{Float64Type, Int64Type, UInt64Type};
use super::*;
fn decimal_array(values: Vec<i128>) -> ColumnarValue {
ColumnarValue::Array(Arc::new(
Decimal128Array::from(values)
.with_precision_and_scale(10, 2)
.unwrap(),
))
}
fn decimal_scalar(value: i128) -> ColumnarValue {
ColumnarValue::Scalar(ScalarValue::Decimal128(Some(value), 10, 2))
}
#[test]
#[allow(deprecated)]
fn modulo_decimal_coercion_executes_as_float64() {
let function = ModuloFunction::default();
let test_cases = [
(
vec![decimal_array(vec![500, 600]), decimal_scalar(200)],
vec![1.0, 0.0],
),
(
vec![decimal_scalar(500), decimal_array(vec![200, 300])],
vec![1.0, 2.0],
),
];
for (args, expected) in test_cases {
let input_types = args
.iter()
.map(ColumnarValue::data_type)
.collect::<Vec<_>>();
let planned_types = datafusion_expr::type_coercion::functions::data_types(
function.name(),
&input_types,
function.signature(),
)
.unwrap();
assert_eq!(vec![DataType::Float64; 2], planned_types);
let args = args
.into_iter()
.zip(planned_types)
.map(|(arg, planned_type)| arg.cast_to(&planned_type, None))
.collect::<datafusion_common::Result<Vec<_>>>()
.unwrap();
let result = function
.invoke_with_args(ScalarFunctionArgs {
args,
arg_fields: vec![],
number_rows: 2,
return_field: Arc::new(Field::new("x", DataType::Float64, false)),
config_options: Arc::new(Default::default()),
})
.unwrap()
.to_array(2)
.unwrap();
let result = result.as_primitive::<Float64Type>();
assert_eq!(&Float64Array::from(expected), result);
}
}
#[test]
fn test_mod_function_signed() {
let function = ModuloFunction::default();
src/common/function/src/scalars/math/rate.rs
+8 -7
View File
@@ -18,11 +18,11 @@ use common_query::error;
use datafusion::arrow::compute::kernels::numeric;
use datafusion_common::arrow::compute::kernels::cast;
use datafusion_common::arrow::datatypes::DataType;
use datafusion_expr::type_coercion::aggregates::NUMERICS;
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility};
use snafu::ResultExt;
use crate::function::{Function, extract_args};
use crate::helper::NUMERICS;
/// generates rates from a sequence of adjacent data points.
#[derive(Clone, Debug)]
@@ -96,12 +96,13 @@ mod tests {
let rate = RateFunction::default();
assert_eq!("rate", rate.name());
assert_eq!(DataType::Float64, rate.return_type(&[]).unwrap());
assert!(matches!(rate.signature(),
Signature {
type_signature: TypeSignature::Uniform(2, valid_types),
volatility: Volatility::Immutable,
..
} if valid_types == NUMERICS
assert!(matches!(
rate.signature(),
Signature {
type_signature: TypeSignature::Uniform(2, valid_types),
volatility: Volatility::Immutable,
..
} if valid_types == NUMERICS
));
let values = vec![1.0, 3.0, 6.0];
let ts = vec![0, 1, 2];
src/common/function/src/scalars/udf.rs
-5
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::fmt::{Debug, Formatter};
use std::hash::{Hash, Hasher};
@@ -49,10 +48,6 @@ impl Hash for ScalarUdf {
}
impl ScalarUDFImpl for ScalarUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
self.function.name()
}
src/common/function/src/system.rs
+6 -1
View File
@@ -21,7 +21,8 @@ mod version;
use build::BuildFunction;
use database::{
ConnectionIdFunction, DatabaseFunction, PgBackendPidFunction, ReadPreferenceFunction,
ConnectionIdFunction, CurrentUserFunction, DatabaseFunction, PgBackendPidFunction,
ReadPreferenceFunction, SchemaFunction, SystemUserFunction, UserFunction,
};
use pg_catalog::PGCatalogFunction;
use procedure_state::ProcedureStateFunction;
@@ -37,6 +38,10 @@ impl SystemFunction {
registry.register_scalar(BuildFunction::default());
registry.register_scalar(VersionFunction::default());
registry.register_scalar(DatabaseFunction::default());
registry.register_scalar(SchemaFunction::default());
registry.register_scalar(UserFunction::default());
registry.register_scalar(CurrentUserFunction::default());
registry.register_scalar(SystemUserFunction::default());
registry.register_scalar(ReadPreferenceFunction::default());
registry.register_scalar(PgBackendPidFunction::default());
registry.register_scalar(ConnectionIdFunction::default());
src/common/function/src/system/database.rs
+70 -18
View File
@@ -20,39 +20,91 @@ use crate::function::{Function, find_function_context};
use crate::system::define_nullary_udf;
define_nullary_udf!(DatabaseFunction);
define_nullary_udf!(SchemaFunction);
define_nullary_udf!(UserFunction);
define_nullary_udf!(CurrentUserFunction);
define_nullary_udf!(SystemUserFunction);
define_nullary_udf!(ReadPreferenceFunction);
define_nullary_udf!(PgBackendPidFunction);
define_nullary_udf!(ConnectionIdFunction);
const DATABASE_FUNCTION_NAME: &str = "database";
const SCHEMA_FUNCTION_NAME: &str = "schema";
const USER_FUNCTION_NAME: &str = "user";
const CURRENT_USER_FUNCTION_NAME: &str = "current_user";
const SYSTEM_USER_FUNCTION_NAME: &str = "system_user";
const READ_PREFERENCE_FUNCTION_NAME: &str = "read_preference";
const PG_BACKEND_PID: &str = "pg_backend_pid";
const CONNECTION_ID: &str = "connection_id";
impl Function for DatabaseFunction {
fn name(&self) -> &str {
DATABASE_FUNCTION_NAME
}
macro_rules! impl_current_schema_function {
($name: ident, $fn_name: expr) => {
impl Function for $name {
fn name(&self) -> &str {
$fn_name
}
fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
Ok(DataType::Utf8View)
}
fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
Ok(DataType::Utf8View)
}
fn signature(&self) -> &Signature {
&self.signature
}
fn signature(&self) -> &Signature {
&self.signature
}
fn invoke_with_args(
&self,
args: ScalarFunctionArgs,
) -> datafusion_common::Result<ColumnarValue> {
let func_ctx = find_function_context(&args)?;
let db = func_ctx.query_ctx.current_schema();
fn invoke_with_args(
&self,
args: ScalarFunctionArgs,
) -> datafusion_common::Result<ColumnarValue> {
let func_ctx = find_function_context(&args)?;
let db = func_ctx.query_ctx.current_schema();
Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(db))))
}
Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(db))))
}
}
};
}
impl_current_schema_function!(DatabaseFunction, DATABASE_FUNCTION_NAME);
// MySQL's `SCHEMA()` is a synonym for `DATABASE()`.
impl_current_schema_function!(SchemaFunction, SCHEMA_FUNCTION_NAME);
macro_rules! impl_current_user_function {
($name: ident, $fn_name: expr) => {
impl Function for $name {
fn name(&self) -> &str {
$fn_name
}
fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
Ok(DataType::Utf8View)
}
fn signature(&self) -> &Signature {
&self.signature
}
fn invoke_with_args(
&self,
args: ScalarFunctionArgs,
) -> datafusion_common::Result<ColumnarValue> {
let func_ctx = find_function_context(&args)?;
let user = func_ctx.query_ctx.current_user();
Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
user.username().to_string(),
))))
}
}
};
}
// GreptimeDB has no notion of a user switching identity mid-session, so `USER()`,
// `CURRENT_USER()` and `SYSTEM_USER()` all report the authenticated user.
impl_current_user_function!(UserFunction, USER_FUNCTION_NAME);
impl_current_user_function!(CurrentUserFunction, CURRENT_USER_FUNCTION_NAME);
impl_current_user_function!(SystemUserFunction, SYSTEM_USER_FUNCTION_NAME);
impl Function for ReadPreferenceFunction {
fn name(&self) -> &str {
READ_PREFERENCE_FUNCTION_NAME
src/common/function/src/system/pg_catalog.rs
+3
View File
@@ -389,6 +389,9 @@ impl PGCatalogFunction {
registry.register(pg_catalog::create_pg_get_partition_ancestors_udf());
registry.register(pg_catalog::quote_ident_udf::create_quote_ident_udf());
registry.register(pg_catalog::quote_ident_udf::create_parse_ident_udf());
// Register array bound UDFs used by pg_catalog views.
registry.register(pg_catalog::array_bounds_udf::create_array_upper_udf());
registry.register(pg_catalog::array_bounds_udf::create_array_lower_udf());
registry.register_scalar(ObjDescriptionFunction::new());
registry.register_scalar(ColDescriptionFunction::new());
registry.register_scalar(ShobjDescriptionFunction::new());
src/common/macro/src/admin_fn.rs
-4
View File
@@ -239,10 +239,6 @@ fn build_struct(
// Implement DataFusion's ScalarUDFImpl trait
impl datafusion::logical_expr::ScalarUDFImpl for #name {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
#display_name
}
src/common/query/src/request.rs
+8 -10
View File
@@ -137,9 +137,7 @@ fn portable_remote_dyn_filter_expr(
bounds_only: bool,
) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
expr.transform_up(|node| {
if node.as_any().is::<HashTableLookupExpr>()
|| (bounds_only && node.as_any().is::<InListExpr>())
{
if node.is::<HashTableLookupExpr>() || (bounds_only && node.is::<InListExpr>()) {
Ok(Transformed::yes(lit(true)))
} else {
Ok(Transformed::no(node))
@@ -183,7 +181,7 @@ fn validate_payload_size(
fn validate_supported_payload_expr(expr: &Arc<dyn PhysicalExpr>) -> DataFusionResult<()> {
expr.apply(|node| {
if node.as_any().is::<HashTableLookupExpr>() {
if node.is::<HashTableLookupExpr>() {
return Err(DataFusionError::Plan(
"HashTableLookupExpr cannot be encoded into DynFilterPayload::Datafusion"
.to_string(),
@@ -207,7 +205,7 @@ fn validate_decoded_payload_expr(
input_schema: &datafusion::arrow::datatypes::Schema,
) -> DataFusionResult<()> {
expr.apply(|node| {
if let Some(column) = node.as_any().downcast_ref::<Column>() {
if let Some(column) = node.downcast_ref::<Column>() {
let Some(field) = input_schema.fields().get(column.index()) else {
return Err(DataFusionError::Plan(format!(
"Decoded Column '{}' references out-of-bounds index {} for input schema of size {}",
@@ -391,8 +389,8 @@ mod tests {
.decode_datafusion_expr(&TaskContext::default(), &schema, 1024)
.unwrap();
let original = expr.as_any().downcast_ref::<Column>().unwrap();
let decoded = decoded.as_any().downcast_ref::<Column>().unwrap();
let original = expr.downcast_ref::<Column>().unwrap();
let decoded = decoded.downcast_ref::<Column>().unwrap();
assert_eq!(decoded.name(), original.name());
assert_eq!(decoded.index(), original.index());
@@ -457,7 +455,7 @@ mod tests {
)) as Arc<dyn PhysicalExpr>;
let lookup = Arc::new(HashTableLookupExpr::new(
vec![Arc::clone(&device_id)],
SeededRandomState::with_seeds(0, 0, 0, 0),
SeededRandomState::with_seed(0),
Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(0)))),
"hash_lookup".to_string(),
)) as Arc<dyn PhysicalExpr>;
@@ -539,10 +537,10 @@ mod tests {
));
}
fn contains_expr<T: 'static>(expr: &Arc<dyn PhysicalExpr>) -> bool {
fn contains_expr<T: PhysicalExpr>(expr: &Arc<dyn PhysicalExpr>) -> bool {
let mut found = false;
expr.apply(|node| {
if node.as_any().is::<T>() {
if node.is::<T>() {
found = true;
Ok(TreeNodeRecursion::Stop)
} else {
src/common/query/src/request/initial_remote_dyn_filter_reg.rs
+1 -1
View File
@@ -371,7 +371,7 @@ mod tests {
let decoded = reg
.decode_children(&TaskContext::default(), &schema, 1024)
.unwrap();
let decoded = decoded[0].as_any().downcast_ref::<Column>().unwrap();
let decoded = decoded[0].downcast_ref::<Column>().unwrap();
assert_eq!(reg.filter_id, "filter-1");
assert_eq!(decoded.name(), "host");
src/common/query/src/stream.rs
+11 -6
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::fmt::{Debug, Formatter};
use std::sync::{Arc, Mutex};
@@ -22,8 +21,11 @@ use datafusion::execution::SendableRecordBatchStream as DfSendableRecordBatchStr
use datafusion::execution::context::TaskContext;
use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalSortExpr};
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, PlanProperties,
};
use datafusion_common::DataFusionError;
use datafusion_common::tree_node::TreeNodeRecursion;
use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef;
use datatypes::schema::SchemaRef;
@@ -99,10 +101,6 @@ impl DisplayAs for StreamScanAdapter {
}
impl ExecutionPlan for StreamScanAdapter {
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> ArrowSchemaRef {
self.arrow_schema.clone()
}
@@ -115,6 +113,13 @@ impl ExecutionPlan for StreamScanAdapter {
vec![]
}
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> datafusion_common::Result<TreeNodeRecursion>,
) -> datafusion_common::Result<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
// DataFusion will swap children unconditionally.
// But since this node is leaf node, it's safe to just return self.
fn with_new_children(
src/common/recordbatch/src/adapter.rs
+18 -7
View File
@@ -29,6 +29,7 @@ use datafusion::arrow::datatypes::SchemaRef as DfSchemaRef;
use datafusion::error::Result as DfResult;
use datafusion::execution::context::ExecutionProps;
use datafusion::logical_expr::Expr;
use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion::logical_expr::utils::conjunction;
use datafusion::physical_expr::create_physical_expr;
use datafusion::physical_plan::metrics::{BaselineMetrics, MetricValue};
@@ -98,8 +99,13 @@ where
.to_dfschema_ref()
.context(error::PhysicalExprSnafu)?;
let filters = create_physical_expr(&expr, &df_schema, &ExecutionProps::new())
.context(error::PhysicalExprSnafu)?;
let filters = create_physical_expr(
&expr,
&df_schema,
&ExecutionProps::new(),
&PhysicalPlanningContext::default(),
)
.context(error::PhysicalExprSnafu)?;
Some(filters)
} else {
None
@@ -931,7 +937,6 @@ fn convert_map_to_json_binary(
#[cfg(test)]
mod test {
use std::any::Any;
use std::time::Duration;
use common_error::ext::BoxedError;
@@ -942,6 +947,7 @@ mod test {
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet};
use datafusion::physical_plan::{DisplayAs, PlanProperties};
use datafusion_common::tree_node::TreeNodeRecursion;
use datatypes::arrow::array::{ArrayRef, MapArray, StringArray, StructArray};
use datatypes::arrow::buffer::OffsetBuffer;
use datatypes::arrow::datatypes::Field;
@@ -1021,10 +1027,6 @@ mod test {
REGION_SCAN_EXEC_NAME
}
fn as_any(&self) -> &dyn Any {
self
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.properties
}
@@ -1033,6 +1035,15 @@ mod test {
vec![]
}
fn apply_expressions(
&self,
_f: &mut dyn FnMut(
&Arc<dyn PhysicalExpr>,
) -> datafusion_common::Result<TreeNodeRecursion>,
) -> datafusion_common::Result<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
fn with_new_children(
self: Arc<Self>,
_children: Vec<Arc<dyn ExecutionPlan>>,
src/common/recordbatch/src/filter.rs
+15 -2
View File
@@ -616,6 +616,7 @@ mod test {
use std::sync::Arc;
use datafusion::execution::context::ExecutionProps;
use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion::logical_expr::{BinaryExpr, col, lit};
use datafusion::physical_expr::create_physical_expr;
use datafusion_common::{Column, DFSchema};
@@ -737,7 +738,13 @@ mod test {
]);
let df_schema = DFSchema::try_from(schema.clone()).unwrap();
let props = ExecutionProps::new();
let physical_expr = create_physical_expr(&expr, &df_schema, &props).unwrap();
let physical_expr = create_physical_expr(
&expr,
&df_schema,
&props,
&PhysicalPlanningContext::default(),
)
.unwrap();
let batch = RecordBatch::try_new(
Arc::new(schema),
vec![
@@ -783,7 +790,13 @@ mod test {
let schema = Schema::new(vec![Field::new("col", DataType::Utf8, false)]);
let df_schema = DFSchema::try_from(schema.clone()).unwrap();
let props = ExecutionProps::new();
let physical_expr = create_physical_expr(&col_or_expr, &df_schema, &props).unwrap();
let physical_expr = create_physical_expr(
&col_or_expr,
&df_schema,
&props,
&PhysicalPlanningContext::default(),
)
.unwrap();
// Create test data
let col_data = Arc::new(datatypes::arrow::array::StringArray::from(vec![
src/common/recordbatch/src/lib.rs
+44 -6
View File
@@ -44,7 +44,9 @@ pub use datatypes::arrow::record_batch::RecordBatch as DfRecordBatch;
use datatypes::arrow::util::display::{
ArrayFormatter, ArrayFormatterFactory, DisplayIndex, FormatOptions, FormatResult,
};
use datatypes::arrow::util::pretty::pretty_format_batches_with_options;
use datatypes::arrow::util::pretty::{
pretty_format_batches_with_options, pretty_format_batches_with_schema,
};
use datatypes::extension::json::is_any_json_extension_type;
use datatypes::prelude::{ConcreteDataType, DataType, VectorRef};
use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
@@ -396,12 +398,19 @@ impl RecordBatches {
.iter()
.map(|x| x.df_record_batch().clone())
.collect::<Vec<_>>();
let options =
FormatOptions::default().with_formatter_factory(Some(&BinaryFormatterFactory));
let result =
pretty_format_batches_with_options(df_batches, &options).context(error::FormatSnafu)?;
let result: String = if df_batches.is_empty() {
pretty_format_batches_with_schema(self.schema.arrow_schema().clone(), df_batches)
.context(error::FormatSnafu)?
.to_string()
} else {
let options =
FormatOptions::default().with_formatter_factory(Some(&BinaryFormatterFactory));
pretty_format_batches_with_options(df_batches, &options)
.context(error::FormatSnafu)?
.to_string()
};
Ok(result.to_string())
Ok(result)
}
pub fn try_new(schema: SchemaRef, batches: Vec<RecordBatch>) -> Result<Self> {
@@ -1082,6 +1091,35 @@ mod tests {
assert_eq!(r.take(), expected);
}
#[tokio::test]
async fn test_recordbatches_pretty_print_empty_batches_preserves_schema() {
let schema = Arc::new(Schema::new(vec![
ColumnSchema::new("unit", ConcreteDataType::string_datatype(), false),
ColumnSchema::new(
"ts",
ConcreteDataType::timestamp_millisecond_datatype(),
false,
),
ColumnSchema::new(
"lhs.degrees(val) + rhs.radians(val)",
ConcreteDataType::float64_datatype(),
false,
),
]));
let batches =
RecordBatches::try_collect(Box::pin(EmptyRecordBatchStream::new(schema.clone())))
.await
.unwrap();
assert_eq!(schema, batches.schema());
let expected = "\
+------+----+-------------------------------------+
| unit | ts | lhs.degrees(val) + rhs.radians(val) |
+------+----+-------------------------------------+
+------+----+-------------------------------------+";
assert_eq!(expected, batches.pretty_print().unwrap());
}
#[test]
fn test_recordbatches_try_new() {
let column_a = ColumnSchema::new("a", ConcreteDataType::int32_datatype(), false);
src/datanode/src/region_server/catalog.rs
-9
View File
@@ -179,9 +179,6 @@ impl NameAwareCatalogList {
}
impl CatalogProviderList for NameAwareCatalogList {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn register_catalog(
&self,
_name: String,
@@ -203,9 +200,6 @@ struct NameAwareCatalogProvider {
}
impl CatalogProvider for NameAwareCatalogProvider {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn schema_names(&self) -> Vec<String> {
vec![]
}
@@ -229,9 +223,6 @@ impl std::fmt::Debug for NameAwareSchemaProvider {
#[async_trait::async_trait]
impl SchemaProvider for NameAwareSchemaProvider {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn table_names(&self) -> Vec<String> {
vec![]
}
src/datatypes/src/value.rs
+2
View File
@@ -1240,6 +1240,8 @@ impl TryFrom<ScalarValue> for Value {
| ScalarValue::Decimal256(_, _, _)
| ScalarValue::FixedSizeList(_)
| ScalarValue::LargeList(_)
| ScalarValue::ListView(_)
| ScalarValue::LargeListView(_)
| ScalarValue::Union(_, _, _)
| ScalarValue::Float16(_)
| ScalarValue::Utf8View(_)
src/datatypes/src/vectors/helper.rs
+2
View File
@@ -150,6 +150,8 @@ impl Helper {
| ScalarValue::Decimal256(_, _, _)
| ScalarValue::FixedSizeList(_)
| ScalarValue::LargeList(_)
| ScalarValue::ListView(_)
| ScalarValue::LargeListView(_)
| ScalarValue::Dictionary(_, _)
| ScalarValue::Union(_, _, _)
| ScalarValue::Utf8View(_)
src/file-engine/src/query/file_stream.rs
+16 -10
View File
@@ -23,7 +23,8 @@ use datafusion::config::CsvOptions;
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::object_store::ObjectStoreUrl;
use datafusion::datasource::physical_plan::{
CsvSource, FileGroup, FileScanConfigBuilder, FileSource, FileStream, JsonSource, ParquetSource,
CsvSource, FileGroup, FileScanConfigBuilder, FileSource, FileStreamBuilder, JsonSource,
ParquetSource,
};
use datafusion::datasource::source::DataSourceExec;
use datafusion::physical_expr::create_physical_expr;
@@ -34,6 +35,7 @@ use datafusion::physical_plan::{
};
use datafusion::prelude::SessionContext;
use datafusion_expr::expr::Expr;
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion_expr::utils::conjunction;
use datatypes::schema::SchemaRef;
use object_store::ObjectStore;
@@ -66,13 +68,12 @@ fn build_record_batch_stream(
));
let file_opener = config.file_source().create_file_opener(store, &config, 0)?;
let stream = FileStream::new(
&config,
0, // partition: hard-code
file_opener,
&ExecutionPlanMetricsSet::new(),
)
.context(error::BuildStreamSnafu)?;
let stream = FileStreamBuilder::new(&config)
.with_partition(0) // partition: hard-code
.with_file_opener(file_opener)
.with_metrics(&ExecutionPlanMetricsSet::new())
.build()
.context(error::BuildStreamSnafu)?;
Ok(Box::pin(stream))
}
@@ -138,8 +139,13 @@ fn new_parquet_stream_with_exec_plan(
.to_dfschema_ref()
.context(error::ParquetScanPlanSnafu)?;
let filters = create_physical_expr(&expr, &df_schema, &ExecutionProps::new())
.context(error::ParquetScanPlanSnafu)?;
let filters = create_physical_expr(
&expr,
&df_schema,
&ExecutionProps::new(),
&PhysicalPlanningContext::default(),
)
.context(error::ParquetScanPlanSnafu)?;
parquet_source = parquet_source.with_predicate(filters);
};
src/flow/src/batching_mode/time_window.rs
+11 -2
View File
@@ -37,6 +37,7 @@ use datafusion_common::tree_node::{
Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter, TreeNodeVisitor,
};
use datafusion_common::{DFSchema, TableReference};
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion_expr::{ColumnarValue, LogicalPlan};
use datafusion_physical_expr::PhysicalExprRef;
use datatypes::prelude::{ConcreteDataType, DataType};
@@ -830,7 +831,15 @@ fn to_phy_expr(
let phy_planner = DefaultPhysicalPlanner::default();
let phy_expr: PhysicalExprRef = phy_planner
.create_physical_expr(expr, df_schema, session)
// Time-window expressions are standalone scalar expressions over the input
// time column, so they cannot contain scalar subqueries or lambda variables
// that would require a plan-level physical planning context.
.create_physical_expr(
expr,
df_schema,
session,
&PhysicalPlanningContext::default(),
)
.with_context(|_e| DatafusionSnafu {
context: format!(
"Failed to create physical expression from {expr:?} using {df_schema:?}"
@@ -993,7 +1002,7 @@ mod test {
Some(Timestamp::new(0, TimeUnit::Millisecond)),
Some(Timestamp::new(300000, TimeUnit::Millisecond)),
),
"SELECT sum(numbers_with_ts.number), numbers_with_ts.number, date_bin('5 minutes', numbers_with_ts.ts) AS time_window, bucket_name FROM (SELECT numbers_with_ts.number, numbers_with_ts.ts, CASE WHEN (numbers_with_ts.number < 5) THEN 'bucket_0_5' WHEN (numbers_with_ts.number >= 5) THEN 'bucket_5_inf' END AS bucket_name FROM numbers_with_ts WHERE ((ts >= CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AND (ts <= CAST('1970-01-01 00:05:00' AS TIMESTAMP)))) GROUP BY numbers_with_ts.number, date_bin('5 minutes', numbers_with_ts.ts), bucket_name",
"SELECT sum(number), number, date_bin('5 minutes', ts) AS time_window, bucket_name FROM (SELECT numbers_with_ts.number, numbers_with_ts.ts, CASE WHEN (numbers_with_ts.number < 5) THEN 'bucket_0_5' WHEN (numbers_with_ts.number >= 5) THEN 'bucket_5_inf' END AS bucket_name FROM numbers_with_ts WHERE ((ts >= CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AND (ts <= CAST('1970-01-01 00:05:00' AS TIMESTAMP)))) GROUP BY number, date_bin('5 minutes', ts), bucket_name",
),
// complex subquery alias
(
src/flow/src/batching_mode/utils.rs
+3 -5
View File
@@ -31,7 +31,7 @@ use datafusion_common::tree_node::{
use datafusion_common::{
Column, DFSchema, DataFusionError, NullEquality, ScalarValue, TableReference,
};
use datafusion_expr::logical_plan::{Aggregate, TableScan};
use datafusion_expr::logical_plan::{Aggregate, TableScanBuilder};
use datafusion_expr::{
Distinct, ExprSchemable, JoinType, LogicalPlan, LogicalPlanBuilder, Operator, Projection, and,
binary_expr, bitwise_and, bitwise_or, bitwise_xor, is_null, or, when,
@@ -698,17 +698,15 @@ pub async fn rewrite_incremental_aggregate_with_sink_merge(
let table_provider = Arc::new(DfTableProviderAdapter::new(sink_table));
let table_source = Arc::new(DefaultTableSource::new(table_provider));
let sink_scan = LogicalPlan::TableScan(
TableScan::try_new(
TableScanBuilder::new(
TableReference::Full {
catalog: sink_table_name[0].clone().into(),
schema: sink_table_name[1].clone().into(),
table: sink_table_name[2].clone().into(),
},
table_source,
None,
vec![],
None,
)
.build()
.with_context(|_| DatafusionSnafu {
context: "Failed to build sink table scan for incremental sink merge".to_string(),
})?,
src/flow/src/batching_mode/utils/test.rs
+4 -6
View File
@@ -20,7 +20,7 @@ use common_recordbatch::recordbatch::merge_record_batches;
use common_recordbatch::{RecordBatch, util};
use common_time::Timestamp;
use datafusion_common::tree_node::TreeNode as _;
use datafusion_expr::GroupingSet;
use datafusion_expr::{GroupingSet, TableScanBuilder};
use datatypes::arrow::array::{Array, AsArray};
use datatypes::arrow::datatypes::{Float64Type, Int64Type, UInt64Type};
use datatypes::prelude::{ConcreteDataType, MutableVector, Scalar, ScalarVectorBuilder, VectorRef};
@@ -96,17 +96,15 @@ fn test_sink_scan(sink_table: TableRef, sink_table_name: &TableName) -> LogicalP
let table_provider = Arc::new(DfTableProviderAdapter::new(sink_table));
let table_source = Arc::new(DefaultTableSource::new(table_provider));
LogicalPlan::TableScan(
TableScan::try_new(
TableScanBuilder::new(
TableReference::Full {
catalog: sink_table_name[0].clone().into(),
schema: sink_table_name[1].clone().into(),
table: sink_table_name[2].clone().into(),
},
table_source,
None,
vec![],
None,
)
.build()
.unwrap(),
)
}
@@ -300,7 +298,7 @@ async fn test_add_filter() {
// complex subquery without alias
(
"SELECT sum(number), number, date_bin('5 minutes', ts) as time_window, bucket_name FROM (SELECT number, ts, case when number < 5 THEN 'bucket_0_5' when number >= 5 THEN 'bucket_5_inf' END as bucket_name FROM numbers_with_ts) GROUP BY number, time_window, bucket_name;",
"SELECT sum(numbers_with_ts.number), numbers_with_ts.number, date_bin('5 minutes', numbers_with_ts.ts) AS time_window, bucket_name FROM (SELECT numbers_with_ts.number, numbers_with_ts.ts, CASE WHEN (numbers_with_ts.number < 5) THEN 'bucket_0_5' WHEN (numbers_with_ts.number >= 5) THEN 'bucket_5_inf' END AS bucket_name FROM numbers_with_ts WHERE (number > 4)) GROUP BY numbers_with_ts.number, date_bin('5 minutes', numbers_with_ts.ts), bucket_name",
"SELECT sum(number), number, date_bin('5 minutes', ts) AS time_window, bucket_name FROM (SELECT numbers_with_ts.number, numbers_with_ts.ts, CASE WHEN (numbers_with_ts.number < 5) THEN 'bucket_0_5' WHEN (numbers_with_ts.number >= 5) THEN 'bucket_5_inf' END AS bucket_name FROM numbers_with_ts WHERE (number > 4)) GROUP BY number, date_bin('5 minutes', ts), bucket_name",
),
// complex subquery alias
(
src/flow/src/transform/aggr.rs
+1 -1
View File
@@ -725,7 +725,7 @@ mod test {
df_scalar_fn: DfScalarFunction::try_from_raw_fn(
RawDfScalarFn {
f: BytesMut::from(
b"\x08\x02\"\x0f\x1a\r\n\x0b\xa2\x02\x08\n\0\x12\x04\x10\x1e \t\"\n\x1a\x08\x12\x06\n\x04\x12\x02\x08\x01".as_ref(),
b"\x08\x02\x1a\x07\x8a\x02\x04\x08\x03\x18\x01\"\x0f\x1a\r\n\x0b\xa2\x02\x08\n\0\x12\x04\x10\x1e \t\"\n\x1a\x08\x12\x06\n\x04\x12\x02\x08\x01".as_ref(),
),
input_schema: RelationType::new(vec![ColumnType::new(
ConcreteDataType::interval_month_day_nano_datatype(),
src/flow/src/transform/expr.rs
+10 -5
View File
@@ -20,6 +20,7 @@ use common_error::ext::BoxedError;
use common_telemetry::debug;
use datafusion::execution::SessionStateBuilder;
use datafusion::functions::all_default_functions;
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion_physical_expr::PhysicalExpr;
use datafusion_substrait::logical_plan::consumer::DefaultSubstraitConsumer;
use datatypes::data_type::ConcreteDataType as CDT;
@@ -101,11 +102,15 @@ pub(crate) async fn from_scalar_fn_to_df_fn_impl(
context: "Failed to convert substrait scalar function to datafusion scalar function",
}
})?;
let phy_expr =
datafusion::physical_expr::create_physical_expr(&expr, &schema, &Default::default())
.context(DatafusionSnafu {
context: "Failed to create physical expression from logical expression",
})?;
let phy_expr = datafusion::physical_expr::create_physical_expr(
&expr,
&schema,
&Default::default(),
&PhysicalPlanningContext::default(),
)
.context(DatafusionSnafu {
context: "Failed to create physical expression from logical expression",
})?;
Ok(phy_expr)
}
src/frontend/src/instance/dashboard.rs
+1 -1
View File
@@ -29,9 +29,9 @@ use common_query::OutputData;
use common_recordbatch::util as record_util;
use common_telemetry::info;
use common_time::FOREVER;
use datafusion::common::TableReference;
use datafusion::datasource::DefaultTableSource;
use datafusion::logical_expr::col;
use datafusion::sql::TableReference;
use datafusion_expr::{DmlStatement, LogicalPlan, lit};
use datatypes::arrow::array::{Array, AsArray};
use servers::error::{
src/meta-srv/src/procedure/region_migration/manager.rs
+152 -18
View File
@@ -510,17 +510,14 @@ impl RegionMigrationManager {
ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context);
let procedure_id = procedure_with_id.id;
info!("Starting region migration procedure {procedure_id} for {task}");
let procedure_manager = self.procedure_manager.clone();
let mut watcher = self
.procedure_manager
.submit(procedure_with_id)
.await
.context(error::SubmitProcedureSnafu)?;
let num_region = task.region_ids.len();
common_runtime::spawn_global(async move {
let watcher = &mut match procedure_manager.submit(procedure_with_id).await {
Ok(watcher) => watcher,
Err(e) => {
error!(e; "Failed to submit region migration procedure {procedure_id} for {task}");
return;
}
};
METRIC_META_REGION_MIGRATION_DATANODES
.with_label_values(&["src", &task.from_peer.id.to_string()])
.inc_by(num_region as u64);
@@ -528,7 +525,7 @@ impl RegionMigrationManager {
.with_label_values(&["desc", &task.to_peer.id.to_string()])
.inc_by(num_region as u64);
if let Err(e) = watcher::wait(watcher).await {
if let Err(e) = watcher::wait(&mut watcher).await {
error!(e; "Failed to wait region migration procedure {procedure_id} for {task}");
METRIC_META_REGION_MIGRATION_FAIL.inc();
return;
@@ -613,15 +610,12 @@ impl RegionMigrationManager {
ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context);
let procedure_id = procedure_with_id.id;
info!("Starting region migration procedure {procedure_id} for {task}");
let procedure_manager = self.procedure_manager.clone();
let mut watcher = self
.procedure_manager
.submit(procedure_with_id)
.await
.context(error::SubmitProcedureSnafu)?;
common_runtime::spawn_global(async move {
let watcher = &mut match procedure_manager.submit(procedure_with_id).await {
Ok(watcher) => watcher,
Err(e) => {
error!(e; "Failed to submit region migration procedure {procedure_id} for {task}");
return;
}
};
METRIC_META_REGION_MIGRATION_DATANODES
.with_label_values(&["src", &task.from_peer.id.to_string()])
.inc();
@@ -629,7 +623,7 @@ impl RegionMigrationManager {
.with_label_values(&["desc", &task.to_peer.id.to_string()])
.inc();
if let Err(e) = watcher::wait(watcher).await {
if let Err(e) = watcher::wait(&mut watcher).await {
error!(e; "Failed to wait region migration procedure {procedure_id} for {task}");
METRIC_META_REGION_MIGRATION_FAIL.inc();
return;
@@ -649,10 +643,148 @@ mod test {
use common_meta::key::table_route::LogicalTableRouteValue;
use common_meta::key::test_utils::new_test_table_info;
use common_meta::rpc::router::Region;
use common_meta::state_store::KvStateStore;
use common_procedure::ProcedureManager;
use common_procedure::local::{LocalManager, ManagerConfig, PauseAware};
use common_procedure::test_util::InMemoryPoisonStore;
use tokio::sync::Notify;
use super::*;
use crate::procedure::region_migration::test_util::TestingEnv;
#[derive(Default)]
struct SubmissionGate {
entered: Notify,
release: Notify,
}
#[async_trait::async_trait]
impl PauseAware for SubmissionGate {
async fn is_paused(&self) -> std::result::Result<bool, common_error::ext::BoxedError> {
self.entered.notify_one();
self.release.notified().await;
Ok(false)
}
}
async fn check_submission_registration(batch: bool) {
let env = TestingEnv::new();
let gate = Arc::new(SubmissionGate::default());
let procedure_manager = Arc::new(LocalManager::new(
ManagerConfig::default(),
Arc::new(KvStateStore::new(env.kv_backend())),
Arc::new(InMemoryPoisonStore::default()),
Some(gate.clone()),
None,
));
let manager = RegionMigrationManager::new(procedure_manager.clone(), env.context_factory());
let region_ids = if batch {
vec![RegionId::new(1024, 1), RegionId::new(1024, 2)]
} else {
vec![RegionId::new(1024, 1)]
};
let region_routes = region_ids
.iter()
.map(|region_id| RegionRoute {
region: Region::new_test(*region_id),
leader_peer: Some(Peer::empty(1)),
..Default::default()
})
.collect();
env.create_physical_table_metadata(new_test_table_info(1024), region_routes)
.await;
// A stopped manager must reject submission and release all migration guards.
// Starting it then allows the same regions to be submitted successfully.
for started in [false, true] {
if started {
procedure_manager.start().await.unwrap();
}
let submission = async {
if batch {
let result = manager
.submit_region_migration_task(RegionMigrationTaskBatch {
region_ids: region_ids.clone(),
from_peer: Peer::empty(1),
to_peer: Peer::empty(2),
timeout: Duration::from_secs(10),
trigger_reason: RegionMigrationTriggerReason::Manual,
})
.await?;
assert_eq!(result.submitted, region_ids);
Ok(result.procedure_id)
} else {
manager
.submit_procedure(
ProcedureContext::default(),
RegionMigrationProcedureTask::new(
region_ids[0],
Peer::empty(1),
Peer::empty(2),
Duration::from_secs(10),
RegionMigrationTriggerReason::Manual,
),
)
.await
}
};
tokio::pin!(submission);
tokio::select! {
biased;
result = &mut submission => panic!("Submission returned before registration: {result:?}"),
_ = gate.entered.notified() => {}
}
assert!(futures::poll!(&mut submission).is_pending());
assert!(
procedure_manager
.list_procedures()
.await
.unwrap()
.is_empty()
);
for region_id in &region_ids {
assert!(manager.tracker.contains(*region_id));
}
gate.release.notify_one();
let result = tokio::time::timeout(Duration::from_secs(5), &mut submission)
.await
.expect("Submission should return after registration");
if started {
let procedure_id = result.unwrap().unwrap();
assert!(
procedure_manager
.procedure_state(procedure_id)
.await
.unwrap()
.is_some()
);
} else {
assert_matches!(
result.unwrap_err(),
error::Error::SubmitProcedure {
source: common_procedure::Error::ManagerNotStart { .. },
..
}
);
for region_id in &region_ids {
assert!(!manager.tracker.contains(*region_id));
}
}
}
procedure_manager.stop().await.unwrap();
}
#[tokio::test]
async fn test_submit_procedure_registration() {
check_submission_registration(false).await;
}
#[tokio::test]
async fn test_submit_region_migration_task_registration() {
check_submission_registration(true).await;
}
#[tokio::test]
async fn test_insert_running_procedure() {
let env = TestingEnv::new();
@@ -1039,6 +1171,7 @@ mod test {
#[tokio::test]
async fn test_running_regions() {
let env = TestingEnv::new();
env.procedure_manager().start().await.unwrap();
let context_factory = env.context_factory();
let manager = RegionMigrationManager::new(env.procedure_manager().clone(), context_factory);
let region_id = RegionId::new(1024, 1);
@@ -1072,5 +1205,6 @@ mod test {
assert_eq!(result.migrating, vec![region_id]);
assert_eq!(result.submitted, vec![RegionId::new(1024, 2)]);
assert!(result.procedure_id.is_some());
env.procedure_manager().stop().await.unwrap();
}
}
src/metric-engine/src/engine/put.rs
+1 -1
View File
@@ -576,7 +576,7 @@ impl MetricEngineInner {
region_id: data_region_id,
})?;
(
physical_state.physical_columns().clone(),
physical_state.physical_columns_snapshot(),
physical_state.time_index_column_name().to_string(),
)
};
src/metric-engine/src/engine/state.rs
+13 -3
View File
@@ -15,6 +15,7 @@
//! Internal states of metric engine
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use api::v1::SemanticType;
use common_time::timestamp::TimeUnit;
@@ -30,7 +31,10 @@ use crate::utils::to_data_region_id;
pub struct PhysicalRegionState {
logical_regions: HashSet<RegionId>,
physical_columns: HashMap<String, ColumnMetadata>,
/// Columns of the physical region, wrapped in an [`Arc`] so that hot read
/// paths (e.g. write request verification) can hold a cheap snapshot
/// instead of deep-cloning the whole map on every row batch.
physical_columns: Arc<HashMap<String, ColumnMetadata>>,
/// Name of the time index column, cached at region load so that the write
/// path doesn't have to scan `physical_columns` for the timestamp on every
/// row batch. The time index is fixed at region creation and never
@@ -58,7 +62,7 @@ impl PhysicalRegionState {
.unwrap_or_default();
Self {
logical_regions: HashSet::new(),
physical_columns,
physical_columns: Arc::new(physical_columns),
time_index_column_name,
primary_key_encoding,
options,
@@ -76,6 +80,12 @@ impl PhysicalRegionState {
&self.physical_columns
}
/// Returns a cheap snapshot of the physical columns that stays valid
/// after releasing the state lock.
pub fn physical_columns_snapshot(&self) -> Arc<HashMap<String, ColumnMetadata>> {
self.physical_columns.clone()
}
/// Returns the cached name of the time index column.
pub fn time_index_column_name(&self) -> &str {
&self.time_index_column_name
@@ -144,7 +154,7 @@ impl MetricEngineState {
SemanticType::Timestamp,
"unexpected time index column {col} added to an existing physical region"
);
state.physical_columns.insert(col, meta);
Arc::make_mut(&mut state.physical_columns).insert(col, meta);
}
}
src/mito2/AGENTS.md
+1 -1
View File
@@ -26,7 +26,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`.
| `compaction` | `src/mito2/src/compaction/` | Compaction scheduler (`scheduler.rs` + `scheduler/`), TWCS picker, strict-window manual picker, compactor, memory control |
| `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store |
| `sst` | `src/mito2/src/sst/` | Parquet format, file metadata, index layout |
| `series_index` | `src/mito2/src/series_index/` | Series index writer/searcher, bucket planning and SST builders, catalogs, immutable snapshots, file lifecycle, and background maintenance |
| `series_index` | `src/mito2/src/series_index/` | Series index writer/searcher, bucket planning, catalogs, worker-owned maintenance and purge tasks, and index lifecycle |
| `read` | `src/mito2/src/read/` | `ScanRegion`, merge, dedup, projection, streaming |
| `manifest` | `src/mito2/src/manifest/` | `RegionManifestManager`, manifest actions/edits |
| `cache` | `src/mito2/src/cache.rs` | Write/file/page caches |
src/mito2/src/access_layer.rs
+2
View File
@@ -568,6 +568,8 @@ pub struct SstWriteRequest {
pub cache_manager: CacheManagerRef,
#[allow(dead_code)]
pub storage: Option<String>,
/// Optional uniform row sequence for writes that do not preserve sequences.
/// Compaction passes `None` to retain the reader's effective input sequences.
pub max_sequence: Option<SequenceNumber>,
pub sst_write_format: FormatType,
src/mito2/src/compaction/compactor.rs
+296 -41
View File
@@ -404,10 +404,33 @@ pub trait SstMerger: Send + Sync + 'static {
#[derive(Clone)]
pub struct DefaultSstMerger;
/// File-level sequence metadata is independent of physical row encoding.
struct OutputSequenceMetadata {
file_sequence: Option<NonZeroU64>,
exact_sequence_trusted: bool,
}
/// Compaction does not admit new data, so it inherits the input sequence bounds
/// without allocating a new admission marker. Retaining physical row sequences
/// does not by itself restore exact-read capability for untrusted inputs.
fn output_sequence_metadata(
region: &CompactionRegion,
inputs: &[FileHandle],
) -> OutputSequenceMetadata {
let exact_sequence_trusted = region.region_options.preserve_row_sequence
&& inputs
.iter()
.all(|f| f.is_effective_target_sequence_trusted(region.region_id));
OutputSequenceMetadata {
file_sequence: known_max_input_sequence(inputs),
exact_sequence_trusted,
}
}
/// Computes the maximum target-domain sequence bound of the output of merging
/// `inputs`. Foreign files are described by their target-local barrier; local
/// trusted files use their physical sequence bound. Unknown bounds never become
/// trusted through a partial maximum.
/// `inputs`. A bound can be a row maximum or an admission marker; compaction
/// must preserve both for sequence-based pruning. Unknown bounds must remain
/// unknown rather than being replaced with a partial maximum.
fn known_max_input_sequence(inputs: &[FileHandle]) -> Option<NonZeroU64> {
let mut max: Option<NonZeroU64> = None;
for input in inputs {
@@ -452,40 +475,7 @@ impl SstMerger for DefaultSstMerger {
.iter()
.map(|f| f.file_id().to_string())
.join(",");
let input_max_sequence = known_max_input_sequence(&output.inputs);
// The output is trusted only when every input is trusted in the target
// domain. Foreign marker values describe the source and are ignored;
// their present FileMeta.sequence is the target-local barrier.
let output_preserves_sequence = compaction_region.region_options.preserve_row_sequence
&& output
.inputs
.iter()
.all(|f| f.is_effective_target_sequence_trusted(region_id));
let output_sequence = if output_preserves_sequence {
input_max_sequence
} else {
// The manifest records the region's latest accepted sequence for
// edit paths and its flushed frontier otherwise. Compaction only
// rewrites the immutable SST snapshot, so this is the same
// region-local admission barrier used when accepting files.
let manifest = compaction_region
.manifest_ctx
.manifest_manager
.read()
.await
.manifest();
NonZeroU64::new(
manifest
.committed_sequence
.unwrap_or(manifest.flushed_sequence)
+ 1,
)
};
// For untrusted output, keep the physical sequence override compatible
// with the main write path: use the known maximum sequence of the
// inputs. The FileMeta sequence remains the admission barrier, but it
// must not be written into the rows (or replaced with zero).
let write_max_sequence = input_max_sequence.map(NonZeroU64::get);
let sequence_metadata = output_sequence_metadata(&compaction_region, &output.inputs);
let builder = CompactionSstReaderBuilder {
metadata: compaction_region.region_metadata.clone(),
sst_layer: compaction_region.access_layer.clone(),
@@ -509,13 +499,16 @@ impl SstMerger for DefaultSstMerger {
source,
cache_manager: compaction_region.cache_manager.clone(),
storage,
max_sequence: write_max_sequence,
// Readers resolve file overrides before merge/dedup. Replacing
// their effective sequences here could promote old rows above
// versions in SSTs that were not part of this merge.
max_sequence: None,
sst_write_format: if flat_format {
FormatType::Flat
} else {
FormatType::PrimaryKey
},
preserve_row_sequence: output_preserves_sequence,
preserve_row_sequence: sequence_metadata.exact_sequence_trusted,
index_options,
index_config,
inverted_index_config,
@@ -564,12 +557,12 @@ impl SstMerger for DefaultSstMerger {
index_version: 0,
num_rows: sst_info.num_rows as u64,
num_row_groups: sst_info.num_row_groups,
sequence: output_sequence,
sequence: sequence_metadata.file_sequence,
partition_expr: partition_expr.clone(),
num_series: sst_info.num_series,
primary_key_min,
primary_key_max,
preserve_row_sequence: output_preserves_sequence,
preserve_row_sequence: sequence_metadata.exact_sequence_trusted,
}
})
.collect::<Vec<_>>();
@@ -888,6 +881,268 @@ mod tests {
assert_eq!(None, known_max_input_sequence(&[]));
}
#[tokio::test]
async fn test_output_sequence_metadata_preserves_trust_and_unknown_bounds() {
let mut region = new_test_compaction_region().await;
let local = region.region_id;
let foreign = RegionId::new(1, 2);
let file = |region_id, sequence: Option<u64>, preserve_row_sequence| {
new_file_handle(FileMeta {
region_id,
sequence: sequence.and_then(NonZeroU64::new),
preserve_row_sequence,
..dummy_file_meta()
})
};
// Bounds are independent of row trust: known, mixed trust, unknown in
// either sequence domain, empty inputs, and the largest possible bound.
let cases = [
(
vec![file(local, Some(3), true), file(foreign, Some(9), false)],
Some(9),
true,
),
(
vec![file(local, Some(3), true), file(local, Some(6), false)],
Some(6),
false,
),
(
vec![file(foreign, Some(9), false), file(foreign, Some(11), true)],
Some(11),
true,
),
(
vec![file(local, Some(3), true), file(local, None, true)],
None,
true,
),
(
vec![file(local, None, false), file(local, Some(6), true)],
None,
false,
),
(
vec![file(local, Some(3), true), file(foreign, None, true)],
None,
false,
),
(vec![], None, true),
(
vec![file(local, Some(u64::MAX), false)],
Some(u64::MAX),
false,
),
];
for committed_sequence in [100, 101] {
region
.manifest_ctx
.manifest_manager
.write()
.await
.update(
RegionMetaActionList::with_action(RegionMetaAction::Edit(RegionEdit {
committed_sequence: Some(committed_sequence),
files_to_add: vec![],
files_to_remove: vec![],
timestamp_ms: None,
compaction_time_window: None,
flushed_entry_id: None,
flushed_sequence: None,
})),
false,
)
.await
.unwrap();
for preserve in [false, true] {
region.region_options.preserve_row_sequence = preserve;
for (inputs, known_bound, inputs_trusted) in &cases {
let metadata = output_sequence_metadata(&region, inputs);
let trusted = preserve && *inputs_trusted;
assert_eq!(trusted, metadata.exact_sequence_trusted);
assert_eq!(
known_bound.and_then(NonZeroU64::new),
metadata.file_sequence
);
}
}
}
}
/// Independent outputs must be safe in either visibility order, including when
/// only one output succeeds. Merely committing both outputs together is not a
/// substitute for preserving version order across their separate merge streams.
#[rstest::rstest]
#[case(true)]
#[case(false)]
#[tokio::test]
async fn test_independent_outputs_preserve_delete_in_every_visibility_state(
#[case] flat_format: bool,
#[values(false, true)] unknown_put_bound: bool,
) {
use datatypes::arrow::array::AsArray;
use datatypes::arrow::datatypes::TimestampMillisecondType;
use store_api::region_engine::RegionEngine;
use store_api::region_request::RegionRequest;
use crate::compaction::CompactionOutput;
use crate::compaction::compactor::{CompactionRegion, Compactor, DefaultCompactor};
use crate::compaction::picker::PickerOutput;
use crate::compaction::reader::CompactionSstReaderBuilder;
use crate::engine::compaction_test::{delete_and_flush, put_and_flush};
use crate::region::options::MergeMode;
use crate::sst::file::FileHandle;
use crate::sst::file_purger::NoopFilePurger;
use crate::test_util::{CreateRequestBuilder, TestEnv, rows_schema};
let mut env = TestEnv::new().await;
let region_id = RegionId::new(1, 1);
let config = MitoConfig {
default_flat_format: flat_format,
min_compaction_interval: Duration::from_secs(3600),
..Default::default()
};
let engine = env.create_engine(config).await;
let request = CreateRequestBuilder::new().build();
let columns = rows_schema(&request);
engine
.handle_request(region_id, RegionRequest::Create(request))
.await
.unwrap();
put_and_flush(&engine, region_id, &columns, 0..1).await;
delete_and_flush(&engine, region_id, &columns, 0..1).await;
put_and_flush(&engine, region_id, &columns, 10..11).await;
let region = engine.get_region(region_id).unwrap();
let version = region.version();
let mut files: Vec<_> = version
.ssts
.levels()
.iter()
.flat_map(|level| level.files())
.cloned()
.collect();
assert_eq!(3, files.len());
files.sort_unstable_by_key(|file| file.meta_ref().sequence);
if unknown_put_bound {
// Model a legacy SST whose physical rows are readable but whose
// file-level boundary is unknown. Merging cannot invent that bound.
let mut meta = files[0].meta_ref().clone();
meta.sequence = None;
files[0] = FileHandle::new(meta, Arc::new(NoopFilePurger));
}
let compaction_region = CompactionRegion {
region_id,
region_options: version.options.clone(),
engine_config: Arc::new(MitoConfig {
default_flat_format: flat_format,
..Default::default()
}),
region_metadata: version.metadata.clone(),
cache_manager: engine.cache_manager(),
access_layer: region.access_layer.clone(),
manifest_ctx: region.manifest_ctx.clone(),
current_version: version.into(),
file_purger: None,
ttl: None,
max_parallelism: 2,
plugins: common_base::Plugins::new(),
};
let puts = vec![files[0].clone(), files[2].clone()];
let deletes = vec![files[1].clone()];
let picker_output = PickerOutput {
outputs: [puts.clone(), deletes.clone()]
.into_iter()
.map(|inputs| CompactionOutput {
output_level: 1,
inputs,
filter_deleted: false,
output_time_range: None,
})
.collect(),
expired_ssts: vec![],
time_window_size: 3600,
max_file_size: None,
};
let merged = DefaultCompactor::with_merger(DefaultSstMerger)
.merge_ssts(&compaction_region, picker_output)
.await
.unwrap();
assert_eq!(2, merged.files_to_add.len());
let outputs: Vec<_> = merged
.files_to_add
.into_iter()
.map(|meta| FileHandle::new(meta, Arc::new(NoopFilePurger)))
.collect();
let put_output = outputs
.iter()
.find(|file| file.meta_ref().num_rows == 2)
.unwrap();
let delete_output = outputs
.iter()
.find(|file| file.meta_ref().num_rows == 1)
.unwrap();
assert_eq!(
if unknown_put_bound {
None
} else {
NonZeroU64::new(3)
},
put_output.meta_ref().sequence,
);
assert_eq!(NonZeroU64::new(2), delete_output.meta_ref().sequence);
assert!(
outputs
.iter()
.all(|file| !file.meta_ref().preserve_row_sequence)
);
// Cover the original snapshot, each partial success/cancellation state,
// and both completed outputs using real persisted SSTs and merge readers.
for (put_done, delete_done) in [(false, false), (true, false), (false, true), (true, true)]
{
let mut visible = if put_done {
vec![put_output.clone()]
} else {
puts.clone()
};
visible.extend(if delete_done {
vec![delete_output.clone()]
} else {
deletes.clone()
});
let mut reader = CompactionSstReaderBuilder {
metadata: compaction_region.region_metadata.clone(),
sst_layer: region.access_layer.clone(),
cache: engine.cache_manager(),
inputs: &visible,
append_mode: false,
filter_deleted: true,
time_range: None,
merge_mode: MergeMode::LastRow,
}
.build_flat_sst_reader()
.await
.unwrap();
let mut timestamps = Vec::new();
while let Some(batch) = reader.next_batch().await.unwrap() {
timestamps.extend(
batch
.column_by_name("ts")
.unwrap()
.as_primitive::<TimestampMillisecondType>()
.values()
.iter()
.copied(),
);
}
assert_eq!(
vec![10_000],
timestamps,
"put_done={put_done}, delete_done={delete_done}"
);
}
}
/// Build a minimal [`CompactionRegion`] suitable for tests where the
/// [`SstMerger`] is mocked and never touches the access layer.
async fn new_test_compaction_region() -> CompactionRegion {
src/mito2/src/config.rs
+40
View File
@@ -95,6 +95,10 @@ pub struct MitoConfig {
/// Interval between series-index maintenance runs (default 5 min). Zero uses the default.
#[serde(with = "humantime_serde")]
pub experimental_series_index_maintenance_interval: Duration,
/// Under development; do not enable. Requested minimum bucket width for series indexes.
/// It is rounded up to an exact multiple of each region's compaction time window.
#[serde(with = "humantime_serde")]
pub experimental_series_index_bucket_width: Duration,
/// Max number of running background flush jobs (default: 1/2 of cpu cores).
pub max_background_flushes: usize,
/// Max number of running background compaction jobs (default: 1/4 of cpu cores).
@@ -215,6 +219,7 @@ impl Default for MitoConfig {
experimental_enable_series_index: false,
experimental_series_index_maintenance_interval:
DEFAULT_SERIES_INDEX_MAINTENANCE_INTERVAL,
experimental_series_index_bucket_width: Duration::from_secs(5 * 24 * 60 * 60),
max_background_flushes: divide_num_cpus(2),
max_background_compactions: divide_num_cpus(4),
max_background_purges: get_total_cpu_cores(),
@@ -421,6 +426,41 @@ mod tests {
assert_eq!(ReadableSize::mb(128), config.prefilter_result_cache_size);
}
#[test]
fn test_series_index_config() {
assert!(!MitoConfig::default().experimental_enable_series_index);
assert_eq!(
MitoConfig::default().experimental_series_index_bucket_width,
Duration::from_secs(5 * 24 * 60 * 60)
);
let mut config: MitoConfig = toml::from_str(
"experimental_enable_series_index = true
experimental_series_index_maintenance_interval = '30s'
experimental_series_index_bucket_width = '2days'",
)
.unwrap();
config.sanitize("/data").unwrap();
assert!(config.experimental_enable_series_index);
assert_eq!(
config.experimental_series_index_maintenance_interval,
Duration::from_secs(30)
);
assert_eq!(
config.experimental_series_index_bucket_width,
Duration::from_secs(2 * 24 * 60 * 60)
);
let restored: MitoConfig = toml::from_str(&toml::to_string(&config).unwrap()).unwrap();
assert_eq!(config, restored);
let mut config: MitoConfig =
toml::from_str("experimental_series_index_maintenance_interval = '0s'").unwrap();
config.sanitize("/data").unwrap();
assert_eq!(
config.experimental_series_index_maintenance_interval,
MitoConfig::default().experimental_series_index_maintenance_interval
);
}
#[test]
fn test_experimental_series_scan_v2_config() {
assert!(MitoConfig::default().experimental_series_scan_v2);
src/mito2/src/engine/basic_test.rs
+1
View File
@@ -1341,6 +1341,7 @@ async fn test_all_index_metas_list_all_types_with_format(flat_format: bool, expe
if let Some(inverted) = value.get_mut("inverted").and_then(|v| v.as_object_mut()) {
inverted.insert("base_offset".to_string(), serde_json::Value::from(0));
}
value.sort_all_objects();
*meta_json = value.to_string();
}
}
src/mito2/src/engine/compaction_test.rs
+291 -2
View File
@@ -12,7 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::ops::Range;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
@@ -39,11 +39,11 @@ use tokio::sync::{Notify, Semaphore};
use crate::config::MitoConfig;
use crate::engine::MitoEngine;
use crate::engine::flush_test::MockTimeProvider;
use crate::engine::listener::{CompactionListener, EventListener};
use crate::test_util::{
CreateRequestBuilder, TestEnv, build_rows_for_key, column_metadata_to_column_schema, put_rows,
};
use crate::time_provider::mock::MockTimeProvider;
pub(crate) async fn put_and_flush(
engine: &MitoEngine,
@@ -137,6 +137,295 @@ async fn collect_stream_ts(stream: SendableRecordBatchStream) -> Vec<i64> {
res
}
/// Flush may collapse versions within one file, but compaction must not promote
/// them, even with no external overlaps or after rewriting a previous output.
#[rstest::rstest]
#[case(false, false)]
#[case(true, false)]
#[case(true, true)]
#[tokio::test]
async fn test_compaction_preserves_flush_sequences_across_generations(
#[values(false, true)] flat_format: bool,
#[case] append_mode: bool,
#[case] preserve_row_sequence: bool,
) {
let mut env = TestEnv::new().await;
let config = MitoConfig {
default_flat_format: flat_format,
min_compaction_interval: Duration::from_secs(3600),
..Default::default()
};
let engine = env.create_engine(config.clone()).await;
let region_id = RegionId::new(1, 1);
let request = CreateRequestBuilder::new()
.insert_option("append_mode", &append_mode.to_string())
.insert_option("preserve_row_sequence", &preserve_row_sequence.to_string())
.insert_option("compaction.type", "twcs")
.insert_option("compaction.twcs.time_window", "1h")
.build();
let table_dir = request.table_dir.clone();
let options = request.options.clone();
let columns = crate::test_util::rows_schema(&request);
engine
.handle_request(region_id, RegionRequest::Create(request))
.await
.unwrap();
async fn read_sequences(engine: &MitoEngine, region_id: RegionId) -> Vec<u64> {
let region = engine.get_region(region_id).unwrap();
let mut sequences = Vec::new();
for file in region
.version()
.ssts
.levels()
.iter()
.flat_map(|l| l.files())
{
let mut reader = region
.access_layer
.read_sst(file.clone())
.build()
.await
.unwrap()
.unwrap();
while let Some(batch) = reader.next_record_batch().await.unwrap() {
sequences.extend_from_slice(
batch
.column(batch.num_columns() - 2)
.as_primitive::<datatypes::arrow::datatypes::UInt64Type>()
.values(),
);
}
}
sequences.sort_unstable();
sequences
}
let mut expected = Vec::new();
for generation in 0..=2 {
// Interleaving time ranges exercise a merge without duplicate rows.
let rows = [generation, generation + 10]
.into_iter()
.flat_map(|ts| build_rows_for_key("a", ts, ts + 1, 0))
.collect();
put_rows(
&engine,
region_id,
Rows {
schema: columns.clone(),
rows,
},
)
.await;
flush(&engine, region_id).await;
let max = (generation * 2 + 2) as u64;
expected.extend(if preserve_row_sequence {
[max - 1, max]
} else {
[max, max]
});
assert_eq!(expected, read_sequences(&engine, region_id).await);
if generation == 0 {
continue;
}
engine
.handle_request(
region_id,
RegionRequest::Compact(RegionCompactRequest {
options: api::v1::region::compact_request::Options::StrictWindow(
api::v1::region::StrictWindow {
window_seconds: 3600,
},
),
..Default::default()
}),
)
.await
.unwrap();
let version = engine.get_region(region_id).unwrap().version();
let files: Vec<_> = version
.ssts
.levels()
.iter()
.flat_map(|l| l.files())
.collect();
assert_eq!(
1,
files.len(),
"generation {generation} must merge all inputs"
);
assert_eq!(
preserve_row_sequence,
files[0].meta_ref().preserve_row_sequence
);
assert_eq!(Some(max), files[0].meta_ref().sequence.map(|s| s.get()));
assert_eq!(
expected,
read_sequences(&engine, region_id).await,
"generation {generation}"
);
}
let engine = env.reopen_engine(engine, config).await;
crate::test_util::reopen_region(&engine, region_id, table_dir, false, options).await;
assert_eq!(expected, read_sequences(&engine, region_id).await);
}
#[tokio::test]
async fn test_partial_compaction_preserves_delete_order_flat() {
assert_partial_compaction_preserves_delete_order(true).await;
}
#[tokio::test]
async fn test_partial_compaction_preserves_delete_order_primary_key() {
assert_partial_compaction_preserves_delete_order(false).await;
}
/// A newer unrelated input must not promote an old Put above an unselected Delete.
/// All files fit in the same one-hour window; no cross-window SST is required.
async fn assert_partial_compaction_preserves_delete_order(flat_format: bool) {
let mut env = TestEnv::new().await;
let region_id = RegionId::new(1, 1);
let (engine, columns) = env_for_manual_compaction(&mut env, region_id, flat_format).await;
put_and_flush(&engine, region_id, &columns, 10..16).await;
let mut deletes = build_rows_for_key("a", 10, 16, 0);
// Distinct keys keep the Delete SST large even after compression, so the
// byte-balance rule excludes it without fabricating FileMeta sizes.
let mut rng = <rand::rngs::StdRng as rand::SeedableRng>::seed_from_u64(0);
for _ in 0..3000 {
deletes.extend(build_rows_for_key(
&format!("{:032x}", rand::Rng::random::<u128>(&mut rng)),
1,
2,
0,
));
}
engine
.handle_request(
region_id,
RegionRequest::Delete(RegionDeleteRequest {
rows: Rows {
schema: columns.clone(),
rows: deletes,
},
hint: None,
partition_expr_version: None,
}),
)
.await
.unwrap();
flush(&engine, region_id).await;
// These newer writes are legitimate, but must not change the version of 10..16.
for i in 10..25 {
put_and_flush(&engine, region_id, &columns, i * 10..i * 10 + 6).await;
}
let region = engine.get_region(region_id).unwrap();
let table_dir = region.table_dir().to_string();
let version = region.version();
let files = version
.ssts
.levels()
.iter()
.flat_map(|level| level.files())
.collect::<Vec<_>>();
assert_eq!(17, files.len());
assert!(files.iter().all(|file| {
file.time_range().0 >= Timestamp::new_millisecond(0)
&& file.time_range().1 < Timestamp::new_millisecond(3_600_000)
}));
let delete_file = files
.iter()
.find(|file| file.meta_ref().num_rows == 3006)
.unwrap();
let delete_id = delete_file.file_id();
let delete_sequence = delete_file.meta_ref().sequence;
let old_ids = files
.iter()
.map(|file| file.file_id())
.collect::<HashSet<_>>();
let before = collect_stream_ts(
engine
.scanner(region_id, ScanRequest::default())
.await
.unwrap()
.scan()
.await
.unwrap(),
)
.await;
let expected = (10..25)
.flat_map(|i| (i * 10..i * 10 + 6).map(|ts| ts * 1000))
.collect::<Vec<i64>>();
assert_eq!(
expected, before,
"the Delete must hide old rows before compaction"
);
compact(&engine, region_id).await;
let scanner = engine
.scanner(region_id, ScanRequest::default())
.await
.unwrap();
let current_ids = scanner.file_ids();
assert_eq!(
2,
current_ids.len(),
"the 16 small inputs must merge, leaving the large Delete SST"
);
assert!(current_ids.contains(&delete_id));
assert_eq!(
1,
current_ids.iter().filter(|id| old_ids.contains(id)).count()
);
let after = collect_stream_ts(scanner.scan().await.unwrap()).await;
let current = region.version();
let output = current
.ssts
.levels()
.iter()
.flat_map(|level| level.files())
.find(|file| !old_ids.contains(&file.file_id()))
.unwrap();
let mut reader = region
.access_layer
.read_sst(output.clone())
.build()
.await
.unwrap()
.unwrap();
let batch = reader.next_record_batch().await.unwrap().unwrap();
let output_sequences = batch
.column(batch.num_columns() - 2)
.as_primitive::<datatypes::arrow::datatypes::UInt64Type>();
let first_output_sequence = output_sequences.value(0);
crate::test_util::reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await;
let reopened = collect_stream_ts(
engine
.scanner(region_id, ScanRequest::default())
.await
.unwrap()
.scan()
.await
.unwrap(),
)
.await;
let resurrected = after
.iter()
.filter(|ts| (10_000..16_000).contains(*ts))
.collect::<Vec<_>>();
assert!(
after == expected && reopened == expected,
"partial compaction changed delete ordering (flat_format={flat_format}, unselected Delete sequence={delete_sequence:?}, first output row sequence={first_output_sequence}); expected rows={}, after rows={}, reopened rows={}, resurrected timestamps={resurrected:?}",
expected.len(),
after.len(),
reopened.len()
);
}
struct CompactionListenerGuard(Option<Arc<CompactionListener>>);
impl CompactionListenerGuard {
src/mito2/src/engine/edit_region_test.rs
+1 -1
View File
@@ -35,12 +35,12 @@ use tokio::sync::{Barrier, mpsc, oneshot};
use crate::config::MitoConfig;
use crate::engine::MitoEngine;
use crate::engine::flush_test::MockTimeProvider;
use crate::engine::listener::EventListener;
use crate::manifest::action::RegionEdit;
use crate::region::{MitoRegionRef, RegionLeaderState, RegionRoleState};
use crate::sst::file::FileMeta;
use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, rows_schema};
use crate::time_provider::mock::MockTimeProvider;
#[tokio::test]
async fn test_edit_region_schedule_compaction() {
src/mito2/src/engine/flush_test.rs
+6 -41
View File
@@ -17,7 +17,7 @@
use std::assert_matches;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicUsize, Ordering};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::Duration;
use api::v1::Rows;
@@ -55,7 +55,7 @@ use crate::test_util::{
prepare_test_for_kafka_log_store, put_rows, raft_engine_log_store_factory, reopen_region,
rows_schema, single_kafka_log_store_factory,
};
use crate::time_provider::TimeProvider;
use crate::time_provider::mock::MockTimeProvider;
use crate::worker::MAX_INITIAL_CHECK_DELAY_SECS;
async fn set_write_buffer_size_to_current_usage(engine: &MitoEngine, region_id: RegionId) {
@@ -1057,8 +1057,10 @@ async fn test_flush_empty_with_format(flat_format: bool) {
let stream = scanner.scan().await.unwrap();
let batches = RecordBatches::try_collect(stream).await.unwrap();
let expected = "\
++
++";
+-------+---------+----+
| tag_0 | field_0 | ts |
+-------+---------+----+
+-------+---------+----+";
assert_eq!(expected, batches.pretty_print().unwrap());
}
@@ -1369,43 +1371,6 @@ fn kafka_wal_options(topic: &Option<String>) -> HashMap<String, String> {
.unwrap_or_default()
}
#[derive(Debug)]
pub(crate) struct MockTimeProvider {
now: AtomicI64,
elapsed: AtomicI64,
}
impl TimeProvider for MockTimeProvider {
fn current_time_millis(&self) -> i64 {
self.now.load(Ordering::Relaxed)
}
fn elapsed_since(&self, _current_millis: i64) -> i64 {
self.elapsed.load(Ordering::Relaxed)
}
fn wait_duration(&self, _duration: Duration) -> Duration {
Duration::from_millis(20)
}
}
impl MockTimeProvider {
pub(crate) fn new(now: i64) -> Self {
Self {
now: AtomicI64::new(now),
elapsed: AtomicI64::new(0),
}
}
pub(crate) fn set_now(&self, now: i64) {
self.now.store(now, Ordering::Relaxed);
}
fn set_elapsed(&self, elapsed: i64) {
self.elapsed.store(elapsed, Ordering::Relaxed);
}
}
#[tokio::test]
async fn test_auto_flush_engine() {
test_auto_flush_engine_with_format(false).await;
src/mito2/src/engine/merge_mode_test.rs
+189
View File
@@ -16,6 +16,9 @@
use api::v1::Rows;
use common_recordbatch::RecordBatches;
use datafusion_expr::{col, lit};
use datatypes::arrow::array::AsArray;
use datatypes::arrow::datatypes::UInt64Type;
use store_api::region_engine::RegionEngine;
use store_api::region_request::{RegionCompactRequest, RegionRequest};
use store_api::storage::{RegionId, ScanRequest};
@@ -27,6 +30,192 @@ use crate::test_util::{
delete_rows_schema, flush_region, put_rows, reopen_region, rows_schema,
};
#[rstest::rstest]
#[tokio::test]
async fn test_partial_compaction_preserves_last_row_put_versions(
#[values(false, true)] flat_format: bool,
) {
let mut env = TestEnv::new().await;
let config = MitoConfig {
default_flat_format: flat_format,
min_compaction_interval: std::time::Duration::from_secs(3600),
..Default::default()
};
let engine = env.create_engine(config.clone()).await;
let region_id = RegionId::new(1, 1);
env.get_schema_metadata_manager()
.register_region_table_info(
region_id.table_id(),
"test_table",
"test_catalog",
"test_schema",
None,
env.get_kv_backend(),
)
.await;
let request = CreateRequestBuilder::new()
.field_num(2)
.insert_option("compaction.type", "twcs")
.insert_option("compaction.twcs.time_window", "1h")
.insert_option("merge_mode", "last_row")
.build();
let table_dir = request.table_dir.clone();
let region_opts = request.options.clone();
let schema = rows_schema(&request);
engine
.handle_request(region_id, RegionRequest::Create(request))
.await
.unwrap();
// A(x=1), B(x=2) share a key and timestamp; C is unrelated.
// Promoting A when merging A+C must not hide the unselected B.
let a = build_rows_with_fields("a", &[10], &[(Some(1), None)]);
put_rows(
&engine,
region_id,
Rows {
schema: schema.clone(),
rows: a,
},
)
.await;
flush_region(&engine, region_id, None).await;
let mut b = build_rows_with_fields("a", &[10], &[(Some(2), None)]);
let mut rng = <rand::rngs::StdRng as rand::SeedableRng>::seed_from_u64(0);
for _ in 0..3000 {
b.extend(build_rows_with_fields(
&format!("{:032x}", rand::Rng::random::<u128>(&mut rng)),
&[1],
&[(Some(2), None)],
));
}
put_rows(
&engine,
region_id,
Rows {
schema: schema.clone(),
rows: b,
},
)
.await;
flush_region(&engine, region_id, None).await;
let b_id = engine
.get_region(region_id)
.unwrap()
.version()
.ssts
.levels()
.iter()
.flat_map(|level| level.files())
.find(|file| file.meta_ref().num_rows == 3001)
.unwrap()
.file_id();
let c = build_rows_with_fields("z", &[10], &[(None, Some(3))]);
put_rows(
&engine,
region_id,
Rows {
schema: schema.clone(),
rows: c,
},
)
.await;
flush_region(&engine, region_id, None).await;
// Sixteen small inputs plus the large middle-version file B. The normal
// byte-balanced candidate excludes B.
for i in 0..14 {
let rows = build_rows_with_fields("z", &[100 + i], &[(Some(4), None)]);
put_rows(
&engine,
region_id,
Rows {
schema: schema.clone(),
rows,
},
)
.await;
flush_region(&engine, region_id, None).await;
}
let scan = || ScanRequest {
filters: vec![col("tag_0").eq(lit("a"))],
..Default::default()
};
let before =
RecordBatches::try_collect(engine.scan_to_stream(region_id, scan()).await.unwrap())
.await
.unwrap()
.pretty_print()
.unwrap();
let expected = "\
+-------+---------+---------+---------------------+
| tag_0 | field_0 | field_1 | ts |
+-------+---------+---------+---------------------+
| a | 2.0 | | 1970-01-01T00:00:10 |
+-------+---------+---------+---------------------+";
assert_eq!(expected, before);
engine
.handle_request(
region_id,
RegionRequest::Compact(RegionCompactRequest::default()),
)
.await
.unwrap();
let after = RecordBatches::try_collect(engine.scan_to_stream(region_id, scan()).await.unwrap())
.await
.unwrap()
.pretty_print()
.unwrap();
assert_eq!(expected, after, "flat_format={flat_format}");
let version = engine.get_region(region_id).unwrap().version();
let files: Vec<_> = version
.ssts
.levels()
.iter()
.flat_map(|level| level.files())
.collect();
assert_eq!(2, files.len());
assert!(files.iter().any(|file| file.file_id() == b_id));
assert!(
files
.iter()
.all(|file| !file.meta_ref().preserve_row_sequence)
);
let region = engine.get_region(region_id).unwrap();
let mut output_sequences = std::collections::HashSet::new();
for file in files.iter().filter(|file| file.file_id() != b_id) {
let mut reader = region
.access_layer
.read_sst((*file).clone())
.build()
.await
.unwrap()
.unwrap();
while let Some(batch) = reader.next_record_batch().await.unwrap() {
output_sequences.extend(
batch
.column(batch.num_columns() - 2)
.as_primitive::<UInt64Type>()
.values()
.iter()
.copied(),
);
}
}
assert!(
output_sequences.len() > 1,
"compaction must retain effective input sequences"
);
let engine = env.reopen_engine(engine, config).await;
reopen_region(&engine, region_id, table_dir, false, region_opts).await;
let reopened =
RecordBatches::try_collect(engine.scan_to_stream(region_id, scan()).await.unwrap())
.await
.unwrap()
.pretty_print()
.unwrap();
assert_eq!(expected, reopened);
}
#[tokio::test]
async fn test_merge_mode_write_query() {
test_merge_mode_write_query_with_format(false).await;
src/mito2/src/engine/scan_test.rs
+107 -52
View File
@@ -1053,8 +1053,10 @@ async fn test_scan_with_min_sst_sequence_with_format(flat_format: bool) {
Some(9),
0,
"\
++
++",
+-------+---------+----+
| tag_0 | field_0 | ts |
+-------+---------+----+
+-------+---------+----+",
)
.await;
}
@@ -2907,9 +2909,8 @@ async fn test_bulk_write_sequence_not_committed_before_install() {
}
}
/// Non-preserving compaction must keep the physical input sequence below the
/// admission barrier. Otherwise a later row at the barrier can collide with
/// the compacted row when ordinary reads deduplicate overlapping SSTs.
/// Compaction must not allocate a new sequence that can collide with the next
/// write. Preserve the input's effective row sequence and file-level bound.
#[tokio::test]
async fn test_non_preserve_compaction_sequence_collision() {
for flat_format in [false, true] {
@@ -2979,8 +2980,7 @@ async fn test_non_preserve_compaction_sequence_collision_with_format(flat_format
);
// Both current SSTs are compacted through the real picker, merger, writer,
// and manifest update. With committed/flushed sequence 2, the untrusted
// output's admission barrier is exactly 3.
// and manifest update. Both the row version and inherited file bound stay 2.
engine
.handle_request(
region_id,
@@ -3003,14 +3003,13 @@ async fn test_non_preserve_compaction_sequence_collision_with_format(flat_format
assert!(!compacted.meta_ref().preserve_row_sequence);
assert_eq!(1, compacted.num_rows());
assert_eq!(
Some(std::num::NonZeroU64::new(3).unwrap()),
Some(std::num::NonZeroU64::new(2).unwrap()),
compacted.meta_ref().sequence,
"admission barrier is input max 2 plus one"
"compaction must inherit the input bound without allocating a sequence"
);
// Hide FileMeta.sequence so this is a direct physical read, not a reader
// admission-barrier override. The output must contain physical sequence 2,
// not zero and not the output barrier 3.
// override. The output must contain physical sequence 2, not zero or 3.
let mut physical_meta = compacted.meta_ref().clone();
physical_meta.sequence = None;
let mut reader = region
@@ -3037,10 +3036,8 @@ async fn test_non_preserve_compaction_sequence_collision_with_format(flat_format
"after compaction: {after_compaction}"
);
// The next write receives physical sequence 3, equal to the compacted
// output's admission barrier. It must still win because the compacted row
// remains physically at sequence 2; an old zero/barrier encoding would
// collide here and incorrectly retain value 2.
// The next write receives sequence 3 and must win. The old zero/barrier
// encoding promoted the compacted row to 3 and could incorrectly retain 2.
test_util::put_rows(
&engine,
region_id,
@@ -3075,7 +3072,7 @@ async fn test_non_preserve_compaction_sequence_collision_with_format(flat_format
.find(|file| file.meta_ref().level == 1)
.expect("compacted L1 SST");
assert_eq!(
Some(std::num::NonZeroU64::new(3).unwrap()),
Some(std::num::NonZeroU64::new(2).unwrap()),
compacted.meta_ref().sequence
);
@@ -3086,24 +3083,27 @@ async fn test_non_preserve_compaction_sequence_collision_with_format(flat_format
);
}
/// Compaction rewrites a legacy (unmarked) input as an untrusted output with
/// the region-local admission barrier. Its physical rows retain the known
/// input maximum, while exact scans skip it once C reaches that barrier and
/// fail closed while C is below it.
/// A legacy input must retain its admission boundary without gaining row trust.
/// A later flush can advance beyond manifest.committed_sequence, so compaction
/// must inherit the maximum input bound rather than sample that stale frontier.
#[rstest::rstest]
#[tokio::test]
async fn test_compaction_output_non_preserve_not_laundered_from_legacy_input() {
async fn test_compaction_output_non_preserve_not_laundered_from_legacy_input(
#[values(false, true)] flat_format: bool,
#[values(false, true)] flush_after_edit: bool,
) {
let mut env =
TestEnv::with_prefix("test_compaction_output_not_laundered_from_legacy_input").await;
// Suppress automatic edit-triggered compactions. The high TWCS trigger below
// also prevents flush-triggered compaction from consuming the inputs, so the
// explicit Compact below is the only compaction in flight (deterministic).
let engine = env
.create_engine(MitoConfig {
min_compaction_interval: std::time::Duration::from_secs(60 * 60),
schedule_compaction_after_edit: false,
..Default::default()
})
.await;
let config = MitoConfig {
default_flat_format: flat_format,
min_compaction_interval: std::time::Duration::from_secs(60 * 60),
schedule_compaction_after_edit: false,
..Default::default()
};
let engine = env.create_engine(config.clone()).await;
let region_id = RegionId::new(1, 1);
let request = CreateRequestBuilder::new()
@@ -3113,6 +3113,8 @@ async fn test_compaction_output_non_preserve_not_laundered_from_legacy_input() {
.insert_option("compaction.twcs.trigger_file_num", "100")
.build();
let column_schemas = test_util::rows_schema(&request);
let table_dir = request.table_dir.clone();
let region_options = request.options.clone();
engine
.handle_request(region_id, RegionRequest::Create(request))
@@ -3172,6 +3174,30 @@ async fn test_compaction_output_non_preserve_not_laundered_from_legacy_input() {
"seeded file should be unmarked"
);
let mut expected_sequences = vec![1, 2, 3, 4, 5, 6];
let expected_bound = if flush_after_edit {
test_util::put_rows(
&engine,
region_id,
Rows {
schema: column_schemas,
rows: test_util::build_rows(6, 9),
},
)
.await;
test_util::flush_region(&engine, region_id, None).await;
expected_sequences.extend([8, 9, 10]);
10
} else {
7
};
let manifest = region.manifest_ctx.manifest().await;
assert_eq!(Some(7), manifest.committed_sequence);
assert_eq!(
if flush_after_edit { 10 } else { 6 },
manifest.flushed_sequence
);
// Compact: the rewritten output must NOT be laundered back to marked.
engine
.handle_request(
@@ -3198,25 +3224,24 @@ async fn test_compaction_output_non_preserve_not_laundered_from_legacy_input() {
assert_eq!(
1,
outputs.len(),
"two inputs should rewrite into one output file"
"all inputs should rewrite into one output file"
);
assert!(
!outputs[0].meta_ref().preserve_row_sequence,
"legacy input must not be laundered into a marked output"
);
// The compacted output is sequence-less and carries the current region's
// admission barrier rather than any source-domain sequence.
// Inherit the admitted legacy bound and, if present, the newer flush bound.
let barrier = outputs[0]
.meta_ref()
.sequence
.expect("compaction output barrier")
.get();
assert_eq!(8, barrier);
assert_eq!(expected_bound, barrier);
// Reinstalling the legacy input assigns it sequence 7, so the physical
// parquet retains that known input maximum rather than encoding either
// zero or the output admission barrier 8. The manifest marker stays false.
// Reinstalling the legacy input assigns its metadata sequence 7, but the
// existing rows keep sequences 1..=6. Compaction must not rewrite them to
// the metadata bound, including when newer flushed rows join the merge.
let mut sequence_meta = outputs[0].meta_ref().clone();
sequence_meta.sequence = None;
let sequence_handle = FileHandle::new(
@@ -3230,22 +3255,34 @@ async fn test_compaction_output_non_preserve_not_laundered_from_legacy_input() {
.await
.unwrap()
.expect("compaction output reader");
let batch = reader
.next_record_batch()
.await
.unwrap()
.expect("compaction output batch");
let sequence = batch
.column(batch.num_columns() - 2)
.as_any()
.downcast_ref::<datatypes::arrow::array::UInt64Array>()
.expect("sequence column");
assert!(
sequence
.values()
.iter()
.all(|sequence| *sequence == barrier - 1)
let mut sequences = Vec::new();
while let Some(batch) = reader.next_record_batch().await.unwrap() {
let sequence = batch
.column(batch.num_columns() - 2)
.as_any()
.downcast_ref::<datatypes::arrow::array::UInt64Array>()
.expect("sequence column");
sequences.extend_from_slice(sequence.values());
}
sequences.sort_unstable();
assert_eq!(expected_sequences, sequences);
// The inherited bound and untrusted marker must survive manifest recovery.
let engine = env.reopen_engine(engine, config).await;
test_util::reopen_region(&engine, region_id, table_dir, false, region_options).await;
let version = engine.get_region(region_id).unwrap().version();
let files: Vec<_> = version
.ssts
.levels()
.iter()
.flat_map(|level| level.files())
.collect();
assert_eq!(1, files.len());
assert_eq!(
Some(expected_bound),
files[0].meta_ref().sequence.map(|s| s.get())
);
assert!(!files[0].meta_ref().preserve_row_sequence);
// A cursor before the barrier must fail closed.
let err = engine
@@ -3261,9 +3298,27 @@ async fn test_compaction_output_non_preserve_not_laundered_from_legacy_input() {
.await
.err()
.expect("newer barrier must disable exact scanning");
assert!(matches!(err, Error::SequenceRangeUnsupported { .. }));
if flush_after_edit {
// Without exact capability, the lower-bound fence is checked first.
assert!(
matches!(
err,
Error::IncrementalQueryStale {
given_seq: 9,
min_readable_seq: 10,
..
}
),
"unexpected error: {err:?}"
);
} else {
assert!(
matches!(err, Error::SequenceRangeUnsupported { .. }),
"unexpected error: {err:?}"
);
}
// Once C reaches the barrier the sequence-less file is skipped, so exact
// Once C reaches the barrier the untrusted file is skipped, so exact
// capability is restored without attempting row-level filtering.
let scanner = engine
.scanner(
src/mito2/src/engine/sync_test.rs
+10 -2
View File
@@ -148,7 +148,11 @@ async fn test_sync_after_flush_region_with_format(flat_format: bool) {
common_telemetry::info!("Scan the region on the follower engine");
// Scan the region on the follower engine
let expected = "++\n++";
let expected = "\
+-------+---------+----+
| tag_0 | field_0 | ts |
+-------+---------+----+
+-------+---------+----+";
scan_check(&follower_engine, region_id, expected, 0, 0).await;
// Returns error since the max manifest is 1
@@ -262,7 +266,11 @@ async fn test_sync_after_alter_region_with_format(flat_format: bool) {
+-------+-------+---------+---------------------+";
scan_check(&engine, region_id, expected, 0, 1).await;
let expected = "++\n++";
let expected = "\
+-------+---------+----+
| tag_0 | field_0 | ts |
+-------+---------+----+
+-------+---------+----+";
scan_check(&follower_engine, region_id, expected, 0, 0).await;
// Sync the region from the leader engine to the follower engine
src/mito2/src/engine/truncate_test.rs
+10 -2
View File
@@ -155,7 +155,11 @@ async fn test_engine_truncate_region_basic_with_format(flat_format: bool) {
let request = ScanRequest::default();
let stream = engine.scan_to_stream(region_id, request).await.unwrap();
let batches = RecordBatches::try_collect(stream).await.unwrap();
let expected = "++\n++";
let expected = "\
+-------+---------+----+
| tag_0 | field_0 | ts |
+-------+---------+----+
+-------+---------+----+";
assert_eq!(expected, batches.pretty_print().unwrap());
}
@@ -401,7 +405,11 @@ async fn test_engine_truncate_reopen_with_format(flat_format: bool) {
let request = ScanRequest::default();
let stream = engine.scan_to_stream(region_id, request).await.unwrap();
let batches = RecordBatches::try_collect(stream).await.unwrap();
let expected = "++\n++";
let expected = "\
+-------+---------+----+
| tag_0 | field_0 | ts |
+-------+---------+----+
+-------+---------+----+";
assert_eq!(expected, batches.pretty_print().unwrap());
}
src/mito2/src/memtable/bulk/part.rs
+1 -1
View File
@@ -1566,7 +1566,7 @@ impl PruningStatistics for BatchPruningStats<'_> {
None
}
fn row_counts(&self, _column: &Column) -> Option<ArrayRef> {
fn row_counts(&self) -> Option<ArrayRef> {
None
}
src/mito2/src/metrics.rs
+13
View File
@@ -319,6 +319,19 @@ lazy_static! {
// Index metrics.
lazy_static! {
// Index metrics.
/// Outcomes of series-index reconciliation passes.
pub static ref SERIES_INDEX_RECONCILE_TOTAL: IntCounterVec = register_int_counter_vec!(
"greptime_mito_series_index_reconcile_total",
"series-index reconciliation passes",
&["result"],
).unwrap();
/// Elapsed time of series-index reconciliation phases.
pub static ref SERIES_INDEX_RECONCILE_ELAPSED: HistogramVec = register_histogram_vec!(
"greptime_mito_series_index_reconcile_elapsed",
"series-index reconciliation elapsed time",
&["phase"],
exponential_buckets(0.01, 10.0, 7).unwrap(),
).unwrap();
/// Series and range index file operations.
pub static ref SERIES_INDEX_FILE_OPERATION_TOTAL: IntCounterVec = register_int_counter_vec!(
"greptime_mito_series_index_file_operation_total",
src/mito2/src/read/scan_region.rs
+63 -4
View File
@@ -1688,6 +1688,27 @@ impl ScanInput {
);
}
/// Whether physical source row counts equal the rows visible to this region,
/// before applying query predicates.
pub(crate) fn total_rows_is_exact(&self) -> bool {
if self.region_partition_expr.is_none() {
return true;
}
// Extension ranges do not expose their partition expressions.
#[cfg(feature = "enterprise")]
if !self.extension_ranges.is_empty() {
return false;
}
// Repartition flushes existing memtables before entering staging. New
// writes are routed under the target partition rules, whereas historical
// SSTs can be shared by regions and still need partition filtering.
self.files
.iter()
.all(|file| self.should_skip_region_partition(file))
}
pub(crate) fn total_rows(&self) -> usize {
let rows_in_files: usize = self.files.iter().map(|f| f.num_rows()).sum();
let rows_in_memtables: usize = self.memtables.iter().map(|m| m.stats().num_rows()).sum();
@@ -1796,7 +1817,7 @@ impl PruningStatistics for FileLevelPruningStats {
}
}
fn row_counts(&self, _column: &Column) -> Option<ArrayRef> {
fn row_counts(&self) -> Option<ArrayRef> {
None
}
@@ -1836,9 +1857,10 @@ fn pre_filter_mode(append_mode: bool, merge_mode: MergeMode) -> PreFilterMode {
/// contribute a row to `(C, H]`.
///
/// Unmarked local files use `FileMeta.sequence` as an admission barrier rather
/// than a row maximum. Compaction and edit assign it as `committed_sequence + 1`;
/// `C >= barrier` proves Flow has already consumed the entire file, so such a
/// file is excluded before the capability check.
/// than a row maximum. Region edits allocate this barrier; compaction inherits
/// the maximum input bound without assigning a new barrier. `C >= barrier`
/// proves Flow has already consumed the entire file, so such a file is excluded
/// before the capability check. An unknown bound cannot prove this exclusion.
///
/// A foreign file is different: the parquet reader virtualizes every row to its
/// target-local `FileMeta.sequence`. Consequently, a present sequence is the
@@ -2324,6 +2346,43 @@ mod tests {
.with_files(vec![file])
}
#[tokio::test]
async fn test_total_rows_is_exact_after_partition_filter() {
let expr = partition_col("k0").gt_eq(Value::String("foo".into()));
let other = partition_col("k0").gt_eq(Value::String("bar".into()));
for (region_expr, file_exprs, exact) in [
(None, vec![None], true),
(None, vec![Some(expr.clone())], true),
(Some(expr.clone()), vec![], true),
(Some(expr.clone()), vec![Some(expr.clone())], true),
(Some(expr.clone()), vec![None], false),
(Some(expr.clone()), vec![Some(other)], false),
(Some(expr.clone()), vec![Some(expr), None], false),
] {
let mut builder =
RegionMetadataBuilder::from_existing(metadata_with_primary_key(vec![0, 1], false));
builder.partition_expr_json(region_expr.map(|expr| expr.as_json_str().unwrap()));
let metadata = Arc::new(builder.build_without_validation().unwrap());
let files = file_exprs
.into_iter()
.map(|partition_expr| {
FileHandle::new(
FileMeta {
partition_expr,
..Default::default()
},
Arc::new(crate::sst::file_purger::NoopFilePurger),
)
})
.collect();
let input = new_scan_input(metadata, vec![])
.await
.with_files(files)
.build();
assert_eq!(input.total_rows_is_exact(), exact);
}
}
/// Helper to create a timestamp millisecond literal.
fn ts_lit(val: i64) -> datafusion_expr::Expr {
lit(ScalarValue::TimestampMillisecond(Some(val), None))
src/mito2/src/read/seq_scan.rs
+2 -1
View File
@@ -76,7 +76,8 @@ impl SeqScan {
pub(crate) fn new(input: ScanInput) -> Self {
let mut properties = ScannerProperties::default()
.with_append_mode(input.append_mode)
.with_total_rows(input.total_rows());
.with_total_rows(input.total_rows())
.with_total_rows_is_exact(input.append_mode && input.total_rows_is_exact());
if let Some(counters) = input.query_stat_counters.clone() {
properties.set_query_stat_counters(counters);
}
src/mito2/src/read/series_scan.rs
+2 -1
View File
@@ -147,7 +147,8 @@ impl SeriesScan {
};
let mut properties = ScannerProperties::default()
.with_append_mode(input.append_mode)
.with_total_rows(input.total_rows());
.with_total_rows(input.total_rows())
.with_total_rows_is_exact(input.append_mode && input.total_rows_is_exact());
if let Some(counters) = input.query_stat_counters.clone() {
properties.set_query_stat_counters(counters);
}
src/mito2/src/read/unordered_scan.rs
+2 -1
View File
@@ -61,7 +61,8 @@ impl UnorderedScan {
pub(crate) fn new(input: ScanInput) -> Self {
let mut properties = ScannerProperties::default()
.with_append_mode(input.append_mode)
.with_total_rows(input.total_rows());
.with_total_rows(input.total_rows())
.with_total_rows_is_exact(input.append_mode && input.total_rows_is_exact());
if let Some(counters) = input.query_stat_counters.clone() {
properties.set_query_stat_counters(counters);
}
src/mito2/src/region.rs
-1
View File
@@ -242,7 +242,6 @@ impl StagingPartitionInfo {
impl MitoRegion {
/// Returns the current immutable series-index snapshot.
#[allow(dead_code)] // Used by the upcoming query integration.
pub(crate) fn series_index_version(&self) -> Arc<SeriesIndexVersion> {
self.series_index_version_control.current()
}
src/mito2/src/region/options.rs
+2 -1
View File
@@ -121,7 +121,8 @@ pub struct RegionOptions {
/// the configured size; zero disables both limits.
#[serde(skip_serializing_if = "Option::is_none")]
pub write_buffer_size: Option<ReadableSize>,
/// Whether to preserve per-row sequence numbers through flush and compaction.
/// Whether to preserve original per-row sequence numbers when flushing.
/// Compaction always retains effective input sequences, regardless of this option.
///
/// Only meaningful for append-only tables (`append_mode = true`): when enabled,
/// every row keeps its exact sequence number in memtables, flushed SSTs and
src/mito2/src/series_index.rs
+9 -14
View File
@@ -12,42 +12,37 @@
// See the License for the specific language governing permissions and
// limitations under the License.
//! Series index writer and searcher.
//! Series index construction, search, and maintenance.
//!
//! Under development. Index files are currently stored on the local filesystem.
// These components are consumed by the upcoming query and maintenance integration.
#[allow(dead_code)]
mod catalog;
// Consumed by the follow-up background-maintenance integration.
#[allow(dead_code)]
mod bucket;
// Consumed by the follow-up background-maintenance integration.
#[allow(dead_code)]
mod builder;
#[allow(dead_code)]
mod catalog;
mod maintenance;
mod purger;
mod searcher;
mod task;
#[cfg(test)]
mod tests;
#[allow(dead_code)]
mod version;
mod writer;
use futures::stream::BoxStream;
pub use searcher::SeriesIndexSearcher;
use store_api::metric_engine_consts::{
DATA_SCHEMA_TABLE_ID_COLUMN_NAME as TABLE_ID_COLUMN,
DATA_SCHEMA_TSID_COLUMN_NAME as TSID_COLUMN,
};
pub use writer::{
SeriesIndexWriter, SeriesIndexWriterMetrics, SeriesIndexWriterOptions, series_index_schema,
};
use crate::error::Result;
pub(crate) use crate::series_index::catalog::{delete_catalogs, load_version_control};
pub(crate) use crate::series_index::purger::{IndexFilePurger, series_index_channel};
pub use crate::series_index::searcher::SeriesIndexSearcher;
pub(crate) use crate::series_index::task::{SeriesIndexTaskState, spawn_series_index_tasks};
pub(crate) use crate::series_index::version::{SeriesIndexVersion, SeriesIndexVersionControl};
pub use crate::series_index::writer::{
SeriesIndexWriter, SeriesIndexWriterMetrics, SeriesIndexWriterOptions, series_index_schema,
};
pub(crate) const MIN_TS_COLUMN: &str = "__series_min_ts";
pub(crate) const MAX_TS_COLUMN: &str = "__series_max_ts";
src/mito2/src/series_index/maintenance.rs
+319
View File
@@ -0,0 +1,319 @@
// Copyright 2023 Greptime Team
//
// Licensed 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.
//! Worker-owned series-index reconciliation and publication.
use std::collections::HashSet;
use std::sync::Arc;
use std::time::{Duration, Instant};
use common_telemetry::{debug, info};
use object_store::ObjectStore;
use store_api::storage::RegionId;
use crate::error::Result;
use crate::metrics::{SERIES_INDEX_RECONCILE_ELAPSED, SERIES_INDEX_RECONCILE_TOTAL};
use crate::read::series_candidate::is_sparse_metric_metadata;
use crate::region::version::VersionRef;
use crate::region::{MitoRegionRef, RegionLeaderState, RegionRoleState};
use crate::series_index::bucket::{
group_files_into_series_buckets, plan_series_indexes, rounded_bucket_width,
};
use crate::series_index::builder::{build_range_index, build_series_index};
use crate::series_index::catalog::{
RangeIndexCatalog, SeriesIndexCatalog, delete_catalogs, range_catalog_path,
series_catalog_path, store_catalog,
};
use crate::series_index::purger::IndexFilePurger;
use crate::series_index::version::{SeriesIndexFileHandle, SeriesIndexVersion};
/// Retires newly completed series files unless their snapshot is published.
#[derive(Default)]
struct UnpublishedSeriesFiles(Vec<SeriesIndexFileHandle>);
impl UnpublishedSeriesFiles {
fn disarm(&mut self) {
self.0.clear();
}
}
impl Drop for UnpublishedSeriesFiles {
fn drop(&mut self) {
for handle in &self.0 {
handle.mark_deleted();
}
}
}
#[derive(Debug, Default)]
pub(crate) struct ReconcileStats {
pub(crate) source_files: usize,
pub(crate) built_range: usize,
pub(crate) built_series: usize,
pub(crate) removed_range: usize,
pub(crate) removed_series: usize,
pub(crate) computed_buckets: usize,
pub(crate) skipped_buckets: usize,
}
impl ReconcileStats {
fn changed(&self) -> bool {
self.built_range + self.built_series + self.removed_range + self.removed_series > 0
}
}
/// Reconciles indexes for one region snapshot, persists catalogs, then atomically publishes it.
pub(crate) async fn reconcile_series_indexes(
worker_id: u32,
store: ObjectStore,
region: MitoRegionRef,
requested_bucket_width: Duration,
now_ms: i64,
purger: IndexFilePurger,
) -> Result<ReconcileStats> {
let total_start = Instant::now();
// Use this snapshot throughout reconciliation, even if the region version advances.
let version = region.version_control.current().version;
if !is_sparse_metric_metadata(&version.metadata) {
SERIES_INDEX_RECONCILE_TOTAL
.with_label_values(&["noop"])
.inc();
return Ok(ReconcileStats::default());
}
let build_start = Instant::now();
let mut unpublished = UnpublishedSeriesFiles::default();
let (next, mut stats) = build_index_version(
worker_id,
&store,
&region,
&version,
requested_bucket_width,
now_ms,
&purger,
&mut unpublished,
)
.await?;
SERIES_INDEX_RECONCILE_ELAPSED
.with_label_values(&["build"])
.observe(build_start.elapsed().as_secs_f64());
// Persist both catalogs before making the new snapshot visible to readers.
let publish_result: Result<()> = async {
if let Some(next) = next {
persist_index_catalogs(&store, region.region_id, &next).await?;
publish_index_version(&region, Arc::new(next));
unpublished.disarm();
}
Ok(())
}
.await;
// Drop may have cleaned up while catalogs were being written, even if a write failed.
// If dropping starts after this check, the normal drop path cleans up our publication.
// This assumes the region ID is not reopened or replaced during cleanup.
if region.state() == RegionRoleState::Leader(RegionLeaderState::Dropping) {
delete_catalogs(&store, region.region_id).await;
region.series_index_version_control.mark_dropped();
stats = ReconcileStats::default();
}
publish_result?;
let result = if stats.changed() { "changed" } else { "noop" };
SERIES_INDEX_RECONCILE_TOTAL
.with_label_values(&[result])
.inc();
SERIES_INDEX_RECONCILE_ELAPSED
.with_label_values(&["total"])
.observe(total_start.elapsed().as_secs_f64());
if stats.changed() {
info!(
"Reconciled series-index snapshot, worker: {worker_id}, region: {}, elapsed: {:?}, stats: {:?}",
region.region_id,
total_start.elapsed(),
stats
);
} else {
debug!(
"Series-index reconciliation made no changes, worker: {worker_id}, region: {}",
region.region_id
);
}
Ok(stats)
}
/// Builds a changed snapshot, retaining reusable indexes and removing obsolete coverage.
/// Returns `None` when no range or series indexes were added or removed.
#[allow(clippy::too_many_arguments)]
async fn build_index_version(
worker_id: u32,
store: &ObjectStore,
region: &MitoRegionRef,
version: &VersionRef,
requested_bucket_width: Duration,
now_ms: i64,
purger: &IndexFilePurger,
unpublished: &mut UnpublishedSeriesFiles,
) -> Result<(Option<SeriesIndexVersion>, ReconcileStats)> {
let mut stats = ReconcileStats::default();
let files = version
.ssts
.levels()
.iter()
.flat_map(|level| level.files())
.cloned()
.collect::<Vec<_>>();
stats.source_files = files.len();
let visible = files
.iter()
.map(|file| file.file_id().file_id())
.collect::<HashSet<_>>();
let current = region.series_index_version();
// The SST purger deletes companion range files after final handle release. Prune
// metadata here using the captured SST snapshot, independently of physical deletion;
// a later region-version change is picked up by the next reconciliation.
stats.removed_range = current.range_indexes.difference(&visible).count();
let buckets = match version.compaction_time_window {
Some(window) => rounded_bucket_width(requested_bucket_width, window)
// Successful rounding guarantees the window fits in i64; subsecond windows
// use the same one-second minimum as rounded_bucket_width.
.map(|width| {
group_files_into_series_buckets(&files, width, (window.as_secs() as i64).max(1))
})
.unwrap_or_default(),
None => {
debug!(
"Deferring series indexes without compaction window, worker: {worker_id}, region: {}",
region.region_id
);
Vec::new()
}
};
let plan = plan_series_indexes(
buckets,
current.index_buckets.clone(),
version.options.ttl,
now_ms,
);
stats.computed_buckets = plan.computed_buckets;
stats.skipped_buckets = plan.skipped_buckets;
if plan.builds.is_empty()
&& plan.expired_index_ids.is_empty()
&& stats.removed_range == 0
&& current.range_indexes.len() == visible.len()
{
return Ok((None, stats));
}
let mut range_indexes = current.range_indexes.clone();
range_indexes.retain(|file_id| visible.contains(file_id));
let mut series_indexes = current.series_indexes.clone();
for id in plan
.expired_index_ids
.iter()
.chain(&plan.superseded_index_ids)
{
series_indexes.remove(id);
}
for (bucket, expected) in plan.builds {
// Complete companion indexes independently so a failed series build preserves them.
for file in &bucket.files {
if !range_indexes.contains(&file.file_id().file_id())
&& let Some(file_id) =
build_range_index(store, region, version, file.clone()).await?
{
stats.built_range += 1;
range_indexes.insert(file_id);
}
}
let series_handle =
build_series_index(store, region, version, &bucket, &expected, purger).await?;
unpublished.0.push(series_handle.clone());
stats.built_series += 1;
series_indexes.insert(expected.index_uuid, series_handle);
}
// Cover SSTs outside planned aggregate builds, including skipped buckets.
for file in files {
let file_id = file.file_id().file_id();
if range_indexes.contains(&file_id) {
continue;
}
if let Some(file_id) = build_range_index(store, region, version, file).await? {
stats.built_range += 1;
range_indexes.insert(file_id);
}
}
stats.removed_series = current
.series_indexes
.keys()
.filter(|id| !series_indexes.contains_key(id))
.count();
// Bucket reconciliation is speculative until indexes change. Recompute its coverage
// next time rather than replacing the published snapshot on a no-op pass.
if !stats.changed() {
return Ok((None, stats));
}
let next = SeriesIndexVersion {
range_indexes,
series_indexes,
index_buckets: plan.index_buckets,
};
Ok((Some(next), stats))
}
/// Writes catalogs in a stable order; the two writes are not atomic together.
async fn persist_index_catalogs(
store: &ObjectStore,
region_id: RegionId,
next: &SeriesIndexVersion,
) -> Result<()> {
let mut range_entries = next.range_indexes.iter().copied().collect::<Vec<_>>();
range_entries.sort_unstable_by(|left, right| left.as_bytes().cmp(right.as_bytes()));
let mut series_entries = next
.series_indexes
.values()
.map(|handle| handle.entry().clone())
.collect::<Vec<_>>();
series_entries.sort_unstable_by_key(|entry| {
(
entry.bucket_start,
entry.bucket_end,
entry.min_file_sequence,
entry.max_file_sequence,
)
});
store_catalog(
store,
&range_catalog_path(region_id),
&RangeIndexCatalog {
indexes: range_entries,
},
)
.await?;
store_catalog(
store,
&series_catalog_path(region_id),
&SeriesIndexCatalog {
indexes: series_entries,
},
)
.await?;
Ok(())
}
/// Publishes the snapshot and retires series files absent from the new version.
fn publish_index_version(region: &MitoRegionRef, next: Arc<SeriesIndexVersion>) {
let previous = region.series_index_version_control.publish(next.clone());
for (id, handle) in &previous.series_indexes {
if !next.series_indexes.contains_key(id) {
// Purge only after readers release their retained handles.
handle.mark_deleted();
}
}
}
src/mito2/src/series_index/purger.rs
+54 -1
View File
@@ -60,7 +60,7 @@ impl IndexFilePurger {
pub(crate) fn purge(&self, request: PurgeRequest) {
if let Err(error) = self.sender.send(request) {
let store = self.store.clone();
common_runtime::spawn_global(async move {
common_runtime::spawn_compact(async move {
purge_file(&store, error.0).await;
});
}
@@ -108,3 +108,56 @@ pub(crate) fn series_index_channel(
let (sender, receiver) = unbounded_channel();
(IndexFilePurger { store, sender }, receiver)
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use object_store::ObjectStore;
use object_store::services::Memory;
use store_api::storage::{FileId, RegionId};
use super::*;
#[tokio::test]
async fn test_purge_drains_queue_after_last_sender_drops() {
let store = ObjectStore::new(Memory::default()).unwrap();
let (purger, receiver) = series_index_channel(store.clone());
let mut paths = Vec::new();
for _ in 0..3 {
let file_id = RegionFileId::new(RegionId::new(1, 1), FileId::random());
let path = series_index_path(file_id.region_id(), file_id.file_id());
store.write(&path, "index").await.unwrap();
paths.push(path);
purger.purge(PurgeRequest { file_id });
}
drop(purger);
tokio::time::timeout(
Duration::from_secs(10),
run_index_purge_task(0, store.clone(), receiver),
)
.await
.unwrap();
for path in paths {
assert!(!store.exists(&path).await.unwrap());
}
}
#[tokio::test]
async fn test_purge_falls_back_after_receiver_drops() {
let store = ObjectStore::new(Memory::default()).unwrap();
let (purger, receiver) = series_index_channel(store.clone());
drop(receiver);
let file_id = RegionFileId::new(RegionId::new(1, 1), FileId::random());
let path = series_index_path(file_id.region_id(), file_id.file_id());
store.write(&path, "index").await.unwrap();
purger.purge(PurgeRequest { file_id });
tokio::time::timeout(Duration::from_secs(10), async {
while store.exists(&path).await.unwrap() {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
}
}
src/mito2/src/series_index/task.rs
+128 -13
View File
@@ -18,14 +18,18 @@ use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use common_telemetry::info;
use common_telemetry::{info, warn};
use object_store::ObjectStore;
use tokio::sync::Notify;
use tokio::sync::mpsc::UnboundedReceiver;
use tokio::task::JoinHandle;
use tokio::time::{Instant, MissedTickBehavior};
use crate::series_index::purger::{PurgeRequest, run_index_purge_task};
use crate::metrics::SERIES_INDEX_RECONCILE_TOTAL;
use crate::region::{RegionLeaderState, RegionMapRef, RegionRoleState};
use crate::series_index::maintenance::reconcile_series_indexes;
use crate::series_index::purger::{IndexFilePurger, PurgeRequest, run_index_purge_task};
use crate::time_provider::TimeProviderRef;
/// Shared lifecycle state for a worker's series-index task.
#[derive(Debug)]
@@ -46,6 +50,10 @@ impl SeriesIndexTaskState {
self.running.load(Ordering::Acquire)
}
pub(crate) fn wake(&self) {
self.notify.notify_one();
}
pub(crate) fn stop(&self) {
self.running.store(false, Ordering::Release);
// Retain a permit if maintenance has not started waiting yet.
@@ -57,21 +65,35 @@ impl SeriesIndexTaskState {
}
}
/// Starts both tasks, detaching purge and returning the maintenance handle.
/// Starts both tasks on the compaction runtime, detaching purge and returning the maintenance handle.
#[allow(clippy::too_many_arguments)]
pub(crate) fn spawn_series_index_tasks(
worker_id: u32,
store: ObjectStore,
regions: RegionMapRef,
state: Arc<SeriesIndexTaskState>,
bucket_width: Duration,
purger: IndexFilePurger,
purge_receiver: UnboundedReceiver<PurgeRequest>,
interval: Duration,
time_provider: TimeProviderRef,
) -> JoinHandle<()> {
// Snapshots may retain senders after the worker stops; purge until all senders drop.
common_runtime::spawn_global(run_index_purge_task(worker_id, store, purge_receiver));
common_runtime::spawn_global(async move {
common_runtime::spawn_compact(run_index_purge_task(
worker_id,
store.clone(),
purge_receiver,
));
common_runtime::spawn_compact(async move {
SeriesIndexTask {
worker_id,
store,
regions,
bucket_width,
purger,
state,
interval,
time_provider,
}
.run()
.await;
@@ -80,9 +102,14 @@ pub(crate) fn spawn_series_index_tasks(
/// Periodic series-index maintenance for one region worker.
struct SeriesIndexTask {
store: ObjectStore,
regions: RegionMapRef,
bucket_width: Duration,
purger: IndexFilePurger,
worker_id: u32,
state: Arc<SeriesIndexTaskState>,
interval: Duration,
time_provider: TimeProviderRef,
}
impl SeriesIndexTask {
@@ -90,17 +117,17 @@ impl SeriesIndexTask {
async fn run(mut self) {
let worker_id = self.worker_id;
info!("Start series-index background task, worker: {worker_id}");
let interval = self.interval;
let interval = self.time_provider.wait_duration(self.interval);
let mut timer = tokio::time::interval_at(Instant::now() + interval, interval);
timer.set_missed_tick_behavior(MissedTickBehavior::Skip);
// Schedule future ticks from a late tick rather than the original cadence.
timer.set_missed_tick_behavior(MissedTickBehavior::Delay);
while self.state.is_running() {
tokio::select! {
_ = self.state.notified() => {}
_ = timer.tick() => {
if self.state.is_running() {
self.maintain().await;
}
}
_ = timer.tick() => {}
}
if self.state.is_running() {
self.maintain().await;
}
}
info!("Stop series-index background task, worker: {worker_id}");
@@ -108,6 +135,94 @@ impl SeriesIndexTask {
/// Runs periodic maintenance independently of incoming deletion requests.
async fn maintain(&mut self) {
// TODO: Reconcile indexes and perform other periodic maintenance here.
for region in self.regions.list_regions() {
if !self.state.is_running() {
break;
}
// Best effort: the region can still change state during reconciliation.
// Local indexes can be built on followers as well as writable leaders.
if !matches!(
region.state(),
RegionRoleState::Follower | RegionRoleState::Leader(RegionLeaderState::Writable)
) {
continue;
}
if let Err(error) = reconcile_series_indexes(
self.worker_id,
self.store.clone(),
region.clone(),
self.bucket_width,
self.time_provider.current_time_millis(),
self.purger.clone(),
)
.await
{
SERIES_INDEX_RECONCILE_TOTAL
.with_label_values(&["failure"])
.inc();
warn!(error; "Failed to reconcile series indexes, worker: {}, region: {}", self.worker_id, region.region_id);
}
}
}
}
#[cfg(test)]
mod tests {
use object_store::services::Memory;
use store_api::region_engine::RegionEngine;
use super::*;
use crate::region::RegionMap;
use crate::series_index::catalog::{range_catalog_path, series_catalog_path};
use crate::series_index::purger::series_index_channel;
use crate::series_index::tests::prepare_region;
use crate::test_util::TestEnv;
#[rstest::rstest]
#[case::follower(RegionRoleState::Follower, true)]
#[case::writable(RegionRoleState::Leader(RegionLeaderState::Writable), true)]
#[case::staging(RegionRoleState::Leader(RegionLeaderState::Staging), false)]
#[case::entering_staging(RegionRoleState::Leader(RegionLeaderState::EnteringStaging), false)]
#[case::altering(RegionRoleState::Leader(RegionLeaderState::Altering), false)]
#[case::dropping(RegionRoleState::Leader(RegionLeaderState::Dropping), false)]
#[case::truncating(RegionRoleState::Leader(RegionLeaderState::Truncating), false)]
#[case::editing(RegionRoleState::Leader(RegionLeaderState::Editing), false)]
#[case::downgrading(RegionRoleState::Leader(RegionLeaderState::Downgrading), false)]
#[tokio::test]
async fn test_maintenance_region_states(#[case] role: RegionRoleState, #[case] builds: bool) {
let mut env = TestEnv::with_prefix("series-maintenance-state").await;
let (engine, region) = prepare_region(&mut env).await;
// Install the desired state without triggering the corresponding DDL.
region.switch_state_to_staging(RegionLeaderState::Writable);
region
.manifest_ctx
.exit_staging(region.region_id, role)
.unwrap();
let store = ObjectStore::new(Memory::default()).unwrap();
let (purger, _receiver) = series_index_channel(store.clone());
let regions = Arc::new(RegionMap::default());
regions.insert_region(region.clone());
let mut task = SeriesIndexTask {
store: store.clone(),
regions,
bucket_width: Duration::from_secs(100),
purger,
worker_id: 0,
state: Arc::new(SeriesIndexTaskState::new()),
interval: Duration::from_secs(3600),
time_provider: Arc::new(crate::time_provider::StdTimeProvider),
};
task.maintain().await;
assert_eq!(
builds,
!region.series_index_version().series_indexes.is_empty()
);
for path in [
range_catalog_path(region.region_id),
series_catalog_path(region.region_id),
] {
assert_eq!(builds, store.exists(&path).await.unwrap());
}
engine.stop().await.unwrap();
}
}
src/mito2/src/series_index/tests.rs
+693 -7
View File
@@ -12,25 +12,36 @@
// See the License for the specific language governing permissions and
// limitations under the License.
//! Shared fixtures for series-index construction tests.
//! Behavioral coverage for series-index reconciliation.
use std::sync::Arc;
use std::time::Duration;
use api::v1::helper::row;
use api::v1::value::ValueData;
use api::v1::{ColumnDataType, Rows, SemanticType, WriteHint};
use object_store::ObjectStore;
use object_store::layers::mock::{self, MockLayerBuilder, oio};
use object_store::services::Memory;
use store_api::codec::PrimaryKeyEncoding;
use store_api::metric_engine_consts::PRIMARY_KEY_ENCODING;
use store_api::region_engine::RegionEngine;
use store_api::region_request::{RegionPutRequest, RegionRequest};
use store_api::storage::RegionId;
use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
use store_api::storage::{FileId, RegionId};
use super::catalog::{
delete_catalogs, load_version_control, range_catalog_path, range_index_path,
series_catalog_path, series_index_path,
};
use super::maintenance::reconcile_series_indexes;
use super::purger::series_index_channel;
use crate::config::MitoConfig;
use crate::engine::MitoEngine;
use crate::region::MitoRegionRef;
use crate::region::{MitoRegionRef, RegionMap};
use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding};
use crate::test_util::{CreateRequestBuilder, TestEnv, flush_region, rows_schema};
use crate::time_provider::StdTimeProvider;
/// Builds real sparse SSTs; background maintenance is disabled so tests control publication.
pub(super) async fn prepare_region(env: &mut TestEnv) -> (MitoEngine, MitoRegionRef) {
@@ -66,10 +77,12 @@ async fn prepare_region_with_timestamps(
"100s".to_string(),
);
// Keep the source SSTs stable while tests reconcile multiple buckets.
request.options.insert(
"compaction.twcs.trigger_file_num".to_string(),
"100".to_string(),
);
for window in ["active_window", "inactive_window"] {
request.options.insert(
format!("compaction.twcs.{window}.trigger_file_num"),
"100".to_string(),
);
}
let full_schema = rows_schema(&request);
let mut pk_column = full_schema[0].clone();
pk_column.column_name = PRIMARY_KEY_COLUMN_NAME.to_string();
@@ -112,3 +125,676 @@ async fn prepare_region_with_timestamps(
let region = engine.get_region(region_id).unwrap();
(engine, region)
}
struct IndexTest {
region: MitoRegionRef,
store: ObjectStore,
purger: super::purger::IndexFilePurger,
}
impl IndexTest {
fn new(
region: MitoRegionRef,
) -> (
Self,
tokio::sync::mpsc::UnboundedReceiver<super::purger::PurgeRequest>,
) {
let store = ObjectStore::new(Memory::default()).unwrap();
let (purger, receiver) = series_index_channel(store.clone());
(
Self {
region,
store,
purger,
},
receiver,
)
}
async fn reconcile(
&self,
store: ObjectStore,
) -> crate::error::Result<super::maintenance::ReconcileStats> {
reconcile_series_indexes(
0,
store,
self.region.clone(),
Duration::from_secs(100),
0,
self.purger.clone(),
)
.await
}
}
#[tokio::test]
async fn test_reconcile_restores_and_reuses_indexes() {
let mut env = TestEnv::with_prefix("series-reconcile").await;
let (engine, region) = prepare_region(&mut env).await;
let (test, _receiver) = IndexTest::new(region.clone());
let store = &test.store;
let purger = &test.purger;
let stats = test.reconcile(store.clone()).await.unwrap();
assert_eq!((4, 1), (stats.built_range, stats.built_series));
let first = region.series_index_version();
let first_id = *first.series_indexes.keys().next().unwrap();
let missing_paths = [
range_index_path(
region.region_id,
*first.range_indexes.iter().next().unwrap(),
),
series_index_path(region.region_id, first_id),
];
for path in &missing_paths {
store.delete(path).await.unwrap();
}
// Restore catalog entries even when their index files are missing.
let restored = load_version_control(store, region.region_id, purger)
.await
.current();
assert_eq!(first.range_indexes, restored.range_indexes);
assert_eq!(first.index_buckets, restored.index_buckets);
assert_eq!(
first.series_indexes[&first_id].entry(),
restored.series_indexes[&first_id].entry()
);
// Catalog changes are not reloaded, and missing index files are not repaired.
for path in [
range_catalog_path(region.region_id),
series_catalog_path(region.region_id),
] {
store.write(&path, "{}").await.unwrap();
}
let stats = test.reconcile(store.clone()).await.unwrap();
assert_eq!((0, 0), (stats.built_range, stats.built_series));
assert!(Arc::ptr_eq(&first, &region.series_index_version()));
for path in [
range_catalog_path(region.region_id),
series_catalog_path(region.region_id),
] {
assert_eq!(
b"{}".as_slice(),
store.read(&path).await.unwrap().to_bytes().as_ref()
);
}
for path in &missing_paths {
assert!(!store.exists(path).await.unwrap());
}
engine.stop().await.unwrap();
}
#[tokio::test]
async fn test_reconcile_range_only_build_and_removal() {
let mut env = TestEnv::with_prefix("series-range-lifecycle").await;
let (engine, region) = prepare_region_with_timestamps(&mut env, &[1000]).await;
let (test, _receiver) = IndexTest::new(region.clone());
let initial = region.series_index_version();
let stats = test.reconcile(test.store.clone()).await.unwrap();
assert_eq!((1, 0), (stats.built_range, stats.built_series));
let built = region.series_index_version();
assert!(!Arc::ptr_eq(&initial, &built));
let files = region
.version()
.ssts
.levels()
.iter()
.flat_map(|level| level.files())
.map(|file| file.meta_ref().clone())
.collect();
region.version_control.apply_edit(
Some(crate::manifest::action::RegionEdit {
files_to_remove: files,
files_to_add: Vec::new(),
timestamp_ms: None,
compaction_time_window: None,
flushed_entry_id: None,
flushed_sequence: None,
committed_sequence: None,
}),
&[],
crate::test_util::new_noop_file_purger(),
);
let stats = test.reconcile(test.store.clone()).await.unwrap();
assert_eq!(
(0, 0, 1),
(stats.built_range, stats.built_series, stats.removed_range)
);
let removed = region.series_index_version();
assert!(removed.range_indexes.is_empty());
assert!(!Arc::ptr_eq(&built, &removed));
let restored = load_version_control(&test.store, region.region_id, &test.purger).await;
assert!(restored.current().range_indexes.is_empty());
test.reconcile(test.store.clone()).await.unwrap();
assert!(Arc::ptr_eq(&removed, &region.series_index_version()));
engine.stop().await.unwrap();
}
struct FailingSeriesWriter {
inner: oio::Writer,
fail: bool,
}
#[rstest::rstest]
#[case::series("series")]
#[case::range_catalog("range-index.json")]
#[case::series_catalog("series-index.json")]
#[tokio::test]
async fn test_reconcile_replaces_whole_bucket_after_failure(#[case] failure: &str) {
let mut env = TestEnv::with_prefix("series-replacement").await;
let (engine, region) =
prepare_region_with_timestamps(&mut env, &[1000, 2000, 3000, 4000, 5000]).await;
// Retain all source handles while hiding the newest SST for the initial build.
let sources = region.version();
let newest = sources
.ssts
.levels()
.iter()
.flat_map(|level| level.files())
.max_by_key(|file| file.meta_ref().sequence)
.unwrap()
.meta_ref()
.clone();
let edit = |add| crate::manifest::action::RegionEdit {
files_to_add: if add { vec![newest.clone()] } else { vec![] },
files_to_remove: if add { vec![] } else { vec![newest.clone()] },
timestamp_ms: None,
compaction_time_window: None,
flushed_entry_id: None,
flushed_sequence: None,
committed_sequence: None,
};
region.version_control.apply_edit(
Some(edit(false)),
&[],
crate::test_util::new_noop_file_purger(),
);
let (test, mut receiver) = IndexTest::new(region.clone());
test.reconcile(test.store.clone()).await.unwrap();
let previous = region.series_index_version();
let old_id = *previous.series_indexes.keys().next().unwrap();
assert_eq!(
4,
previous.series_indexes[&old_id]
.entry()
.source_file_ids
.len()
);
region.version_control.apply_edit(
Some(edit(true)),
&[],
crate::test_util::new_noop_file_purger(),
);
let failure = failure.to_string();
let layer = MockLayerBuilder::default()
.writer_factory(Arc::new(move |path, _, inner| {
Box::new(FailingSeriesWriter {
inner,
fail: if failure == "series" {
path.contains("/series/")
} else {
path.ends_with(&failure)
},
})
}))
.build()
.unwrap();
assert!(
test.reconcile(test.store.clone().layer(layer))
.await
.is_err()
);
assert!(Arc::ptr_eq(&previous, &region.series_index_version()));
// Failed catalog publication may retire a completed replacement, never its predecessor.
while let Ok(request) = receiver.try_recv() {
assert_ne!(old_id, request.file_id.file_id());
}
let stats = test.reconcile(test.store.clone()).await.unwrap();
assert_eq!((1, 1), (stats.built_series, stats.removed_series));
let replacement = region.series_index_version();
assert_eq!(1, replacement.series_indexes.len());
assert!(!replacement.series_indexes.contains_key(&old_id));
let entry = replacement.series_indexes.values().next().unwrap().entry();
assert_eq!(5, entry.source_file_ids.len());
assert!(entry.source_file_ids.contains(&newest.file_id));
let restored = load_version_control(&test.store, region.region_id, &test.purger).await;
assert_eq!(replacement.index_buckets, restored.current().index_buckets);
assert_eq!(1, restored.current().series_indexes.len());
test.reconcile(test.store.clone()).await.unwrap();
assert!(Arc::ptr_eq(&replacement, &region.series_index_version()));
// Publication retires the old file only after the last reader releases its snapshot.
let old_path = series_index_path(region.region_id, old_id);
assert!(test.store.exists(&old_path).await.unwrap());
assert!(receiver.try_recv().is_err());
drop(previous);
assert_eq!(old_id, receiver.try_recv().unwrap().file_id.file_id());
assert!(receiver.try_recv().is_err());
engine.stop().await.unwrap();
}
impl mock::Write for FailingSeriesWriter {
async fn write(&mut self, buffer: mock::Buffer) -> mock::Result<()> {
self.inner.write(buffer).await
}
async fn close(&mut self) -> mock::Result<mock::Metadata> {
if self.fail {
return Err(mock::Error::new(
mock::ErrorKind::Unexpected,
"injected series write failure",
));
}
self.inner.close().await
}
async fn abort(&mut self) -> mock::Result<()> {
self.inner.abort().await
}
}
#[tokio::test]
async fn test_failed_series_build_keeps_completed_sst_range_indexes() {
let mut env = TestEnv::with_prefix("series-build-failure").await;
let (engine, region) = prepare_region(&mut env).await;
let (test, mut receiver) = IndexTest::new(region.clone());
let store = &test.store;
let layer = MockLayerBuilder::default()
.writer_factory(Arc::new(|path, _, inner| {
Box::new(FailingSeriesWriter {
inner,
fail: path.contains("/series/"),
})
}))
.build()
.unwrap();
let failing_store = store.clone().layer(layer);
assert!(test.reconcile(failing_store).await.is_err());
assert!(region.series_index_version().range_indexes.is_empty());
assert!(receiver.try_recv().is_err());
for file in region
.version()
.ssts
.levels()
.iter()
.flat_map(|level| level.files())
{
let path = range_index_path(region.region_id, file.file_id().file_id());
assert!(store.exists(&path).await.unwrap());
}
// A retry can rebuild and publish the complete snapshot.
let stats = test.reconcile(store.clone()).await.unwrap();
assert_eq!((4, 1), (stats.built_range, stats.built_series));
engine.stop().await.unwrap();
}
#[tokio::test]
async fn test_reconcile_publishes_after_region_version_changes() {
for during_catalog_write in [false, true] {
let mut env = TestEnv::with_prefix("series-version-change").await;
let (engine, region) = prepare_region(&mut env).await;
let initial_version = region.version();
let (test, mut receiver) = IndexTest::new(region.clone());
let store = &test.store;
let purger = &test.purger;
let target_region = region.clone();
let catalog_path = range_catalog_path(region.region_id);
let layer = MockLayerBuilder::default()
.writer_factory(Arc::new(move |path, _, inner| {
let should_advance = if during_catalog_write {
path == catalog_path
} else {
path.contains("/series/")
};
if should_advance {
// Replace the version while reconciliation is writing its captured snapshot.
target_region
.version_control
.alter_options(target_region.version().options.clone());
}
inner
}))
.build()
.unwrap();
let stats = test.reconcile(store.clone().layer(layer)).await.unwrap();
assert!(!Arc::ptr_eq(&initial_version, &region.version()));
assert_eq!((4, 1), (stats.built_range, stats.built_series));
let published = region.series_index_version();
assert_eq!(4, published.range_indexes.len());
assert_eq!(1, published.series_indexes.len());
let restored = load_version_control(store, region.region_id, purger).await;
assert_eq!(published.range_indexes, restored.current().range_indexes);
let handle = published.series_indexes.values().next().unwrap();
assert_eq!(
handle.entry(),
restored.current().series_indexes[&handle.entry().index_uuid].entry()
);
assert!(receiver.try_recv().is_err());
engine.stop().await.unwrap();
}
}
#[tokio::test]
async fn test_reconcile_cleans_up_after_region_drop() {
let mut env = TestEnv::with_prefix("series-reconcile-drop").await;
let (engine, region) = prepare_region(&mut env).await;
let (test, mut receiver) = IndexTest::new(region.clone());
let store = test.store.clone();
let target_region = region.clone();
let catalog_path = range_catalog_path(region.region_id);
let layer = MockLayerBuilder::default()
.writer_factory(Arc::new(move |path, _, inner| {
if path == catalog_path {
// Simulate drop cleanup after building indexes but before persisting the
// first catalog. Neither catalog exists yet, so deletion is a no-op.
target_region
.set_dropping(crate::region::RegionLeaderState::Writable)
.unwrap();
target_region.version_control.mark_dropped();
target_region.series_index_version_control.mark_dropped();
}
inner
}))
.build()
.unwrap();
let stats = test.reconcile(store.clone().layer(layer)).await.unwrap();
assert_eq!((0, 0), (stats.built_range, stats.built_series));
for path in [
range_catalog_path(region.region_id),
series_catalog_path(region.region_id),
] {
assert!(!store.exists(&path).await.unwrap());
}
let published = region.series_index_version();
assert!(published.range_indexes.is_empty());
assert!(published.series_indexes.is_empty());
let retired = receiver.try_recv().unwrap();
assert!(receiver.try_recv().is_err());
let series_path = series_index_path(region.region_id, retired.file_id.file_id());
test.purger.purge(retired);
drop(test);
super::purger::run_index_purge_task(0, store.clone(), receiver).await;
assert!(!store.exists(&series_path).await.unwrap());
engine.stop().await.unwrap();
}
#[tokio::test]
async fn test_failed_reconcile_retires_only_unpublished_series() {
use std::collections::{HashMap, HashSet};
use std::sync::Mutex;
use super::bucket::{group_files_into_series_buckets, plan_series_indexes};
use super::builder::{build_range_index, build_series_index};
use super::purger::run_index_purge_task;
use super::version::SeriesIndexVersion;
for failure in ["series", "range", "range-index.json", "series-index.json"] {
let mut env = TestEnv::with_prefix("series-unpublished-cleanup").await;
let mut timestamps = vec![
-99000, -98000, -97000, -96000, 1000, 2000, 3000, 4000, 101000,
];
if failure != "range" {
timestamps.extend([102000, 103000, 104000]);
}
let (engine, region) = prepare_region_with_timestamps(&mut env, &timestamps).await;
let (test, mut receiver) = IndexTest::new(region.clone());
let store = &test.store;
let purger = &test.purger;
// Publish the first bucket so failed attempts must preserve a reused handle.
let version = region.version();
let files = version
.ssts
.levels()
.iter()
.flat_map(|level| level.files())
.cloned()
.collect::<Vec<_>>();
let plan = plan_series_indexes(
group_files_into_series_buckets(
&files,
100,
version.compaction_time_window.unwrap().as_secs() as i64,
),
Default::default(),
None,
0,
);
assert_eq!(if failure == "range" { 2 } else { 3 }, plan.builds.len());
let (bucket, entry) = &plan.builds[0];
let mut ranges = HashSet::new();
for file in &bucket.files {
let id = build_range_index(store, &region, &version, file.clone())
.await
.unwrap()
.unwrap();
ranges.insert(id);
}
let handle = build_series_index(store, &region, &version, bucket, entry, purger)
.await
.unwrap();
let reused_path = series_index_path(region.region_id, entry.index_uuid);
region
.series_index_version_control
.publish(Arc::new(SeriesIndexVersion::new(
ranges,
HashMap::from([(entry.index_uuid, handle)]),
)));
let previous = region.series_index_version();
let standalone_path = files
.iter()
.find(|file| file.time_range().0 == common_time::Timestamp::new_millisecond(101000))
.map(|file| range_index_path(region.region_id, file.file_id().file_id()))
.unwrap();
// Every completed unpublished output must be retired.
let attempted = Arc::new(Mutex::new(Vec::<String>::new()));
let recorded = attempted.clone();
let layer = MockLayerBuilder::default()
.writer_factory(Arc::new(move |path, _, inner| {
let mut paths = recorded.lock().unwrap();
if path.contains("/series/") {
paths.push(path.to_string());
}
let fail = match failure {
"series" => path.contains("/series/") && paths.len() == 2,
"range" => path == standalone_path,
catalog => path.ends_with(catalog),
};
Box::new(FailingSeriesWriter { inner, fail })
}))
.build()
.unwrap();
let result = test.reconcile(store.clone().layer(layer)).await;
assert!(
result.is_err(),
"expected {failure} failure, got {result:?}; attempted series paths: {:?}",
attempted.lock().unwrap()
);
assert!(Arc::ptr_eq(&previous, &region.series_index_version()));
let completed = attempted
.lock()
.unwrap()
.iter()
.take(if failure.ends_with(".json") { 2 } else { 1 })
.cloned()
.collect::<HashSet<_>>();
assert_eq!(
if failure.ends_with(".json") { 2 } else { 1 },
completed.len()
);
let mut retired = HashSet::new();
// Drain through the real purge task using a finite channel.
let (drain_purger, drain_receiver) = series_index_channel(store.clone());
while let Ok(request) = receiver.try_recv() {
let path = series_index_path(request.file_id.region_id(), request.file_id.file_id());
assert!(store.exists(&path).await.unwrap());
assert!(retired.insert(path));
drain_purger.purge(request);
}
assert_eq!(completed, retired);
drop(drain_purger);
run_index_purge_task(0, store.clone(), drain_receiver).await;
let attempted_paths = attempted.lock().unwrap().clone();
for path in attempted_paths {
assert!(!store.exists(&path).await.unwrap());
}
assert!(store.exists(&reused_path).await.unwrap());
let stats = test.reconcile(store.clone()).await.unwrap();
assert_eq!(if failure == "range" { 1 } else { 2 }, stats.built_series);
let published = region.series_index_version();
for handle in published.series_indexes.values() {
assert!(
store
.exists(&series_index_path(
region.region_id,
handle.entry().index_uuid
))
.await
.unwrap()
);
}
// Release all snapshots without retiring them: neither reused nor published files
// should have been marked deleted by the cleanup guard.
region.series_index_version_control.publish(Arc::default());
drop(published);
drop(previous);
assert!(receiver.try_recv().is_err());
engine.stop().await.unwrap();
}
}
/// Waits for a background state change without relying on exact scheduling delays.
async fn wait_for(mut ready: impl FnMut() -> bool) {
tokio::time::timeout(Duration::from_secs(10), async {
while !ready() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.unwrap();
}
#[tokio::test]
async fn test_maintenance_wakeup_and_timer() {
let mut env = TestEnv::with_prefix("series-task").await;
let (engine, region) = prepare_region(&mut env).await;
let mut options = region.version().options.clone();
options.ttl = Some(common_time::TimeToLive::Duration(Duration::from_secs(100)));
region.version_control.alter_options(options);
let store = ObjectStore::new(Memory::default()).unwrap();
let (purger, receiver) = series_index_channel(store.clone());
let state = Arc::new(super::task::SeriesIndexTaskState::new());
let clock = Arc::new(crate::time_provider::mock::MockTimeProvider::new(0));
// A notification issued before the task starts must also trigger maintenance.
state.wake();
let regions = Arc::new(RegionMap::default());
regions.insert_region(region.clone());
let task = super::task::spawn_series_index_tasks(
0,
store,
regions,
state.clone(),
Duration::from_secs(100),
purger,
receiver,
Duration::from_secs(3600),
clock.clone(),
);
wait_for(|| !region.series_index_version().series_indexes.is_empty()).await;
// Advance event time; periodic maintenance must expire coverage without a wakeup.
clock.set_now(201_000);
wait_for(|| region.series_index_version().series_indexes.is_empty()).await;
state.stop();
task.await.unwrap();
engine.stop().await.unwrap();
}
#[tokio::test]
async fn test_drop_catalogs_and_retained_snapshot_after_task_stop() {
use std::collections::BTreeMap;
use common_time::Timestamp;
use super::catalog::{SeriesIndexCatalog, SeriesIndexEntry, WindowSequence, store_catalog};
let store = ObjectStore::new(Memory::default()).unwrap();
let region_id = RegionId::new(1, 1);
let entry = SeriesIndexEntry {
index_uuid: FileId::random(),
bucket_start: Timestamp::new_second(0),
bucket_end: Timestamp::new_second(100),
source_file_ids: vec![FileId::random()],
min_file_sequence: 1,
max_file_sequence: 2,
compaction_window_secs: 10,
window_sequences: BTreeMap::from([(
0,
WindowSequence {
start: 0,
end: 100,
max_sequence: 2,
},
)]),
};
let path = series_index_path(region_id, entry.index_uuid);
store.write(&path, "index").await.unwrap();
store_catalog(
&store,
&series_catalog_path(region_id),
&SeriesIndexCatalog {
indexes: vec![entry.clone()],
},
)
.await
.unwrap();
store_catalog(
&store,
&range_catalog_path(region_id),
&super::catalog::RangeIndexCatalog::default(),
)
.await
.unwrap();
let (purger, receiver) = series_index_channel(store.clone());
let control = load_version_control(&store, region_id, &purger).await;
let snapshot = control.current();
assert_eq!(&entry, snapshot.series_indexes[&entry.index_uuid].entry());
let state = Arc::new(super::task::SeriesIndexTaskState::new());
state.stop();
super::task::spawn_series_index_tasks(
0,
store.clone(),
Arc::new(RegionMap::default()),
state,
Duration::from_secs(100),
purger,
receiver,
Duration::from_secs(3600),
Arc::new(StdTimeProvider),
)
.await
.unwrap();
delete_catalogs(&store, region_id).await;
// Removing already absent catalogs is harmless.
delete_catalogs(&store, region_id).await;
assert!(!store.exists(&range_catalog_path(region_id)).await.unwrap());
assert!(!store.exists(&series_catalog_path(region_id)).await.unwrap());
control.mark_dropped();
assert!(store.exists(&path).await.unwrap());
drop(snapshot);
tokio::time::timeout(Duration::from_secs(10), async {
while store.exists(&path).await.unwrap() {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
}
src/mito2/src/sst.rs
+1 -3
View File
@@ -1079,9 +1079,7 @@ mod tests {
&builder.parquet_schema().root_schema().get_fields()[0].get_fields()[0];
assert_eq!(
parquet_remainder.get_basic_info().logical_type_ref(),
Some(&LogicalType::Variant {
specification_version: None,
})
Some(&LogicalType::variant(None))
);
let ArrowDataType::Struct(children) = builder.schema().field_with_name("data")?.data_type()
src/mito2/src/sst/file.rs
+8 -3
View File
@@ -243,10 +243,13 @@ pub struct FileMeta {
/// the default value `0` doesn't means the file doesn't contains any rows,
/// but instead means the number of rows is unknown.
pub num_row_groups: u64,
/// Sequence in this file.
/// File-level sequence bound or admission marker in the target region.
///
/// This sequence is the only sequence in this file. And it's retrieved from the max
/// sequence of the rows on generating this file.
/// Flush records the maximum input row sequence. Compaction outputs inherit
/// the maximum input bound, or remain unknown if any input bound is unknown,
/// independently of per-row sequence trust. This does not imply that every
/// physical row has this sequence.
/// Readers also use it to normalize foreign and legacy all-zero files.
pub sequence: Option<NonZeroU64>,
/// Partition expression from the region metadata when the file is created.
///
@@ -284,6 +287,8 @@ pub struct FileMeta {
pub primary_key_max: Option<Bytes>,
/// Whether the file preserves per-row sequence numbers usable for exact
/// row-level sequence filtering.
/// Merely retaining physical input sequences during compaction does not
/// restore this capability for untrusted inputs.
#[serde(default, skip_serializing_if = "is_false")]
pub preserve_row_sequence: bool,
}
src/mito2/src/sst/file_purger.rs
+32 -5
View File
@@ -288,8 +288,14 @@ mod tests {
use crate::sst::index::puffin_manager::PuffinManagerFactory;
use crate::sst::location;
#[rstest::rstest]
#[case(false)]
#[case(true)]
#[tokio::test]
async fn test_file_purge() {
async fn test_range_index_purge_on_handle_release(
#[case] gc_enabled: bool,
#[values(false, true)] is_delete: bool,
) {
common_telemetry::init_default_ut_logging();
let dir = create_temp_dir("file-purge");
@@ -320,7 +326,19 @@ mod tests {
let scheduler = Arc::new(LocalScheduler::new(3));
let file_purger = Arc::new(LocalFilePurger::new(scheduler.clone(), layer, None));
let index_store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
let owner = RegionId::new(9, 1);
let index_path = crate::sst::range_index::range_index_path(owner, sst_file_id.file_id());
index_store.write(&index_path, "range index").await.unwrap();
let file_purger = create_file_purger(
gc_enabled,
PathType::Bare,
scheduler.clone(),
layer,
None,
Arc::new(crate::sst::file_ref::FileReferenceManager::new(None)),
Some(RangeIndexDeleter::new(index_store.clone(), owner)),
);
{
let handle = FileHandle::new(
@@ -344,13 +362,22 @@ mod tests {
},
file_purger,
);
// mark file as deleted and drop the handle, we expect the file is deleted.
handle.mark_deleted();
if is_delete {
handle.mark_deleted();
}
let reader = handle.clone();
drop(handle);
assert!(index_store.exists(&index_path).await.unwrap());
drop(reader);
}
scheduler.stop(true).await.unwrap();
assert!(!object_store.exists(&path).await.unwrap());
assert_eq!(
object_store.exists(&path).await.unwrap(),
gc_enabled || !is_delete
);
assert_eq!(index_store.exists(&index_path).await.unwrap(), !is_delete);
}
#[tokio::test]
src/mito2/src/sst/parquet/index_reader.rs
+1 -1
View File
@@ -172,7 +172,7 @@ impl PruningStatistics for IndexRowGroupPruningStats<'_> {
column_null_counts(self.row_groups, column_index)
}
fn row_counts(&self, _column: &Column) -> Option<ArrayRef> {
fn row_counts(&self) -> Option<ArrayRef> {
None
}
src/mito2/src/sst/parquet/json_align/stream.rs
+1 -1
View File
@@ -247,7 +247,7 @@ fn align_array(
return Ok(array.clone());
}
cast_column(array, field.as_ref(), &DEFAULT_CAST_OPTIONS).context(CastColumnSnafu)
cast_column(array, field.data_type(), &DEFAULT_CAST_OPTIONS).context(CastColumnSnafu)
}
#[cfg(test)]
src/mito2/src/sst/parquet/read_columns.rs
+3 -3
View File
@@ -386,7 +386,7 @@ fn is_variant_leaf(leaf_col: &ColumnDescriptor) -> bool {
mod tests {
use std::sync::Arc;
use parquet::basic::{ConvertedType, LogicalType, Repetition};
use parquet::basic::{ConvertedType, LogicalType, Repetition, VariantType};
use parquet::errors::ParquetError;
use parquet::schema::types::Type;
@@ -816,9 +816,9 @@ mod tests {
let remainder = Arc::new(
Type::group_type_builder(JSON2_REMAINDER_FIELD_NAME)
.with_repetition(Repetition::OPTIONAL)
.with_logical_type(Some(LogicalType::Variant {
.with_logical_type(Some(LogicalType::Variant(VariantType {
specification_version: None,
}))
})))
.with_fields(vec![metadata, value])
.build()?,
);
src/mito2/src/sst/parquet/reader.rs
-9
View File
@@ -2556,7 +2556,6 @@ impl FlatRowGroupReader {
#[cfg(test)]
mod tests {
use std::any::Any;
use std::collections::HashMap;
use std::fmt::{Debug, Formatter};
use std::sync::{Arc, LazyLock};
@@ -3048,10 +3047,6 @@ mod tests {
}
impl ScalarUDFImpl for PanicDebugUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"panic_debug_udf"
}
@@ -3150,10 +3145,6 @@ mod tests {
}
impl ScalarUDFImpl for TestVolatilityUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
&self.name
}
src/mito2/src/sst/parquet/stats.rs
+1 -1
View File
@@ -174,7 +174,7 @@ impl<T: Borrow<RowGroupMetaData>> PruningStatistics for RowGroupPruningStats<'_,
}
}
fn row_counts(&self, _column: &Column) -> Option<ArrayRef> {
fn row_counts(&self) -> Option<ArrayRef> {
// TODO(LFC): Impl it.
None
}
src/mito2/src/time_provider.rs
+46
View File
@@ -50,3 +50,49 @@ impl TimeProvider for StdTimeProvider {
current_time_millis() - current_millis
}
}
/// Controllable event time with short waits for background-task tests.
#[cfg(test)]
pub(crate) mod mock {
use std::sync::atomic::{AtomicI64, Ordering};
use std::time::Duration;
use crate::time_provider::TimeProvider;
#[derive(Debug)]
pub(crate) struct MockTimeProvider {
now: AtomicI64,
elapsed: AtomicI64,
}
impl TimeProvider for MockTimeProvider {
fn current_time_millis(&self) -> i64 {
self.now.load(Ordering::Relaxed)
}
fn elapsed_since(&self, _current_millis: i64) -> i64 {
self.elapsed.load(Ordering::Relaxed)
}
fn wait_duration(&self, _duration: Duration) -> Duration {
Duration::from_millis(20)
}
}
impl MockTimeProvider {
pub(crate) fn new(now: i64) -> Self {
Self {
now: AtomicI64::new(now),
elapsed: AtomicI64::new(0),
}
}
pub(crate) fn set_now(&self, now: i64) {
self.now.store(now, Ordering::Relaxed);
}
pub(crate) fn set_elapsed(&self, elapsed: i64) {
self.elapsed.store(elapsed, Ordering::Relaxed);
}
}
}
src/mito2/src/worker.rs
+11 -3
View File
@@ -611,13 +611,17 @@ impl<S: LogStore> WorkerStarter<S> {
.zip(series_index_task_state.clone())
.map(|(store, state)| {
let (purger, purge_receiver) = series_index_channel(store.clone());
series_index_purger = Some(purger);
series_index_purger = Some(purger.clone());
spawn_series_index_tasks(
self.id,
store,
regions.clone(),
state,
self.config.experimental_series_index_bucket_width,
purger,
purge_receiver,
self.config.experimental_series_index_maintenance_interval,
self.time_provider.clone(),
)
});
let now = self.time_provider.current_time_millis();
@@ -644,6 +648,7 @@ impl<S: LogStore> WorkerStarter<S> {
self.index_build_job_pool,
self.config.max_background_index_builds,
),
series_index_task_state: series_index_task_state.clone(),
series_index_store: self.series_index_store,
series_index_purger,
flush_scheduler: FlushScheduler::new(self.flush_job_pool),
@@ -708,9 +713,9 @@ pub(crate) struct RegionWorker {
sender: Sender<WorkerRequestWithTime>,
/// Handle to the worker thread.
handle: Mutex<Option<JoinHandle<()>>>,
/// Handle to the series-index maintenance task.
/// Handle to the sequential series-index task.
series_index_handle: Mutex<Option<JoinHandle<()>>>,
/// Controls the worker-owned series-index maintenance task.
/// Controls the worker-owned series-index task.
series_index_task_state: Option<Arc<SeriesIndexTaskState>>,
/// Whether to run the worker thread.
running: Arc<AtomicBool>,
@@ -943,6 +948,8 @@ struct RegionWorkerLoop<S> {
write_buffer_manager: WriteBufferManagerRef,
/// Scheduler for index build task.
index_build_scheduler: IndexBuildScheduler,
/// Controls the worker-owned series-index task.
series_index_task_state: Option<Arc<SeriesIndexTaskState>>,
/// Store for companion range indexes deleted by the region SST purger.
series_index_store: Option<ObjectStore>,
series_index_purger: Option<IndexFilePurger>,
@@ -1642,6 +1649,7 @@ mod tests {
.create_worker_group(MitoConfig {
num_workers: 4,
experimental_enable_series_index: enabled,
experimental_series_index_maintenance_interval: Duration::from_secs(3600),
..Default::default()
})
.await;
src/mito2/src/worker/handle_catchup.rs
+3
View File
@@ -134,6 +134,9 @@ impl<S: LogStore> RegionWorkerLoop<S> {
.await?;
debug_assert!(!reopened_region.is_writable());
self.regions.insert_region(reopened_region.clone());
if let Some(state) = &self.series_index_task_state {
state.wake();
}
Ok(reopened_region)
}

Some files were not shown because too many files have changed in this diff Show More