diff --git a/.github/actions/release-cn-artifacts/action.yaml b/.github/actions/release-cn-artifacts/action.yaml index 2825d3f5d08..fe78d5a760a 100644 --- a/.github/actions/release-cn-artifacts/action.yaml +++ b/.github/actions/release-cn-artifacts/action.yaml @@ -37,17 +37,14 @@ inputs: description: Whether to push the latest tag of the image required: false default: 'true' - aws-cn-s3-bucket: - description: S3 bucket to store released artifacts in CN region + proxy-url: + description: The url of the S3 proxy server required: true - aws-cn-access-key-id: - description: AWS access key id in CN region + proxy-username: + description: The username of the S3 proxy required: true - aws-cn-secret-access-key: - description: AWS secret access key in CN region - required: true - aws-cn-region: - description: AWS region in CN + proxy-password: + description: The password of the S3 proxy required: true upload-to-s3: description: Upload to S3 @@ -77,21 +74,13 @@ runs: with: path: ${{ inputs.artifacts-dir }} - - name: Install s5cmd - shell: bash - run: | - wget https://github.com/peak/s5cmd/releases/download/v2.3.0/s5cmd_2.3.0_Linux-64bit.tar.gz - tar -xzf s5cmd_2.3.0_Linux-64bit.tar.gz - sudo mv s5cmd /usr/local/bin/ - sudo chmod +x /usr/local/bin/s5cmd - - name: Release artifacts to cn region uses: nick-invision/retry@v2 if: ${{ inputs.upload-to-s3 == 'true' }} env: - AWS_ACCESS_KEY_ID: ${{ inputs.aws-cn-access-key-id }} - AWS_SECRET_ACCESS_KEY: ${{ inputs.aws-cn-secret-access-key }} - AWS_REGION: ${{ inputs.aws-cn-region }} + PROXY_URL: ${{ inputs.proxy-url }} + PROXY_USERNAME: ${{ inputs.proxy-username }} + PROXY_PASSWORD: ${{ inputs.proxy-password }} UPDATE_VERSION_INFO: ${{ inputs.update-version-info }} with: max_attempts: ${{ inputs.upload-max-retry-times }} @@ -99,8 +88,7 @@ runs: command: | ./.github/scripts/upload-artifacts-to-s3.sh \ ${{ inputs.artifacts-dir }} \ - ${{ inputs.version }} \ - ${{ inputs.aws-cn-s3-bucket }} + ${{ inputs.version }} - name: Push greptimedb image from Dockerhub to ACR shell: bash diff --git a/.github/scripts/upload-artifacts-to-s3.sh b/.github/scripts/upload-artifacts-to-s3.sh index 75c8f8d9329..310575c069f 100755 --- a/.github/scripts/upload-artifacts-to-s3.sh +++ b/.github/scripts/upload-artifacts-to-s3.sh @@ -5,16 +5,15 @@ set -o pipefail ARTIFACTS_DIR=$1 VERSION=$2 -AWS_S3_BUCKET=$3 RELEASE_DIRS="releases/greptimedb" GREPTIMEDB_REPO="GreptimeTeam/greptimedb" # Check if necessary variables are set. function check_vars() { - for var in AWS_S3_BUCKET VERSION ARTIFACTS_DIR; do + for var in VERSION ARTIFACTS_DIR; do if [ -z "${!var}" ]; then echo "$var is not set or empty." - echo "Usage: $0 " + echo "Usage: $0 " exit 1 fi done @@ -33,8 +32,13 @@ function upload_artifacts() { # ├── greptime-darwin-amd64-v0.2.0.sha256sum # └── greptime-darwin-amd64-v0.2.0.tar.gz find "$ARTIFACTS_DIR" -type f \( -name "*.tar.gz" -o -name "*.sha256sum" \) | while IFS= read -r file; do - s5cmd cp \ - "$file" "s3://$AWS_S3_BUCKET/$RELEASE_DIRS/$VERSION/$(basename "$file")" + filename=$(basename "$file") + TARGET_URL="$PROXY_URL/$RELEASE_DIRS/$VERSION/$filename" + + curl -X PUT \ + -u "$PROXY_USERNAME:$PROXY_PASSWORD" \ + -F "file=@$file" \ + "$TARGET_URL" done } @@ -45,16 +49,24 @@ function update_version_info() { if [[ "$VERSION" =~ ^v[0-9]+\.[0-9]+\.[0-9]+$ ]]; then echo "Updating latest-version.txt" echo "$VERSION" > latest-version.txt - s5cmd cp \ - latest-version.txt "s3://$AWS_S3_BUCKET/$RELEASE_DIRS/latest-version.txt" + TARGET_URL="$PROXY_URL/$RELEASE_DIRS/latest-version.txt" + + curl -X PUT \ + -u "$PROXY_USERNAME:$PROXY_PASSWORD" \ + -F "file=@latest-version.txt" \ + "$TARGET_URL" fi # If it's the nightly release, update latest-nightly-version.txt. if [[ "$VERSION" == *"nightly"* ]]; then echo "Updating latest-nightly-version.txt" echo "$VERSION" > latest-nightly-version.txt - s5cmd cp \ - latest-nightly-version.txt "s3://$AWS_S3_BUCKET/$RELEASE_DIRS/latest-nightly-version.txt" + + TARGET_URL="$PROXY_URL/$RELEASE_DIRS/latest-nightly-version.txt" + curl -X PUT \ + -u "$PROXY_USERNAME:$PROXY_PASSWORD" \ + -F "file=@latest-nightly-version.txt" \ + "$TARGET_URL" fi fi } @@ -93,10 +105,10 @@ function main() { } # Usage example: -# AWS_ACCESS_KEY_ID= \ -# AWS_SECRET_ACCESS_KEY= \ -# AWS_DEFAULT_REGION= \ +# PROXY_URL= \ +# PROXY_USERNAME= \ +# PROXY_PASSWORD= \ # UPDATE_VERSION_INFO=true \ # DOWNLOAD_ARTIFACTS_FROM_GITHUB=false \ -# ./upload-artifacts-to-s3.sh +# ./upload-artifacts-to-s3.sh main diff --git a/.github/workflows/dev-build.yml b/.github/workflows/dev-build.yml index 021867e4edd..d03fbeff14a 100644 --- a/.github/workflows/dev-build.yml +++ b/.github/workflows/dev-build.yml @@ -285,10 +285,9 @@ jobs: dst-image-registry: ${{ vars.ACR_IMAGE_REGISTRY }} dst-image-namespace: ${{ vars.IMAGE_NAMESPACE }} version: ${{ needs.allocate-runners.outputs.version }} - aws-cn-s3-bucket: ${{ vars.AWS_RELEASE_BUCKET }} - aws-cn-access-key-id: ${{ secrets.AWS_CN_ACCESS_KEY_ID }} - aws-cn-secret-access-key: ${{ secrets.AWS_CN_SECRET_ACCESS_KEY }} - aws-cn-region: ${{ vars.AWS_RELEASE_BUCKET_REGION }} + proxy-url: ${{ secrets.PROXY_URL }} + proxy-username: ${{ secrets.PROXY_USERNAME }} + proxy-password: ${{ secrets.PROXY_PASSWORD }} upload-to-s3: ${{ inputs.upload_artifacts_to_s3 }} dev-mode: true # Only build the standard images(exclude centos images). push-latest-tag: false # Don't push the latest tag to registry. diff --git a/.github/workflows/develop.yml b/.github/workflows/develop.yml index 0238e92c8d7..b6ab0f8926a 100644 --- a/.github/workflows/develop.yml +++ b/.github/workflows/develop.yml @@ -319,7 +319,13 @@ jobs: include: - target: "fuzz_repartition_table" mode: - name: "Local WAL Repartition GC" + name: "Local WAL mito table repartition" + minio: true + kafka: false + values: "with-minio-repartition-gc.yaml" + - target: "fuzz_repartition_metric_table" + mode: + name: "Local WAL metric table repartition" minio: true kafka: false values: "with-minio-repartition-gc.yaml" @@ -455,6 +461,14 @@ jobs: path: /tmp/fuzz-monitor-dumps if-no-files-found: warn retention-days: 3 + - name: Upload CSV dumps + if: failure() + uses: actions/upload-artifact@v4 + with: + name: fuzz-tests-csv-dumps-${{ matrix.mode.name }}-${{ matrix.target }} + path: /tmp/greptime-fuzz-dumps + if-no-files-found: warn + retention-days: 3 - name: Delete cluster if: success() shell: bash diff --git a/.github/workflows/nightly-build.yml b/.github/workflows/nightly-build.yml index 9eaa38c7896..14ebb6e715e 100644 --- a/.github/workflows/nightly-build.yml +++ b/.github/workflows/nightly-build.yml @@ -236,10 +236,9 @@ jobs: dst-image-registry: ${{ vars.ACR_IMAGE_REGISTRY }} dst-image-namespace: ${{ vars.IMAGE_NAMESPACE }} version: ${{ needs.allocate-runners.outputs.version }} - aws-cn-s3-bucket: ${{ vars.AWS_RELEASE_BUCKET }} - aws-cn-access-key-id: ${{ secrets.AWS_CN_ACCESS_KEY_ID }} - aws-cn-secret-access-key: ${{ secrets.AWS_CN_SECRET_ACCESS_KEY }} - aws-cn-region: ${{ vars.AWS_RELEASE_BUCKET_REGION }} + proxy-url: ${{ secrets.PROXY_URL }} + proxy-username: ${{ secrets.PROXY_USERNAME }} + proxy-password: ${{ secrets.PROXY_PASSWORD }} upload-to-s3: false dev-mode: false update-version-info: false # Don't update version info in S3. diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 3b0eb2d68ce..9f8f2d9703a 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -358,10 +358,9 @@ jobs: dst-image-registry: ${{ vars.ACR_IMAGE_REGISTRY }} dst-image-namespace: ${{ vars.IMAGE_NAMESPACE }} version: ${{ needs.allocate-runners.outputs.version }} - aws-cn-s3-bucket: ${{ vars.AWS_RELEASE_BUCKET }} - aws-cn-access-key-id: ${{ secrets.AWS_CN_ACCESS_KEY_ID }} - aws-cn-secret-access-key: ${{ secrets.AWS_CN_SECRET_ACCESS_KEY }} - aws-cn-region: ${{ vars.AWS_RELEASE_BUCKET_REGION }} + proxy-url: ${{ secrets.PROXY_URL }} + proxy-username: ${{ secrets.PROXY_USERNAME }} + proxy-password: ${{ secrets.PROXY_PASSWORD }} dev-mode: false upload-to-s3: true update-version-info: true diff --git a/Cargo.lock b/Cargo.lock index 85c2b1ed2df..605b037fc9c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2488,7 +2488,6 @@ version = "1.0.0-rc.2" dependencies = [ "common-error", "common-macro", - "common-telemetry", "humantime", "serde", "snafu 0.8.6", @@ -7301,7 +7300,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.52.6", + "windows-targets 0.48.5", ] [[package]] @@ -7887,6 +7886,7 @@ dependencies = [ "common-base", "common-error", "common-function", + "common-grpc", "common-macro", "common-meta", "common-query", @@ -10771,9 +10771,9 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.12" +version = "0.11.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49df843a9161c85bb8aae55f101bc0bac8bcafd637a620d9122fd7e0b2f7422e" +checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098" dependencies = [ "bytes", "getrandom 0.3.3", diff --git a/docs/rfcs/2026-03-16-flow-inc-query.md b/docs/rfcs/2026-03-16-flow-inc-query.md new file mode 100644 index 00000000000..8041d37d2b5 --- /dev/null +++ b/docs/rfcs/2026-03-16-flow-inc-query.md @@ -0,0 +1,190 @@ +--- +Feature Name: Flow Batching Sequence-Based Incremental Query Plan (Lite) +Tracking Issue: TBD +Date: 2026-03-16 +Author: @discord9 +--- + +# Summary + +This RFC proposes a correctness-first incremental query mode for Flow batching. +Flow queries can read only `seq > checkpoint` and advance checkpoints using per-region correctness watermarks. +When incremental reads are stale or correctness cannot be proven, Flow falls back to full recomputation. + +# Motivation + +Flow batching still needs to repeatedly compute old data in the same time window, so incremental query can improve Flow performance. + +# Goals + +1. Add opt-in incremental reads (`seq > given_seq`) for Flow. +2. Return per-region correctness watermarks for checkpoint advancement. +3. Keep existing query behavior unchanged unless explicitly enabled. +4. Define deterministic fallback for stale or unprovable incremental reads. + +# Non-Goals + +1. No business-schema changes (no synthetic watermark columns in result rows). +2. No global throughput optimization in v1 (correctness first). +3. No observational watermark output when correctness is unprovable. + +# Proposal + +## 1) Query options + +Introduce three `QueryContext` extension keys: + +- `flow.incremental_after_seqs` +- `flow.incremental_mode` +- `flow.return_region_seq` + +These options are opt-in and only affect Flow incremental execution paths. + +## 2) Scan mapping + +When incremental mode is enabled: + +- map `after_seq` to `memtable_min_sequence` (exclusive lower bound) +- keep existing snapshot upper-bound behavior (`memtable_max_sequence`) + +Important limitation in v1: + +- incremental filtering is correctness-proven only for memtable rows +- SST files do not preserve detailed row-level sequence metadata; they only expose coarser file-level sequence information +- therefore `seq > checkpoint` must not assume precise incremental pruning across memtable->SST flush boundaries + +If required incremental parameters are missing or invalid, return argument error. + +## 3) Stale protection + +Add dedicated stale error: + +- `IncrementalQueryStale { region_id, given_seq, min_readable_seq }` + +Behavior: + +- if `given_seq < min_readable_seq`, return stale error +- if `given_seq == min_readable_seq`, query is valid and reads `seq > given_seq` +- if `given_seq > min_readable_seq`, query is also valid and reads `seq > given_seq` + +`IncrementalQueryStale` also covers the case where rows newer than the checkpoint have crossed a memtable->SST flush boundary and sequence-precise incremental exclusion can no longer be proven. +In other words, the flush-boundary case is not a separate fallback category in v1; it is one concrete way an incremental cursor becomes stale. + +## 4) Watermark return + +Extend query metrics with optional per-region watermark map: + +- `region_latest_sequences: Vec<(region_id: u64, latest_sequence: u64)>` + +Rules: + +- only terminal metrics of successful query can advance checkpoints +- for multi-region query, watermark must be complete map or absent +- if correctness is unprovable, business rows may return but watermark is absent + +## 5) Flow state machine + +Checkpoint and watermark state are kept only in flownode memory in v1; they are not persisted as durable flow metadata. +Cold start or flownode restart therefore always re-enters through a full snapshot read. +Only after that full query succeeds with a complete correctness watermark may Flow switch back to incremental mode. + +Flow starts in full mode, then transitions: + +1. Full query succeeds with correctness watermark -> enter incremental mode +2. Incremental query succeeds with correctness watermark -> advance checkpoint +3. Incremental stale/failure -> fallback to full mode +4. Full query without correctness watermark -> remain in full mode + +```mermaid +stateDiagram-v2 + [*] --> FullSnapshot: Flow starts + + state FullSnapshot { + [*] --> RunFull + RunFull --> RunFull: Full query succeeds but watermark is unprovable
no region_latest_sequences returned + } + + FullSnapshot --> Incremental: Full query succeeds and correctness watermark is returned
(checkpoint updated) + + state Incremental { + [*] --> RunInc + RunInc --> RunInc: Incremental succeeds
(checkpoint advances) + } + + Incremental --> FullSnapshot: IncrementalQueryStale
(cursor too old, fallback required) + Incremental --> FullSnapshot: Incremental fails
and fallback policy is triggered + + FullSnapshot --> [*]: Flow stops + Incremental --> [*]: Flow stops +``` + +### Fallback Policy + +Fallback to full mode is deterministic and is triggered by any of the following: + +1. `IncrementalQueryStale` is returned. +2. Incremental query fails with execution errors. +3. Incremental query succeeds but watermark is absent or incomplete for participating regions. + +Policy behavior: + +1. Do not advance any checkpoint in the failed/incomplete round. +2. Switch to full mode for the affected flow/window in the next round. +3. Return to incremental mode only after a full query succeeds with a complete correctness watermark map. + +### Persistence and recovery model + +The v1 design is intentionally correctness-first and keeps the progress cursor lightweight: + +1. Watermarks/checkpoints live only in flownode memory; v1 does not persist them separately. +2. On cold start, the flow re-establishes progress by running a successful full-query snapshot read, then resumes incremental mode only after that round returns a complete correctness watermark map. +3. Sequence-precise incremental correctness is currently limited to rows still visible in memtables. +4. Once relevant rows have been flushed into SST, the system cannot use `seq > checkpoint` alone to prove precise incremental exclusion, because SST lacks detailed row-level sequence metadata. +5. In that case the correct behavior is to fall back to full recomputation, not to continue a best-effort incremental scan. + +# Distributed and Compatibility Requirements + +1. Distributed path must preserve region-level snapshot/read-bound semantics end-to-end. +2. `snapshot_seqs` transport and `flow.*` options must both be carried correctly. + - `snapshot_seqs` means the per-region snapshot upper-bound map: `region_id -> sequence`. +3. New metrics fields must be backward-compatible (old clients ignore unknown fields). + +# Rollout Plan + +## Phase 1 (MVP, correctness first) + +1. Add extension constants and parsing. +2. Add incremental scan mapping and stale detection. +3. Add watermark metrics field and terminal-watermark checkpoint update path. +4. Complete standalone and distributed passthrough. + +## Phase 2 (performance and observability) + +1. Improve batching key strategy with sequence/watermark context. +2. Optimize watermark serialization overhead. +3. Add metrics: incremental hit rate, fallback rate, fallback window size. + +# Testing Plan + +1. Unit tests for incremental bounds and stale detection. +2. Query-path tests for extension mapping and watermark semantics. +3. Flow integration tests for full->incremental->fallback transitions. +4. Distributed tests for end-to-end snapshot/watermark propagation. +5. Compatibility tests for old/new client-server combinations. + +# Risks + +1. Boundary semantic mismatch (`<` vs `<=`) may cause correctness bugs. +2. Incomplete distributed propagation can silently invalidate watermark safety. +3. Frequent fallback can reduce throughput before phase-2 optimizations. +4. Memtable->SST flushes may force more full recomputation than expected until finer-grained SST sequence tracking exists. + +# Alternatives + +1. Put watermark into business rows (rejected: schema pollution). +2. Add new dedicated Flight message type in v1 (deferred to reduce scope). + +# Conclusion + +This plan enables a practical, correctness-first incremental path for Flow batching. +It reuses existing sequence scan capability, adds strict stale handling, and advances checkpoints only from correctness-proven per-region watermarks. diff --git a/src/cmd/src/datanode/objbench.rs b/src/cmd/src/datanode/objbench.rs index d8f53b9d710..f6d8674d4c2 100644 --- a/src/cmd/src/datanode/objbench.rs +++ b/src/cmd/src/datanode/objbench.rs @@ -20,13 +20,14 @@ use clap::Parser; use colored::Colorize; use datanode::config::RegionEngineConfig; use datanode::store; -use either::Either; +use futures::stream; use mito2::access_layer::{ AccessLayer, AccessLayerRef, Metrics, OperationType, SstWriteRequest, WriteType, }; use mito2::cache::{CacheManager, CacheManagerRef}; use mito2::config::{FulltextIndexConfig, MitoConfig, Mode}; -use mito2::read::Source; +use mito2::read::FlatSource; +use mito2::sst::FormatType; use mito2::sst::file::{FileHandle, FileMeta}; use mito2::sst::file_purger::{FilePurger, FilePurgerRef}; use mito2::sst::index::intermediate::IntermediateManager; @@ -210,6 +211,7 @@ impl ObjbenchCommand { object_store.clone(), ) .expected_metadata(Some(region_meta.clone())) + .flat_format(true) .build() .await .map_err(|e| { @@ -231,6 +233,10 @@ impl ObjbenchCommand { let reader_build_elapsed = reader_build_start.elapsed(); let total_rows = reader.parquet_metadata().file_metadata().num_rows(); println!("{} Reader built in {:?}", "✓".green(), reader_build_elapsed); + let reader_stream = Box::pin(stream::try_unfold(reader, |mut reader| async move { + let batch = reader.next_record_batch().await?; + Ok(batch.map(|batch| (batch, reader))) + })); // Build write request let fulltext_index_config = FulltextIndexConfig { @@ -241,10 +247,11 @@ impl ObjbenchCommand { let write_req = SstWriteRequest { op_type: OperationType::Flush, metadata: region_meta, - source: Either::Left(Source::Reader(Box::new(reader))), + source: FlatSource::Stream(reader_stream), cache_manager, storage: None, max_sequence: None, + sst_write_format: FormatType::PrimaryKey, index_options: Default::default(), index_config: mito_engine_config.index.clone(), inverted_index_config: MitoConfig::default().inverted_index, diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 92638d3c4a7..215bea0ec56 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -32,14 +32,15 @@ use common_meta::cache::LayeredCacheRegistryBuilder; use common_meta::ddl::flow_meta::FlowMetadataAllocator; use common_meta::ddl::table_meta::TableMetadataAllocator; use common_meta::ddl::{DdlContext, NoopRegionFailureDetectorControl}; -use common_meta::ddl_manager::{DdlManager, DdlManagerConfiguratorRef}; +use common_meta::ddl_manager::{DdlManager, DdlManagerConfiguratorRef, DdlManagerRef}; use common_meta::key::flow::FlowMetadataManager; use common_meta::key::{TableMetadataManager, TableMetadataManagerRef}; use common_meta::kv_backend::KvBackendRef; -use common_meta::procedure_executor::LocalProcedureExecutor; +use common_meta::node_manager::{FlownodeRef, NodeManagerRef}; +use common_meta::procedure_executor::{LocalProcedureExecutor, ProcedureExecutorRef}; use common_meta::region_keeper::MemoryRegionKeeper; use common_meta::region_registry::LeaderRegionRegistry; -use common_meta::sequence::SequenceBuilder; +use common_meta::sequence::{Sequence, SequenceBuilder}; use common_meta::wal_provider::{WalProviderRef, build_wal_provider}; use common_procedure::ProcedureManagerRef; use common_query::prelude::set_default_prefix; @@ -49,6 +50,7 @@ use common_time::timezone::set_default_timezone; use common_version::{short_version, verbose_version}; use datanode::config::DatanodeOptions; use datanode::datanode::{Datanode, DatanodeBuilder}; +use datanode::region_server::RegionServer; use flow::{ FlownodeBuilder, FlownodeInstance, FlownodeOptions, FrontendClient, FrontendInvoker, GrpcQueryHandlerWithBoxedError, @@ -58,6 +60,7 @@ use frontend::instance::StandaloneDatanodeManager; use frontend::instance::builder::FrontendBuilder; use frontend::server::Services; use meta_srv::metasrv::{FLOW_ID_SEQ, TABLE_ID_SEQ}; +use plugins::PluginOptions; use plugins::frontend::context::{ CatalogManagerConfigureContext, StandaloneCatalogManagerConfigureContext, }; @@ -130,6 +133,18 @@ impl Instance { pub fn server_addr(&self, name: &str) -> Option { self.frontend.server_handlers().addr(name) } + + /// Get the mutable Frontend component of this Standalone instance for externally modification + /// by others (might not be in this code base, so don't delete this function). + pub fn mut_frontend(&mut self) -> &mut Frontend { + &mut self.frontend + } + + /// Get the Datanode component of this Standalone instance for externally usage + /// by others (might not be in this code base, so don't delete this function). + pub fn datanode(&self) -> &Datanode { + &self.datanode + } } #[async_trait] @@ -342,9 +357,18 @@ impl StartCommand { info!("Standalone start command: {:#?}", self); info!("Standalone options: {opts:#?}"); + let (mut instance, _) = + Self::build_with(opts.component, opts.plugins, InstanceCreator::default()).await?; + instance._guard.extend(guard); + Ok(instance) + } + + pub async fn build_with( + mut opts: StandaloneOptions, + plugin_opts: Vec, + creator: InstanceCreator, + ) -> Result<(Instance, InstanceCreatorResult)> { let mut plugins = Plugins::new(); - let plugin_opts = opts.plugins; - let mut opts = opts.component; set_default_prefix(opts.default_column_prefix.as_deref()) .map_err(BoxedError::new) .context(error::BuildCliSnafu)?; @@ -462,17 +486,16 @@ impl StartCommand { .await; } - let node_manager = Arc::new(StandaloneDatanodeManager { - region_server: datanode.region_server(), - flow_server: flownode.flow_engine(), - }); + let node_manager = creator + .node_manager_creator + .create( + &kv_backend, + datanode.region_server(), + flownode.flow_engine(), + ) + .await?; - let table_id_allocator = Arc::new( - SequenceBuilder::new(TABLE_ID_SEQ, kv_backend.clone()) - .initial(MIN_USER_TABLE_ID as u64) - .step(10) - .build(), - ); + let table_id_allocator = creator.table_id_allocator_creator.create(&kv_backend); let flow_id_sequence = Arc::new( SequenceBuilder::new(FLOW_ID_SEQ, kv_backend.clone()) .initial(MIN_USER_FLOW_ID as u64) @@ -489,7 +512,7 @@ impl StartCommand { .context(error::BuildWalProviderSnafu)?; let wal_provider = Arc::new(wal_provider); let table_metadata_allocator = Arc::new(TableMetadataAllocator::new( - table_id_allocator, + table_id_allocator.clone(), wal_provider.clone(), )); let flow_metadata_allocator = Arc::new(FlowMetadataAllocator::with_noop_peer_allocator( @@ -532,10 +555,10 @@ impl StartCommand { ddl_manager }; - let procedure_executor = Arc::new(LocalProcedureExecutor::new( - Arc::new(ddl_manager), - procedure_manager.clone(), - )); + let procedure_executor = creator + .procedure_executor_creator + .create(Arc::new(ddl_manager), procedure_manager.clone()) + .await?; let fe_instance = FrontendBuilder::new( fe_opts.clone(), @@ -568,7 +591,7 @@ impl StartCommand { kv_backend.clone(), layered_cache_registry.clone(), procedure_executor, - node_manager, + node_manager.clone(), ) .await .context(StartFlownodeSnafu)?; @@ -584,14 +607,20 @@ impl StartCommand { heartbeat_task: None, }; - Ok(Instance { + let instance = Instance { datanode, frontend, flownode, procedure_manager, wal_provider, - _guard: guard, - }) + _guard: vec![], + }; + let result = InstanceCreatorResult { + kv_backend, + node_manager, + table_id_allocator, + }; + Ok((instance, result)) } pub async fn create_table_metadata_manager( @@ -608,6 +637,115 @@ impl StartCommand { } } +#[async_trait] +pub trait NodeManagerCreator { + async fn create( + &self, + kv_backend: &KvBackendRef, + region_server: RegionServer, + flow_server: FlownodeRef, + ) -> Result; +} + +pub struct DefaultNodeManagerCreator; + +#[async_trait] +impl NodeManagerCreator for DefaultNodeManagerCreator { + async fn create( + &self, + _: &KvBackendRef, + region_server: RegionServer, + flow_server: FlownodeRef, + ) -> Result { + Ok(Arc::new(StandaloneDatanodeManager { + region_server, + flow_server, + })) + } +} + +pub trait TableIdAllocatorCreator { + fn create(&self, kv_backend: &KvBackendRef) -> Arc; +} + +struct DefaultTableIdAllocatorCreator; + +impl TableIdAllocatorCreator for DefaultTableIdAllocatorCreator { + fn create(&self, kv_backend: &KvBackendRef) -> Arc { + Arc::new( + SequenceBuilder::new(TABLE_ID_SEQ, kv_backend.clone()) + .initial(MIN_USER_TABLE_ID as u64) + .step(10) + .build(), + ) + } +} + +#[async_trait] +pub trait ProcedureExecutorCreator { + async fn create( + &self, + ddl_manager: DdlManagerRef, + procedure_manager: ProcedureManagerRef, + ) -> Result; +} + +pub struct DefaultProcedureExecutorCreator; + +#[async_trait] +impl ProcedureExecutorCreator for DefaultProcedureExecutorCreator { + async fn create( + &self, + ddl_manager: DdlManagerRef, + procedure_manager: ProcedureManagerRef, + ) -> Result { + Ok(Arc::new(LocalProcedureExecutor::new( + ddl_manager, + procedure_manager, + ))) + } +} + +/// `InstanceCreator` is used for grouping various component creators for building the +/// Standalone instance, suitable for customizing how the instance can be built. +pub struct InstanceCreator { + node_manager_creator: Box, + table_id_allocator_creator: Box, + procedure_executor_creator: Box, +} + +impl InstanceCreator { + pub fn new( + node_manager_creator: Box, + table_id_allocator_creator: Box, + procedure_executor_creator: Box, + ) -> Self { + Self { + node_manager_creator, + table_id_allocator_creator, + procedure_executor_creator, + } + } +} + +impl Default for InstanceCreator { + fn default() -> Self { + Self { + node_manager_creator: Box::new(DefaultNodeManagerCreator), + table_id_allocator_creator: Box::new(DefaultTableIdAllocatorCreator), + procedure_executor_creator: Box::new(DefaultProcedureExecutorCreator), + } + } +} + +/// `InstanceCreatorResult` is expected to be used paired with [InstanceCreator]. +/// It stores the created and other important components for further reusing. +pub struct InstanceCreatorResult { + pub kv_backend: KvBackendRef, + pub node_manager: NodeManagerRef, + pub table_id_allocator: Arc, +} + #[cfg(test)] mod tests { use std::default::Default; diff --git a/src/common/config/src/config.rs b/src/common/config/src/config.rs index e25c46a0c04..85ce3d206f6 100644 --- a/src/common/config/src/config.rs +++ b/src/common/config/src/config.rs @@ -53,7 +53,7 @@ pub trait Configurable: Serialize + DeserializeOwned + Default + Sized { env.try_parsing(true) .separator(ENV_VAR_SEP) - .ignore_empty(true) + .ignore_empty(false) }; // Workaround: Replacement for `Config::try_from(&default_opts)` due to @@ -237,4 +237,31 @@ mod tests { }, ); } + + #[derive(Debug, Serialize, Deserialize, Default)] + struct SimpleConfig { + name: Option, + prefix: Option, + } + + impl Configurable for SimpleConfig {} + + #[test] + fn test_empty_env_var_is_not_ignored() { + let env_prefix = "SIMPLE_CFG_UT"; + temp_env::with_vars( + [( + [env_prefix.to_string(), "PREFIX".to_string()].join(ENV_VAR_SEP), + Some(""), + )], + || { + let opts = SimpleConfig::load_layered_options(None, env_prefix).unwrap(); + // With ignore_empty(false), an empty env var should yield Some("") + // rather than None (which was the previous behavior with ignore_empty(true)). + assert_eq!(opts.prefix, Some("".to_string())); + // Unset env var should remain None. + assert_eq!(opts.name, None); + }, + ); + } } diff --git a/src/common/function/src/scalars/json/json_to_string.rs b/src/common/function/src/scalars/json/json_to_string.rs index 6c0cc260b29..6364dff4de0 100644 --- a/src/common/function/src/scalars/json/json_to_string.rs +++ b/src/common/function/src/scalars/json/json_to_string.rs @@ -19,6 +19,7 @@ use datafusion_common::DataFusionError; use datafusion_common::arrow::array::{Array, AsArray, StringViewBuilder}; use datafusion_common::arrow::datatypes::DataType; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility}; +use datatypes::types::jsonb_to_string; use crate::function::{Function, extract_args}; @@ -74,7 +75,7 @@ impl Function for JsonToStringFunction { for i in 0..size { let json = jsons.is_valid(i).then(|| jsons.value(i)); let result = json - .map(|json| jsonb::from_slice(json).map(|x| x.to_string())) + .map(jsonb_to_string) .transpose() .map_err(|e| DataFusionError::Execution(format!("invalid json binary: {e}")))?; diff --git a/src/common/memory-manager/Cargo.toml b/src/common/memory-manager/Cargo.toml index a6be50f774b..6686c981671 100644 --- a/src/common/memory-manager/Cargo.toml +++ b/src/common/memory-manager/Cargo.toml @@ -10,7 +10,6 @@ workspace = true [dependencies] common-error = { workspace = true } common-macro = { workspace = true } -common-telemetry = { workspace = true } humantime = { workspace = true } serde = { workspace = true } snafu = { workspace = true } diff --git a/src/common/memory-manager/src/guard.rs b/src/common/memory-manager/src/guard.rs index 770b6dec24a..ad3111581b6 100644 --- a/src/common/memory-manager/src/guard.rs +++ b/src/common/memory-manager/src/guard.rs @@ -14,14 +14,13 @@ use std::{fmt, mem}; -use common_telemetry::debug; use snafu::ensure; use tokio::sync::{OwnedSemaphorePermit, TryAcquireError}; use crate::error::{ MemoryAcquireTimeoutSnafu, MemoryLimitExceededSnafu, MemorySemaphoreClosedSnafu, Result, }; -use crate::manager::{MemoryMetrics, MemoryQuota}; +use crate::manager::{MemoryMetrics, MemoryQuota, UnlimitedMemoryQuota}; use crate::policy::OnExhaustedPolicy; /// Guard representing a slice of reserved memory. @@ -30,31 +29,57 @@ pub struct MemoryGuard { } pub(crate) enum GuardState { - Unlimited, + Released, + Unlimited { + quota: UnlimitedMemoryQuota, + granted_bytes: u64, + }, Limited { - permit: OwnedSemaphorePermit, quota: MemoryQuota, + permit: OwnedSemaphorePermit, }, } +impl GuardState { + fn release(self) { + match self { + GuardState::Released => {} + GuardState::Unlimited { + quota, + granted_bytes, + } => { + quota.sub_in_use(granted_bytes); + } + GuardState::Limited { quota, permit } => { + quota.release_permit(permit); + } + } + } +} + impl MemoryGuard { - pub(crate) fn unlimited() -> Self { + pub(crate) fn unlimited(quota: UnlimitedMemoryQuota, bytes: u64) -> Self { + quota.add_in_use(bytes); Self { - state: GuardState::Unlimited, + state: GuardState::Unlimited { + quota, + granted_bytes: bytes, + }, } } - pub(crate) fn limited(permit: OwnedSemaphorePermit, quota: MemoryQuota) -> Self { + pub(crate) fn limited(quota: MemoryQuota, permit: OwnedSemaphorePermit) -> Self { Self { - state: GuardState::Limited { permit, quota }, + state: GuardState::Limited { quota, permit }, } } /// Returns granted quota in bytes. pub fn granted_bytes(&self) -> u64 { match &self.state { - GuardState::Unlimited => 0, - GuardState::Limited { permit, quota } => { + GuardState::Released => 0, + GuardState::Unlimited { granted_bytes, .. } => *granted_bytes, + GuardState::Limited { quota, permit } => { quota.permits_to_bytes(permit.num_permits() as u32) } } @@ -68,13 +93,24 @@ impl MemoryGuard { /// - Returns error if requested bytes would exceed the manager's total limit /// - Returns error if the semaphore is unexpectedly closed pub async fn acquire_additional(&mut self, bytes: u64) -> Result<()> { - match &mut self.state { - GuardState::Unlimited => Ok(()), - GuardState::Limited { permit, quota } => { - if bytes == 0 { - return Ok(()); - } + if bytes == 0 { + return Ok(()); + } + match &mut self.state { + GuardState::Released => { + debug_assert!(false, "released memory guard state should not be reused"); + Ok(()) + } + GuardState::Unlimited { + quota, + granted_bytes, + } => { + quota.add_in_use(bytes); + *granted_bytes = granted_bytes.saturating_add(bytes); + Ok(()) + } + GuardState::Limited { quota, permit } => { let additional_permits = quota.bytes_to_permits(bytes); let current_permits = permit.num_permits() as u32; @@ -95,7 +131,6 @@ impl MemoryGuard { permit.merge(additional_permit); quota.update_in_use_metric(); - debug!("Acquired additional {} bytes", bytes); Ok(()) } } @@ -106,13 +141,24 @@ impl MemoryGuard { /// On success, merges the new memory into this guard and returns true. /// On failure, returns false and leaves this guard unchanged. pub fn try_acquire_additional(&mut self, bytes: u64) -> bool { - match &mut self.state { - GuardState::Unlimited => true, - GuardState::Limited { permit, quota } => { - if bytes == 0 { - return true; - } + if bytes == 0 { + return true; + } + match &mut self.state { + GuardState::Released => { + debug_assert!(false, "released memory guard state should not be reused"); + false + } + GuardState::Unlimited { + quota, + granted_bytes, + } => { + quota.add_in_use(bytes); + *granted_bytes = granted_bytes.saturating_add(bytes); + true + } + GuardState::Limited { quota, permit } => { let additional_permits = quota.bytes_to_permits(bytes); match quota @@ -123,7 +169,6 @@ impl MemoryGuard { Ok(additional_permit) => { permit.merge(additional_permit); quota.update_in_use_metric(); - debug!("Acquired additional {} bytes", bytes); true } Err(TryAcquireError::NoPermits) | Err(TryAcquireError::Closed) => { @@ -168,7 +213,8 @@ impl MemoryGuard { MemoryLimitExceededSnafu { requested_bytes: bytes, limit_bytes: match &self.state { - GuardState::Unlimited => 0, // unreachable: unlimited mode always succeeds + GuardState::Released => 0, + GuardState::Unlimited { .. } => 0, // unreachable: unlimited mode always succeeds GuardState::Limited { quota, .. } => { quota.permits_to_bytes(quota.limit_permits) } @@ -184,22 +230,30 @@ impl MemoryGuard { /// /// Returns true if the release succeeds or is a no-op; false if the request exceeds granted. pub fn release_partial(&mut self, bytes: u64) -> bool { + if bytes == 0 { + return true; + } + match &mut self.state { - GuardState::Unlimited => true, - GuardState::Limited { permit, quota } => { - if bytes == 0 { - return true; + GuardState::Released => true, + GuardState::Unlimited { + quota, + granted_bytes, + } => { + if bytes > *granted_bytes { + return false; } + quota.sub_in_use(bytes); + *granted_bytes = granted_bytes.saturating_sub(bytes); + true + } + GuardState::Limited { quota, permit } => { let release_permits = quota.bytes_to_permits(bytes); match permit.split(release_permits as usize) { Some(released_permit) => { - let released_bytes = - quota.permits_to_bytes(released_permit.num_permits() as u32); - drop(released_permit); - quota.update_in_use_metric(); - debug!("Released {} bytes from memory guard", released_bytes); + quota.release_permit(released_permit); true } None => false, @@ -211,14 +265,7 @@ impl MemoryGuard { impl Drop for MemoryGuard { fn drop(&mut self) { - if let GuardState::Limited { permit, quota } = - mem::replace(&mut self.state, GuardState::Unlimited) - { - let bytes = quota.permits_to_bytes(permit.num_permits() as u32); - drop(permit); - quota.update_in_use_metric(); - debug!("Released memory: {} bytes", bytes); - } + mem::replace(&mut self.state, GuardState::Released).release(); } } diff --git a/src/common/memory-manager/src/manager.rs b/src/common/memory-manager/src/manager.rs index 50360d2a319..8cca5f220c6 100644 --- a/src/common/memory-manager/src/manager.rs +++ b/src/common/memory-manager/src/manager.rs @@ -13,9 +13,10 @@ // limitations under the License. use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use snafu::ensure; -use tokio::sync::{Semaphore, TryAcquireError}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError}; use crate::error::{ MemoryAcquireTimeoutSnafu, MemoryLimitExceededSnafu, MemorySemaphoreClosedSnafu, Result, @@ -34,7 +35,7 @@ pub trait MemoryMetrics: Clone + Send + Sync + 'static { /// Generic memory manager for quota-controlled operations. #[derive(Clone)] pub struct MemoryManager { - quota: Option>, + quota: MemoryQuotaState, } impl Default for MemoryManager { @@ -51,6 +52,18 @@ pub(crate) struct MemoryQuota { pub(crate) metrics: M, } +#[derive(Clone)] +pub(crate) struct UnlimitedMemoryQuota { + pub(crate) current_bytes: Arc, + pub(crate) metrics: M, +} + +#[derive(Clone)] +pub(crate) enum MemoryQuotaState { + Unlimited(UnlimitedMemoryQuota), + Limited(MemoryQuota), +} + impl MemoryManager { /// Creates a new memory manager with the given limit in bytes. /// `limit_bytes = 0` disables the limit. @@ -62,7 +75,12 @@ impl MemoryManager { pub fn with_granularity(limit_bytes: u64, granularity: PermitGranularity, metrics: M) -> Self { if limit_bytes == 0 { metrics.set_limit(0); - return Self { quota: None }; + return Self { + quota: MemoryQuotaState::Unlimited(UnlimitedMemoryQuota { + current_bytes: Arc::new(AtomicU64::new(0)), + metrics, + }), + }; } let limit_permits = granularity.bytes_to_permits(limit_bytes); @@ -70,7 +88,7 @@ impl MemoryManager { metrics.set_limit(limit_aligned_bytes as i64); Self { - quota: Some(MemoryQuota { + quota: MemoryQuotaState::Limited(MemoryQuota { semaphore: Arc::new(Semaphore::new(limit_permits as usize)), limit_permits, granularity, @@ -81,26 +99,30 @@ impl MemoryManager { /// Returns the configured limit in bytes (0 if unlimited). pub fn limit_bytes(&self) -> u64 { - self.quota - .as_ref() - .map(|quota| quota.permits_to_bytes(quota.limit_permits)) - .unwrap_or(0) + match &self.quota { + MemoryQuotaState::Unlimited(_) => 0, + MemoryQuotaState::Limited(quota) => quota.permits_to_bytes(quota.limit_permits), + } } /// Returns currently used bytes. pub fn used_bytes(&self) -> u64 { - self.quota - .as_ref() - .map(|quota| quota.permits_to_bytes(quota.used_permits())) - .unwrap_or(0) + match &self.quota { + MemoryQuotaState::Unlimited(quota) => quota.current_bytes.load(Ordering::Acquire), + MemoryQuotaState::Limited(quota) => quota.permits_to_bytes(quota.used_permits()), + } } /// Returns available bytes. + /// + /// Unlimited managers report `u64::MAX`. pub fn available_bytes(&self) -> u64 { - self.quota - .as_ref() - .map(|quota| quota.permits_to_bytes(quota.available_permits_clamped())) - .unwrap_or(0) + match &self.quota { + MemoryQuotaState::Unlimited(_) => u64::MAX, + MemoryQuotaState::Limited(quota) => { + quota.permits_to_bytes(quota.available_permits_clamped()) + } + } } /// Acquires memory, waiting if necessary until enough is available. @@ -110,8 +132,8 @@ impl MemoryManager { /// - Returns error if the semaphore is unexpectedly closed pub async fn acquire(&self, bytes: u64) -> Result> { match &self.quota { - None => Ok(MemoryGuard::unlimited()), - Some(quota) => { + MemoryQuotaState::Unlimited(quota) => Ok(MemoryGuard::unlimited(quota.clone(), bytes)), + MemoryQuotaState::Limited(quota) => { let permits = quota.bytes_to_permits(bytes); ensure!( @@ -129,7 +151,7 @@ impl MemoryManager { .await .map_err(|_| MemorySemaphoreClosedSnafu.build())?; quota.update_in_use_metric(); - Ok(MemoryGuard::limited(permit, quota.clone())) + Ok(MemoryGuard::limited(quota.clone(), permit)) } } } @@ -137,14 +159,16 @@ impl MemoryManager { /// Tries to acquire memory. Returns Some(guard) on success, None if insufficient. pub fn try_acquire(&self, bytes: u64) -> Option> { match &self.quota { - None => Some(MemoryGuard::unlimited()), - Some(quota) => { + MemoryQuotaState::Unlimited(quota) => { + Some(MemoryGuard::unlimited(quota.clone(), bytes)) + } + MemoryQuotaState::Limited(quota) => { let permits = quota.bytes_to_permits(bytes); match quota.semaphore.clone().try_acquire_many_owned(permits) { Ok(permit) => { quota.update_in_use_metric(); - Some(MemoryGuard::limited(permit, quota.clone())) + Some(MemoryGuard::limited(quota.clone(), permit)) } Err(TryAcquireError::NoPermits) | Err(TryAcquireError::Closed) => { quota.metrics.inc_rejected("try_acquire"); @@ -219,4 +243,49 @@ impl MemoryQuota { let bytes = self.permits_to_bytes(self.used_permits()); self.metrics.set_in_use(bytes as i64); } + + pub(crate) fn release_permit(&self, permit: OwnedSemaphorePermit) { + drop(permit); + self.update_in_use_metric(); + } +} + +impl UnlimitedMemoryQuota { + pub(crate) fn add_in_use(&self, bytes: u64) { + if bytes == 0 { + return; + } + + let previous = self + .current_bytes + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + Some(current.saturating_add(bytes)) + }) + .unwrap(); + let new_total = previous.saturating_add(bytes); + debug_assert!( + new_total >= previous, + "unlimited memory usage counter overflowed" + ); + self.metrics.set_in_use(new_total as i64); + } + + pub(crate) fn sub_in_use(&self, bytes: u64) { + if bytes == 0 { + return; + } + + let previous = self + .current_bytes + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + Some(current.saturating_sub(bytes)) + }) + .unwrap(); + debug_assert!( + previous >= bytes, + "unlimited memory usage counter underflowed: current={previous}, release={bytes}" + ); + let new_total = previous.saturating_sub(bytes); + self.metrics.set_in_use(new_total as i64); + } } diff --git a/src/common/memory-manager/src/tests.rs b/src/common/memory-manager/src/tests.rs index 886eef9dac4..fe02703f0b5 100644 --- a/src/common/memory-manager/src/tests.rs +++ b/src/common/memory-manager/src/tests.rs @@ -24,7 +24,9 @@ fn test_try_acquire_unlimited() { let manager = MemoryManager::new(0, NoOpMetrics); let guard = manager.try_acquire(10 * PERMIT_GRANULARITY_BYTES).unwrap(); assert_eq!(manager.limit_bytes(), 0); - assert_eq!(guard.granted_bytes(), 0); + assert_eq!(manager.available_bytes(), u64::MAX); + assert_eq!(guard.granted_bytes(), 10 * PERMIT_GRANULARITY_BYTES); + assert_eq!(manager.used_bytes(), 10 * PERMIT_GRANULARITY_BYTES); } #[test] @@ -136,7 +138,10 @@ fn test_request_additional_unlimited() { // Should always succeed with unlimited manager assert!(guard.try_acquire_additional(100 * PERMIT_GRANULARITY_BYTES)); - assert_eq!(guard.granted_bytes(), 0); + assert_eq!(guard.granted_bytes(), 105 * PERMIT_GRANULARITY_BYTES); + assert_eq!(manager.used_bytes(), 105 * PERMIT_GRANULARITY_BYTES); + + drop(guard); assert_eq!(manager.used_bytes(), 0); } @@ -187,9 +192,10 @@ fn test_early_release_partial_unlimited() { let manager = MemoryManager::new(0, NoOpMetrics); let mut guard = manager.try_acquire(100 * PERMIT_GRANULARITY_BYTES).unwrap(); - // Unlimited guard - release should succeed (no-op) + // Unlimited guard should track and release exact bytes. assert!(guard.release_partial(50 * PERMIT_GRANULARITY_BYTES)); - assert_eq!(guard.granted_bytes(), 0); + assert_eq!(guard.granted_bytes(), 50 * PERMIT_GRANULARITY_BYTES); + assert_eq!(manager.used_bytes(), 50 * PERMIT_GRANULARITY_BYTES); } #[test] @@ -406,6 +412,6 @@ async fn test_acquire_additional_unlimited() { .acquire_additional(1000 * PERMIT_GRANULARITY_BYTES) .await .unwrap(); - assert_eq!(guard.granted_bytes(), 0); - assert_eq!(manager.used_bytes(), 0); + assert_eq!(guard.granted_bytes(), 1000 * PERMIT_GRANULARITY_BYTES); + assert_eq!(manager.used_bytes(), 1000 * PERMIT_GRANULARITY_BYTES); } diff --git a/src/common/query/src/prelude.rs b/src/common/query/src/prelude.rs index c27b94294ea..50668bbbb19 100644 --- a/src/common/query/src/prelude.rs +++ b/src/common/query/src/prelude.rs @@ -27,7 +27,16 @@ static GREPTIME_TIMESTAMP_CELL: OnceCell = OnceCell::new(); static GREPTIME_VALUE_CELL: OnceCell = OnceCell::new(); pub fn set_default_prefix(prefix: Option<&str>) -> Result<()> { - match prefix { + // Strip surrounding double quotes as a defensive measure against upstream + // sources (scripts, CI, template engines, incorrect shell escaping) that may + // pass literal `""` as the value instead of an empty string. + let stripped = prefix.map(|s| { + s.strip_prefix('"') + .and_then(|s| s.strip_suffix('"')) + .unwrap_or(s) + }); + + match stripped { None => { // use default greptime prefix GREPTIME_TIMESTAMP_CELL.get_or_init(|| GREPTIME_TIMESTAMP.to_string()); @@ -70,3 +79,45 @@ const GREPTIME_VALUE: &str = "greptime_value"; pub const GREPTIME_COUNT: &str = "greptime_count"; /// Default physical table name pub const GREPTIME_PHYSICAL_TABLE: &str = "greptime_physical_table"; + +#[cfg(test)] +mod tests { + use super::*; + + // Each test runs in a separate process via `cargo nextest`, so OnceCell + // state does not leak between tests. + + #[test] + fn test_set_default_prefix_none() { + set_default_prefix(None).unwrap(); + assert_eq!(greptime_timestamp(), "greptime_timestamp"); + assert_eq!(greptime_value(), "greptime_value"); + } + + #[test] + fn test_set_default_prefix_empty_string() { + set_default_prefix(Some("")).unwrap(); + assert_eq!(greptime_timestamp(), "timestamp"); + assert_eq!(greptime_value(), "value"); + } + + #[test] + fn test_set_default_prefix_quoted_empty() { + // Handles upstream sources that pass literal `""` instead of an empty string + set_default_prefix(Some("\"\"")).unwrap(); + assert_eq!(greptime_timestamp(), "timestamp"); + assert_eq!(greptime_value(), "value"); + } + + #[test] + fn test_set_default_prefix_custom() { + set_default_prefix(Some("mydb")).unwrap(); + assert_eq!(greptime_timestamp(), "mydb_timestamp"); + assert_eq!(greptime_value(), "mydb_value"); + } + + #[test] + fn test_set_default_prefix_invalid() { + assert!(set_default_prefix(Some("invalid prefix!")).is_err()); + } +} diff --git a/src/datatypes/src/schema.rs b/src/datatypes/src/schema.rs index 9070e2babee..50f2dba2709 100644 --- a/src/datatypes/src/schema.rs +++ b/src/datatypes/src/schema.rs @@ -16,8 +16,8 @@ mod column_schema; pub mod constraint; use std::collections::HashMap; -use std::fmt; use std::sync::Arc; +use std::{fmt, mem}; use arrow::datatypes::{Field, Schema as ArrowSchema}; use datafusion_common::DFSchemaRef; @@ -177,6 +177,26 @@ impl Schema { &self.arrow_schema.metadata } + /// Returns the estimated memory footprint of this schema. + pub fn estimated_size(&self) -> usize { + mem::size_of_val(self) + + mem::size_of::() * self.column_schemas.capacity() + + self + .column_schemas + .iter() + .map(|column_schema| { + column_schema.estimated_size() - mem::size_of::() + }) + .sum::() + + mem::size_of::<(String, usize)>() * self.name_to_index.capacity() + + self + .name_to_index + .keys() + .map(|name| name.capacity()) + .sum::() + + arrow_schema_size(self.arrow_schema.as_ref()) + } + /// Generate a new projected schema /// /// # Panic @@ -213,6 +233,17 @@ impl Schema { } } +fn arrow_schema_size(schema: &ArrowSchema) -> usize { + mem::size_of_val(schema) + + schema.fields.size() + + mem::size_of::<(String, String)>() * schema.metadata.capacity() + + schema + .metadata + .iter() + .map(|(key, value)| key.capacity() + value.capacity()) + .sum::() +} + #[derive(Default)] pub struct SchemaBuilder { column_schemas: Vec, diff --git a/src/datatypes/src/schema/column_schema.rs b/src/datatypes/src/schema/column_schema.rs index 183cf05da8f..2479f4fc418 100644 --- a/src/datatypes/src/schema/column_schema.rs +++ b/src/datatypes/src/schema/column_schema.rs @@ -13,8 +13,8 @@ // limitations under the License. use std::collections::HashMap; -use std::fmt; use std::str::FromStr; +use std::{fmt, mem}; use arrow::datatypes::Field; use arrow_schema::extension::{ @@ -178,6 +178,19 @@ impl ColumnSchema { self } + /// Returns the estimated memory footprint of this schema. + pub fn estimated_size(&self) -> usize { + mem::size_of_val(self) - mem::size_of_val(&self.data_type) + + self.data_type.as_arrow_type().size() + + self.name.capacity() + + self + .default_constraint + .as_ref() + .map(column_default_constraint_size) + .unwrap_or_default() + + metadata_size(&self.metadata) + } + /// Set the inverted index for the column. /// Similar to [with_inverted_index] but don't take the ownership. /// @@ -493,6 +506,21 @@ impl ColumnSchema { } } +fn metadata_size(metadata: &Metadata) -> usize { + mem::size_of::<(String, String)>() * metadata.capacity() + + metadata + .iter() + .map(|(key, value)| key.capacity() + value.capacity()) + .sum::() +} + +fn column_default_constraint_size(default_constraint: &ColumnDefaultConstraint) -> usize { + match default_constraint { + ColumnDefaultConstraint::Function(expr) => expr.capacity(), + ColumnDefaultConstraint::Value(value) => value.as_value_ref().data_size(), + } +} + /// Column extended type set in column schema's metadata. #[derive(Debug, Clone, PartialEq, Eq)] pub enum ColumnExtType { diff --git a/src/datatypes/src/types/json_type.rs b/src/datatypes/src/types/json_type.rs index 61586fc460f..912bbfca54b 100644 --- a/src/datatypes/src/types/json_type.rs +++ b/src/datatypes/src/types/json_type.rs @@ -396,7 +396,7 @@ pub fn jsonb_to_string(val: &[u8]) -> Result { match jsonb::from_slice(val) { Ok(jsonb_value) => { let serialized = jsonb_value.to_string(); - Ok(serialized) + fix_unicode_point(&serialized) } Err(e) => InvalidJsonbSnafu { error: e }.fail(), } @@ -405,18 +405,12 @@ pub fn jsonb_to_string(val: &[u8]) -> Result { /// Converts a json type value to serde_json::Value pub fn jsonb_to_serde_json(val: &[u8]) -> Result { let json_string = jsonb_to_string(val)?; - jsonb_string_to_serde_value(&json_string) + serde_json::Value::from_str(&json_string).context(DeserializeSnafu { json: json_string }) } -/// Attempts to deserialize a JSON text into `serde_json::Value`, with a best-effort -/// fallback for Rust-style Unicode escape sequences. +/// Normalizes a JSON string by converting Rust-style Unicode escape sequences to JSON-compatible format. /// -/// This function is intended to be used on JSON strings produced from the internal -/// JSONB representation (e.g. via [`jsonb_to_string`]). It first calls -/// `serde_json::Value::from_str` directly. If that succeeds, the parsed value is -/// returned as-is. -/// -/// If the initial parse fails, the input is scanned for Rust-style Unicode code +/// The input is scanned for Rust-style Unicode code /// point escapes of the form `\\u{H...}` (a backslash, `u`, an opening brace, /// followed by 1–6 hexadecimal digits, and a closing brace). Each such escape is /// converted into JSON-compatible UTF‑16 escape sequences: @@ -427,59 +421,44 @@ pub fn jsonb_to_serde_json(val: &[u8]) -> Result { /// the code point is encoded as a UTF‑16 surrogate pair and emitted as two consecutive /// `\\uXXXX` sequences (as JSON format required). /// -/// After this normalization, the function retries parsing the resulting string as -/// JSON and returns the deserialized value or a `DeserializeSnafu` error if it -/// still cannot be parsed. -fn jsonb_string_to_serde_value(json: &str) -> Result { - match serde_json::Value::from_str(json) { - Ok(v) => Ok(v), - Err(e) => { - // If above deserialization is failed, the JSON string might contain some Rust chars - // that are somehow incorrectly represented as Unicode code point literal. For example, - // "\u{fe0f}". We have to convert them to JSON compatible format, like "\uFE0F", then - // try to deserialize the JSON string again. - if !e.is_syntax() || !e.to_string().contains("invalid escape") { - return Err(e).context(DeserializeSnafu { json }); - } +/// After this normalization, the function returns the normalized string +fn fix_unicode_point(json: &str) -> Result { + static UNICODE_CODE_POINT_PATTERN: LazyLock = LazyLock::new(|| { + // Match literal "\u{...}" sequences, capturing 1–6 (code point range) hex digits + // inside braces. + Regex::new(r"\\u\{([0-9a-fA-F]{1,6})}").unwrap_or_else(|e| panic!("{}", e)) + }); - static UNICODE_CODE_POINT_PATTERN: LazyLock = LazyLock::new(|| { - // Match literal "\u{...}" sequences, capturing 1–6 (code point range) hex digits - // inside braces. - Regex::new(r"\\u\{([0-9a-fA-F]{1,6})}").unwrap_or_else(|e| panic!("{}", e)) - }); + let v = UNICODE_CODE_POINT_PATTERN.replace_all(json, |caps: &Captures| { + // Extract the hex payload (without braces) and parse to a code point. + let hex = &caps[1]; + let Ok(code) = u32::from_str_radix(hex, 16) else { + // On parse failure, leave the original escape sequence unchanged. + return caps[0].to_string(); + }; - let v = UNICODE_CODE_POINT_PATTERN.replace_all(json, |caps: &Captures| { - // Extract the hex payload (without braces) and parse to a code point. - let hex = &caps[1]; - let Ok(code) = u32::from_str_radix(hex, 16) else { - // On parse failure, leave the original escape sequence unchanged. - return caps[0].to_string(); - }; + if code <= 0xFFFF { + // Basic Multilingual Plane: JSON can represent this directly as \uXXXX. + format!("\\u{:04X}", code) + } else if code > 0x10FFFF { + // Beyond max Unicode code point + caps[0].to_string() + } else { + // Supplementary planes: JSON needs UTF-16 surrogate pairs. + // Convert the code point to a 20-bit value. + let code = code - 0x10000; - if code <= 0xFFFF { - // Basic Multilingual Plane: JSON can represent this directly as \uXXXX. - format!("\\u{:04X}", code) - } else if code > 0x10FFFF { - // Beyond max Unicode code point - caps[0].to_string() - } else { - // Supplementary planes: JSON needs UTF-16 surrogate pairs. - // Convert the code point to a 20-bit value. - let code = code - 0x10000; + // High surrogate: top 10 bits, offset by 0xD800. + let high = 0xD800 + ((code >> 10) & 0x3FF); - // High surrogate: top 10 bits, offset by 0xD800. - let high = 0xD800 + ((code >> 10) & 0x3FF); + // Low surrogate: bottom 10 bits, offset by 0xDC00. + let low = 0xDC00 + (code & 0x3FF); - // Low surrogate: bottom 10 bits, offset by 0xDC00. - let low = 0xDC00 + (code & 0x3FF); - - // Emit two \uXXXX escapes in sequence. - format!("\\u{:04X}\\u{:04X}", high, low) - } - }); - serde_json::Value::from_str(&v).context(DeserializeSnafu { json }) + // Emit two \uXXXX escapes in sequence. + format!("\\u{:04X}\\u{:04X}", high, low) } - } + }); + Ok(v.to_string()) } /// Parses a string to a json type value @@ -495,45 +474,54 @@ mod tests { use crate::json::JsonStructureSettings; #[test] - fn test_jsonb_string_to_serde_value() -> Result<()> { + fn test_fix_unicode_point() -> Result<()> { let valid_cases = vec![ - (r#"{"data": "simple ascii"}"#, r#"{"data":"simple ascii"}"#), + (r#"{"data": "simple ascii"}"#, r#"{"data": "simple ascii"}"#), ( - r#"{"data": "Greek sigma: \u{03a3}"}"#, - r#"{"data":"Greek sigma: Σ"}"#, + r#"{"data":"Greek sigma: \u{03a3}"}"#, + r#"{"data":"Greek sigma: \u03A3"}"#, ), ( - r#"{"data": "Joker card: \u{1f0df}"}"#, - r#"{"data":"Joker card: 🃟"}"#, + r#"{"data":"Joker card: \u{1f0df}"}"#, + r#"{"data":"Joker card: \uD83C\uDCDF"}"#, ), ( - r#"{"data": "BMP boundary: \u{ffff}"}"#, - r#"{"data":"BMP boundary: ￿"}"#, + r#"{"data":"BMP boundary: \u{ffff}"}"#, + r#"{"data":"BMP boundary: \uFFFF"}"#, ), ( - r#"{"data": "Supplementary min: \u{10000}"}"#, - r#"{"data":"Supplementary min: 𐀀"}"#, + r#"{"data":"Supplementary min: \u{10000}"}"#, + r#"{"data":"Supplementary min: \uD800\uDC00"}"#, ), ( - r#"{"data": "Supplementary max: \u{10ffff}"}"#, - r#"{"data":"Supplementary max: 􏿿"}"#, + r#"{"data":"Supplementary max: \u{10ffff}"}"#, + r#"{"data":"Supplementary max: \uDBFF\uDFFF"}"#, ), ]; for (input, expect) in valid_cases { - let v = jsonb_string_to_serde_value(input)?; - assert_eq!(v.to_string(), expect); + let v = fix_unicode_point(input)?; + assert_eq!(v, expect); } - let invalid_cases = vec![ - r#"{"data": "Invalid hex: \u{gggg}"}"#, - r#"{"data": "Beyond max Unicode code point: \u{110000}"}"#, - r#"{"data": "Out of range: \u{1100000}"}"#, // 7 digit - r#"{"data": "Empty braces: \u{}"}"#, + let invalid_escape_cases = vec![ + ( + r#"{"data": "Invalid hex: \u{gggg}"}"#, + r#"{"data": "Invalid hex: \u{gggg}"}"#, + ), + ( + r#"{"data": "Empty braces: \u{}"}"#, + r#"{"data": "Empty braces: \u{}"}"#, + ), + ( + r#"{"data": "Out of range: \u{1100000}"}"#, + r#"{"data": "Out of range: \u{1100000}"}"#, + ), ]; - for input in invalid_cases { - let result = jsonb_string_to_serde_value(input); - assert!(result.is_err()); + for (input, expect) in invalid_escape_cases { + let v = fix_unicode_point(input)?; + assert_eq!(v, expect); } + Ok(()) } diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index fa8a74cad2b..ce589bb6774 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -13,6 +13,7 @@ // limitations under the License. pub mod builder; +mod dashboard; mod grpc; mod influxdb; mod jaeger; diff --git a/src/frontend/src/instance/dashboard.rs b/src/frontend/src/instance/dashboard.rs new file mode 100644 index 00000000000..373961dbfaf --- /dev/null +++ b/src/frontend/src/instance/dashboard.rs @@ -0,0 +1,405 @@ +// 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. + +use std::collections::HashMap; +use std::sync::Arc; + +use api::v1::value::ValueData; +use api::v1::{ + ColumnDataType, ColumnDef, ColumnSchema as PbColumnSchema, Row, RowInsertRequest, + RowInsertRequests, Rows, SemanticType, +}; +use async_trait::async_trait; +use common_catalog::consts::{DEFAULT_PRIVATE_SCHEMA_NAME, default_engine}; +use common_error::ext::BoxedError; +use common_query::OutputData; +use common_recordbatch::util as record_util; +use common_telemetry::info; +use common_time::FOREVER; +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::{ + CatalogSnafu, CollectRecordbatchSnafu, DataFusionSnafu, ExecuteQuerySnafu, NotSupportedSnafu, + TableNotFoundSnafu, +}; +use servers::query_handler::DashboardDefinition; +use session::context::{QueryContextBuilder, QueryContextRef}; +use snafu::{OptionExt, ResultExt}; +use table::TableRef; +use table::metadata::TableInfo; +use table::requests::TTL_KEY; +use table::table::adapter::DfTableProviderAdapter; + +use crate::instance::Instance; + +pub const DASHBOARD_TABLE_NAME: &str = "dashboard"; +pub const DASHBOARD_TABLE_NAME_COLUMN_NAME: &str = "name"; +pub const DASHBOARD_TABLE_DEFINITION_COLUMN_NAME: &str = "definition"; +pub const DASHBOARD_TABLE_CREATED_AT_COLUMN_NAME: &str = "created_at"; + +impl Instance { + /// Build a schema for dashboard table. + /// Returns the (time index, primary keys, column) definitions. + fn build_dashboard_schema() -> (String, Vec, Vec) { + ( + DASHBOARD_TABLE_CREATED_AT_COLUMN_NAME.to_string(), + vec![DASHBOARD_TABLE_NAME_COLUMN_NAME.to_string()], + vec![ + ColumnDef { + name: DASHBOARD_TABLE_NAME_COLUMN_NAME.to_string(), + data_type: ColumnDataType::String as i32, + is_nullable: false, + default_constraint: vec![], + semantic_type: SemanticType::Tag as i32, + comment: String::new(), + datatype_extension: None, + options: None, + }, + ColumnDef { + name: DASHBOARD_TABLE_DEFINITION_COLUMN_NAME.to_string(), + data_type: ColumnDataType::String as i32, + is_nullable: false, + default_constraint: vec![], + semantic_type: SemanticType::Field as i32, + comment: String::new(), + datatype_extension: None, + options: None, + }, + ColumnDef { + name: DASHBOARD_TABLE_CREATED_AT_COLUMN_NAME.to_string(), + data_type: ColumnDataType::TimestampNanosecond as i32, + is_nullable: false, + default_constraint: vec![], + semantic_type: SemanticType::Timestamp as i32, + comment: String::new(), + datatype_extension: None, + options: None, + }, + ], + ) + } + + /// Build a column schemas for inserting a row into the dashboard table. + fn build_dashboard_insert_column_schemas() -> Vec { + vec![ + PbColumnSchema { + column_name: DASHBOARD_TABLE_NAME_COLUMN_NAME.to_string(), + datatype: ColumnDataType::String.into(), + semantic_type: SemanticType::Tag.into(), + ..Default::default() + }, + PbColumnSchema { + column_name: DASHBOARD_TABLE_DEFINITION_COLUMN_NAME.to_string(), + datatype: ColumnDataType::String.into(), + semantic_type: SemanticType::Field.into(), + ..Default::default() + }, + PbColumnSchema { + column_name: DASHBOARD_TABLE_CREATED_AT_COLUMN_NAME.to_string(), + datatype: ColumnDataType::TimestampNanosecond.into(), + semantic_type: SemanticType::Timestamp.into(), + ..Default::default() + }, + ] + } + + fn dashboard_query_ctx(table_info: &TableInfo) -> QueryContextRef { + QueryContextBuilder::default() + .current_catalog(table_info.catalog_name.clone()) + .current_schema(table_info.schema_name.clone()) + .build() + .into() + } + + async fn create_dashboard_table_if_not_exists( + &self, + ctx: QueryContextRef, + ) -> servers::error::Result { + let catalog = ctx.current_catalog(); + + if let Some(table) = self + .catalog_manager + .table( + catalog, + DEFAULT_PRIVATE_SCHEMA_NAME, + DASHBOARD_TABLE_NAME, + Some(&ctx), + ) + .await + .context(CatalogSnafu)? + { + return Ok(table); + } + + let (time_index, primary_keys, column_defs) = Self::build_dashboard_schema(); + + let mut table_options = HashMap::new(); + table_options.insert(TTL_KEY.to_string(), FOREVER.to_string()); + + let mut create_table_expr = api::v1::CreateTableExpr { + catalog_name: catalog.to_string(), + schema_name: DEFAULT_PRIVATE_SCHEMA_NAME.to_string(), + table_name: DASHBOARD_TABLE_NAME.to_string(), + desc: "GreptimeDB dashboard table".to_string(), + column_defs, + time_index, + primary_keys, + create_if_not_exists: true, + table_options, + table_id: None, + engine: default_engine().to_string(), + }; + + self.statement_executor + .create_table_inner(&mut create_table_expr, None, ctx.clone()) + .await + .map_err(BoxedError::new) + .context(ExecuteQuerySnafu)?; + + let table = self + .catalog_manager + .table( + catalog, + DEFAULT_PRIVATE_SCHEMA_NAME, + DASHBOARD_TABLE_NAME, + Some(&ctx), + ) + .await + .context(CatalogSnafu)? + .context(TableNotFoundSnafu { + catalog: catalog.to_string(), + schema: DEFAULT_PRIVATE_SCHEMA_NAME.to_string(), + table: DASHBOARD_TABLE_NAME.to_string(), + })?; + + Ok(table) + } + + /// Insert a dashboard into the dashboard table. + async fn insert_dashboard( + &self, + name: &str, + definition: &str, + query_ctx: QueryContextRef, + ) -> servers::error::Result<()> { + let table = self + .create_dashboard_table_if_not_exists(query_ctx.clone()) + .await?; + let table_info = table.table_info(); + + let insert = RowInsertRequest { + table_name: DASHBOARD_TABLE_NAME.to_string(), + rows: Some(Rows { + schema: Self::build_dashboard_insert_column_schemas(), + rows: vec![Row { + values: vec![ + ValueData::StringValue(name.to_string()).into(), + ValueData::StringValue(definition.to_string()).into(), + ValueData::TimestampNanosecondValue(0).into(), + ], + }], + }), + }; + + let requests = RowInsertRequests { + inserts: vec![insert], + }; + + let output = self + .inserter + .handle_row_inserts( + requests, + Self::dashboard_query_ctx(&table_info), + &self.statement_executor, + false, + false, + ) + .await + .map_err(BoxedError::new) + .context(ExecuteQuerySnafu)?; + + info!( + "Insert dashboard success, name: {}, table: {}, output: {:?}", + name, + table_info.full_table_name(), + output + ); + + Ok(()) + } + + /// List all dashboards. + async fn list_dashboards( + &self, + query_ctx: QueryContextRef, + ) -> servers::error::Result> { + let table = if let Some(table) = self + .catalog_manager + .table( + query_ctx.current_catalog(), + DEFAULT_PRIVATE_SCHEMA_NAME, + DASHBOARD_TABLE_NAME, + Some(&query_ctx), + ) + .await + .context(CatalogSnafu)? + { + table + } else { + return Ok(vec![]); + }; + + let table_info = table.table_info(); + + let dataframe = self + .query_engine + .read_table(table.clone()) + .map_err(BoxedError::new) + .context(ExecuteQuerySnafu)?; + + let dataframe = dataframe + .select_columns(&[ + DASHBOARD_TABLE_NAME_COLUMN_NAME, + DASHBOARD_TABLE_DEFINITION_COLUMN_NAME, + ]) + .context(DataFusionSnafu)?; + + let plan = dataframe.into_parts().1; + + let output = self + .query_engine + .execute(plan, Self::dashboard_query_ctx(&table_info)) + .await + .map_err(BoxedError::new) + .context(ExecuteQuerySnafu)?; + + let stream = match output.data { + OutputData::Stream(stream) => stream, + OutputData::RecordBatches(record_batches) => record_batches.as_stream(), + _ => unreachable!(), + }; + + let records = record_util::collect(stream) + .await + .context(CollectRecordbatchSnafu)?; + + let mut dashboards = Vec::new(); + + for r in &records { + let name_column = r.column(0); + let definition_column = r.column(1); + + let name = name_column + .as_string_opt::() + .context(NotSupportedSnafu { + feat: "Invalid data type for greptime_private.dashboard.name", + })?; + + let definition = + definition_column + .as_string_opt::() + .context(NotSupportedSnafu { + feat: "Invalid data type for greptime_private.dashboard.definition", + })?; + + for i in 0..name.len() { + dashboards.push(DashboardDefinition { + name: name.value(i).to_string(), + definition: definition.value(i).to_string(), + }); + } + } + + Ok(dashboards) + } + + /// Delete a dashboard by name. + async fn delete_dashboard( + &self, + name: &str, + query_ctx: QueryContextRef, + ) -> servers::error::Result<()> { + let table = self + .create_dashboard_table_if_not_exists(query_ctx.clone()) + .await?; + let table_info = table.table_info(); + + let dataframe = self + .query_engine + .read_table(table.clone()) + .map_err(BoxedError::new) + .context(ExecuteQuerySnafu)?; + + let name_condition = col(DASHBOARD_TABLE_NAME_COLUMN_NAME).eq(lit(name)); + + let dataframe = dataframe.filter(name_condition).context(DataFusionSnafu)?; + + let table_name = TableReference::full( + table_info.catalog_name.clone(), + table_info.schema_name.clone(), + table_info.name.clone(), + ); + + let table_provider = Arc::new(DfTableProviderAdapter::new(table.clone())); + let table_source = Arc::new(DefaultTableSource::new(table_provider)); + + let stmt = DmlStatement::new( + table_name, + table_source, + datafusion_expr::WriteOp::Delete, + Arc::new(dataframe.into_parts().1), + ); + + let plan = LogicalPlan::Dml(stmt); + + let output = self + .query_engine + .execute(plan, Self::dashboard_query_ctx(&table_info)) + .await + .map_err(BoxedError::new) + .context(ExecuteQuerySnafu)?; + + info!( + "Delete dashboard success, name: {}, table: {}, output: {:?}", + name, + table_info.full_table_name(), + output + ); + + Ok(()) + } +} + +#[async_trait] +impl servers::query_handler::DashboardHandler for Instance { + async fn save( + &self, + name: &str, + definition: &str, + ctx: QueryContextRef, + ) -> servers::error::Result<()> { + self.insert_dashboard(name, definition, ctx).await + } + + async fn list(&self, ctx: QueryContextRef) -> servers::error::Result> { + self.list_dashboards(ctx).await + } + + async fn delete(&self, name: &str, ctx: QueryContextRef) -> servers::error::Result<()> { + self.delete_dashboard(name, ctx).await + } +} diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 45c3ec3649e..4b51efbd336 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -143,6 +143,8 @@ where builder = builder.with_jaeger_handler(self.instance.clone()); } + builder = builder.with_dashboard_handler(self.instance.clone()); + if let Some(configurator) = self.plugins.get::() { info!("Adding extra router from plugins"); builder = builder.with_extra_router(configurator.router()); diff --git a/src/metric-engine/Cargo.toml b/src/metric-engine/Cargo.toml index 567210b9523..5b561997abf 100644 --- a/src/metric-engine/Cargo.toml +++ b/src/metric-engine/Cargo.toml @@ -17,6 +17,7 @@ bytes.workspace = true fxhash = "0.2" common-base.workspace = true common-error.workspace = true +common-grpc.workspace = true common-macro.workspace = true common-query.workspace = true common-recordbatch.workspace = true diff --git a/src/metric-engine/src/batch_modifier.rs b/src/metric-engine/src/batch_modifier.rs new file mode 100644 index 00000000000..8a5774889b2 --- /dev/null +++ b/src/metric-engine/src/batch_modifier.rs @@ -0,0 +1,426 @@ +// 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. + +use std::hash::Hasher; +use std::sync::Arc; + +use datatypes::arrow::array::{Array, BinaryBuilder, StringArray, UInt64Array}; +use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::value::ValueRef; +use fxhash::FxHasher; +use mito_codec::row_converter::SparsePrimaryKeyCodec; +use snafu::ResultExt; +use store_api::storage::ColumnId; +use store_api::storage::consts::{PRIMARY_KEY_COLUMN_NAME, ReservedColumnId}; + +use crate::error::{EncodePrimaryKeySnafu, Result, UnexpectedRequestSnafu}; + +/// Info about a tag column for TSID computation and sparse primary key encoding. +#[allow(dead_code)] +pub(crate) struct TagColumnInfo { + /// Column name (used for label-name hash). + pub name: String, + /// Column index in the RecordBatch. + pub index: usize, + /// Column ID in the physical region. + pub column_id: ColumnId, +} + +/// Computes `__tsid` values for each row. +#[allow(dead_code)] +pub(crate) fn compute_tsid_array( + batch: &RecordBatch, + sorted_tag_columns: &[TagColumnInfo], + tag_arrays: &[&StringArray], +) -> UInt64Array { + let num_rows = batch.num_rows(); + + let label_name_hash = { + let mut hasher = FxHasher::default(); + for tag_col in sorted_tag_columns { + hasher.write(tag_col.name.as_bytes()); + hasher.write_u8(0xff); + } + hasher.finish() + }; + + let mut tsid_values = Vec::with_capacity(num_rows); + for row in 0..num_rows { + let has_null = tag_arrays.iter().any(|arr| arr.is_null(row)); + + let tsid = if !has_null { + let mut hasher = FxHasher::default(); + hasher.write_u64(label_name_hash); + for arr in tag_arrays { + hasher.write(arr.value(row).as_bytes()); + hasher.write_u8(0xff); + } + hasher.finish() + } else { + let mut name_hasher = FxHasher::default(); + for (tc, arr) in sorted_tag_columns.iter().zip(tag_arrays.iter()) { + if !arr.is_null(row) { + name_hasher.write(tc.name.as_bytes()); + name_hasher.write_u8(0xff); + } + } + let row_label_hash = name_hasher.finish(); + + let mut val_hasher = FxHasher::default(); + val_hasher.write_u64(row_label_hash); + for arr in tag_arrays { + if !arr.is_null(row) { + val_hasher.write(arr.value(row).as_bytes()); + val_hasher.write_u8(0xff); + } + } + val_hasher.finish() + }; + + tsid_values.push(tsid); + } + + UInt64Array::from(tsid_values) +} + +fn build_tag_arrays<'a>( + batch: &'a RecordBatch, + sorted_tag_columns: &[TagColumnInfo], +) -> Vec<&'a StringArray> { + sorted_tag_columns + .iter() + .map(|tc| { + batch + .column(tc.index) + .as_any() + .downcast_ref::() + .expect("tag column must be utf8") + }) + .collect() +} + +/// Modifies a RecordBatch for sparse primary key encoding. +#[allow(dead_code)] +pub(crate) fn modify_batch_sparse( + batch: RecordBatch, + table_id: u32, + sorted_tag_columns: &[TagColumnInfo], + non_tag_column_indices: &[usize], +) -> Result { + let num_rows = batch.num_rows(); + let codec = SparsePrimaryKeyCodec::schemaless(); + let tag_arrays: Vec<&StringArray> = build_tag_arrays(&batch, sorted_tag_columns); + let tsid_array = compute_tsid_array(&batch, sorted_tag_columns, &tag_arrays); + + let mut pk_builder = BinaryBuilder::with_capacity(num_rows, 0); + let mut buffer = Vec::new(); + for row in 0..num_rows { + buffer.clear(); + let internal = [ + (ReservedColumnId::table_id(), ValueRef::UInt32(table_id)), + ( + ReservedColumnId::tsid(), + ValueRef::UInt64(tsid_array.value(row)), + ), + ]; + codec + .encode_to_vec(internal.into_iter(), &mut buffer) + .context(EncodePrimaryKeySnafu)?; + + let tags = sorted_tag_columns + .iter() + .zip(tag_arrays.iter()) + .filter(|(_, arr)| !arr.is_null(row)) + .map(|(tc, arr)| (tc.column_id, ValueRef::String(arr.value(row)))); + codec + .encode_to_vec(tags, &mut buffer) + .context(EncodePrimaryKeySnafu)?; + + pk_builder.append_value(&buffer); + } + + let pk_array = pk_builder.finish(); + + let mut fields = vec![Arc::new(Field::new( + PRIMARY_KEY_COLUMN_NAME, + DataType::Binary, + false, + ))]; + let mut columns: Vec> = vec![Arc::new(pk_array)]; + + for &idx in non_tag_column_indices { + fields.push(batch.schema().fields()[idx].clone()); + columns.push(batch.column(idx).clone()); + } + + let new_schema = Arc::new(ArrowSchema::new(fields)); + RecordBatch::try_new(new_schema, columns).map_err(|e| { + UnexpectedRequestSnafu { + reason: format!("Failed to build modified sparse RecordBatch: {e}"), + } + .build() + }) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, ColumnSchema, Row, Rows, SemanticType, Value}; + use datatypes::arrow::array::{BinaryArray, Int64Array, StringArray}; + use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; + use datatypes::arrow::record_batch::RecordBatch; + use store_api::codec::PrimaryKeyEncoding; + use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; + + use super::*; + use crate::row_modifier::{RowModifier, RowsIter, TableIdInput}; + + fn build_sparse_test_batch() -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("greptime_timestamp", DataType::Int64, false), + Field::new("greptime_value", DataType::Float64, true), + Field::new("namespace", DataType::Utf8, true), + Field::new("host", DataType::Utf8, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1000])), + Arc::new(datatypes::arrow::array::Float64Array::from(vec![42.0])), + Arc::new(StringArray::from(vec!["greptimedb"])), + Arc::new(StringArray::from(vec!["127.0.0.1"])), + ], + ) + .unwrap() + } + + fn sparse_tag_columns() -> Vec { + vec![ + TagColumnInfo { + name: "host".to_string(), + index: 3, + column_id: 3, + }, + TagColumnInfo { + name: "namespace".to_string(), + index: 2, + column_id: 2, + }, + ] + } + + #[test] + fn test_compute_tsid_basic() { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("namespace", DataType::Utf8, true), + Field::new("host", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(vec!["greptimedb"])), + Arc::new(StringArray::from(vec!["127.0.0.1"])), + ], + ) + .unwrap(); + + let tag_columns: Vec = vec![ + TagColumnInfo { + name: "host".to_string(), + index: 1, + column_id: 2, + }, + TagColumnInfo { + name: "namespace".to_string(), + index: 0, + column_id: 1, + }, + ]; + let tag_arrays = build_tag_arrays(&batch, &tag_columns); + let tsid_array = compute_tsid_array(&batch, &tag_columns, &tag_arrays); + + assert_eq!(tsid_array.value(0), 2721566936019240841); + } + + #[test] + fn test_compute_tsid_with_nulls() { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("a", DataType::Utf8, true), + Field::new("b", DataType::Utf8, true), + ])); + let batch_no_null = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["A"])), + Arc::new(StringArray::from(vec!["B"])), + ], + ) + .unwrap(); + let tag_cols_2: Vec = vec![ + TagColumnInfo { + name: "a".to_string(), + index: 0, + column_id: 1, + }, + TagColumnInfo { + name: "b".to_string(), + index: 1, + column_id: 2, + }, + ]; + let tag_arrays_2 = build_tag_arrays(&batch_no_null, &tag_cols_2); + let tsid_no_null = compute_tsid_array(&batch_no_null, &tag_cols_2, &tag_arrays_2); + + let schema3 = Arc::new(ArrowSchema::new(vec![ + Field::new("a", DataType::Utf8, true), + Field::new("b", DataType::Utf8, true), + Field::new("c", DataType::Utf8, true), + ])); + let batch_with_null = RecordBatch::try_new( + schema3, + vec![ + Arc::new(StringArray::from(vec!["A"])), + Arc::new(StringArray::from(vec!["B"])), + Arc::new(StringArray::from(vec![None as Option<&str>])), + ], + ) + .unwrap(); + let tag_cols_3: Vec = vec![ + TagColumnInfo { + name: "a".to_string(), + index: 0, + column_id: 1, + }, + TagColumnInfo { + name: "b".to_string(), + index: 1, + column_id: 2, + }, + TagColumnInfo { + name: "c".to_string(), + index: 2, + column_id: 3, + }, + ]; + let tag_arrays_3 = build_tag_arrays(&batch_with_null, &tag_cols_3); + let tsid_with_null = compute_tsid_array(&batch_with_null, &tag_cols_3, &tag_arrays_3); + + assert_eq!(tsid_no_null.value(0), tsid_with_null.value(0)); + } + + #[test] + fn test_modify_batch_sparse() { + let batch = build_sparse_test_batch(); + let tag_columns = sparse_tag_columns(); + let non_tag_indices = vec![0, 1]; + let table_id: u32 = 1025; + + let modified = + modify_batch_sparse(batch, table_id, &tag_columns, &non_tag_indices).unwrap(); + + assert_eq!(modified.num_columns(), 3); + assert_eq!(modified.schema().field(0).name(), PRIMARY_KEY_COLUMN_NAME); + assert_eq!(modified.schema().field(1).name(), "greptime_timestamp"); + assert_eq!(modified.schema().field(2).name(), "greptime_value"); + } + + #[test] + fn test_modify_batch_sparse_matches_row_modifier() { + let batch = build_sparse_test_batch(); + let tag_columns = sparse_tag_columns(); + let non_tag_indices = vec![0, 1]; + let table_id: u32 = 1025; + let modified = + modify_batch_sparse(batch, table_id, &tag_columns, &non_tag_indices).unwrap(); + + let name_to_column_id: HashMap = [ + ("greptime_timestamp".to_string(), 0), + ("greptime_value".to_string(), 1), + ("namespace".to_string(), 2), + ("host".to_string(), 3), + ] + .into_iter() + .collect(); + + let rows = Rows { + schema: vec![ + ColumnSchema { + column_name: "greptime_timestamp".to_string(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "greptime_value".to_string(), + datatype: ColumnDataType::Float64 as i32, + semantic_type: SemanticType::Field as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "namespace".to_string(), + datatype: ColumnDataType::String as i32, + semantic_type: SemanticType::Tag as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "host".to_string(), + datatype: ColumnDataType::String as i32, + semantic_type: SemanticType::Tag as i32, + ..Default::default() + }, + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMillisecondValue(1000)), + }, + Value { + value_data: Some(ValueData::F64Value(42.0)), + }, + Value { + value_data: Some(ValueData::StringValue("greptimedb".to_string())), + }, + Value { + value_data: Some(ValueData::StringValue("127.0.0.1".to_string())), + }, + ], + }], + }; + + let row_iter = RowsIter::new(rows, &name_to_column_id); + let rows = RowModifier::default() + .modify_rows( + row_iter, + TableIdInput::Single(table_id), + PrimaryKeyEncoding::Sparse, + ) + .unwrap(); + let ValueData::BinaryValue(expected_pk) = + rows.rows[0].values[0].value_data.clone().unwrap() + else { + panic!("expected binary primary key"); + }; + + let actual_array = modified + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(actual_array.value(0), expected_pk.as_slice()); + } +} diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index 7a1efedac4e..ba90ca960df 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -13,6 +13,7 @@ // limitations under the License. mod alter; +mod bulk_insert; mod catchup; mod close; mod create; @@ -288,9 +289,8 @@ impl RegionEngine for MetricEngine { debug_assert_eq!(region_id, resp_region_id); return response; } - RegionRequest::BulkInserts(_) => { - // todo(hl): find a way to support bulk inserts in metric engine. - UnsupportedRegionRequestSnafu { request }.fail() + RegionRequest::BulkInserts(bulk) => { + self.inner.bulk_insert_region(region_id, bulk).await } }; diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs new file mode 100644 index 00000000000..2a3c26c80ca --- /dev/null +++ b/src/metric-engine/src/engine/bulk_insert.rs @@ -0,0 +1,783 @@ +// 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. + +use std::collections::HashSet; + +use api::v1::{ArrowIpc, ColumnDataType, SemanticType}; +use bytes::Bytes; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; +use common_grpc::flight::{FlightEncoder, FlightMessage}; +use common_query::prelude::{greptime_timestamp, greptime_value}; +use datatypes::arrow::array::{Array, Float64Array, StringArray, TimestampMillisecondArray}; +use datatypes::arrow::record_batch::RecordBatch; +use snafu::{OptionExt, ensure}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metadata::RegionMetadataRef; +use store_api::region_request::{ + AffectedRows, RegionBulkInsertsRequest, RegionPutRequest, RegionRequest, +}; +use store_api::storage::RegionId; + +use crate::batch_modifier::{TagColumnInfo, modify_batch_sparse}; +use crate::engine::MetricEngineInner; +use crate::error; +use crate::error::Result; + +impl MetricEngineInner { + /// Bulk-inserts logical rows into a metric region. + /// + /// This method accepts a `RegionBulkInsertsRequest` whose payload is a logical + /// `RecordBatch` (timestamp, value and tag columns) for the given logical `region_id`. + /// + /// The transformed batch is encoded to Arrow IPC and forwarded as a `BulkInserts` + /// request to the data region, along with the original `partition_expr_version`. + /// If the data region reports `StatusCode::Unsupported` for bulk inserts, the request + /// is transparently retried as a `Put` by converting the original logical batch into + /// `api::v1::Rows`, so callers observe the same semantics as `put_region`. + /// + /// Returns the number of affected rows, or `0` if the input batch is empty. + pub async fn bulk_insert_region( + &self, + region_id: RegionId, + request: RegionBulkInsertsRequest, + ) -> Result { + ensure!( + !self.is_physical_region(region_id), + error::UnsupportedRegionRequestSnafu { + request: RegionRequest::BulkInserts(request), + } + ); + + let (physical_region_id, data_region_id, primary_key_encoding) = + self.find_data_region_meta(region_id)?; + + if primary_key_encoding != PrimaryKeyEncoding::Sparse { + return error::UnsupportedRegionRequestSnafu { + request: RegionRequest::BulkInserts(request), + } + .fail(); + } + + let batch = request.payload; + if batch.num_rows() == 0 { + return Ok(0); + } + + let logical_metadata = self + .logical_region_metadata(physical_region_id, region_id) + .await?; + let (tag_columns, non_tag_indices) = self.resolve_tag_columns_from_metadata( + region_id, + data_region_id, + &batch, + &logical_metadata, + )?; + let modified_batch = modify_batch_sparse( + batch.clone(), + region_id.table_id(), + &tag_columns, + &non_tag_indices, + )?; + let (schema, data_header, payload) = record_batch_to_ipc(&modified_batch)?; + + let partition_expr_version = request.partition_expr_version; + let request = RegionBulkInsertsRequest { + region_id: data_region_id, + payload: modified_batch, + raw_data: ArrowIpc { + schema, + data_header, + payload, + }, + partition_expr_version, + }; + match self + .data_region + .write_data(data_region_id, RegionRequest::BulkInserts(request)) + .await + { + Ok(affected_rows) => Ok(affected_rows), + Err(err) if err.status_code() == StatusCode::Unsupported => { + // todo(hl): fallback path for PartitionTreeMemtable, remove this once we remove it + let rows = record_batch_to_rows(&batch, region_id)?; + self.put_region( + region_id, + RegionPutRequest { + rows, + hint: None, + partition_expr_version, + }, + ) + .await + } + Err(err) => Err(err), + } + } + + fn resolve_tag_columns_from_metadata( + &self, + logical_region_id: RegionId, + data_region_id: RegionId, + batch: &RecordBatch, + logical_metadata: &RegionMetadataRef, + ) -> Result<(Vec, Vec)> { + let tag_names: HashSet<&str> = logical_metadata + .column_metadatas + .iter() + .filter_map(|column| { + if column.semantic_type == SemanticType::Tag { + Some(column.column_schema.name.as_str()) + } else { + None + } + }) + .collect(); + + let mut tag_columns = Vec::new(); + let mut non_tag_indices = Vec::new(); + { + let state = self.state.read().unwrap(); + let physical_columns = state + .physical_region_states() + .get(&data_region_id) + .context(error::PhysicalRegionNotFoundSnafu { + region_id: data_region_id, + })? + .physical_columns(); + + for (index, field) in batch.schema().fields().iter().enumerate() { + let name = field.name(); + let column_id = + *physical_columns + .get(name) + .with_context(|| error::ColumnNotFoundSnafu { + name: name.clone(), + region_id: logical_region_id, + })?; + if tag_names.contains(name.as_str()) { + tag_columns.push(TagColumnInfo { + name: name.clone(), + index, + column_id, + }); + } else { + non_tag_indices.push(index); + } + } + } + + tag_columns.sort_by(|a, b| a.name.cmp(&b.name)); + Ok((tag_columns, non_tag_indices)) + } +} + +fn record_batch_to_rows(batch: &RecordBatch, logical_region_id: RegionId) -> Result { + let schema_ref = batch.schema(); + let fields = schema_ref.fields(); + + let mut ts_idx = None; + let mut val_idx = None; + let mut tag_indices = Vec::new(); + + for (idx, field) in fields.iter().enumerate() { + if field.name() == greptime_timestamp() { + ts_idx = Some(idx); + if !matches!( + field.data_type(), + datatypes::arrow::datatypes::DataType::Timestamp( + datatypes::arrow::datatypes::TimeUnit::Millisecond, + _ + ) + ) { + return error::UnexpectedRequestSnafu { + reason: format!( + "Timestamp column '{}' in region {:?} has incompatible type: {:?}", + field.name(), + logical_region_id, + field.data_type() + ), + } + .fail(); + } + } else if field.name() == greptime_value() { + val_idx = Some(idx); + if !matches!( + field.data_type(), + datatypes::arrow::datatypes::DataType::Float64 + ) { + return error::UnexpectedRequestSnafu { + reason: format!( + "Value column '{}' in region {:?} has incompatible type: {:?}", + field.name(), + logical_region_id, + field.data_type() + ), + } + .fail(); + } + } else { + if !matches!( + field.data_type(), + datatypes::arrow::datatypes::DataType::Utf8 + ) { + return error::UnexpectedRequestSnafu { + reason: format!( + "Tag column '{}' in region {:?} must be Utf8, found: {:?}", + field.name(), + logical_region_id, + field.data_type() + ), + } + .fail(); + } + tag_indices.push(idx); + } + } + + let ts_idx = ts_idx.with_context(|| error::UnexpectedRequestSnafu { + reason: format!( + "Timestamp column '{}' not found in RecordBatch for region {:?}", + greptime_timestamp(), + logical_region_id + ), + })?; + let val_idx = val_idx.with_context(|| error::UnexpectedRequestSnafu { + reason: format!( + "Value column '{}' not found in RecordBatch for region {:?}", + greptime_value(), + logical_region_id + ), + })?; + + let mut schema = Vec::with_capacity(2 + tag_indices.len()); + schema.push(api::v1::ColumnSchema { + column_name: greptime_timestamp().to_string(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + datatype_extension: None, + options: None, + }); + schema.push(api::v1::ColumnSchema { + column_name: greptime_value().to_string(), + datatype: ColumnDataType::Float64 as i32, + semantic_type: SemanticType::Field as i32, + datatype_extension: None, + options: None, + }); + for &idx in &tag_indices { + let field = &fields[idx]; + schema.push(api::v1::ColumnSchema { + column_name: field.name().clone(), + datatype: ColumnDataType::String as i32, + semantic_type: SemanticType::Tag as i32, + datatype_extension: None, + options: None, + }); + } + + let ts_array = batch + .column(ts_idx) + .as_any() + .downcast_ref::() + .expect("validated as TimestampMillisecond"); + let val_array = batch + .column(val_idx) + .as_any() + .downcast_ref::() + .expect("validated as Float64"); + let tag_arrays: Vec<&StringArray> = tag_indices + .iter() + .map(|&idx| { + batch + .column(idx) + .as_any() + .downcast_ref::() + .expect("validated as Utf8") + }) + .collect(); + + let num_rows = batch.num_rows(); + let mut rows = Vec::with_capacity(num_rows); + for row_idx in 0..num_rows { + let mut values = Vec::with_capacity(2 + tag_arrays.len()); + + if ts_array.is_null(row_idx) { + values.push(api::v1::Value { value_data: None }); + } else { + values.push(api::v1::Value { + value_data: Some(api::v1::value::ValueData::TimestampMillisecondValue( + ts_array.value(row_idx), + )), + }); + } + + if val_array.is_null(row_idx) { + values.push(api::v1::Value { value_data: None }); + } else { + values.push(api::v1::Value { + value_data: Some(api::v1::value::ValueData::F64Value( + val_array.value(row_idx), + )), + }); + } + + for arr in &tag_arrays { + if arr.is_null(row_idx) { + values.push(api::v1::Value { value_data: None }); + } else { + values.push(api::v1::Value { + value_data: Some(api::v1::value::ValueData::StringValue( + arr.value(row_idx).to_string(), + )), + }); + } + } + + rows.push(api::v1::Row { values }); + } + + Ok(api::v1::Rows { schema, rows }) +} + +fn record_batch_to_ipc(record_batch: &RecordBatch) -> Result<(Bytes, Bytes, Bytes)> { + let mut encoder = FlightEncoder::default(); + let schema = encoder.encode_schema(record_batch.schema().as_ref()); + let mut iter = encoder + .encode(FlightMessage::RecordBatch(record_batch.clone())) + .into_iter(); + + let Some(flight_data) = iter.next() else { + return error::UnexpectedRequestSnafu { + reason: "Failed to encode empty flight data", + } + .fail(); + }; + ensure!( + iter.next().is_none(), + error::UnexpectedRequestSnafu { + reason: "Bulk insert RecordBatch with dictionary arrays is unsupported".to_string(), + } + ); + + Ok(( + schema.data_header, + flight_data.data_header, + flight_data.data_body, + )) +} + +#[cfg(test)] +mod tests { + use std::assert_matches::assert_matches; + use std::sync::Arc; + + use api::v1::ArrowIpc; + use common_error::ext::ErrorExt; + use common_query::prelude::{greptime_timestamp, greptime_value}; + use common_recordbatch::RecordBatches; + use datatypes::arrow::array::{Float64Array, StringArray, TimestampMillisecondArray}; + use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema, TimeUnit}; + use datatypes::arrow::record_batch::RecordBatch; + use store_api::metric_engine_consts::MEMTABLE_PARTITION_TREE_PRIMARY_KEY_ENCODING; + use store_api::path_utils::table_dir; + use store_api::region_engine::RegionEngine; + use store_api::region_request::{RegionBulkInsertsRequest, RegionPutRequest, RegionRequest}; + use store_api::storage::{RegionId, ScanRequest}; + + use super::record_batch_to_ipc; + use crate::error::Error; + use crate::test_util::{self, TestEnv}; + + fn build_logical_batch(start: usize, rows: usize) -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + greptime_timestamp(), + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new(greptime_value(), DataType::Float64, true), + Field::new("job", DataType::Utf8, true), + ])); + + let mut ts = Vec::with_capacity(rows); + let mut values = Vec::with_capacity(rows); + let mut tags = Vec::with_capacity(rows); + for i in start..start + rows { + ts.push(i as i64); + values.push(i as f64); + tags.push("tag_0".to_string()); + } + + RecordBatch::try_new( + schema, + vec![ + Arc::new(TimestampMillisecondArray::from(ts)), + Arc::new(Float64Array::from(values)), + Arc::new(StringArray::from(tags)), + ], + ) + .unwrap() + } + + fn build_bulk_request(logical_region_id: RegionId, batch: RecordBatch) -> RegionRequest { + let (schema, data_header, payload) = record_batch_to_ipc(&batch).unwrap(); + RegionRequest::BulkInserts(RegionBulkInsertsRequest { + region_id: logical_region_id, + payload: batch, + raw_data: ArrowIpc { + schema, + data_header, + payload, + }, + partition_expr_version: None, + }) + } + + async fn init_dense_metric_region(env: &TestEnv) -> RegionId { + let physical_region_id = env.default_physical_region_id(); + env.create_physical_region( + physical_region_id, + &TestEnv::default_table_dir(), + vec![( + MEMTABLE_PARTITION_TREE_PRIMARY_KEY_ENCODING.to_string(), + "dense".to_string(), + )], + ) + .await; + + let logical_region_id = env.default_logical_region_id(); + let request = test_util::create_logical_region_request( + &["job"], + physical_region_id, + &table_dir("test", logical_region_id.table_id()), + ); + env.metric() + .handle_request(logical_region_id, RegionRequest::Create(request)) + .await + .unwrap(); + logical_region_id + } + + #[tokio::test] + async fn test_bulk_insert_empty_batch_returns_zero() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let logical_region_id = env.default_logical_region_id(); + + let batch = build_logical_batch(0, 0); + let request = RegionRequest::BulkInserts(RegionBulkInsertsRequest { + region_id: logical_region_id, + payload: batch, + raw_data: ArrowIpc::default(), + partition_expr_version: None, + }); + let response = env + .metric() + .handle_request(logical_region_id, request) + .await + .unwrap(); + assert_eq!(response.affected_rows, 0); + } + + #[tokio::test] + async fn test_bulk_insert_physical_region_rejected() { + let env = TestEnv::new().await; + env.init_metric_region().await; + + let physical_region_id = env.default_physical_region_id(); + let batch = build_logical_batch(0, 2); + let request = build_bulk_request(physical_region_id, batch); + + let err = env + .metric() + .handle_request(physical_region_id, request) + .await + .unwrap_err(); + let Some(err) = err.as_any().downcast_ref::() else { + panic!("unexpected error type"); + }; + assert_matches!(err, Error::UnsupportedRegionRequest { .. }); + } + + #[tokio::test] + async fn test_bulk_insert_unknown_column_errors() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let logical_region_id = env.default_logical_region_id(); + + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + greptime_timestamp(), + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new(greptime_value(), DataType::Float64, true), + Field::new("nonexistent_column", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(TimestampMillisecondArray::from(vec![0i64])), + Arc::new(Float64Array::from(vec![1.0])), + Arc::new(StringArray::from(vec!["val"])), + ], + ) + .unwrap(); + + let request = build_bulk_request(logical_region_id, batch); + let err = env + .metric() + .handle_request(logical_region_id, request) + .await + .unwrap_err(); + let Some(err) = err.as_any().downcast_ref::() else { + panic!("unexpected error type"); + }; + assert_matches!(err, Error::ColumnNotFound { .. }); + } + + #[tokio::test] + async fn test_bulk_insert_multiple_tag_columns() { + let env = TestEnv::new().await; + let physical_region_id = env.default_physical_region_id(); + env.create_physical_region(physical_region_id, &TestEnv::default_table_dir(), vec![]) + .await; + let logical_region_id = env.default_logical_region_id(); + let request = test_util::create_logical_region_request( + &["host", "region"], + physical_region_id, + &table_dir("test", logical_region_id.table_id()), + ); + env.metric() + .handle_request(logical_region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + greptime_timestamp(), + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new(greptime_value(), DataType::Float64, true), + Field::new("host", DataType::Utf8, true), + Field::new("region", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(TimestampMillisecondArray::from(vec![0i64, 1, 2])), + Arc::new(Float64Array::from(vec![10.0, 20.0, 30.0])), + Arc::new(StringArray::from(vec!["h1", "h2", "h1"])), + Arc::new(StringArray::from(vec!["us-east", "us-west", "eu-west"])), + ], + ) + .unwrap(); + + let request = build_bulk_request(logical_region_id, batch); + let response = env + .metric() + .handle_request(logical_region_id, request) + .await + .unwrap(); + assert_eq!(response.affected_rows, 3); + + let stream = env + .metric() + .scan_to_stream(logical_region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 3); + } + + #[tokio::test] + async fn test_bulk_insert_accumulates_rows() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let logical_region_id = env.default_logical_region_id(); + + let request = build_bulk_request(logical_region_id, build_logical_batch(0, 3)); + let response = env + .metric() + .handle_request(logical_region_id, request) + .await + .unwrap(); + assert_eq!(response.affected_rows, 3); + + let request = build_bulk_request(logical_region_id, build_logical_batch(3, 5)); + let response = env + .metric() + .handle_request(logical_region_id, request) + .await + .unwrap(); + assert_eq!(response.affected_rows, 5); + + let stream = env + .metric() + .scan_to_stream(logical_region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 8); + } + + #[tokio::test] + async fn test_bulk_insert_sparse_encoding() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let logical_region_id = env.default_logical_region_id(); + + let request = build_bulk_request(logical_region_id, build_logical_batch(0, 4)); + let response = env + .metric() + .handle_request(logical_region_id, request) + .await + .unwrap(); + assert_eq!(response.affected_rows, 4); + + let stream = env + .metric() + .scan_to_stream(logical_region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 4); + } + + #[tokio::test] + async fn test_bulk_insert_dense_encoding_rejected() { + let env = TestEnv::new().await; + let logical_region_id = init_dense_metric_region(&env).await; + + let request = build_bulk_request(logical_region_id, build_logical_batch(0, 2)); + let err = env + .metric() + .handle_request(logical_region_id, request) + .await + .unwrap_err(); + let Some(err) = err.as_any().downcast_ref::() else { + panic!("unexpected error type"); + }; + assert_matches!(err, Error::UnsupportedRegionRequest { .. }); + } + + #[tokio::test] + async fn test_bulk_insert_matches_put() { + let env_put = TestEnv::new().await; + env_put.init_metric_region().await; + let logical_region_id = env_put.default_logical_region_id(); + let schema = test_util::row_schema_with_tags(&["job"]); + let rows = test_util::build_rows(1, 5); + env_put + .metric() + .handle_request( + logical_region_id, + RegionRequest::Put(RegionPutRequest { + rows: api::v1::Rows { schema, rows }, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + let put_stream = env_put + .metric() + .scan_to_stream(logical_region_id, ScanRequest::default()) + .await + .unwrap(); + let put_batches = RecordBatches::try_collect(put_stream).await.unwrap(); + let put_output = put_batches.pretty_print().unwrap(); + + let env_bulk = TestEnv::new().await; + env_bulk.init_metric_region().await; + let request = build_bulk_request(logical_region_id, build_logical_batch(0, 5)); + env_bulk + .metric() + .handle_request(logical_region_id, request) + .await + .unwrap(); + let bulk_stream = env_bulk + .metric() + .scan_to_stream(logical_region_id, ScanRequest::default()) + .await + .unwrap(); + let bulk_batches = RecordBatches::try_collect(bulk_stream).await.unwrap(); + let bulk_output = bulk_batches.pretty_print().unwrap(); + + assert_eq!(put_output, bulk_output); + } + + #[test] + fn test_record_batch_to_rows_with_null_values() { + use datatypes::arrow::array::{Float64Array, StringArray, TimestampMillisecondArray}; + use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema, TimeUnit}; + use datatypes::arrow::record_batch::RecordBatch; + use store_api::storage::RegionId; + + use crate::engine::bulk_insert::record_batch_to_rows; + + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + greptime_timestamp(), + DataType::Timestamp(TimeUnit::Millisecond, None), + true, + ), + Field::new(greptime_value(), DataType::Float64, true), + Field::new("job", DataType::Utf8, true), + Field::new("host", DataType::Utf8, true), + ])); + + let ts_array = TimestampMillisecondArray::from(vec![Some(1000), None, Some(3000)]); + let val_array = Float64Array::from(vec![Some(1.0), Some(2.0), None]); + let job_array = StringArray::from(vec![Some("job1"), None, Some("job3")]); + let host_array = StringArray::from(vec![None, Some("host2"), Some("host3")]); + + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(ts_array), + Arc::new(val_array), + Arc::new(job_array), + Arc::new(host_array), + ], + ) + .unwrap(); + + let region_id = RegionId::new(1, 1); + let rows = record_batch_to_rows(&batch, region_id).unwrap(); + + assert_eq!(rows.rows.len(), 3); + assert_eq!(rows.schema.len(), 4); + + // Row 0: all non-null except host + assert!(rows.rows[0].values[0].value_data.is_some()); + assert!(rows.rows[0].values[1].value_data.is_some()); + assert!(rows.rows[0].values[2].value_data.is_some()); + assert!(rows.rows[0].values[3].value_data.is_none()); + + // Row 1: null timestamp, null job + assert!(rows.rows[1].values[0].value_data.is_none()); + assert!(rows.rows[1].values[1].value_data.is_some()); + assert!(rows.rows[1].values[2].value_data.is_none()); + assert!(rows.rows[1].values[3].value_data.is_some()); + + // Row 2: null value + assert!(rows.rows[2].values[0].value_data.is_some()); + assert!(rows.rows[2].values[1].value_data.is_none()); + assert!(rows.rows[2].values[2].value_data.is_some()); + assert!(rows.rows[2].values[3].value_data.is_some()); + } +} diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index 9251605aea6..edae0d2bb46 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/src/metric-engine/src/engine/put.rs @@ -460,7 +460,7 @@ impl MetricEngineInner { .await } - fn find_data_region_meta( + pub(crate) fn find_data_region_meta( &self, logical_region_id: RegionId, ) -> Result<(RegionId, RegionId, PrimaryKeyEncoding)> { diff --git a/src/metric-engine/src/lib.rs b/src/metric-engine/src/lib.rs index 30daa80b91f..b93029f2f41 100644 --- a/src/metric-engine/src/lib.rs +++ b/src/metric-engine/src/lib.rs @@ -52,6 +52,7 @@ #![feature(assert_matches)] +mod batch_modifier; pub mod config; mod data_region; pub mod engine; diff --git a/src/mito2/benches/memtable_bench.rs b/src/mito2/benches/memtable_bench.rs index ebe994f8618..df991f6f925 100644 --- a/src/mito2/benches/memtable_bench.rs +++ b/src/mito2/benches/memtable_bench.rs @@ -28,7 +28,7 @@ use mito2::memtable::bulk::part_reader::BulkPartBatchIter; use mito2::memtable::bulk::{BulkMemtable, BulkMemtableConfig}; use mito2::memtable::partition_tree::{PartitionTreeConfig, PartitionTreeMemtable}; use mito2::memtable::time_series::TimeSeriesMemtable; -use mito2::memtable::{KeyValues, Memtable, RangesOptions}; +use mito2::memtable::{IterBuilder, KeyValues, Memtable, RangesOptions}; use mito2::read::flat_merge::FlatMergeIterator; use mito2::read::scan_region::PredicateGroup; use mito2::region::options::MergeMode; @@ -105,7 +105,11 @@ fn full_scan(c: &mut Criterion) { } b.iter(|| { - let iter = memtable.iter(None, None, None).unwrap(); + let iter = memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); for batch in iter { let _batch = batch.unwrap(); } @@ -145,7 +149,17 @@ fn filter_1_host(c: &mut Criterion) { let predicate = generator.random_host_filter(); b.iter(|| { - let iter = memtable.iter(None, Some(predicate.clone()), None).unwrap(); + let iter = memtable + .ranges( + None, + RangesOptions { + predicate: PredicateGroup::new(&metadata, predicate.exprs()).unwrap(), + ..Default::default() + }, + ) + .unwrap() + .build(None) + .unwrap(); for batch in iter { let _batch = batch.unwrap(); } diff --git a/src/mito2/benches/simple_bulk_memtable.rs b/src/mito2/benches/simple_bulk_memtable.rs index 02773977686..05035734de6 100644 --- a/src/mito2/benches/simple_bulk_memtable.rs +++ b/src/mito2/benches/simple_bulk_memtable.rs @@ -21,7 +21,7 @@ use criterion::{Criterion, criterion_group, criterion_main}; use datatypes::data_type::ConcreteDataType; use datatypes::schema::ColumnSchema; use mito2::memtable::simple_bulk_memtable::SimpleBulkMemtable; -use mito2::memtable::{KeyValues, Memtable, MemtableRanges, RangesOptions}; +use mito2::memtable::{IterBuilder, KeyValues, Memtable, MemtableRanges, RangesOptions}; use mito2::read; use mito2::read::Source; use mito2::read::dedup::DedupReader; @@ -156,7 +156,11 @@ async fn flush(mem: &SimpleBulkMemtable) { } async fn flush_original(mem: &SimpleBulkMemtable) { - let iter = mem.iter(None, None, None).unwrap(); + let iter = mem + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); for b in iter { black_box(b.unwrap()); } diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index 92c8a3bc368..33180ebf466 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -17,7 +17,6 @@ use std::time::{Duration, Instant}; use async_stream::try_stream; use common_time::Timestamp; -use either::Either; use futures::{Stream, TryStreamExt}; use object_store::services::Fs; use object_store::util::{join_dir, with_instrument_layers}; @@ -37,7 +36,7 @@ use crate::error::{ CleanDirSnafu, DeleteIndexSnafu, DeleteIndexesSnafu, DeleteSstsSnafu, OpenDalSnafu, Result, }; use crate::metrics::{COMPACTION_STAGE_ELAPSED, FLUSH_ELAPSED}; -use crate::read::{FlatSource, Source}; +use crate::read::FlatSource; use crate::region::options::IndexOptions; use crate::sst::file::{FileHandle, RegionFileId, RegionIndexId}; use crate::sst::index::IndexerBuilderImpl; @@ -47,7 +46,7 @@ use crate::sst::location::{self, region_dir_from_table_dir}; use crate::sst::parquet::reader::ParquetReaderBuilder; use crate::sst::parquet::writer::ParquetWriter; use crate::sst::parquet::{SstInfo, WriteOptions}; -use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY}; +use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY, FormatType}; pub type AccessLayerRef = Arc; /// SST write results. @@ -339,6 +338,7 @@ impl AccessLayer { metrics: &mut Metrics, ) -> Result { let region_id = request.metadata.region_id; + let region_metadata = request.metadata.clone(); let cache_manager = request.cache_manager.clone(); let sst_info = if let Some(write_cache) = cache_manager.write_cache() { @@ -391,15 +391,19 @@ impl AccessLayer { ) .await .with_file_cleaner(cleaner); - match request.source { - Either::Left(source) => { + match request.sst_write_format { + FormatType::PrimaryKey => { writer - .write_all(source, request.max_sequence, write_opts) + .write_all_flat_as_primary_key( + request.source, + request.max_sequence, + write_opts, + ) .await? } - Either::Right(flat_source) => { + FormatType::Flat => { writer - .write_all_flat(flat_source, request.max_sequence, write_opts) + .write_all_flat(request.source, request.max_sequence, write_opts) .await? } } @@ -412,6 +416,7 @@ impl AccessLayer { cache_manager.put_parquet_meta_data( RegionFileId::new(region_id, sst.file_id), parquet_metadata.clone(), + Some(region_metadata.clone()), ) } } @@ -520,11 +525,12 @@ pub enum OperationType { pub struct SstWriteRequest { pub op_type: OperationType, pub metadata: RegionMetadataRef, - pub source: Either, + pub source: FlatSource, pub cache_manager: CacheManagerRef, #[allow(dead_code)] pub storage: Option, pub max_sequence: Option, + pub sst_write_format: FormatType, /// Configs for index pub index_options: IndexOptions, diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 3ad71d2a612..c9a8b99166d 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -28,6 +28,7 @@ use std::ops::Range; use std::sync::Arc; use bytes::Bytes; +use common_telemetry::warn; use datatypes::arrow::record_batch::RecordBatch; use datatypes::value::Value; use datatypes::vectors::VectorRef; @@ -36,8 +37,10 @@ use index::result_cache::IndexResultCache; use moka::notification::RemovalCause; use moka::sync::Cache; use object_store::ObjectStore; -use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; +use parquet::file::metadata::{FileMetaData, PageIndexPolicy, ParquetMetaData}; use puffin::puffin_manager::cache::{PuffinMetadataCache, PuffinMetadataCacheRef}; +use snafu::{OptionExt, ResultExt}; +use store_api::metadata::RegionMetadataRef; use store_api::storage::{ConcreteDataType, FileId, RegionId, TimeSeriesRowSelector}; use crate::cache::cache_size::parquet_meta_size; @@ -46,10 +49,13 @@ use crate::cache::index::inverted_index::{InvertedIndexCache, InvertedIndexCache #[cfg(feature = "vector_index")] use crate::cache::index::vector_index::{VectorIndexCache, VectorIndexCacheRef}; use crate::cache::write_cache::WriteCacheRef; +use crate::error::{InvalidMetadataSnafu, InvalidParquetSnafu, Result}; use crate::memtable::record_batch_estimated_size; use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, CACHE_HIT, CACHE_MISS}; use crate::read::Batch; +use crate::read::range_cache::{RangeScanCacheKey, RangeScanCacheValue}; use crate::sst::file::{RegionFileId, RegionIndexId}; +use crate::sst::parquet::PARQUET_METADATA_KEY; use crate::sst::parquet::reader::MetadataCacheMetrics; /// Metrics type key for sst meta. @@ -64,6 +70,108 @@ const FILE_TYPE: &str = "file"; const INDEX_TYPE: &str = "index"; /// Metrics type key for selector result cache. const SELECTOR_RESULT_TYPE: &str = "selector_result"; +/// Metrics type key for range scan result cache. +const RANGE_RESULT_TYPE: &str = "range_result"; + +/// Cached SST metadata combines the parquet footer with the decoded region metadata. +/// +/// The cached parquet footer strips the `greptime:metadata` JSON payload and stores the decoded +/// [RegionMetadata] separately so readers can skip repeated deserialization work. +#[derive(Debug)] +pub(crate) struct CachedSstMeta { + parquet_metadata: Arc, + region_metadata: RegionMetadataRef, + region_metadata_weight: usize, +} + +impl CachedSstMeta { + pub(crate) fn try_new(file_path: &str, parquet_metadata: ParquetMetaData) -> Result { + Self::try_new_with_region_metadata(file_path, parquet_metadata, None) + } + + pub(crate) fn try_new_with_region_metadata( + file_path: &str, + parquet_metadata: ParquetMetaData, + region_metadata: Option, + ) -> Result { + let file_metadata = parquet_metadata.file_metadata(); + let key_values = file_metadata + .key_value_metadata() + .context(InvalidParquetSnafu { + file: file_path, + reason: "missing key value meta", + })?; + let meta_value = key_values + .iter() + .find(|kv| kv.key == PARQUET_METADATA_KEY) + .with_context(|| InvalidParquetSnafu { + file: file_path, + reason: format!("key {} not found", PARQUET_METADATA_KEY), + })?; + let json = meta_value + .value + .as_ref() + .with_context(|| InvalidParquetSnafu { + file: file_path, + reason: format!("No value for key {}", PARQUET_METADATA_KEY), + })?; + let region_metadata = match region_metadata { + Some(region_metadata) => region_metadata, + None => Arc::new( + store_api::metadata::RegionMetadata::from_json(json) + .context(InvalidMetadataSnafu)?, + ), + }; + // Keep the previous JSON-byte floor and charge the decoded structures as well. + let region_metadata_weight = region_metadata.estimated_size().max(json.len()); + let parquet_metadata = Arc::new(strip_region_metadata_from_parquet(parquet_metadata)); + + Ok(Self { + parquet_metadata, + region_metadata, + region_metadata_weight, + }) + } + + pub(crate) fn parquet_metadata(&self) -> Arc { + self.parquet_metadata.clone() + } + + pub(crate) fn region_metadata(&self) -> RegionMetadataRef { + self.region_metadata.clone() + } +} + +fn strip_region_metadata_from_parquet(parquet_metadata: ParquetMetaData) -> ParquetMetaData { + let file_metadata = parquet_metadata.file_metadata(); + let filtered_key_values = file_metadata.key_value_metadata().and_then(|key_values| { + let filtered = key_values + .iter() + .filter(|kv| kv.key != PARQUET_METADATA_KEY) + .cloned() + .collect::>(); + (!filtered.is_empty()).then_some(filtered) + }); + let stripped_file_metadata = FileMetaData::new( + file_metadata.version(), + file_metadata.num_rows(), + file_metadata.created_by().map(ToString::to_string), + filtered_key_values, + file_metadata.schema_descr_ptr(), + file_metadata.column_orders().cloned(), + ); + + let mut builder = parquet_metadata.into_builder(); + let row_groups = builder.take_row_groups(); + let column_index = builder.take_column_index(); + let offset_index = builder.take_offset_index(); + + parquet::file::metadata::ParquetMetaDataBuilder::new(stripped_file_metadata) + .set_row_groups(row_groups) + .set_column_index(column_index) + .set_offset_index(offset_index) + .build() +} /// Cache strategies that may only enable a subset of caches. #[derive(Clone)] @@ -81,18 +189,17 @@ pub enum CacheStrategy { } impl CacheStrategy { - /// Gets parquet metadata with cache metrics tracking. - /// Returns the metadata and updates the provided metrics. - pub(crate) async fn get_parquet_meta_data( + /// Gets fused SST metadata with cache metrics tracking. + pub(crate) async fn get_sst_meta_data( &self, file_id: RegionFileId, metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, - ) -> Option> { + ) -> Option> { match self { CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { cache_manager - .get_parquet_meta_data(file_id, metrics, page_index_policy) + .get_sst_meta_data(file_id, metrics, page_index_policy) .await } CacheStrategy::Disabled => { @@ -102,30 +209,48 @@ impl CacheStrategy { } } - /// Calls [CacheManager::get_parquet_meta_data_from_mem_cache()]. - pub fn get_parquet_meta_data_from_mem_cache( + /// Calls [CacheManager::get_sst_meta_data_from_mem_cache()]. + pub(crate) fn get_sst_meta_data_from_mem_cache( &self, file_id: RegionFileId, - ) -> Option> { + ) -> Option> { match self { - CacheStrategy::EnableAll(cache_manager) => { - cache_manager.get_parquet_meta_data_from_mem_cache(file_id) - } - CacheStrategy::Compaction(cache_manager) => { - cache_manager.get_parquet_meta_data_from_mem_cache(file_id) + CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { + cache_manager.get_sst_meta_data_from_mem_cache(file_id) } CacheStrategy::Disabled => None, } } - /// Calls [CacheManager::put_parquet_meta_data()]. - pub fn put_parquet_meta_data(&self, file_id: RegionFileId, metadata: Arc) { + /// Calls [CacheManager::get_parquet_meta_data_from_mem_cache()]. + pub fn get_parquet_meta_data_from_mem_cache( + &self, + file_id: RegionFileId, + ) -> Option> { + self.get_sst_meta_data_from_mem_cache(file_id) + .map(|metadata| metadata.parquet_metadata()) + } + + /// Calls [CacheManager::put_sst_meta_data()]. + pub(crate) fn put_sst_meta_data(&self, file_id: RegionFileId, metadata: Arc) { match self { - CacheStrategy::EnableAll(cache_manager) => { - cache_manager.put_parquet_meta_data(file_id, metadata); + CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { + cache_manager.put_sst_meta_data(file_id, metadata); } - CacheStrategy::Compaction(cache_manager) => { - cache_manager.put_parquet_meta_data(file_id, metadata); + CacheStrategy::Disabled => {} + } + } + + /// Calls [CacheManager::put_parquet_meta_data()]. + pub fn put_parquet_meta_data( + &self, + file_id: RegionFileId, + metadata: Arc, + region_metadata: Option, + ) { + match self { + CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { + cache_manager.put_parquet_meta_data(file_id, metadata, region_metadata); } CacheStrategy::Disabled => {} } @@ -223,6 +348,32 @@ impl CacheStrategy { } } + /// Calls [CacheManager::get_range_result()]. + /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled]. + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn get_range_result( + &self, + key: &RangeScanCacheKey, + ) -> Option> { + match self { + CacheStrategy::EnableAll(cache_manager) => cache_manager.get_range_result(key), + CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None, + } + } + + /// Calls [CacheManager::put_range_result()]. + /// It does nothing if the strategy isn't [CacheStrategy::EnableAll]. + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn put_range_result( + &self, + key: RangeScanCacheKey, + result: Arc, + ) { + if let CacheStrategy::EnableAll(cache_manager) = self { + cache_manager.put_range_result(key, result); + } + } + /// Calls [CacheManager::write_cache()]. /// It returns None if the strategy is [CacheStrategy::Disabled]. pub fn write_cache(&self) -> Option<&WriteCacheRef> { @@ -324,6 +475,9 @@ pub struct CacheManager { puffin_metadata_cache: Option, /// Cache for time series selectors. selector_result_cache: Option, + /// Cache for range scan outputs in flat format. + #[cfg_attr(not(test), allow(dead_code))] + range_result_cache: Option, /// Cache for index result. index_result_cache: Option, } @@ -336,6 +490,35 @@ impl CacheManager { CacheManagerBuilder::default() } + /// Gets fused SST metadata with metrics tracking. + /// Tries in-memory cache first, then file cache, updating metrics accordingly. + pub(crate) async fn get_sst_meta_data( + &self, + file_id: RegionFileId, + metrics: &mut MetadataCacheMetrics, + page_index_policy: PageIndexPolicy, + ) -> Option> { + if let Some(metadata) = self.get_sst_meta_data_from_mem_cache(file_id) { + metrics.mem_cache_hit += 1; + return Some(metadata); + } + + let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet); + if let Some(write_cache) = &self.write_cache + && let Some(metadata) = write_cache + .file_cache() + .get_sst_meta_data(key, metrics, page_index_policy) + .await + { + metrics.file_cache_hit += 1; + self.put_sst_meta_data(file_id, metadata.clone()); + return Some(metadata); + } + + metrics.cache_miss += 1; + None + } + /// Gets cached [ParquetMetaData] with metrics tracking. /// Tries in-memory cache first, then file cache, updating metrics accordingly. pub(crate) async fn get_parquet_meta_data( @@ -344,29 +527,21 @@ impl CacheManager { metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, ) -> Option> { - // Try to get metadata from sst meta cache - if let Some(metadata) = self.get_parquet_meta_data_from_mem_cache(file_id) { - metrics.mem_cache_hit += 1; - return Some(metadata); - } + self.get_sst_meta_data(file_id, metrics, page_index_policy) + .await + .map(|metadata| metadata.parquet_metadata()) + } - // Try to get metadata from write cache - let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet); - if let Some(write_cache) = &self.write_cache - && let Some(metadata) = write_cache - .file_cache() - .get_parquet_meta_data(key, metrics, page_index_policy) - .await - { - metrics.file_cache_hit += 1; - let metadata = Arc::new(metadata); - // Put metadata into sst meta cache - self.put_parquet_meta_data(file_id, metadata.clone()); - return Some(metadata); - }; - metrics.cache_miss += 1; - - None + /// Gets cached fused SST metadata from in-memory cache. + /// This method does not perform I/O. + pub(crate) fn get_sst_meta_data_from_mem_cache( + &self, + file_id: RegionFileId, + ) -> Option> { + self.sst_meta_cache.as_ref().and_then(|sst_meta_cache| { + let value = sst_meta_cache.get(&SstMetaKey(file_id.region_id(), file_id.file_id())); + update_hit_miss(value, SST_META_TYPE) + }) } /// Gets cached [ParquetMetaData] from in-memory cache. @@ -375,15 +550,12 @@ impl CacheManager { &self, file_id: RegionFileId, ) -> Option> { - // Try to get metadata from sst meta cache - self.sst_meta_cache.as_ref().and_then(|sst_meta_cache| { - let value = sst_meta_cache.get(&SstMetaKey(file_id.region_id(), file_id.file_id())); - update_hit_miss(value, SST_META_TYPE) - }) + self.get_sst_meta_data_from_mem_cache(file_id) + .map(|metadata| metadata.parquet_metadata()) } - /// Puts [ParquetMetaData] into the cache. - pub fn put_parquet_meta_data(&self, file_id: RegionFileId, metadata: Arc) { + /// Puts fused SST metadata into the cache. + pub(crate) fn put_sst_meta_data(&self, file_id: RegionFileId, metadata: Arc) { if let Some(cache) = &self.sst_meta_cache { let key = SstMetaKey(file_id.region_id(), file_id.file_id()); CACHE_BYTES @@ -393,6 +565,34 @@ impl CacheManager { } } + /// Puts [ParquetMetaData] into the cache. + pub fn put_parquet_meta_data( + &self, + file_id: RegionFileId, + metadata: Arc, + region_metadata: Option, + ) { + if self.sst_meta_cache.is_some() { + let file_path = format!( + "region_id={}, file_id={}", + file_id.region_id(), + file_id.file_id() + ); + match CachedSstMeta::try_new_with_region_metadata( + &file_path, + Arc::unwrap_or_clone(metadata), + region_metadata, + ) { + Ok(metadata) => self.put_sst_meta_data(file_id, Arc::new(metadata)), + Err(err) => warn!( + err; "Failed to decode region metadata while caching parquet metadata, region_id: {}, file_id: {}", + file_id.region_id(), + file_id.file_id() + ), + } + } + } + /// Removes [ParquetMetaData] from the cache. pub fn remove_parquet_meta_data(&self, file_id: RegionFileId) { if let Some(cache) = &self.sst_meta_cache { @@ -512,6 +712,32 @@ impl CacheManager { } } + /// Gets cached result for range scan. + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn get_range_result( + &self, + key: &RangeScanCacheKey, + ) -> Option> { + self.range_result_cache + .as_ref() + .and_then(|cache| update_hit_miss(cache.get(key), RANGE_RESULT_TYPE)) + } + + /// Puts range scan result into the cache. + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn put_range_result( + &self, + key: RangeScanCacheKey, + result: Arc, + ) { + if let Some(cache) = &self.range_result_cache { + CACHE_BYTES + .with_label_values(&[RANGE_RESULT_TYPE]) + .add(range_result_cache_weight(&key, &result).into()); + cache.insert(key, result); + } + } + /// Gets the write cache. pub(crate) fn write_cache(&self) -> Option<&WriteCacheRef> { self.write_cache.as_ref() @@ -562,6 +788,7 @@ pub struct CacheManagerBuilder { puffin_metadata_size: u64, write_cache: Option, selector_result_cache_size: u64, + range_result_cache_size: u64, } impl CacheManagerBuilder { @@ -625,6 +852,12 @@ impl CacheManagerBuilder { self } + /// Sets range result cache size. + pub fn range_result_cache_size(mut self, bytes: u64) -> Self { + self.range_result_cache_size = bytes; + self + } + /// Builds the [CacheManager]. pub fn build(self) -> CacheManager { fn to_str(cause: RemovalCause) -> &'static str { @@ -712,6 +945,21 @@ impl CacheManagerBuilder { }) .build() }); + let range_result_cache = (self.range_result_cache_size != 0).then(|| { + Cache::builder() + .max_capacity(self.range_result_cache_size) + .weigher(range_result_cache_weight) + .eviction_listener(|k, v, cause| { + let size = range_result_cache_weight(&k, &v); + CACHE_BYTES + .with_label_values(&[RANGE_RESULT_TYPE]) + .sub(size.into()); + CACHE_EVICTION + .with_label_values(&[RANGE_RESULT_TYPE, to_str(cause)]) + .inc(); + }) + .build() + }); CacheManager { sst_meta_cache, vector_cache, @@ -723,14 +971,15 @@ impl CacheManagerBuilder { vector_index_cache, puffin_metadata_cache: Some(Arc::new(puffin_metadata_cache)), selector_result_cache, + range_result_cache, index_result_cache, } } } -fn meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { +fn meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { // We ignore the size of `Arc`. - (k.estimated_size() + parquet_meta_size(v)) as u32 + (k.estimated_size() + parquet_meta_size(&v.parquet_metadata) + v.region_metadata_weight) as u32 } fn vector_cache_weight(_k: &(ConcreteDataType, Value), v: &VectorRef) -> u32 { @@ -746,6 +995,10 @@ fn selector_result_cache_weight(k: &SelectorResultKey, v: &Arc) -> u32 { + (k.estimated_size() + v.estimated_size()) as u32 +} + /// Updates cache hit/miss metrics. fn update_hit_miss(value: Option, cache_type: &str) -> Option { if value.is_some() { @@ -892,8 +1145,8 @@ impl SelectorResultValue { } } -/// Maps (region id, file id) to [ParquetMetaData]. -type SstMetaCache = Cache>; +/// Maps (region id, file id) to fused SST metadata. +type SstMetaCache = Cache>; /// Maps [Value] to a vector that holds this value repeatedly. /// /// e.g. `"hello" => ["hello", "hello", "hello"]` @@ -902,20 +1155,30 @@ type VectorCache = Cache<(ConcreteDataType, Value), VectorRef>; type PageCache = Cache>; /// Maps (file id, row group id, time series row selector) to [SelectorResultValue]. type SelectorResultCache = Cache>; +/// Maps partition-range scan key to cached flat batches. +type RangeResultCache = Cache>; #[cfg(test)] mod tests { use std::sync::Arc; + use api::v1::SemanticType; use api::v1::index::{BloomFilterMeta, InvertedIndexMetas}; + use datatypes::schema::ColumnSchema; use datatypes::vectors::Int64Vector; use puffin::file_metadata::FileMetadata; + use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::storage::ColumnId; use super::*; use crate::cache::index::bloom_filter_index::Tag; use crate::cache::index::result_cache::PredicateKey; - use crate::cache::test_util::parquet_meta; + use crate::cache::test_util::{ + parquet_meta, sst_parquet_meta, sst_parquet_meta_with_region_metadata, + }; + use crate::read::range_cache::{ + RangeScanCacheKey, RangeScanCacheValue, ScanRequestFingerprintBuilder, + }; use crate::sst::parquet::row_selection::RowGroupSelection; #[tokio::test] @@ -929,7 +1192,7 @@ mod tests { let file_id = RegionFileId::new(region_id, FileId::random()); let metadata = parquet_meta(); let mut metrics = MetadataCacheMetrics::default(); - cache.put_parquet_meta_data(file_id, metadata); + cache.put_parquet_meta_data(file_id, metadata, None); assert!( cache .get_parquet_meta_data(file_id, &mut metrics, Default::default()) @@ -966,13 +1229,23 @@ mod tests { .await .is_none() ); - let metadata = parquet_meta(); - cache.put_parquet_meta_data(file_id, metadata); + let (metadata, region_metadata) = sst_parquet_meta(); + cache.put_parquet_meta_data(file_id, metadata, None); + let cached = cache + .get_sst_meta_data(file_id, &mut metrics, Default::default()) + .await + .unwrap(); + assert_eq!(region_metadata, cached.region_metadata()); assert!( - cache - .get_parquet_meta_data(file_id, &mut metrics, Default::default()) - .await - .is_some() + cached + .parquet_metadata() + .file_metadata() + .key_value_metadata() + .is_none_or(|key_values| { + key_values + .iter() + .all(|key_value| key_value.key != PARQUET_METADATA_KEY) + }) ); cache.remove_parquet_meta_data(file_id); assert!( @@ -983,6 +1256,42 @@ mod tests { ); } + #[tokio::test] + async fn test_parquet_meta_cache_with_provided_region_metadata() { + let cache = CacheManager::builder().sst_meta_cache_size(2000).build(); + let mut metrics = MetadataCacheMetrics::default(); + let region_id = RegionId::new(1, 1); + let file_id = RegionFileId::new(region_id, FileId::random()); + let (metadata, region_metadata) = sst_parquet_meta(); + + cache.put_parquet_meta_data(file_id, metadata, Some(region_metadata.clone())); + + let cached = cache + .get_sst_meta_data(file_id, &mut metrics, Default::default()) + .await + .unwrap(); + assert!(Arc::ptr_eq(®ion_metadata, &cached.region_metadata())); + } + + #[test] + fn test_meta_cache_weight_accounts_for_decoded_region_metadata() { + let region_metadata = Arc::new(wide_region_metadata(128)); + let json_len = region_metadata.to_json().unwrap().len(); + let metadata = sst_parquet_meta_with_region_metadata(region_metadata.clone()); + let cached = Arc::new( + CachedSstMeta::try_new("test.parquet", Arc::unwrap_or_clone(metadata)).unwrap(), + ); + let key = SstMetaKey(region_metadata.region_id, FileId::random()); + + assert!(cached.region_metadata_weight > json_len); + assert_eq!( + meta_cache_weight(&key, &cached) as usize, + key.estimated_size() + + parquet_meta_size(&cached.parquet_metadata) + + cached.region_metadata_weight + ); + } + #[test] fn test_repeated_vector_cache() { let cache = CacheManager::builder().vector_cache_size(4096).build(); @@ -1028,6 +1337,50 @@ mod tests { assert!(cache.get_selector_result(&key).is_some()); } + #[test] + fn test_range_result_cache() { + let cache = Arc::new( + CacheManager::builder() + .range_result_cache_size(1024 * 1024) + .build(), + ); + + let key = RangeScanCacheKey { + region_id: RegionId::new(1, 1), + row_groups: vec![(FileId::random(), 0)], + scan: ScanRequestFingerprintBuilder { + read_column_ids: vec![], + read_column_types: vec![], + filters: vec!["tag_0 = 1".to_string()], + time_filters: vec![], + series_row_selector: None, + append_mode: false, + filter_deleted: true, + merge_mode: crate::region::options::MergeMode::LastRow, + partition_expr_version: 0, + } + .build(), + }; + let value = Arc::new(RangeScanCacheValue::new(Vec::new())); + + assert!(cache.get_range_result(&key).is_none()); + cache.put_range_result(key.clone(), value.clone()); + assert!(cache.get_range_result(&key).is_some()); + + let enable_all = CacheStrategy::EnableAll(cache.clone()); + assert!(enable_all.get_range_result(&key).is_some()); + + let compaction = CacheStrategy::Compaction(cache.clone()); + assert!(compaction.get_range_result(&key).is_none()); + compaction.put_range_result(key.clone(), value.clone()); + assert!(cache.get_range_result(&key).is_some()); + + let disabled = CacheStrategy::Disabled; + assert!(disabled.get_range_result(&key).is_none()); + disabled.put_range_result(key.clone(), value); + assert!(cache.get_range_result(&key).is_some()); + } + #[tokio::test] async fn test_evict_puffin_cache_clears_all_entries() { use std::collections::{BTreeMap, HashMap}; @@ -1122,4 +1475,45 @@ mod tests { assert!(result_cache.get(&predicate, index_id.file_id()).is_none()); assert!(puffin_metadata_cache.get_metadata(&file_id_str).is_none()); } + + fn wide_region_metadata(column_count: u32) -> RegionMetadata { + let region_id = RegionId::new(1024, 7); + let mut builder = RegionMetadataBuilder::new(region_id); + let mut primary_key = Vec::new(); + + for column_id in 0..column_count { + let semantic_type = if column_id < 32 { + primary_key.push(column_id); + SemanticType::Tag + } else { + SemanticType::Field + }; + let mut column_schema = ColumnSchema::new( + format!("wide_column_{column_id}"), + ConcreteDataType::string_datatype(), + true, + ); + column_schema + .mut_metadata() + .insert(format!("cache_key_{column_id}"), "cache_value".repeat(4)); + builder.push_column_metadata(ColumnMetadata { + column_schema, + semantic_type, + column_id, + }); + } + + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: column_count, + }); + builder.primary_key(primary_key); + + builder.build().unwrap() + } } diff --git a/src/mito2/src/cache/file_cache.rs b/src/mito2/src/cache/file_cache.rs index 32a276d0e49..278838b3694 100644 --- a/src/mito2/src/cache/file_cache.rs +++ b/src/mito2/src/cache/file_cache.rs @@ -34,7 +34,7 @@ use store_api::storage::{FileId, RegionId}; use tokio::sync::mpsc::{Sender, UnboundedReceiver}; use crate::access_layer::TempFileCleaner; -use crate::cache::{FILE_TYPE, INDEX_TYPE}; +use crate::cache::{CachedSstMeta, FILE_TYPE, INDEX_TYPE}; use crate::error::{self, OpenDalSnafu, Result}; use crate::metrics::{ CACHE_BYTES, CACHE_HIT, CACHE_MISS, WRITE_CACHE_DOWNLOAD_BYTES_TOTAL, @@ -612,6 +612,34 @@ impl FileCache { } } + /// Get fused SST metadata from the file cache. + /// If the file is not in the cache, or metadata loading/decoding fails, return None. + pub(crate) async fn get_sst_meta_data( + &self, + key: IndexKey, + cache_metrics: &mut MetadataCacheMetrics, + page_index_policy: PageIndexPolicy, + ) -> Option> { + let file_path = self.inner.cache_file_path(key); + self.get_parquet_meta_data(key, cache_metrics, page_index_policy) + .await + .and_then( + |metadata| match CachedSstMeta::try_new(&file_path, metadata) { + Ok(metadata) => Some(Arc::new(metadata)), + Err(err) => { + CACHE_MISS + .with_label_values(&[key.file_type.metric_label()]) + .inc(); + warn!( + err; "Failed to decode cached parquet metadata for key {:?}", + key + ); + None + } + }, + ) + } + async fn get_reader(&self, file_path: &str) -> object_store::Result> { if self.inner.local_store.exists(file_path).await? { Ok(Some(self.inner.local_store.reader(file_path).await?)) diff --git a/src/mito2/src/cache/test_util.rs b/src/mito2/src/cache/test_util.rs index 65ad9d87eb2..ef3d8e93158 100644 --- a/src/mito2/src/cache/test_util.rs +++ b/src/mito2/src/cache/test_util.rs @@ -23,8 +23,13 @@ use object_store::ObjectStore; use object_store::services::Fs; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; -use parquet::file::metadata::ParquetMetaData; +use parquet::file::metadata::{KeyValue, ParquetMetaData}; +use parquet::file::properties::WriterProperties; use parquet::file::statistics::Statistics; +use store_api::metadata::RegionMetadataRef; + +use crate::sst::parquet::PARQUET_METADATA_KEY; +use crate::test_util::sst_util::sst_region_metadata; /// Returns a parquet meta data. pub(crate) fn parquet_meta() -> Arc { @@ -33,13 +38,43 @@ pub(crate) fn parquet_meta() -> Arc { builder.metadata().clone() } +/// Returns parquet metadata for an SST parquet file and its decoded region metadata. +pub(crate) fn sst_parquet_meta() -> (Arc, RegionMetadataRef) { + let region_metadata = Arc::new(sst_region_metadata()); + let file_data = parquet_file_data_with_region_metadata(®ion_metadata); + let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(file_data)).unwrap(); + (builder.metadata().clone(), region_metadata) +} + +/// Returns parquet metadata for an SST parquet file with custom region metadata. +pub(crate) fn sst_parquet_meta_with_region_metadata( + region_metadata: RegionMetadataRef, +) -> Arc { + let file_data = parquet_file_data_with_region_metadata(®ion_metadata); + let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(file_data)).unwrap(); + builder.metadata().clone() +} + /// Write a test parquet file to a buffer fn parquet_file_data() -> Vec { + parquet_file_data_inner(None) +} + +fn parquet_file_data_with_region_metadata(region_metadata: &RegionMetadataRef) -> Vec { + let json = region_metadata.to_json().unwrap(); + let key_value = KeyValue::new(PARQUET_METADATA_KEY.to_string(), json); + parquet_file_data_inner(Some(vec![key_value])) +} + +fn parquet_file_data_inner(key_value_metadata: Option>) -> Vec { let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef; let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap(); let mut buffer = Vec::new(); - let mut writer = ArrowWriter::try_new(&mut buffer, to_write.schema(), None).unwrap(); + let props = WriterProperties::builder() + .set_key_value_metadata(key_value_metadata) + .build(); + let mut writer = ArrowWriter::try_new(&mut buffer, to_write.schema(), Some(props)).unwrap(); writer.write(&to_write).unwrap(); writer.close().unwrap(); diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index a28df3f54cf..e2483ed4e49 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -244,15 +244,19 @@ impl WriteCache { .await .with_file_cleaner(cleaner); - let sst_info = match write_request.source { - either::Left(source) => { + let sst_info = match write_request.sst_write_format { + crate::sst::FormatType::PrimaryKey => { writer - .write_all(source, write_request.max_sequence, write_opts) + .write_all_flat_as_primary_key( + write_request.source, + write_request.max_sequence, + write_opts, + ) .await? } - either::Right(flat_source) => { + crate::sst::FormatType::Flat => { writer - .write_all_flat(flat_source, write_request.max_sequence, write_opts) + .write_all_flat(write_request.source, write_request.max_sequence, write_opts) .await? } }; @@ -509,12 +513,13 @@ mod tests { use crate::cache::test_util::{assert_parquet_metadata_equal, new_fs_store}; use crate::cache::{CacheManager, CacheStrategy}; use crate::error::InvalidBatchSnafu; - use crate::read::Source; + use crate::read::FlatSource; use crate::region::options::IndexOptions; use crate::sst::parquet::reader::ParquetReaderBuilder; use crate::test_util::TestEnv; use crate::test_util::sst_util::{ - new_batch_by_range, new_source, sst_file_handle_with_file_id, sst_region_metadata, + new_flat_source_from_record_batches, new_record_batch_by_range, + sst_file_handle_with_file_id, sst_region_metadata, }; #[tokio::test] @@ -532,21 +537,22 @@ mod tests { .create_write_cache(local_store.clone(), ReadableSize::mb(10)) .await; - // Create Source + // Create source. let metadata = Arc::new(sst_region_metadata()); let region_id = metadata.region_id; - let source = new_source(&[ - new_batch_by_range(&["a", "d"], 0, 60), - new_batch_by_range(&["b", "f"], 0, 40), - new_batch_by_range(&["b", "h"], 100, 200), + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_by_range(&["a", "d"], 0, 60), + new_record_batch_by_range(&["b", "f"], 0, 40), + new_record_batch_by_range(&["b", "h"], 100, 200), ]); let write_request = SstWriteRequest { op_type: OperationType::Flush, metadata, - source: either::Left(source), + source, storage: None, max_sequence: None, + sst_write_format: Default::default(), cache_manager: Default::default(), index_options: IndexOptions::default(), index_config: Default::default(), @@ -636,19 +642,20 @@ mod tests { // Create source let metadata = Arc::new(sst_region_metadata()); - let source = new_source(&[ - new_batch_by_range(&["a", "d"], 0, 60), - new_batch_by_range(&["b", "f"], 0, 40), - new_batch_by_range(&["b", "h"], 100, 200), + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_by_range(&["a", "d"], 0, 60), + new_record_batch_by_range(&["b", "f"], 0, 40), + new_record_batch_by_range(&["b", "h"], 100, 200), ]); // Write to local cache and upload sst to mock remote store let write_request = SstWriteRequest { op_type: OperationType::Flush, metadata, - source: either::Left(source), + source, storage: None, max_sequence: None, + sst_write_format: Default::default(), cache_manager: cache_manager.clone(), index_options: IndexOptions::default(), index_config: Default::default(), @@ -686,9 +693,15 @@ mod tests { .cache(CacheStrategy::EnableAll(cache_manager.clone())) .page_index_policy(PageIndexPolicy::Optional); let reader = builder.build().await.unwrap().unwrap(); + let cached_write_parquet_metadata = crate::cache::CachedSstMeta::try_new( + "test.sst", + Arc::unwrap_or_clone(write_parquet_metadata), + ) + .unwrap() + .parquet_metadata(); // Check parquet metadata - assert_parquet_metadata_equal(write_parquet_metadata, reader.parquet_metadata()); + assert_parquet_metadata_equal(cached_write_parquet_metadata, reader.parquet_metadata()); } #[tokio::test] @@ -715,9 +728,9 @@ mod tests { let metadata = Arc::new(sst_region_metadata()); // Creates a source that can return an error to abort the writer. - let source = Source::Iter(Box::new( + let source = FlatSource::Iter(Box::new( [ - Ok(new_batch_by_range(&["a", "d"], 0, 60)), + Ok(new_record_batch_by_range(&["a", "d"], 0, 60)), InvalidBatchSnafu { reason: "Abort the writer", } @@ -730,9 +743,10 @@ mod tests { let write_request = SstWriteRequest { op_type: OperationType::Flush, metadata, - source: either::Left(source), + source, storage: None, max_sequence: None, + sst_write_format: Default::default(), cache_manager: cache_manager.clone(), index_options: IndexOptions::default(), index_config: Default::default(), diff --git a/src/mito2/src/compaction.rs b/src/mito2/src/compaction.rs index 6d51d1dd590..ba6957fdaed 100644 --- a/src/mito2/src/compaction.rs +++ b/src/mito2/src/compaction.rs @@ -58,10 +58,10 @@ use crate::error::{ TimeRangePredicateOverflowSnafu, TimeoutSnafu, }; use crate::metrics::{COMPACTION_STAGE_ELAPSED, INFLIGHT_COMPACTION_COUNT}; +use crate::read::BoxedRecordBatchStream; use crate::read::projection::ProjectionMapper; use crate::read::scan_region::{PredicateGroup, ScanInput}; use crate::read::seq_scan::SeqScan; -use crate::read::{BoxedBatchReader, BoxedRecordBatchStream}; use crate::region::options::{MergeMode, RegionOptions}; use crate::region::version::VersionControlRef; use crate::region::{ManifestContextRef, RegionLeaderState, RegionRoleState}; @@ -828,7 +828,7 @@ pub struct SerializedCompactionOutput { output_time_range: Option, } -/// Builders to create [BoxedBatchReader] for compaction. +/// Builders to create [BoxedRecordBatchStream] for compaction. struct CompactionSstReaderBuilder<'a> { metadata: RegionMetadataRef, sst_layer: AccessLayerRef, @@ -841,24 +841,17 @@ struct CompactionSstReaderBuilder<'a> { } impl CompactionSstReaderBuilder<'_> { - /// Builds [BoxedBatchReader] that reads all SST files and yields batches in primary key order. - async fn build_sst_reader(self) -> Result { - let scan_input = self.build_scan_input(false)?.with_compaction(true); - - SeqScan::new(scan_input).build_reader_for_compaction().await - } - /// Builds [BoxedRecordBatchStream] that reads all SST files and yields batches in flat format for compaction. async fn build_flat_sst_reader(self) -> Result { - let scan_input = self.build_scan_input(true)?.with_compaction(true); + let scan_input = self.build_scan_input()?.with_compaction(true); SeqScan::new(scan_input) .build_flat_reader_for_compaction() .await } - fn build_scan_input(self, flat_format: bool) -> Result { - let mapper = ProjectionMapper::all(&self.metadata, flat_format)?; + fn build_scan_input(self) -> Result { + let mapper = ProjectionMapper::all(&self.metadata, true)?; let mut scan_input = ScanInput::new(self.sst_layer, mapper) .with_files(self.inputs.to_vec()) .with_append_mode(self.append_mode) @@ -868,7 +861,7 @@ impl CompactionSstReaderBuilder<'_> { // We ignore file not found error during compaction. .with_ignore_file_not_found(true) .with_merge_mode(self.merge_mode) - .with_flat_format(flat_format); + .with_flat_format(true); // This serves as a workaround of https://github.com/GreptimeTeam/greptimedb/issues/3944 // by converting time ranges into predicate. diff --git a/src/mito2/src/compaction/compactor.rs b/src/mito2/src/compaction/compactor.rs index 1876972b0d6..b03e6415e81 100644 --- a/src/mito2/src/compaction/compactor.rs +++ b/src/mito2/src/compaction/compactor.rs @@ -43,7 +43,7 @@ use crate::error::{ use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; use crate::metrics; -use crate::read::{FlatSource, Source}; +use crate::read::FlatSource; use crate::region::options::RegionOptions; use crate::region::version::VersionRef; use crate::region::{ManifestContext, RegionLeaderState, RegionRoleState}; @@ -356,13 +356,8 @@ impl DefaultCompactor { time_range: output.output_time_range, merge_mode, }; - let source = if flat_format { - let reader = builder.build_flat_sst_reader().await?; - Either::Right(FlatSource::Stream(reader)) - } else { - let reader = builder.build_sst_reader().await?; - Either::Left(Source::Reader(reader)) - }; + let reader = builder.build_flat_sst_reader().await?; + let source = FlatSource::Stream(reader); let mut metrics = Metrics::new(WriteType::Compaction); let region_metadata = compaction_region.region_metadata.clone(); let sst_infos = compaction_region @@ -375,6 +370,11 @@ impl DefaultCompactor { cache_manager: compaction_region.cache_manager.clone(), storage, max_sequence: max_sequence.map(NonZero::get), + sst_write_format: if flat_format { + FormatType::Flat + } else { + FormatType::PrimaryKey + }, index_options, index_config, inverted_index_config, diff --git a/src/mito2/src/engine/row_selector_test.rs b/src/mito2/src/engine/row_selector_test.rs index 317ede5a97d..d79152e57fa 100644 --- a/src/mito2/src/engine/row_selector_test.rs +++ b/src/mito2/src/engine/row_selector_test.rs @@ -24,7 +24,7 @@ use crate::test_util::{ CreateRequestBuilder, TestEnv, build_rows_for_key, flush_region, put_rows, rows_schema, }; -async fn test_last_row(append_mode: bool) { +async fn test_last_row(append_mode: bool, flat_format: bool) { let mut env = TestEnv::new().await; let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -39,9 +39,12 @@ async fn test_last_row(append_mode: bool) { env.get_kv_backend(), ) .await; - let request = CreateRequestBuilder::new() - .insert_option("append_mode", &append_mode.to_string()) - .build(); + let mut request_builder = + CreateRequestBuilder::new().insert_option("append_mode", &append_mode.to_string()); + if flat_format { + request_builder = request_builder.insert_option("sst_format", "flat"); + } + let request = request_builder.build(); let column_schemas = rows_schema(&request); engine .handle_request(region_id, RegionRequest::Create(request)) @@ -106,10 +109,20 @@ async fn test_last_row(append_mode: bool) { #[tokio::test] async fn test_last_row_append_mode_disabled() { - test_last_row(false).await; + test_last_row(false, false).await; } #[tokio::test] async fn test_last_row_append_mode_enabled() { - test_last_row(true).await; + test_last_row(true, false).await; +} + +#[tokio::test] +async fn test_last_row_flat_format_append_mode_disabled() { + test_last_row(false, true).await; +} + +#[tokio::test] +async fn test_last_row_flat_format_append_mode_enabled() { + test_last_row(true, true).await; } diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index 0c16544b6e9..fedac95d279 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -22,7 +22,6 @@ use std::time::Instant; use common_telemetry::{debug, error, info}; use datatypes::arrow::datatypes::SchemaRef; -use either::Either; use partition::expr::PartitionExpr; use smallvec::{SmallVec, smallvec}; use snafu::ResultExt; @@ -41,18 +40,14 @@ use crate::error::{ }; use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; use crate::memtable::bulk::ENCODE_ROW_THRESHOLD; -use crate::memtable::{ - BoxedRecordBatchIterator, EncodedRange, IterBuilder, MemtableRanges, RangesOptions, -}; +use crate::memtable::{BoxedRecordBatchIterator, EncodedRange, MemtableRanges, RangesOptions}; use crate::metrics::{ FLUSH_BYTES_TOTAL, FLUSH_ELAPSED, FLUSH_FAILURE_TOTAL, FLUSH_FILE_TOTAL, FLUSH_REQUESTS_TOTAL, INFLIGHT_FLUSH_COUNT, }; -use crate::read::dedup::{DedupReader, LastNonNull, LastRow}; +use crate::read::FlatSource; use crate::read::flat_dedup::{FlatDedupIterator, FlatLastNonNull, FlatLastRow}; use crate::read::flat_merge::FlatMergeIterator; -use crate::read::merge::MergeReaderBuilder; -use crate::read::{FlatSource, Source}; use crate::region::options::{IndexOptions, MergeMode, RegionOptions}; use crate::region::version::{VersionControlData, VersionControlRef, VersionRef}; use crate::region::{ManifestContextRef, RegionLeaderState, RegionRoleState, parse_partition_expr}; @@ -62,8 +57,10 @@ use crate::request::{ }; use crate::schedule::scheduler::{Job, SchedulerRef}; use crate::sst::file::FileMeta; -use crate::sst::parquet::{DEFAULT_READ_BATCH_SIZE, DEFAULT_ROW_GROUP_SIZE, SstInfo, WriteOptions}; -use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; +use crate::sst::parquet::{ + DEFAULT_READ_BATCH_SIZE, DEFAULT_ROW_GROUP_SIZE, SstInfo, WriteOptions, flat_format, +}; +use crate::sst::{FlatSchemaOptions, FormatType, to_flat_sst_arrow_schema}; use crate::worker::WorkerListener; /// Global write buffer (memtable) manager. @@ -480,78 +477,29 @@ impl RegionFlushTask { // the counter may have more series than the actual series count. series_count += memtable_series_count; - if mem_ranges.is_record_batch() { - let flush_start = Instant::now(); - let FlushFlatMemResult { - num_encoded, - num_sources, - results, - } = self - .flush_flat_mem_ranges(version, &write_opts, mem_ranges) - .await?; - encoded_part_count += num_encoded; - for (source_idx, result) in results.into_iter().enumerate() { - let (max_sequence, ssts_written, metrics) = result?; - if ssts_written.is_empty() { - // No data written. - continue; - } - - common_telemetry::debug!( - "Region {} flush one memtable {} {}/{}, metrics: {:?}", - self.region_id, - memtable_id, - source_idx, - num_sources, - metrics - ); - - flush_metrics = flush_metrics.merge(metrics); - - file_metas.extend(ssts_written.into_iter().map(|sst_info| { - flushed_bytes += sst_info.file_size; - Self::new_file_meta( - self.region_id, - max_sequence, - sst_info, - partition_expr.clone(), - ) - })); - } - - common_telemetry::debug!( - "Region {} flush {} memtables for {}, num_mem_ranges: {}, num_encoded: {}, num_rows: {}, flush_cost: {:?}, compact_cost: {:?}", - self.region_id, - num_sources, - memtable_id, - num_mem_ranges, - num_encoded, - num_mem_rows, - flush_start.elapsed(), - compact_cost, - ); - } else { - let max_sequence = mem_ranges.max_sequence(); - let source = memtable_source(mem_ranges, &version.options).await?; - - // Flush to level 0. - let source = Either::Left(source); - let write_request = self.new_write_request(version, max_sequence, source); - - let mut metrics = Metrics::new(WriteType::Flush); - let ssts_written = self - .access_layer - .write_sst(write_request, &write_opts, &mut metrics) - .await?; - FLUSH_FILE_TOTAL.inc_by(ssts_written.len() as u64); + let flush_start = Instant::now(); + let FlushFlatMemResult { + num_encoded, + num_sources, + results, + } = self + .flush_flat_mem_ranges(version, &write_opts, mem_ranges) + .await?; + encoded_part_count += num_encoded; + for (source_idx, result) in results.into_iter().enumerate() { + let (max_sequence, ssts_written, metrics) = result?; if ssts_written.is_empty() { // No data written. continue; } - debug!( - "Region {} flush one memtable, num_mem_ranges: {}, num_rows: {}, metrics: {:?}", - self.region_id, num_mem_ranges, num_mem_rows, metrics + common_telemetry::debug!( + "Region {} flush one memtable {} {}/{}, metrics: {:?}", + self.region_id, + memtable_id, + source_idx, + num_sources, + metrics ); flush_metrics = flush_metrics.merge(metrics); @@ -565,7 +513,19 @@ impl RegionFlushTask { partition_expr.clone(), ) })); - }; + } + + common_telemetry::debug!( + "Region {} flush {} memtables for {}, num_mem_ranges: {}, num_encoded: {}, num_rows: {}, flush_cost: {:?}, compact_cost: {:?}", + self.region_id, + num_sources, + memtable_id, + num_mem_ranges, + num_encoded, + num_mem_rows, + flush_start.elapsed(), + compact_cost, + ); } Ok(DoFlushMemtablesResult { @@ -587,16 +547,17 @@ impl RegionFlushTask { &version.metadata, &FlatSchemaOptions::from_encoding(version.metadata.primary_key_encoding), ); + let field_column_start = + flat_format::field_column_start(&version.metadata, batch_schema.fields().len()); let flat_sources = memtable_flat_sources( batch_schema, mem_ranges, &version.options, - version.metadata.primary_key.len(), + field_column_start, )?; let mut tasks = Vec::with_capacity(flat_sources.encoded.len() + flat_sources.sources.len()); let num_encoded = flat_sources.encoded.len(); for (source, max_sequence) in flat_sources.sources { - let source = Either::Right(source); let write_request = self.new_write_request(version, max_sequence, source); let access_layer = self.access_layer.clone(); let write_opts = write_opts.clone(); @@ -667,8 +628,13 @@ impl RegionFlushTask { &self, version: &VersionRef, max_sequence: u64, - source: Either, + source: FlatSource, ) -> SstWriteRequest { + let flat_format = version + .options + .sst_format + .map(|f| f == FormatType::Flat) + .unwrap_or(self.engine_config.default_experimental_flat_format); SstWriteRequest { op_type: OperationType::Flush, metadata: version.metadata.clone(), @@ -676,6 +642,11 @@ impl RegionFlushTask { cache_manager: self.cache_manager.clone(), storage: version.options.storage.clone(), max_sequence: Some(max_sequence), + sst_write_format: if flat_format { + FormatType::Flat + } else { + FormatType::PrimaryKey + }, index_options: self.index_options.clone(), index_config: self.engine_config.index.clone(), inverted_index_config: self.engine_config.inverted_index.clone(), @@ -722,41 +693,6 @@ struct DoFlushMemtablesResult { flush_metrics: Metrics, } -/// Returns a [Source] for the given memtable. -async fn memtable_source(mem_ranges: MemtableRanges, options: &RegionOptions) -> Result { - let source = if mem_ranges.ranges.len() == 1 { - let only_range = mem_ranges.ranges.into_values().next().unwrap(); - let iter = only_range.build_iter()?; - Source::Iter(iter) - } else { - // todo(hl): a workaround since sync version of MergeReader is wip. - let sources = mem_ranges - .ranges - .into_values() - .map(|r| r.build_iter().map(Source::Iter)) - .collect::>>()?; - let merge_reader = MergeReaderBuilder::from_sources(sources).build().await?; - let maybe_dedup = if options.append_mode { - // no dedup in append mode - Box::new(merge_reader) as _ - } else { - // dedup according to merge mode - match options.merge_mode.unwrap_or(MergeMode::LastRow) { - MergeMode::LastRow => { - Box::new(DedupReader::new(merge_reader, LastRow::new(false), None)) as _ - } - MergeMode::LastNonNull => Box::new(DedupReader::new( - merge_reader, - LastNonNull::new(false), - None, - )) as _, - } - }; - Source::Reader(maybe_dedup) - }; - Ok(source) -} - struct FlatSources { sources: SmallVec<[(FlatSource, SequenceNumber); 4]>, encoded: SmallVec<[(EncodedRange, SequenceNumber); 4]>, diff --git a/src/mito2/src/memtable.rs b/src/mito2/src/memtable.rs index c39bbfa3469..7494ec68ede 100644 --- a/src/mito2/src/memtable.rs +++ b/src/mito2/src/memtable.rs @@ -28,6 +28,7 @@ use mito_codec::key_values::KeyValue; pub use mito_codec::key_values::KeyValues; use mito_codec::row_converter::{PrimaryKeyCodec, build_primary_key_codec}; use serde::{Deserialize, Serialize}; +use snafu::ensure; use store_api::metadata::RegionMetadataRef; use store_api::storage::{ColumnId, SequenceNumber, SequenceRange}; @@ -231,10 +232,17 @@ impl MemtableRanges { impl IterBuilder for MemtableRanges { fn build(&self, _metrics: Option) -> Result { - UnsupportedOperationSnafu { - err_msg: "MemtableRanges does not support build iterator", - } - .fail() + ensure!( + self.ranges.len() == 1, + UnsupportedOperationSnafu { + err_msg: format!( + "Building an iterator from MemtableRanges expects 1 range, but got {}", + self.ranges.len() + ), + } + ); + + self.ranges.values().next().unwrap().build_iter() } fn is_record_batch(&self) -> bool { @@ -256,20 +264,6 @@ pub trait Memtable: Send + Sync + fmt::Debug { /// Writes an encoded batch of into memtable. fn write_bulk(&self, part: crate::memtable::bulk::part::BulkPart) -> Result<()>; - /// Scans the memtable. - /// `projection` selects columns to read, `None` means reading all columns. - /// `filters` are the predicates to be pushed down to memtable. - /// - /// # Note - /// This method should only be used for tests. - #[cfg(any(test, feature = "test"))] - fn iter( - &self, - projection: Option<&[ColumnId]>, - predicate: Option, - sequence: Option, - ) -> Result; - /// Returns the ranges in the memtable. /// /// The returned map contains the range id and the range after applying the predicate. diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index cf2ced06fef..4dad4fb8851 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -57,7 +57,7 @@ use crate::memtable::{ use crate::read::flat_dedup::{FlatDedupIterator, FlatLastNonNull, FlatLastRow}; use crate::read::flat_merge::FlatMergeIterator; use crate::region::options::MergeMode; -use crate::sst::parquet::format::FIXED_POS_COLUMN_NUM; +use crate::sst::parquet::flat_format::field_column_start; use crate::sst::parquet::{DEFAULT_READ_BATCH_SIZE, DEFAULT_ROW_GROUP_SIZE}; use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; @@ -462,16 +462,6 @@ impl Memtable for BulkMemtable { Ok(()) } - #[cfg(any(test, feature = "test"))] - fn iter( - &self, - _projection: Option<&[ColumnId]>, - _predicate: Option, - _sequence: Option, - ) -> Result { - todo!() - } - fn ranges( &self, projection: Option<&[ColumnId]>, @@ -1186,13 +1176,8 @@ impl MemtableCompactor { Box::new(dedup_iter) } MergeMode::LastNonNull => { - // Calculates field column start: total columns - fixed columns - field columns - // Field column count = total metadata columns - time index column - primary key columns - let field_column_count = - metadata.column_metadatas.len() - 1 - metadata.primary_key.len(); - let total_columns = arrow_schema.fields().len(); let field_column_start = - total_columns - FIXED_POS_COLUMN_NUM - field_column_count; + field_column_start(metadata, arrow_schema.fields().len()); let dedup_iter = FlatDedupIterator::new( merged_iter, diff --git a/src/mito2/src/memtable/partition_tree.rs b/src/mito2/src/memtable/partition_tree.rs index febae467849..662bfd99f64 100644 --- a/src/mito2/src/memtable/partition_tree.rs +++ b/src/mito2/src/memtable/partition_tree.rs @@ -177,16 +177,6 @@ impl Memtable for PartitionTreeMemtable { .fail() } - #[cfg(any(test, feature = "test"))] - fn iter( - &self, - projection: Option<&[ColumnId]>, - predicate: Option, - sequence: Option, - ) -> Result { - self.tree.read(projection, predicate, sequence, None) - } - fn ranges( &self, projection: Option<&[ColumnId]>, @@ -396,8 +386,6 @@ mod tests { use api::v1::{Mutation, OpType, Rows, SemanticType}; use common_query::prelude::{greptime_timestamp, greptime_value}; use common_time::Timestamp; - use datafusion_common::Column; - use datafusion_expr::{BinaryExpr, Expr, Literal, Operator}; use datatypes::data_type::ConcreteDataType; use datatypes::prelude::Vector; use datatypes::scalars::ScalarVector; @@ -548,7 +536,10 @@ mod tests { let expect = (0..100).collect::>(); let kvs = memtable_util::build_key_values(&metadata, "hello".to_string(), 10, &expect, 1); memtable.write(&kvs).unwrap(); - let iter = memtable.iter(Some(&[3]), None, None).unwrap(); + let ranges = memtable + .ranges(Some(&[3]), RangesOptions::default()) + .unwrap(); + let iter = ranges.build(None).unwrap(); let mut v0_all = vec![]; for res in iter { @@ -625,41 +616,6 @@ mod tests { assert_eq!(expect, read); } - #[test] - fn test_memtable_filter() { - let metadata = Arc::new(memtable_util::metadata_with_primary_key(vec![0, 1], false)); - // Try to build a memtable via the builder. - let memtable = PartitionTreeMemtableBuilder::new( - PartitionTreeConfig { - index_max_keys_per_shard: 40, - ..Default::default() - }, - None, - ) - .build(1, &metadata); - - for i in 0..100 { - let timestamps: Vec<_> = (0..10).map(|v| i as i64 * 1000 + v).collect(); - let kvs = - memtable_util::build_key_values(&metadata, "hello".to_string(), i, ×tamps, 1); - memtable.write(&kvs).unwrap(); - } - - for i in 0..100 { - let timestamps: Vec<_> = (0..10).map(|v| i as i64 * 1000 + v).collect(); - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("k1"))), - op: Operator::Eq, - right: Box::new((i as u32).lit()), - }); - let iter = memtable - .iter(None, Some(Predicate::new(vec![expr])), None) - .unwrap(); - let read = collect_iter_timestamps(iter); - assert_eq!(timestamps, read); - } - } - #[test] fn test_deserialize_config() { let config = PartitionTreeConfig { @@ -811,7 +767,11 @@ mod tests { )) .unwrap(); - let mut reader = new_memtable.iter(None, None, None).unwrap(); + let mut reader = new_memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let batch = reader.next().unwrap().unwrap(); let pk = codec.decode(batch.primary_key()).unwrap().into_dense(); if let Value::String(s) = &pk[2] { @@ -916,7 +876,14 @@ mod tests { .unwrap(); memtable.freeze().unwrap(); assert_eq!( - collect_kvs(memtable.iter(None, None, None).unwrap(), &metadata), + collect_kvs( + memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(), + &metadata + ), ('a'..'h').map(|c| (c.to_string(), c.to_string())).collect() ); let forked = memtable.fork(2, &metadata); @@ -925,7 +892,14 @@ mod tests { forked.write(&key_values(&metadata, keys.iter())).unwrap(); forked.freeze().unwrap(); assert_eq!( - collect_kvs(forked.iter(None, None, None).unwrap(), &metadata), + collect_kvs( + forked + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(), + &metadata + ), keys.iter() .map(|c| (c.to_string(), c.to_string())) .collect() @@ -936,7 +910,14 @@ mod tests { let keys = ["g", "e", "a", "f", "b", "c", "h"]; forked2.write(&key_values(&metadata, keys.iter())).unwrap(); - let kvs = collect_kvs(forked2.iter(None, None, None).unwrap(), &metadata); + let kvs = collect_kvs( + forked2 + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(), + &metadata, + ); let expected = keys .iter() .map(|c| (c.to_string(), c.to_string())) diff --git a/src/mito2/src/memtable/simple_bulk_memtable.rs b/src/mito2/src/memtable/simple_bulk_memtable.rs index 4dcaa2bac0e..6d91f003619 100644 --- a/src/mito2/src/memtable/simple_bulk_memtable.rs +++ b/src/mito2/src/memtable/simple_bulk_memtable.rs @@ -213,22 +213,6 @@ impl Memtable for SimpleBulkMemtable { Ok(()) } - #[cfg(any(test, feature = "test"))] - fn iter( - &self, - projection: Option<&[ColumnId]>, - _predicate: Option, - sequence: Option, - ) -> error::Result { - let iter = self.create_iter(projection, sequence)?.build(None)?; - if self.merge_mode == MergeMode::LastNonNull { - let iter = LastNonNullIter::new(iter); - Ok(Box::new(iter)) - } else { - Ok(Box::new(iter)) - } - } - fn ranges( &self, projection: Option<&[ColumnId]>, @@ -526,7 +510,11 @@ mod tests { )) .unwrap(); - let mut iter = memtable.iter(None, None, None).unwrap(); + let mut iter = memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let batch = iter.next().unwrap().unwrap(); assert_eq!(2, batch.num_rows()); assert_eq!(2, batch.fields().len()); @@ -551,7 +539,11 @@ mod tests { )) .unwrap(); - let mut iter = memtable.iter(None, None, None).unwrap(); + let mut iter = memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let batch = iter.next().unwrap().unwrap(); assert_eq!(1, batch.num_rows()); assert_eq!(2, batch.fields().len()); @@ -565,7 +557,11 @@ mod tests { // Only project column 2 (f1) let projection = vec![2]; - let mut iter = memtable.iter(Some(&projection), None, None).unwrap(); + let mut iter = memtable + .ranges(Some(&projection), RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let batch = iter.next().unwrap().unwrap(); assert_eq!(1, batch.num_rows()); @@ -592,7 +588,11 @@ mod tests { OpType::Put, )) .unwrap(); - let mut iter = memtable.iter(None, None, None).unwrap(); + let mut iter = memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let batch = iter.next().unwrap().unwrap(); assert_eq!(1, batch.num_rows()); // deduped to 1 row @@ -611,7 +611,11 @@ mod tests { let kv = kvs.iter().next().unwrap(); memtable.write_one(kv).unwrap(); - let mut iter = memtable.iter(None, None, None).unwrap(); + let mut iter = memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let batch = iter.next().unwrap().unwrap(); assert_eq!(1, batch.num_rows()); } @@ -745,7 +749,11 @@ mod tests { }; memtable.write_bulk(part).unwrap(); - let mut iter = memtable.iter(None, None, None).unwrap(); + let mut iter = memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let batch = iter.next().unwrap().unwrap(); assert_eq!(2, batch.num_rows()); @@ -764,7 +772,11 @@ mod tests { OpType::Put, ); memtable.write(&kvs).unwrap(); - let mut iter = memtable.iter(None, None, None).unwrap(); + let mut iter = memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let batch = iter.next().unwrap().unwrap(); assert_eq!(3, batch.num_rows()); assert_eq!( @@ -854,7 +866,15 @@ mod tests { // Filter with sequence 0 should only return first write let mut iter = memtable - .iter(None, None, Some(SequenceRange::LtEq { max: 0 })) + .ranges( + None, + RangesOptions { + sequence: Some(SequenceRange::LtEq { max: 0 }), + ..Default::default() + }, + ) + .unwrap() + .build(None) .unwrap(); let batch = iter.next().unwrap().unwrap(); assert_eq!(1, batch.num_rows()); diff --git a/src/mito2/src/memtable/simple_bulk_memtable/test_only.rs b/src/mito2/src/memtable/simple_bulk_memtable/test_only.rs index b71a86c5540..08edebdbb2a 100644 --- a/src/mito2/src/memtable/simple_bulk_memtable/test_only.rs +++ b/src/mito2/src/memtable/simple_bulk_memtable/test_only.rs @@ -12,98 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashSet; -use std::time::Instant; - use store_api::metadata::RegionMetadataRef; -use store_api::storage::{ColumnId, SequenceRange}; -use crate::error; -use crate::memtable::simple_bulk_memtable::{Iter, SimpleBulkMemtable}; -use crate::memtable::time_series::Values; -use crate::memtable::{BoxedBatchIterator, IterBuilder, MemScanMetrics}; -use crate::read::dedup::LastNonNullIter; -use crate::region::options::MergeMode; +use crate::memtable::simple_bulk_memtable::SimpleBulkMemtable; impl SimpleBulkMemtable { pub fn region_metadata(&self) -> RegionMetadataRef { self.region_metadata.clone() } - - pub(crate) fn create_iter( - &self, - projection: Option<&[ColumnId]>, - sequence: Option, - ) -> error::Result { - let mut series = self.series.write().unwrap(); - - let values = if series.is_empty() { - None - } else { - Some(series.compact(&self.region_metadata)?.clone()) - }; - let projection = self.build_projection(projection); - Ok(BatchIterBuilderDeprecated { - region_metadata: self.region_metadata.clone(), - values, - projection, - dedup: self.dedup, - sequence, - merge_mode: self.merge_mode, - }) - } -} - -#[derive(Clone)] -pub(crate) struct BatchIterBuilderDeprecated { - region_metadata: RegionMetadataRef, - values: Option, - projection: HashSet, - sequence: Option, - dedup: bool, - merge_mode: MergeMode, -} - -impl IterBuilder for BatchIterBuilderDeprecated { - fn build(&self, metrics: Option) -> error::Result { - let start_time = Instant::now(); - let Some(values) = self.values.clone() else { - return Ok(Box::new(Iter { batch: None })); - }; - - let maybe_batch = values - .to_batch( - &[], - &self.region_metadata, - &self.projection, - self.sequence, - self.dedup, - self.merge_mode, - ) - .map(Some) - .transpose(); - - // Collect metrics from the batch - if let Some(metrics) = metrics { - let (num_rows, num_batches) = match &maybe_batch { - Some(Ok(batch)) => (batch.num_rows(), 1), - _ => (0, 0), - }; - let inner = crate::memtable::MemScanMetricsData { - total_series: 1, - num_rows, - num_batches, - scan_cost: start_time.elapsed(), - }; - metrics.merge_inner(&inner); - } - - let iter = Iter { batch: maybe_batch }; - - if self.merge_mode == MergeMode::LastNonNull { - Ok(Box::new(LastNonNullIter::new(iter))) - } else { - Ok(Box::new(iter)) - } - } } diff --git a/src/mito2/src/memtable/time_partition.rs b/src/mito2/src/memtable/time_partition.rs index 6f11c813cba..ee695aceb87 100644 --- a/src/mito2/src/memtable/time_partition.rs +++ b/src/mito2/src/memtable/time_partition.rs @@ -827,6 +827,7 @@ mod tests { use super::*; use crate::memtable::partition_tree::PartitionTreeMemtableBuilder; use crate::memtable::time_series::TimeSeriesMemtableBuilder; + use crate::memtable::{IterBuilder, RangesOptions}; use crate::test_util::memtable_util::{self, collect_iter_timestamps}; #[test] @@ -852,7 +853,11 @@ mod tests { partitions.list_memtables(&mut memtables); assert_eq!(0, memtables[0].id()); - let iter = memtables[0].iter(None, None, None).unwrap(); + let iter = memtables[0] + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let timestamps = collect_iter_timestamps(iter); assert_eq!(&[1000, 3000, 5000, 6000, 7000], ×tamps[..]); } @@ -890,7 +895,11 @@ mod tests { let mut memtables = Vec::new(); partitions.list_memtables(&mut memtables); - let iter = memtables[0].iter(None, None, None).unwrap(); + let iter = memtables[0] + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let timestamps = collect_iter_timestamps(iter); assert_eq!(&[0, 2000, 3000, 4000, 5000, 7000], ×tamps[..]); let parts = partitions.list_partitions(); @@ -943,7 +952,12 @@ mod tests { let partitions = new_multi_partitions(&metadata); let parts = partitions.list_partitions(); - let iter = parts[0].memtable.iter(None, None, None).unwrap(); + let iter = parts[0] + .memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let timestamps = collect_iter_timestamps(iter); assert_eq!(0, parts[0].memtable.id()); assert_eq!( @@ -955,7 +969,12 @@ mod tests { parts[0].time_range.max_timestamp ); assert_eq!(&[0, 2000, 3000, 4000], ×tamps[..]); - let iter = parts[1].memtable.iter(None, None, None).unwrap(); + let iter = parts[1] + .memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); assert_eq!(1, parts[1].memtable.id()); let timestamps = collect_iter_timestamps(iter); assert_eq!(&[5000, 7000], ×tamps[..]); @@ -1273,7 +1292,12 @@ mod tests { let parts = partitions.list_partitions(); assert_eq!(1, parts.len()); - let iter = parts[0].memtable.iter(None, None, None).unwrap(); + let iter = parts[0] + .memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let timestamps = collect_iter_timestamps(iter); assert_eq!(&[1000, 2000, 3000], ×tamps[..]); @@ -1284,11 +1308,21 @@ mod tests { let parts = partitions.list_partitions(); assert_eq!(2, parts.len()); // Check first partition [0, 5000) - let iter = parts[0].memtable.iter(None, None, None).unwrap(); + let iter = parts[0] + .memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let timestamps = collect_iter_timestamps(iter); assert_eq!(&[1000, 2000, 3000, 4000], ×tamps[..]); // Check second partition [5000, 10000) - let iter = parts[1].memtable.iter(None, None, None).unwrap(); + let iter = parts[1] + .memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let timestamps = collect_iter_timestamps(iter); assert_eq!(&[5000, 6000], ×tamps[..]); @@ -1301,7 +1335,12 @@ mod tests { assert_eq!(3, parts.len()); // Check new partition [10000, 15000) - let iter = parts[2].memtable.iter(None, None, None).unwrap(); + let iter = parts[2] + .memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let timestamps = collect_iter_timestamps(iter); assert_eq!(&[11000, 12000], ×tamps[..]); @@ -1314,7 +1353,12 @@ mod tests { let parts = partitions.list_partitions(); assert_eq!(1, parts.len()); - let iter = parts[0].memtable.iter(None, None, None).unwrap(); + let iter = parts[0] + .memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let timestamps = collect_iter_timestamps(iter); assert_eq!(&[1000, 5000, 9000], ×tamps[..]); } diff --git a/src/mito2/src/memtable/time_series.rs b/src/mito2/src/memtable/time_series.rs index 271a9343eb2..97f5f3c9ceb 100644 --- a/src/mito2/src/memtable/time_series.rs +++ b/src/mito2/src/memtable/time_series.rs @@ -267,39 +267,6 @@ impl Memtable for TimeSeriesMemtable { Ok(()) } - #[cfg(any(test, feature = "test"))] - fn iter( - &self, - projection: Option<&[ColumnId]>, - filters: Option, - sequence: Option, - ) -> Result { - let projection = if let Some(projection) = projection { - projection.iter().copied().collect() - } else { - self.region_metadata - .field_columns() - .map(|c| c.column_id) - .collect() - }; - - let iter = self.series_set.iter_series( - projection, - filters, - self.dedup, - self.merge_mode, - sequence, - None, - )?; - - if self.merge_mode == MergeMode::LastNonNull { - let iter = LastNonNullIter::new(iter); - Ok(Box::new(iter)) - } else { - Ok(Box::new(iter)) - } - } - fn ranges( &self, projection: Option<&[ColumnId]>, @@ -1798,7 +1765,9 @@ mod tests { *expected_ts.entry(ts).or_default() += if dedup { 1 } else { 2 }; } - let iter = memtable.iter(None, None, None).unwrap(); + let ranges = memtable.ranges(None, RangesOptions::default()).unwrap(); + let range = ranges.ranges.into_values().next().unwrap(); + let iter = range.build_iter().unwrap(); let mut read = HashMap::new(); for ts in iter @@ -1838,7 +1807,11 @@ mod tests { let memtable = TimeSeriesMemtable::new(schema, 42, None, true, MergeMode::LastRow); memtable.write(&kvs).unwrap(); - let iter = memtable.iter(Some(&[3]), None, None).unwrap(); + let iter = memtable + .ranges(Some(&[3]), RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let mut v0_all = vec![]; @@ -1917,7 +1890,11 @@ mod tests { barrier.wait(); for _ in 0..10 { - let iter = memtable.iter(None, None, None).unwrap(); + let iter = memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); for batch_result in iter { let _ = batch_result.unwrap(); } @@ -1936,7 +1913,11 @@ mod tests { handle.join().unwrap(); } - let iter = memtable.iter(None, None, None).unwrap(); + let iter = memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .build(None) + .unwrap(); let mut series_count = 0; let mut row_count = 0; diff --git a/src/mito2/src/read.rs b/src/mito2/src/read.rs index 5fbd63ce8be..240a99c247d 100644 --- a/src/mito2/src/read.rs +++ b/src/mito2/src/read.rs @@ -27,6 +27,7 @@ pub mod projection; pub(crate) mod prune; pub(crate) mod pruner; pub mod range; +pub(crate) mod range_cache; pub mod scan_region; pub mod scan_util; pub(crate) mod seq_scan; diff --git a/src/mito2/src/read/batch_adapter.rs b/src/mito2/src/read/batch_adapter.rs index 461dbeba696..4698229c5b2 100644 --- a/src/mito2/src/read/batch_adapter.rs +++ b/src/mito2/src/read/batch_adapter.rs @@ -59,7 +59,7 @@ impl BatchToRecordBatchAdapter { /// - `metadata`: region metadata describing the schema. /// - `codec`: codec for decoding the encoded primary key bytes. /// - `read_column_ids`: projected column ids to read. - pub(crate) fn new( + pub fn new( iter: BoxedBatchIterator, metadata: RegionMetadataRef, codec: Arc, diff --git a/src/mito2/src/read/last_row.rs b/src/mito2/src/read/last_row.rs index 69bd02ae60f..b3c5a30f5cf 100644 --- a/src/mito2/src/read/last_row.rs +++ b/src/mito2/src/read/last_row.rs @@ -21,6 +21,7 @@ use datatypes::arrow::array::{Array, BinaryArray}; use datatypes::arrow::compute::concat_batches; use datatypes::arrow::record_batch::RecordBatch; use datatypes::vectors::UInt32Vector; +use futures::{Stream, TryStreamExt}; use mito_codec::row_converter::PrimaryKeyFilter; use snafu::ResultExt; use store_api::storage::{FileId, TimeSeriesRowSelector}; @@ -31,7 +32,7 @@ use crate::cache::{ }; use crate::error::{ComputeArrowSnafu, Result}; use crate::memtable::partition_tree::data::timestamp_array_to_i64_slice; -use crate::read::{Batch, BatchReader, BoxedBatchReader}; +use crate::read::{Batch, BatchReader, BoxedBatchReader, BoxedRecordBatchStream}; use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; use crate::sst::parquet::flat_format::{primary_key_column_index, time_index_column_index}; use crate::sst::parquet::format::{PrimaryKeyArray, primary_key_offsets}; @@ -637,6 +638,41 @@ impl FlatLastTimestampSelector { } } +/// Reader that keeps only the last row of each time series from a flat RecordBatch stream. +/// Assumes input is sorted, deduped, and contains no delete operations. +pub(crate) struct FlatLastRowReader { + stream: BoxedRecordBatchStream, + selector: FlatLastTimestampSelector, + pending: BatchBuffer, +} + +impl FlatLastRowReader { + /// Creates a new `FlatLastRowReader`. + pub(crate) fn new(stream: BoxedRecordBatchStream) -> Self { + Self { + stream, + selector: FlatLastTimestampSelector::default(), + pending: BatchBuffer::new(), + } + } + + /// Converts the reader into a stream of RecordBatches. + pub(crate) fn into_stream(mut self) -> impl Stream> { + async_stream::try_stream! { + while let Some(batch) = self.stream.try_next().await? { + self.selector.on_next(batch, &mut self.pending)?; + if self.pending.is_full() { + yield self.pending.concat()?; + } + } + self.selector.finish(&mut self.pending)?; + if !self.pending.is_empty() { + yield self.pending.concat()?; + } + } + } +} + /// Gets the primary key bytes at `index` from the primary key dictionary column. fn primary_key_bytes_at(batch: &RecordBatch, pk_col_idx: usize, index: usize) -> &[u8] { let pk_dict = batch diff --git a/src/mito2/src/read/prune.rs b/src/mito2/src/read/prune.rs index b08f6c55d8b..4e731a4a6a2 100644 --- a/src/mito2/src/read/prune.rs +++ b/src/mito2/src/read/prune.rs @@ -81,11 +81,6 @@ impl PruneReader { } } - pub(crate) fn reset_source(&mut self, source: Source, skip_fields: bool) { - self.source = source; - self.skip_fields = skip_fields; - } - /// Merge metrics with the inner reader and return the merged metrics. pub(crate) fn metrics(&self) -> ReaderMetrics { let mut metrics = self.metrics.clone(); diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs new file mode 100644 index 00000000000..5b90e68baef --- /dev/null +++ b/src/mito2/src/read/range_cache.rs @@ -0,0 +1,252 @@ +// 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. + +//! Utilities for the partition range scan result cache. + +use std::mem; +use std::sync::Arc; + +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::prelude::ConcreteDataType; +use store_api::storage::{ColumnId, FileId, RegionId, TimeSeriesRowSelector}; + +use crate::memtable::record_batch_estimated_size; +use crate::region::options::MergeMode; + +/// Fingerprint of the scan request fields that affect partition range cache reuse. +/// +/// It records a normalized view of the projected columns and filters, plus +/// scan options that can change the returned rows. Schema-dependent metadata +/// and the partition expression version are included so cached results are not +/// reused across incompatible schema or partitioning changes. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct ScanRequestFingerprint { + /// Projection and filters without the time index and partition exprs. + inner: Arc, + /// Filters with the time index column. + time_filters: Option>>, + series_row_selector: Option, + append_mode: bool, + filter_deleted: bool, + merge_mode: MergeMode, + /// We keep the partition expr version to ensure we won't reuse the fingerprint after we change the partition expr. + /// We store the version instead of the whole partition expr or partition expr filters. + partition_expr_version: u64, +} + +#[derive(Debug)] +pub(crate) struct ScanRequestFingerprintBuilder { + pub(crate) read_column_ids: Vec, + pub(crate) read_column_types: Vec>, + pub(crate) filters: Vec, + pub(crate) time_filters: Vec, + pub(crate) series_row_selector: Option, + pub(crate) append_mode: bool, + pub(crate) filter_deleted: bool, + pub(crate) merge_mode: MergeMode, + pub(crate) partition_expr_version: u64, +} + +impl ScanRequestFingerprintBuilder { + pub(crate) fn build(self) -> ScanRequestFingerprint { + let Self { + read_column_ids, + read_column_types, + filters, + time_filters, + series_row_selector, + append_mode, + filter_deleted, + merge_mode, + partition_expr_version, + } = self; + + ScanRequestFingerprint { + inner: Arc::new(SharedScanRequestFingerprint { + read_column_ids, + read_column_types, + filters, + }), + time_filters: (!time_filters.is_empty()).then(|| Arc::new(time_filters)), + series_row_selector, + append_mode, + filter_deleted, + merge_mode, + partition_expr_version, + } + } +} + +/// Non-copiable struct of the fingerprint. +#[derive(Debug, PartialEq, Eq, Hash)] +struct SharedScanRequestFingerprint { + /// Column ids of the projection. + read_column_ids: Vec, + /// Column types of the projection. + /// We keep this to ensure we won't reuse the fingerprint after a schema change. + read_column_types: Vec>, + /// Filters without the time index column and region partition exprs. + filters: Vec, +} + +impl ScanRequestFingerprint { + #[cfg(test)] + pub(crate) fn read_column_ids(&self) -> &[ColumnId] { + &self.inner.read_column_ids + } + + #[cfg(test)] + pub(crate) fn read_column_types(&self) -> &[Option] { + &self.inner.read_column_types + } + + #[cfg(test)] + pub(crate) fn filters(&self) -> &[String] { + &self.inner.filters + } + + #[cfg(test)] + pub(crate) fn time_filters(&self) -> &[String] { + self.time_filters + .as_deref() + .map(Vec::as_slice) + .unwrap_or(&[]) + } + + #[cfg(test)] + pub(crate) fn without_time_filters(&self) -> Self { + Self { + inner: Arc::clone(&self.inner), + time_filters: None, + series_row_selector: self.series_row_selector, + append_mode: self.append_mode, + filter_deleted: self.filter_deleted, + merge_mode: self.merge_mode, + partition_expr_version: self.partition_expr_version, + } + } + + pub(crate) fn estimated_size(&self) -> usize { + mem::size_of::() + + self.inner.read_column_ids.capacity() * mem::size_of::() + + self.inner.read_column_types.capacity() * mem::size_of::>() + + self.inner.filters.capacity() * mem::size_of::() + + self + .inner + .filters + .iter() + .map(|filter| filter.capacity()) + .sum::() + + self.time_filters.as_ref().map_or(0, |filters| { + mem::size_of::>() + + filters.capacity() * mem::size_of::() + + filters + .iter() + .map(|filter| filter.capacity()) + .sum::() + }) + } +} + +/// Cache key for range scan outputs. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct RangeScanCacheKey { + pub(crate) region_id: RegionId, + /// Sorted (file_id, row_group_index) pairs that uniquely identify the covered data. + pub(crate) row_groups: Vec<(FileId, i64)>, + pub(crate) scan: ScanRequestFingerprint, +} + +impl RangeScanCacheKey { + pub(crate) fn estimated_size(&self) -> usize { + mem::size_of::() + + self.row_groups.capacity() * mem::size_of::<(FileId, i64)>() + + self.scan.estimated_size() + } +} + +/// Cached result for one range scan. +pub(crate) struct RangeScanCacheValue { + pub(crate) batches: Vec, +} + +impl RangeScanCacheValue { + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn new(batches: Vec) -> Self { + Self { batches } + } + + pub(crate) fn estimated_size(&self) -> usize { + mem::size_of::() + + self.batches.capacity() * mem::size_of::() + + self + .batches + .iter() + .map(record_batch_estimated_size) + .sum::() + } +} + +#[cfg(test)] +mod tests { + use store_api::storage::TimeSeriesRowSelector; + + use super::*; + + #[test] + fn normalizes_and_clears_time_filters() { + let normalized = ScanRequestFingerprintBuilder { + read_column_ids: vec![1, 2], + read_column_types: vec![None, None], + filters: vec!["k0 = 'foo'".to_string()], + time_filters: vec![], + series_row_selector: None, + append_mode: false, + filter_deleted: true, + merge_mode: MergeMode::LastRow, + partition_expr_version: 0, + } + .build(); + + assert!(normalized.time_filters().is_empty()); + + let fingerprint = ScanRequestFingerprintBuilder { + read_column_ids: vec![1, 2], + read_column_types: vec![None, None], + filters: vec!["k0 = 'foo'".to_string()], + time_filters: vec!["ts >= 1000".to_string()], + series_row_selector: Some(TimeSeriesRowSelector::LastRow), + append_mode: false, + filter_deleted: true, + merge_mode: MergeMode::LastRow, + partition_expr_version: 7, + } + .build(); + + let reset = fingerprint.without_time_filters(); + + assert_eq!(reset.read_column_ids(), fingerprint.read_column_ids()); + assert_eq!(reset.read_column_types(), fingerprint.read_column_types()); + assert_eq!(reset.filters(), fingerprint.filters()); + assert!(reset.time_filters().is_empty()); + assert_eq!(reset.series_row_selector, fingerprint.series_row_selector); + assert_eq!(reset.append_mode, fingerprint.append_mode); + assert_eq!(reset.filter_deleted, fingerprint.filter_deleted); + assert_eq!(reset.merge_mode, fingerprint.merge_mode); + assert_eq!( + reset.partition_expr_version, + fingerprint.partition_expr_version + ); + } +} diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 5d934afd2d9..5cb2d75e25e 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -55,6 +55,7 @@ use crate::metrics::READ_SST_COUNT; use crate::read::compat::{self, CompatBatch, FlatCompatBatch, PrimaryKeyCompatBatch}; use crate::read::projection::ProjectionMapper; use crate::read::range::{FileRangeBuilder, MemRangeBuilder, RangeMeta, RowGroupIndex}; +use crate::read::range_cache::ScanRequestFingerprint; use crate::read::seq_scan::SeqScan; use crate::read::series_scan::SeriesScan; use crate::read::stream::ScanBatchStream; @@ -815,7 +816,7 @@ pub struct ScanInput { /// But this read columns might also include non-projected columns needed for filtering. pub(crate) read_column_ids: Vec, /// Time range filter for time index. - time_range: Option, + pub(crate) time_range: Option, /// Predicate to push down. pub(crate) predicate: PredicateGroup, /// Region partition expr applied at read time. @@ -1417,6 +1418,92 @@ fn pre_filter_mode(append_mode: bool, merge_mode: MergeMode) -> PreFilterMode { } } +/// Builds a [ScanRequestFingerprint] from a [ScanInput] if the scan is eligible +/// for partition range caching. +#[cfg_attr(not(test), allow(dead_code))] +pub(crate) fn build_scan_fingerprint(input: &ScanInput) -> Option { + let eligible = input.flat_format + && !input.compaction + && !input.files.is_empty() + && matches!(input.cache_strategy, CacheStrategy::EnableAll(_)); + + if !eligible { + return None; + } + + let metadata = input.region_metadata(); + let tag_names: HashSet<&str> = metadata + .column_metadatas + .iter() + .filter(|col| col.semantic_type == SemanticType::Tag) + .map(|col| col.column_schema.name.as_str()) + .collect(); + + let time_index_name = metadata.time_index_column().column_schema.name.clone(); + + let exprs = input + .predicate_group() + .predicate_without_region() + .map(|predicate| predicate.exprs()) + .unwrap_or_default(); + + let mut filters = Vec::new(); + let mut time_filters = Vec::new(); + let mut has_tag_filter = false; + let mut columns = HashSet::new(); + + for expr in exprs { + columns.clear(); + let is_time_only = match expr_to_columns(expr, &mut columns) { + Ok(()) if !columns.is_empty() => { + has_tag_filter |= columns + .iter() + .any(|col| tag_names.contains(col.name.as_str())); + columns.iter().all(|col| col.name == time_index_name) + } + _ => false, + }; + + if is_time_only { + time_filters.push(expr.to_string()); + } else { + filters.push(expr.to_string()); + } + } + + if !has_tag_filter { + // We only cache requests that have tag filters to avoid caching all series. + return None; + } + + // Ensure the filters are sorted for consistent fingerprinting. + filters.sort_unstable(); + time_filters.sort_unstable(); + + Some( + crate::read::range_cache::ScanRequestFingerprintBuilder { + read_column_ids: input.read_column_ids.clone(), + read_column_types: input + .read_column_ids + .iter() + .map(|id| { + metadata + .column_by_id(*id) + .map(|col| col.column_schema.data_type.clone()) + }) + .collect(), + filters, + time_filters, + series_row_selector: input.series_row_selector, + append_mode: input.append_mode, + filter_deleted: input.filter_deleted, + merge_mode: input.merge_mode, + partition_expr_version: metadata.partition_expr_version, + } + .build(), + ) +} + /// Context shared by different streams from a scanner. /// It contains the input and ranges to scan. pub struct StreamContext { @@ -1763,10 +1850,15 @@ mod tests { use datafusion::physical_plan::expressions::lit as physical_lit; use datafusion_expr::{col, lit}; - use store_api::storage::ScanRequest; + use datatypes::value::Value; + use partition::expr::col as partition_col; + use store_api::metadata::RegionMetadataBuilder; + use store_api::storage::{ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector}; use super::*; + use crate::cache::CacheManager; use crate::memtable::time_partition::TimePartitions; + use crate::read::range_cache::ScanRequestFingerprintBuilder; use crate::region::options::RegionOptions; use crate::region::version::VersionBuilder; use crate::sst::FormatType; @@ -1804,6 +1896,26 @@ mod tests { ) } + async fn new_scan_input(metadata: RegionMetadataRef, filters: Vec) -> ScanInput { + let env = SchedulerEnv::new().await; + let mapper = ProjectionMapper::new(&metadata, [0, 2, 3].into_iter(), true).unwrap(); + let predicate = PredicateGroup::new(metadata.as_ref(), &filters).unwrap(); + let file = FileHandle::new( + crate::sst::file::FileMeta::default(), + Arc::new(crate::sst::file_purger::NoopFilePurger), + ); + + ScanInput::new(env.access_layer.clone(), mapper) + .with_predicate(predicate) + .with_cache(CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .range_result_cache_size(1024) + .build(), + ))) + .with_flat_format(true) + .with_files(vec![file]) + } + #[tokio::test] async fn test_build_read_column_ids_includes_filters() { let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); @@ -1923,6 +2035,133 @@ mod tests { assert!(scan_region.use_flat_format()); } + #[tokio::test] + async fn test_build_scan_fingerprint_for_eligible_scan() { + let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); + let input = new_scan_input( + metadata.clone(), + vec![ + col("ts").gt_eq(lit(1000)), + col("k0").eq(lit("foo")), + col("v0").gt(lit(1)), + ], + ) + .await + .with_distribution(Some(TimeSeriesDistribution::PerSeries)) + .with_series_row_selector(Some(TimeSeriesRowSelector::LastRow)) + .with_merge_mode(MergeMode::LastNonNull) + .with_filter_deleted(false); + + let fingerprint = build_scan_fingerprint(&input).unwrap(); + + let expected = ScanRequestFingerprintBuilder { + read_column_ids: input.read_column_ids.clone(), + read_column_types: vec![ + metadata + .column_by_id(0) + .map(|col| col.column_schema.data_type.clone()), + metadata + .column_by_id(2) + .map(|col| col.column_schema.data_type.clone()), + metadata + .column_by_id(3) + .map(|col| col.column_schema.data_type.clone()), + ], + filters: vec![ + col("k0").eq(lit("foo")).to_string(), + col("v0").gt(lit(1)).to_string(), + ], + time_filters: vec![col("ts").gt_eq(lit(1000)).to_string()], + series_row_selector: Some(TimeSeriesRowSelector::LastRow), + append_mode: false, + filter_deleted: false, + merge_mode: MergeMode::LastNonNull, + partition_expr_version: 0, + } + .build(); + assert_eq!(expected, fingerprint); + } + + #[tokio::test] + async fn test_build_scan_fingerprint_requires_tag_filter() { + let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); + let input = new_scan_input( + metadata, + vec![col("ts").gt_eq(lit(1000)), col("v0").gt(lit(1))], + ) + .await; + + assert!(build_scan_fingerprint(&input).is_none()); + } + + #[tokio::test] + async fn test_build_scan_fingerprint_respects_scan_eligibility() { + let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); + let filters = vec![col("k0").eq(lit("foo"))]; + + let disabled = ScanInput::new( + SchedulerEnv::new().await.access_layer.clone(), + ProjectionMapper::new(&metadata, [0, 2, 3].into_iter(), true).unwrap(), + ) + .with_predicate(PredicateGroup::new(metadata.as_ref(), &filters).unwrap()) + .with_flat_format(true); + assert!(build_scan_fingerprint(&disabled).is_none()); + + let non_flat = new_scan_input(metadata.clone(), filters.clone()) + .await + .with_flat_format(false); + assert!(build_scan_fingerprint(&non_flat).is_none()); + + let compaction = new_scan_input(metadata.clone(), filters.clone()) + .await + .with_compaction(true); + assert!(build_scan_fingerprint(&compaction).is_none()); + + // No files to read. + let no_files = new_scan_input(metadata, filters).await.with_files(vec![]); + assert!(build_scan_fingerprint(&no_files).is_none()); + } + + #[tokio::test] + async fn test_build_scan_fingerprint_tracks_schema_and_partition_expr_changes() { + let base = metadata_with_primary_key(vec![0, 1], false); + let mut builder = RegionMetadataBuilder::from_existing(base); + let partition_expr = partition_col("k0") + .gt_eq(Value::String("foo".into())) + .as_json_str() + .unwrap(); + builder.partition_expr_json(Some(partition_expr)); + let metadata = Arc::new(builder.build_without_validation().unwrap()); + + let input = new_scan_input(metadata.clone(), vec![col("k0").eq(lit("foo"))]).await; + let fingerprint = build_scan_fingerprint(&input).unwrap(); + + let expected = ScanRequestFingerprintBuilder { + read_column_ids: input.read_column_ids.clone(), + read_column_types: vec![ + metadata + .column_by_id(0) + .map(|col| col.column_schema.data_type.clone()), + metadata + .column_by_id(2) + .map(|col| col.column_schema.data_type.clone()), + metadata + .column_by_id(3) + .map(|col| col.column_schema.data_type.clone()), + ], + filters: vec![col("k0").eq(lit("foo")).to_string()], + time_filters: vec![], + series_row_selector: None, + append_mode: false, + filter_deleted: true, + merge_mode: MergeMode::LastRow, + partition_expr_version: metadata.partition_expr_version, + } + .build(); + assert_eq!(expected, fingerprint); + assert_ne!(0, metadata.partition_expr_version); + } + #[test] fn test_update_dyn_filters_with_empty_base_predicates() { let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index c13b40d111a..a1b3b8f350a 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -39,7 +39,7 @@ use crate::error::{PartitionOutOfRangeSnafu, Result, TooManyFilesToReadSnafu, Un use crate::read::dedup::{DedupReader, LastNonNull, LastRow}; use crate::read::flat_dedup::{FlatDedupReader, FlatLastNonNull, FlatLastRow}; use crate::read::flat_merge::FlatMergeReader; -use crate::read::last_row::LastRowReader; +use crate::read::last_row::{FlatLastRowReader, LastRowReader}; use crate::read::merge::MergeReaderBuilder; use crate::read::pruner::{PartitionPruner, Pruner}; use crate::read::range::RangeMeta; @@ -128,28 +128,6 @@ impl SeqScan { Ok(Box::pin(futures::stream::iter(streams).flatten())) } - /// Builds a [BoxedBatchReader] from sequential scan for compaction. - /// - /// # Panics - /// Panics if the compaction flag is not set. - pub async fn build_reader_for_compaction(&self) -> Result { - assert!(self.stream_ctx.input.compaction); - - let metrics_set = ExecutionPlanMetricsSet::new(); - let part_metrics = self.new_partition_metrics(false, &metrics_set, 0); - debug_assert_eq!(1, self.properties.partitions.len()); - let partition_ranges = &self.properties.partitions[0]; - - let reader = Self::merge_all_ranges_for_compaction( - &self.stream_ctx, - partition_ranges, - &part_metrics, - self.pruner.clone(), - ) - .await?; - Ok(Box::new(reader)) - } - /// Builds a [BoxedRecordBatchStream] from sequential scan for flat format compaction. /// /// # Panics @@ -172,40 +150,6 @@ impl SeqScan { Ok(reader) } - /// Builds a merge reader that reads all ranges. - /// Callers MUST not split ranges before calling this method. - async fn merge_all_ranges_for_compaction( - stream_ctx: &Arc, - partition_ranges: &[PartitionRange], - part_metrics: &PartitionMetrics, - pruner: Arc, - ) -> Result { - pruner.add_partition_ranges(partition_ranges); - let partition_pruner = Arc::new(PartitionPruner::new(pruner, partition_ranges)); - - let mut sources = Vec::new(); - for part_range in partition_ranges { - build_sources( - stream_ctx, - part_range, - true, - part_metrics, - partition_pruner.clone(), - &mut sources, - None, - ) - .await?; - } - - common_telemetry::debug!( - "Build reader to read all parts, region_id: {}, num_part_ranges: {}, num_sources: {}", - stream_ctx.input.mapper.metadata().region_id, - partition_ranges.len(), - sources.len() - ); - Self::build_reader_from_sources(stream_ctx, sources, None, None).await - } - /// Builds a merge reader that reads all flat ranges. /// Callers MUST not split ranges before calling this method. async fn merge_all_flat_ranges_for_compaction( @@ -345,6 +289,13 @@ impl SeqScan { Box::pin(reader.into_stream()) as _ }; + let reader = match &stream_ctx.input.series_row_selector { + Some(TimeSeriesRowSelector::LastRow) => { + Box::pin(FlatLastRowReader::new(reader).into_stream()) as _ + } + None => reader, + }; + Ok(reader) } diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 014c50820fd..d089493f81f 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -1043,7 +1043,7 @@ async fn preload_parquet_meta_cache_for_files( let loader = MetadataLoader::new(object_store.clone(), &file_path, file_size); match loader.load(&mut cache_metrics).await { Ok(metadata) => { - cache_manager.put_parquet_meta_data(file_id, Arc::new(metadata)); + cache_manager.put_parquet_meta_data(file_id, Arc::new(metadata), None); loaded += 1; } Err(err) => { @@ -1153,6 +1153,8 @@ mod tests { use object_store::ObjectStore; use object_store::services::{Fs, Memory}; use parquet::arrow::ArrowWriter; + use parquet::file::metadata::KeyValue; + use parquet::file::properties::WriterProperties; use store_api::region_request::PathType; use store_api::storage::{FileId, RegionId}; @@ -1161,7 +1163,27 @@ mod tests { use crate::cache::file_cache::{FileType, IndexKey}; use crate::sst::file::{FileHandle, FileMeta}; use crate::sst::file_purger::NoopFilePurger; + use crate::sst::parquet::PARQUET_METADATA_KEY; use crate::test_util::TestEnv; + use crate::test_util::sst_util::sst_region_metadata; + + fn sst_parquet_bytes(batch: &RecordBatch) -> Vec { + let key_value_meta = KeyValue::new( + PARQUET_METADATA_KEY.to_string(), + sst_region_metadata().to_json().unwrap(), + ); + let props = WriterProperties::builder() + .set_key_value_metadata(Some(vec![key_value_meta])) + .build(); + + let mut parquet_bytes = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), Some(props)).unwrap(); + writer.write(batch).unwrap(); + writer.close().unwrap(); + + parquet_bytes + } #[tokio::test] async fn test_preload_parquet_meta_cache_uses_file_cache() { @@ -1183,10 +1205,7 @@ mod tests { let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef; let batch = RecordBatch::try_from_iter([("col", col)]).unwrap(); - let mut parquet_bytes = Vec::new(); - let mut writer = ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), None).unwrap(); - writer.write(&batch).unwrap(); - writer.close().unwrap(); + let parquet_bytes = sst_parquet_bytes(&batch); let file_size = parquet_bytes.len() as u64; let file_meta = FileMeta { @@ -1334,10 +1353,7 @@ mod tests { let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef; let batch = RecordBatch::try_from_iter([("col", col)]).unwrap(); - let mut parquet_bytes = Vec::new(); - let mut writer = ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), None).unwrap(); - writer.write(&batch).unwrap(); - writer.close().unwrap(); + let parquet_bytes = sst_parquet_bytes(&batch); // file_size is 0 when it's missing/defaulted in manifests; MetadataLoader::load will stat // the local filesystem to retrieve it. diff --git a/src/mito2/src/region/options.rs b/src/mito2/src/region/options.rs index 0fe0a8f12a3..fcf68a92166 100644 --- a/src/mito2/src/region/options.rs +++ b/src/mito2/src/region/options.rs @@ -50,7 +50,7 @@ pub(crate) fn parse_wal_options( } /// Mode to handle duplicate rows while merging. -#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, EnumString)] +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, EnumString)] #[serde(rename_all = "snake_case")] #[strum(serialize_all = "snake_case")] pub enum MergeMode { diff --git a/src/mito2/src/sst.rs b/src/mito2/src/sst.rs index 78e4c563b1a..94bc1feea83 100644 --- a/src/mito2/src/sst.rs +++ b/src/mito2/src/sst.rs @@ -31,7 +31,6 @@ use store_api::storage::consts::{ OP_TYPE_COLUMN_NAME, PRIMARY_KEY_COLUMN_NAME, SEQUENCE_COLUMN_NAME, }; -use crate::read::Batch; use crate::sst::parquet::flat_format::time_index_column_index; pub mod file; @@ -260,33 +259,6 @@ pub(crate) struct SeriesEstimator { } impl SeriesEstimator { - /// Updates the estimator with a new Batch. - /// - /// Since each Batch contains only one series, this increments the series count - /// and updates the last timestamp. - pub(crate) fn update(&mut self, batch: &Batch) { - let Some(last_ts) = batch.last_timestamp() else { - return; - }; - - // Checks if there's a boundary between the last batch and this batch - if let Some(prev_last_ts) = self.last_timestamp { - // If the first timestamp of this batch is less than the last timestamp - // we've seen, it indicates a new series - if let Some(first_ts) = batch.first_timestamp() - && first_ts.value() <= prev_last_ts - { - self.series_count += 1; - } - } else { - // First batch, counts as first series - self.series_count = 1; - } - - // Updates the last timestamp - self.last_timestamp = Some(last_ts.value()); - } - /// Updates the estimator with a new record batch in flat format. /// /// This method examines the time index column to detect series boundaries. @@ -340,43 +312,14 @@ impl SeriesEstimator { mod tests { use std::sync::Arc; - use api::v1::OpType; use datatypes::arrow::array::{ - BinaryArray, DictionaryArray, TimestampMillisecondArray, UInt8Array, UInt8Builder, - UInt32Array, UInt64Array, + BinaryArray, DictionaryArray, TimestampMillisecondArray, UInt8Array, UInt32Array, + UInt64Array, }; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema, TimeUnit}; use datatypes::arrow::record_batch::RecordBatch; use super::*; - use crate::read::{Batch, BatchBuilder}; - - fn new_batch( - primary_key: &[u8], - timestamps: &[i64], - sequences: &[u64], - op_types: &[OpType], - ) -> Batch { - let timestamps = Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())); - let sequences = Arc::new(UInt64Array::from(sequences.to_vec())); - let mut op_type_builder = UInt8Builder::with_capacity(op_types.len()); - for op_type in op_types { - op_type_builder.append_value(*op_type as u8); - } - let op_types = Arc::new(UInt8Array::from( - op_types.iter().map(|op| *op as u8).collect::>(), - )); - - let mut builder = BatchBuilder::new(primary_key.to_vec()); - builder - .timestamps_array(timestamps) - .unwrap() - .sequences_array(sequences) - .unwrap() - .op_types_array(op_types) - .unwrap(); - builder.build().unwrap() - } fn new_flat_record_batch(timestamps: &[i64]) -> RecordBatch { // Flat format has: [fields..., time_index, __primary_key, __sequence, __op_type] @@ -411,128 +354,6 @@ mod tests { RecordBatch::try_new(schema, vec![time_array, pk_array, seq_array, op_array]).unwrap() } - #[test] - fn test_series_estimator_empty_batch() { - let mut estimator = SeriesEstimator::default(); - let batch = new_batch(b"test", &[], &[], &[]); - estimator.update(&batch); - assert_eq!(0, estimator.finish()); - } - - #[test] - fn test_series_estimator_single_batch() { - let mut estimator = SeriesEstimator::default(); - let batch = new_batch( - b"test", - &[1, 2, 3], - &[1, 2, 3], - &[OpType::Put, OpType::Put, OpType::Put], - ); - estimator.update(&batch); - assert_eq!(1, estimator.finish()); - } - - #[test] - fn test_series_estimator_multiple_batches_same_series() { - let mut estimator = SeriesEstimator::default(); - - // First batch with timestamps 1, 2, 3 - let batch1 = new_batch( - b"test", - &[1, 2, 3], - &[1, 2, 3], - &[OpType::Put, OpType::Put, OpType::Put], - ); - estimator.update(&batch1); - - // Second batch with timestamps 4, 5, 6 (continuation) - let batch2 = new_batch( - b"test", - &[4, 5, 6], - &[4, 5, 6], - &[OpType::Put, OpType::Put, OpType::Put], - ); - estimator.update(&batch2); - - assert_eq!(1, estimator.finish()); - } - - #[test] - fn test_series_estimator_new_series_detected() { - let mut estimator = SeriesEstimator::default(); - - // First batch with timestamps 1, 2, 3 - let batch1 = new_batch( - b"pk0", - &[1, 2, 3], - &[1, 2, 3], - &[OpType::Put, OpType::Put, OpType::Put], - ); - estimator.update(&batch1); - - // Second batch with timestamps 2, 3, 4 (timestamp goes back, new series) - let batch2 = new_batch( - b"pk1", - &[2, 3, 4], - &[4, 5, 6], - &[OpType::Put, OpType::Put, OpType::Put], - ); - estimator.update(&batch2); - - assert_eq!(2, estimator.finish()); - } - - #[test] - fn test_series_estimator_equal_timestamp_boundary() { - let mut estimator = SeriesEstimator::default(); - - // First batch ending at timestamp 5 - let batch1 = new_batch( - b"test", - &[1, 2, 5], - &[1, 2, 3], - &[OpType::Put, OpType::Put, OpType::Put], - ); - estimator.update(&batch1); - - // Second batch starting at timestamp 5 (equal, indicates new series) - let batch2 = new_batch( - b"test", - &[5, 6, 7], - &[4, 5, 6], - &[OpType::Put, OpType::Put, OpType::Put], - ); - estimator.update(&batch2); - - assert_eq!(2, estimator.finish()); - } - - #[test] - fn test_series_estimator_finish_resets_state() { - let mut estimator = SeriesEstimator::default(); - - let batch1 = new_batch( - b"test", - &[1, 2, 3], - &[1, 2, 3], - &[OpType::Put, OpType::Put, OpType::Put], - ); - estimator.update(&batch1); - - assert_eq!(1, estimator.finish()); - - // After finish, state should be reset - let batch2 = new_batch( - b"test", - &[4, 5, 6], - &[4, 5, 6], - &[OpType::Put, OpType::Put, OpType::Put], - ); - estimator.update(&batch2); - - assert_eq!(1, estimator.finish()); - } - #[test] fn test_series_estimator_flat_empty_batch() { let mut estimator = SeriesEstimator::default(); diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 0df3229e9cc..88aebfc0013 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -58,7 +58,7 @@ use crate::error::{ }; use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; use crate::metrics::INDEX_CREATE_MEMORY_USAGE; -use crate::read::{Batch, BatchReader}; +use crate::read::Batch; use crate::region::options::IndexOptions; use crate::region::version::VersionControlRef; use crate::region::{ManifestContextRef, RegionLeaderState}; @@ -802,9 +802,9 @@ impl IndexBuildTask { if let Some(mut parquet_reader) = parquet_reader { // TODO(SNC123): optimize index batch loop { - match parquet_reader.next_batch().await { - Ok(Some(mut batch)) => { - indexer.update(&mut batch).await; + match parquet_reader.next_record_batch().await { + Ok(Some(batch)) => { + indexer.update_flat(&batch).await; } Ok(None) => break, Err(e) => { @@ -1227,7 +1227,9 @@ mod tests { use crate::sst::parquet::WriteOptions; use crate::test_util::memtable_util::EmptyMemtableBuilder; use crate::test_util::scheduler_util::SchedulerEnv; - use crate::test_util::sst_util::{new_batch_by_range, new_source, sst_region_metadata}; + use crate::test_util::sst_util::{ + new_flat_source_from_record_batches, new_record_batch_by_range, sst_region_metadata, + }; struct MetaConfig { with_inverted: bool, @@ -1358,19 +1360,20 @@ mod tests { env: &SchedulerEnv, build_mode: IndexBuildMode, ) -> SstInfo { - let source = new_source(&[ - new_batch_by_range(&["a", "d"], 0, 60), - new_batch_by_range(&["b", "f"], 0, 40), - new_batch_by_range(&["b", "h"], 100, 200), + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_by_range(&["a", "d"], 0, 60), + new_record_batch_by_range(&["b", "f"], 0, 40), + new_record_batch_by_range(&["b", "h"], 100, 200), ]); let mut index_config = MitoConfig::default().index; index_config.build_mode = build_mode; let write_request = SstWriteRequest { op_type: OperationType::Flush, metadata: metadata.clone(), - source: either::Left(source), + source, storage: None, max_sequence: None, + sst_write_format: Default::default(), cache_manager: Default::default(), index_options: IndexOptions::default(), index_config, diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index aa98b69176c..26bed76fd60 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -110,6 +110,7 @@ mod tests { TimestampMillisecondArray, UInt8Array, UInt64Array, }; use datatypes::arrow::datatypes::{DataType, Field, Schema, UInt32Type}; + use datatypes::arrow::util::pretty::pretty_format_batches; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{FulltextAnalyzer, FulltextBackend, FulltextOptions}; use object_store::ObjectStore; @@ -129,7 +130,7 @@ mod tests { use crate::cache::test_util::assert_parquet_metadata_equal; use crate::cache::{CacheManager, CacheStrategy, PageKey}; use crate::config::IndexConfig; - use crate::read::{BatchBuilder, BatchReader, FlatSource}; + use crate::read::FlatSource; use crate::region::options::{IndexOptions, InvertedIndexOptions}; use crate::sst::file::{FileHandle, FileMeta, RegionFileId, RegionIndexId}; use crate::sst::file_purger::NoopFilePurger; @@ -137,19 +138,19 @@ mod tests { use crate::sst::index::fulltext_index::applier::builder::FulltextIndexApplierBuilder; use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder; use crate::sst::index::{IndexBuildType, Indexer, IndexerBuilder, IndexerBuilderImpl}; - use crate::sst::parquet::format::PrimaryKeyWriteFormat; + use crate::sst::parquet::flat_format::FlatWriteFormat; use crate::sst::parquet::reader::{ParquetReader, ParquetReaderBuilder, ReaderMetrics}; use crate::sst::parquet::writer::ParquetWriter; use crate::sst::{ DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, location, to_flat_sst_arrow_schema, }; + use crate::test_util::TestEnv; use crate::test_util::sst_util::{ - build_test_binary_test_region_metadata, new_batch_by_range, new_batch_with_binary, - new_batch_with_custom_sequence, new_primary_key, new_source, new_sparse_primary_key, - sst_file_handle, sst_file_handle_with_file_id, sst_region_metadata, + build_test_binary_test_region_metadata, new_flat_source_from_record_batches, + new_primary_key, new_record_batch_by_range, new_record_batch_with_custom_sequence, + new_sparse_primary_key, sst_file_handle, sst_file_handle_with_file_id, sst_region_metadata, sst_region_metadata_with_encoding, }; - use crate::test_util::{TestEnv, check_reader_result}; const FILE_DIR: &str = "/"; const REGION_ID: RegionId = RegionId::new(0, 0); @@ -191,10 +192,10 @@ mod tests { region_file_id: handle.file_id(), }; let metadata = Arc::new(sst_region_metadata()); - let source = new_source(&[ - new_batch_by_range(&["a", "d"], 0, 60), - new_batch_by_range(&["b", "f"], 0, 40), - new_batch_by_range(&["b", "h"], 100, 200), + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_by_range(&["a", "d"], 0, 60), + new_record_batch_by_range(&["b", "f"], 0, 40), + new_record_batch_by_range(&["b", "h"], 100, 200), ]); // Use a small row group size for test. let write_opts = WriteOptions { @@ -214,7 +215,7 @@ mod tests { .await; let info = writer - .write_all(source, None, &write_opts) + .write_all_flat_as_primary_key(source, None, &write_opts) .await .unwrap() .remove(0); @@ -235,14 +236,14 @@ mod tests { object_store, ); let mut reader = builder.build().await.unwrap().unwrap(); - check_reader_result( + check_record_batch_reader_result( &mut reader, &[ - new_batch_by_range(&["a", "d"], 0, 50), - new_batch_by_range(&["a", "d"], 50, 60), - new_batch_by_range(&["b", "f"], 0, 40), - new_batch_by_range(&["b", "h"], 100, 150), - new_batch_by_range(&["b", "h"], 150, 200), + new_record_batch_by_range(&["a", "d"], 0, 50), + new_record_batch_by_range(&["a", "d"], 50, 60), + new_record_batch_by_range(&["b", "f"], 0, 40), + new_record_batch_by_range(&["b", "h"], 100, 150), + new_record_batch_by_range(&["b", "h"], 150, 200), ], ) .await; @@ -254,10 +255,10 @@ mod tests { let object_store = env.init_object_store_manager(); let handle = sst_file_handle(0, 1000); let metadata = Arc::new(sst_region_metadata()); - let source = new_source(&[ - new_batch_by_range(&["a", "d"], 0, 60), - new_batch_by_range(&["b", "f"], 0, 40), - new_batch_by_range(&["b", "h"], 100, 200), + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_by_range(&["a", "d"], 0, 60), + new_record_batch_by_range(&["b", "f"], 0, 40), + new_record_batch_by_range(&["b", "h"], 100, 200), ]); // Use a small row group size for test. let write_opts = WriteOptions { @@ -279,7 +280,7 @@ mod tests { .await; let sst_info = writer - .write_all(source, None, &write_opts) + .write_all_flat_as_primary_key(source, None, &write_opts) .await .unwrap() .remove(0); @@ -299,14 +300,14 @@ mod tests { .cache(cache.clone()); for _ in 0..3 { let mut reader = builder.build().await.unwrap().unwrap(); - check_reader_result( + check_record_batch_reader_result( &mut reader, &[ - new_batch_by_range(&["a", "d"], 0, 50), - new_batch_by_range(&["a", "d"], 50, 60), - new_batch_by_range(&["b", "f"], 0, 40), - new_batch_by_range(&["b", "h"], 100, 150), - new_batch_by_range(&["b", "h"], 150, 200), + new_record_batch_by_range(&["a", "d"], 0, 50), + new_record_batch_by_range(&["a", "d"], 50, 60), + new_record_batch_by_range(&["b", "f"], 0, 40), + new_record_batch_by_range(&["b", "h"], 100, 150), + new_record_batch_by_range(&["b", "h"], 150, 200), ], ) .await; @@ -340,10 +341,10 @@ mod tests { let object_store = env.init_object_store_manager(); let handle = sst_file_handle(0, 1000); let metadata = Arc::new(sst_region_metadata()); - let source = new_source(&[ - new_batch_by_range(&["a", "d"], 0, 60), - new_batch_by_range(&["b", "f"], 0, 40), - new_batch_by_range(&["b", "h"], 100, 200), + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_by_range(&["a", "d"], 0, 60), + new_record_batch_by_range(&["b", "f"], 0, 40), + new_record_batch_by_range(&["b", "h"], 100, 200), ]); let write_opts = WriteOptions { row_group_size: 50, @@ -366,7 +367,7 @@ mod tests { .await; let sst_info = writer - .write_all(source, None, &write_opts) + .write_all_flat_as_primary_key(source, None, &write_opts) .await .unwrap() .remove(0); @@ -382,8 +383,12 @@ mod tests { .page_index_policy(PageIndexPolicy::Optional); let reader = builder.build().await.unwrap().unwrap(); let reader_metadata = reader.parquet_metadata(); + let cached_writer_metadata = + crate::cache::CachedSstMeta::try_new("test.sst", Arc::unwrap_or_clone(writer_metadata)) + .unwrap() + .parquet_metadata(); - assert_parquet_metadata_equal(writer_metadata, reader_metadata); + assert_parquet_metadata_equal(cached_writer_metadata, reader_metadata); } #[tokio::test] @@ -392,10 +397,10 @@ mod tests { let object_store = env.init_object_store_manager(); let handle = sst_file_handle(0, 1000); let metadata = Arc::new(sst_region_metadata()); - let source = new_source(&[ - new_batch_by_range(&["a", "d"], 0, 60), - new_batch_by_range(&["b", "f"], 0, 40), - new_batch_by_range(&["b", "h"], 100, 200), + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_by_range(&["a", "d"], 0, 60), + new_record_batch_by_range(&["b", "f"], 0, 40), + new_record_batch_by_range(&["b", "h"], 100, 200), ]); // Use a small row group size for test. let write_opts = WriteOptions { @@ -416,7 +421,7 @@ mod tests { ) .await; writer - .write_all(source, None, &write_opts) + .write_all_flat_as_primary_key(source, None, &write_opts) .await .unwrap() .remove(0); @@ -436,11 +441,11 @@ mod tests { ) .predicate(predicate); let mut reader = builder.build().await.unwrap().unwrap(); - check_reader_result( + check_record_batch_reader_result( &mut reader, &[ - new_batch_by_range(&["a", "d"], 0, 50), - new_batch_by_range(&["a", "d"], 50, 60), + new_record_batch_by_range(&["a", "d"], 0, 50), + new_record_batch_by_range(&["a", "d"], 50, 60), ], ) .await; @@ -452,10 +457,10 @@ mod tests { let object_store = env.init_object_store_manager(); let handle = sst_file_handle(0, 1000); let metadata = Arc::new(sst_region_metadata()); - let source = new_source(&[ - new_batch_by_range(&["a", "z"], 0, 0), - new_batch_by_range(&["a", "z"], 100, 100), - new_batch_by_range(&["a", "z"], 200, 230), + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_by_range(&["a", "z"], 0, 0), + new_record_batch_by_range(&["a", "z"], 100, 100), + new_record_batch_by_range(&["a", "z"], 200, 230), ]); // Use a small row group size for test. let write_opts = WriteOptions { @@ -476,7 +481,7 @@ mod tests { ) .await; writer - .write_all(source, None, &write_opts) + .write_all_flat_as_primary_key(source, None, &write_opts) .await .unwrap() .remove(0); @@ -488,7 +493,11 @@ mod tests { object_store, ); let mut reader = builder.build().await.unwrap().unwrap(); - check_reader_result(&mut reader, &[new_batch_by_range(&["a", "z"], 200, 230)]).await; + check_record_batch_reader_result( + &mut reader, + &[new_record_batch_by_range(&["a", "z"], 200, 230)], + ) + .await; } #[tokio::test] @@ -497,10 +506,10 @@ mod tests { let object_store = env.init_object_store_manager(); let handle = sst_file_handle(0, 1000); let metadata = Arc::new(sst_region_metadata()); - let source = new_source(&[ - new_batch_by_range(&["a", "d"], 0, 60), - new_batch_by_range(&["b", "f"], 0, 40), - new_batch_by_range(&["b", "h"], 100, 200), + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_by_range(&["a", "d"], 0, 60), + new_record_batch_by_range(&["b", "f"], 0, 40), + new_record_batch_by_range(&["b", "h"], 100, 200), ]); // Use a small row group size for test. let write_opts = WriteOptions { @@ -522,7 +531,7 @@ mod tests { .await; writer - .write_all(source, None, &write_opts) + .write_all_flat_as_primary_key(source, None, &write_opts) .await .unwrap() .remove(0); @@ -542,7 +551,11 @@ mod tests { ) .predicate(predicate); let mut reader = builder.build().await.unwrap().unwrap(); - check_reader_result(&mut reader, &[new_batch_by_range(&["b", "h"], 150, 200)]).await; + check_record_batch_reader_result( + &mut reader, + &[new_record_batch_by_range(&["b", "h"], 150, 200)], + ) + .await; } #[tokio::test] @@ -569,7 +582,7 @@ mod tests { let writer_props = props_builder.build(); - let write_format = PrimaryKeyWriteFormat::new(metadata); + let write_format = FlatWriteFormat::new(metadata, &FlatSchemaOptions::default()); let fields: Vec<_> = write_format .arrow_schema() .fields() @@ -603,9 +616,8 @@ mod tests { ) .unwrap(); - let batch = new_batch_with_binary(&["a"], 0, 60); - let arrow_batch = write_format.convert_batch(&batch).unwrap(); - let arrays: Vec<_> = arrow_batch + let batch = new_record_batch_with_binary(&["a"], 0, 60); + let arrays: Vec<_> = batch .columns() .iter() .map(|array| { @@ -629,11 +641,11 @@ mod tests { object_store, ); let mut reader = builder.build().await.unwrap().unwrap(); - check_reader_result( + check_record_batch_reader_result( &mut reader, &[ - new_batch_with_binary(&["a"], 0, 50), - new_batch_with_binary(&["a"], 50, 60), + new_record_batch_with_binary(&["a"], 0, 50), + new_record_batch_with_binary(&["a"], 50, 60), ], ) .await; @@ -646,17 +658,17 @@ mod tests { let mut env = TestEnv::new().await; let object_store = env.init_object_store_manager(); let metadata = Arc::new(sst_region_metadata()); - let batches = &[ - new_batch_by_range(&["a", "d"], 0, 1000), - new_batch_by_range(&["b", "f"], 0, 1000), - new_batch_by_range(&["c", "g"], 0, 1000), - new_batch_by_range(&["b", "h"], 100, 200), - new_batch_by_range(&["b", "h"], 200, 300), - new_batch_by_range(&["b", "h"], 300, 1000), + let batches = vec![ + new_record_batch_by_range(&["a", "d"], 0, 1000), + new_record_batch_by_range(&["b", "f"], 0, 1000), + new_record_batch_by_range(&["c", "g"], 0, 1000), + new_record_batch_by_range(&["b", "h"], 100, 200), + new_record_batch_by_range(&["b", "h"], 200, 300), + new_record_batch_by_range(&["b", "h"], 300, 1000), ]; let total_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum(); - let source = new_source(batches); + let source = new_flat_source_from_record_batches(batches); let write_opts = WriteOptions { row_group_size: 50, max_file_size: Some(1024 * 16), @@ -678,7 +690,10 @@ mod tests { ) .await; - let files = writer.write_all(source, None, &write_opts).await.unwrap(); + let files = writer + .write_all_flat_as_primary_key(source, None, &write_opts) + .await + .unwrap(); assert_eq!(2, files.len()); let mut rows_read = 0; @@ -695,7 +710,7 @@ mod tests { object_store.clone(), ); let mut reader = builder.build().await.unwrap().unwrap(); - while let Some(batch) = reader.next_batch().await.unwrap() { + while let Some(batch) = reader.next_record_batch().await.unwrap() { rows_read += batch.num_rows(); } } @@ -710,12 +725,12 @@ mod tests { let metadata = Arc::new(sst_region_metadata()); let row_group_size = 50; - let source = new_source(&[ - new_batch_by_range(&["a", "d"], 0, 20), - new_batch_by_range(&["b", "d"], 0, 20), - new_batch_by_range(&["c", "d"], 0, 20), - new_batch_by_range(&["c", "f"], 0, 40), - new_batch_by_range(&["c", "h"], 100, 200), + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_by_range(&["a", "d"], 0, 20), + new_record_batch_by_range(&["b", "d"], 0, 20), + new_record_batch_by_range(&["c", "d"], 0, 20), + new_record_batch_by_range(&["c", "f"], 0, 40), + new_record_batch_by_range(&["c", "h"], 100, 200), ]); // Use a small row group size for test. let write_opts = WriteOptions { @@ -760,7 +775,7 @@ mod tests { .await; let info = writer - .write_all(source, None, &write_opts) + .write_all_flat_as_primary_key(source, None, &write_opts) .await .unwrap() .remove(0); @@ -877,6 +892,7 @@ mod tests { handle.clone(), object_store.clone(), ) + .flat_format(true) .predicate(Some(Predicate::new(preds))) .inverted_index_appliers([inverted_index_applier.clone(), None]) .bloom_filter_index_appliers([bloom_filter_applier.clone(), None]) @@ -891,7 +907,11 @@ mod tests { let mut reader = ParquetReader::new(Arc::new(context), selection) .await .unwrap(); - check_reader_result(&mut reader, &[new_batch_by_range(&["b", "d"], 0, 20)]).await; + check_record_batch_reader_result( + &mut reader, + &[new_record_batch_by_range(&["b", "d"], 0, 20)], + ) + .await; assert_eq!(metrics.filter_metrics.rg_total, 4); assert_eq!(metrics.filter_metrics.rg_minmax_filtered, 3); @@ -937,6 +957,7 @@ mod tests { handle.clone(), object_store.clone(), ) + .flat_format(true) .predicate(Some(Predicate::new(preds))) .inverted_index_appliers([inverted_index_applier.clone(), None]) .bloom_filter_index_appliers([bloom_filter_applier.clone(), None]) @@ -991,6 +1012,7 @@ mod tests { handle.clone(), object_store.clone(), ) + .flat_format(true) .predicate(Some(Predicate::new(preds))) .inverted_index_appliers([inverted_index_applier.clone(), None]) .bloom_filter_index_appliers([bloom_filter_applier.clone(), None]) @@ -1005,13 +1027,13 @@ mod tests { let mut reader = ParquetReader::new(Arc::new(context), selection) .await .unwrap(); - check_reader_result( + check_record_batch_reader_result( &mut reader, &[ - new_batch_by_range(&["a", "d"], 0, 20), - new_batch_by_range(&["b", "d"], 0, 20), - new_batch_by_range(&["c", "d"], 0, 10), - new_batch_by_range(&["c", "d"], 10, 20), + new_record_batch_by_range(&["a", "d"], 0, 20), + new_record_batch_by_range(&["b", "d"], 0, 20), + new_record_batch_by_range(&["c", "d"], 0, 10), + new_record_batch_by_range(&["c", "d"], 10, 20), ], ) .await; @@ -1032,37 +1054,32 @@ mod tests { assert!(cached.contains_row_group(3)); } - /// Creates a flat format RecordBatch for testing. - /// Similar to `new_batch_by_range` but returns a RecordBatch in flat format. - fn new_record_batch_by_range(tags: &[&str], start: usize, end: usize) -> RecordBatch { + fn new_record_batch_with_binary(tags: &[&str], start: usize, end: usize) -> RecordBatch { assert!(end >= start); - let metadata = Arc::new(sst_region_metadata()); + let metadata = build_test_binary_test_region_metadata(); let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default()); let num_rows = end - start; let mut columns = Vec::new(); - // Add primary key columns (tag_0, tag_1) as dictionary arrays let mut tag_0_builder = StringDictionaryBuilder::::new(); - let mut tag_1_builder = StringDictionaryBuilder::::new(); - for _ in 0..num_rows { tag_0_builder.append_value(tags[0]); - tag_1_builder.append_value(tags[1]); } - columns.push(Arc::new(tag_0_builder.finish()) as ArrayRef); - columns.push(Arc::new(tag_1_builder.finish()) as ArrayRef); - // Add field column (field_0) - let field_values: Vec = (start..end).map(|v| v as u64).collect(); - columns.push(Arc::new(UInt64Array::from(field_values))); + let values = (0..num_rows) + .map(|_| "some data".as_bytes()) + .collect::>(); + columns.push( + Arc::new(datatypes::arrow::array::BinaryArray::from_iter_values( + values, + )) as ArrayRef, + ); - // Add time index column (ts) let timestamps: Vec = (start..end).map(|v| v as i64).collect(); columns.push(Arc::new(TimestampMillisecondArray::from(timestamps))); - // Add encoded primary key column let pk = new_primary_key(tags); let mut pk_builder = BinaryDictionaryBuilder::::new(); for _ in 0..num_rows { @@ -1070,10 +1087,7 @@ mod tests { } columns.push(Arc::new(pk_builder.finish())); - // Add sequence column columns.push(Arc::new(UInt64Array::from_value(1000, num_rows))); - - // Add op_type column columns.push(Arc::new(UInt8Array::from_value( OpType::Put as u8, num_rows, @@ -1082,9 +1096,19 @@ mod tests { RecordBatch::try_new(flat_schema, columns).unwrap() } - /// Creates a FlatSource from flat format RecordBatches. - fn new_flat_source_from_record_batches(batches: Vec) -> FlatSource { - FlatSource::Iter(Box::new(batches.into_iter().map(Ok))) + async fn check_record_batch_reader_result( + reader: &mut ParquetReader, + expected: &[RecordBatch], + ) { + let mut actual = Vec::new(); + while let Some(batch) = reader.next_record_batch().await.unwrap() { + actual.push(batch); + } + assert_eq!( + pretty_format_batches(expected).unwrap().to_string(), + pretty_format_batches(&actual).unwrap().to_string() + ); + assert!(reader.next_record_batch().await.unwrap().is_none()); } /// Creates a flat format RecordBatch for testing with sparse primary key encoding. @@ -1333,10 +1357,11 @@ mod tests { }; let metadata = Arc::new(sst_region_metadata()); - // Create batches with sequence 0 to trigger override functionality - let batch1 = new_batch_with_custom_sequence(&["a", "d"], 0, 60, 0); - let batch2 = new_batch_with_custom_sequence(&["b", "f"], 0, 40, 0); - let source = new_source(&[batch1, batch2]); + // Create batches with sequence 0 to trigger override functionality. + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_with_custom_sequence(&["a", "d"], 0, 60, 0), + new_record_batch_with_custom_sequence(&["b", "f"], 0, 40, 0), + ]); let write_opts = WriteOptions { row_group_size: 50, @@ -1355,7 +1380,7 @@ mod tests { .await; writer - .write_all(source, None, &write_opts) + .write_all_flat_as_primary_key(source, None, &write_opts) .await .unwrap() .remove(0); @@ -1369,7 +1394,7 @@ mod tests { ); let mut reader = builder.build().await.unwrap().unwrap(); let mut normal_batches = Vec::new(); - while let Some(batch) = reader.next_batch().await.unwrap() { + while let Some(batch) = reader.next_record_batch().await.unwrap() { normal_batches.push(batch); } @@ -1391,22 +1416,19 @@ mod tests { ); let mut reader = builder.build().await.unwrap().unwrap(); let mut override_batches = Vec::new(); - while let Some(batch) = reader.next_batch().await.unwrap() { + while let Some(batch) = reader.next_record_batch().await.unwrap() { override_batches.push(batch); } // Compare the results assert_eq!(normal_batches.len(), override_batches.len()); for (normal, override_batch) in normal_batches.into_iter().zip(override_batches.iter()) { - // Create expected batch with override sequence let expected_batch = { - let num_rows = normal.num_rows(); - let mut builder = BatchBuilder::from(normal); - builder - .sequences_array(Arc::new(UInt64Array::from_value(custom_sequence, num_rows))) - .unwrap(); - - builder.build().unwrap() + let mut columns = normal.columns().to_vec(); + let num_cols = columns.len(); + columns[num_cols - 2] = + Arc::new(UInt64Array::from_value(custom_sequence, normal.num_rows())); + RecordBatch::try_new(normal.schema(), columns).unwrap() }; // Override batch should match expected batch diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs index ad9ddf0280a..5281f96d8d4 100644 --- a/src/mito2/src/sst/parquet/flat_format.rs +++ b/src/mito2/src/sst/parquet/flat_format.rs @@ -52,8 +52,8 @@ use crate::error::{ NewRecordBatchSnafu, Result, }; use crate::sst::parquet::format::{ - FormatProjection, INTERNAL_COLUMN_NUM, PrimaryKeyArray, PrimaryKeyReadFormat, ReadFormat, - StatValues, + FIXED_POS_COLUMN_NUM, FormatProjection, INTERNAL_COLUMN_NUM, PrimaryKeyArray, + PrimaryKeyReadFormat, ReadFormat, StatValues, }; use crate::sst::{ FlatSchemaOptions, flat_sst_arrow_schema_column_num, tag_maybe_to_dictionary_field, @@ -127,6 +127,21 @@ pub(crate) fn op_type_column_index(num_columns: usize) -> usize { num_columns - 1 } +/// Returns the start index of field columns in a flat batch. +/// +/// `num_columns` is the total number of columns in the flat batch schema, +/// including tag columns (if present), field columns, and fixed position columns +/// (time index, primary key, sequence, op type). +/// +/// For Dense encoding (raw PK columns included): field_column_start = primary_key.len() +/// For Sparse encoding (no raw PK columns): field_column_start = 0 +pub(crate) fn field_column_start(metadata: &RegionMetadata, num_columns: usize) -> usize { + // Calculates field column start: total columns - fixed columns - field columns + // Field column count = total metadata columns - time index column - primary key columns + let field_column_count = metadata.column_metadatas.len() - 1 - metadata.primary_key.len(); + num_columns - FIXED_POS_COLUMN_NUM - field_column_count +} + // TODO(yingwen): Add an option to skip reading internal columns if the region is // append only and doesn't use sparse encoding (We need to check the table id under // sparse encoding). @@ -774,3 +789,89 @@ impl FlatReadFormat { .unwrap() } } + +#[cfg(test)] +mod tests { + use api::v1::SemanticType; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::ColumnSchema; + use store_api::codec::PrimaryKeyEncoding; + use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; + use store_api::storage::RegionId; + + use super::field_column_start; + use crate::sst::{FlatSchemaOptions, flat_sst_arrow_schema_column_num}; + + /// Builds a `RegionMetadata` with the given number of tags and fields. + fn build_metadata( + num_tags: usize, + num_fields: usize, + encoding: PrimaryKeyEncoding, + ) -> RegionMetadata { + let mut builder = RegionMetadataBuilder::new(RegionId::new(0, 0)); + let mut col_id = 0u32; + + for i in 0..num_tags { + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + format!("tag_{i}"), + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id: col_id, + }); + col_id += 1; + } + + for i in 0..num_fields { + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + format!("field_{i}"), + ConcreteDataType::uint64_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: col_id, + }); + col_id += 1; + } + + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: col_id, + }); + + let primary_key: Vec = (0..num_tags as u32).collect(); + builder.primary_key(primary_key); + builder.primary_key_encoding(encoding); + builder.build().unwrap() + } + + #[test] + fn test_field_column_start() { + // (num_tags, num_fields, encoding, expected) + let cases = [ + (1, 1, PrimaryKeyEncoding::Dense, 1), + (2, 2, PrimaryKeyEncoding::Dense, 2), + (0, 2, PrimaryKeyEncoding::Dense, 0), + (2, 2, PrimaryKeyEncoding::Sparse, 0), + ]; + + for (num_tags, num_fields, encoding, expected) in cases { + let metadata = build_metadata(num_tags, num_fields, encoding); + let options = FlatSchemaOptions::from_encoding(encoding); + let num_columns = flat_sst_arrow_schema_column_num(&metadata, &options); + let result = field_column_start(&metadata, num_columns); + assert_eq!( + result, expected, + "num_tags={num_tags}, num_fields={num_fields}, encoding={encoding:?}" + ); + } + } +} diff --git a/src/mito2/src/sst/parquet/format.rs b/src/mito2/src/sst/parquet/format.rs index 70d026e6db9..ba64eac78b3 100644 --- a/src/mito2/src/sst/parquet/format.rs +++ b/src/mito2/src/sst/parquet/format.rs @@ -34,12 +34,12 @@ use api::v1::SemanticType; use common_time::Timestamp; use datafusion_common::ScalarValue; use datatypes::arrow::array::{ - ArrayRef, BinaryArray, BinaryDictionaryBuilder, DictionaryArray, UInt32Array, UInt64Array, + ArrayRef, BinaryArray, BinaryDictionaryBuilder, DictionaryArray, UInt64Array, }; use datatypes::arrow::datatypes::{SchemaRef, UInt32Type}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::DataType; -use datatypes::vectors::{Helper, Vector}; +use datatypes::vectors::Helper; use mito_codec::row_converter::{ CompositeValues, PrimaryKeyCodec, SortField, build_primary_key_codec, build_primary_key_codec_with_fields, @@ -51,8 +51,7 @@ use store_api::metadata::{ColumnMetadata, RegionMetadataRef}; use store_api::storage::{ColumnId, SequenceNumber}; use crate::error::{ - ConvertVectorSnafu, DecodeSnafu, InvalidBatchSnafu, InvalidRecordBatchSnafu, - NewRecordBatchSnafu, Result, + ConvertVectorSnafu, DecodeSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result, }; use crate::read::{Batch, BatchBuilder, BatchColumn}; use crate::sst::file::{FileMeta, FileTimeRange}; @@ -73,7 +72,6 @@ pub(crate) const INTERNAL_COLUMN_NUM: usize = 3; /// Helper for writing the SST format with primary key. pub(crate) struct PrimaryKeyWriteFormat { - metadata: RegionMetadataRef, /// SST file schema. arrow_schema: SchemaRef, override_sequence: Option, @@ -84,7 +82,6 @@ impl PrimaryKeyWriteFormat { pub(crate) fn new(metadata: RegionMetadataRef) -> PrimaryKeyWriteFormat { let arrow_schema = to_sst_arrow_schema(&metadata); PrimaryKeyWriteFormat { - metadata, arrow_schema, override_sequence: None, } @@ -104,40 +101,25 @@ impl PrimaryKeyWriteFormat { &self.arrow_schema } - /// Convert `batch` to a arrow record batch to store in parquet. - pub(crate) fn convert_batch(&self, batch: &Batch) -> Result { - debug_assert_eq!( - batch.fields().len() + FIXED_POS_COLUMN_NUM, - self.arrow_schema.fields().len() - ); - let mut columns = Vec::with_capacity(batch.fields().len() + FIXED_POS_COLUMN_NUM); - // Store all fields first. - for (column, column_metadata) in batch.fields().iter().zip(self.metadata.field_columns()) { - ensure!( - column.column_id == column_metadata.column_id, - InvalidBatchSnafu { - reason: format!( - "Batch has column {} but metadata has column {}", - column.column_id, column_metadata.column_id - ), - } - ); - - columns.push(column.data.to_arrow_array()); - } - // Add time index column. - columns.push(batch.timestamps().to_arrow_array()); - // Add internal columns: primary key, sequences, op types. - columns.push(new_primary_key_array(batch.primary_key(), batch.num_rows())); + /// Convert a flat `RecordBatch` to primary-key format, retaining only + /// field columns, time index, and internal columns. + /// + /// `num_fields` is the number of field columns. The method strips + /// leading tag columns: `num_tag_columns = batch.num_columns() - num_fields - FIXED_POS_COLUMN_NUM`. + pub(crate) fn convert_flat_batch( + &self, + batch: &RecordBatch, + num_fields: usize, + ) -> Result { + let num_tag_columns = batch.num_columns() - num_fields - FIXED_POS_COLUMN_NUM; + let mut columns: Vec = batch.columns()[num_tag_columns..].to_vec(); if let Some(override_sequence) = self.override_sequence { - let sequence_array = + let num_cols = columns.len(); + // sequence is at num_cols - 2 (before op_type) + columns[num_cols - 2] = Arc::new(UInt64Array::from(vec![override_sequence; batch.num_rows()])); - columns.push(sequence_array); - } else { - columns.push(batch.sequences().to_arrow_array()); } - columns.push(batch.op_types().to_arrow_array()); RecordBatch::try_new(self.arrow_schema.clone(), columns).context(NewRecordBatchSnafu) } @@ -926,15 +908,6 @@ pub(crate) fn primary_key_offsets(pk_dict_array: &PrimaryKeyArray) -> Result ArrayRef { - let values = Arc::new(BinaryArray::from_iter_values([primary_key])); - let keys = UInt32Array::from_value(0, num_rows); - - // Safety: The key index is valid. - Arc::new(DictionaryArray::new(keys, values)) -} - /// Gets the min/max time index of the row group from the parquet meta. /// It assumes the parquet is created by the mito engine. pub(crate) fn parquet_row_group_time_range( @@ -1017,7 +990,7 @@ mod tests { use api::v1::OpType; use datatypes::arrow::array::{ - Int64Array, StringArray, TimestampMillisecondArray, UInt8Array, UInt64Array, + Int64Array, StringArray, TimestampMillisecondArray, UInt8Array, UInt32Array, UInt64Array, }; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema, TimeUnit}; use datatypes::prelude::ConcreteDataType; @@ -1145,13 +1118,6 @@ mod tests { assert_eq!(&build_test_arrow_schema(), write_format.arrow_schema()); } - #[test] - fn test_new_primary_key_array() { - let array = new_primary_key_array(b"test", 3); - let expect = build_test_pk_array(&[(b"test".to_vec(), 3)]) as ArrayRef; - assert_eq!(&expect, &array); - } - fn build_test_pk_array(pk_row_nums: &[(Vec, usize)]) -> Arc { let values = Arc::new(BinaryArray::from_iter_values( pk_row_nums.iter().map(|v| &v.0), @@ -1164,49 +1130,6 @@ mod tests { Arc::new(DictionaryArray::new(keys, values)) } - #[test] - fn test_convert_batch() { - let metadata = build_test_region_metadata(); - let write_format = PrimaryKeyWriteFormat::new(metadata); - - let num_rows = 4; - let batch = new_batch(b"test", 1, 2, num_rows); - let columns: Vec = vec![ - Arc::new(Int64Array::from(vec![2; num_rows])), // field1 - Arc::new(Int64Array::from(vec![3; num_rows])), // field0 - Arc::new(TimestampMillisecondArray::from(vec![1, 2, 3, 4])), // ts - build_test_pk_array(&[(b"test".to_vec(), num_rows)]), // primary key - Arc::new(UInt64Array::from(vec![TEST_SEQUENCE; num_rows])), // sequence - Arc::new(UInt8Array::from(vec![TEST_OP_TYPE; num_rows])), // op type - ]; - let expect_record = RecordBatch::try_new(build_test_arrow_schema(), columns).unwrap(); - - let actual = write_format.convert_batch(&batch).unwrap(); - assert_eq!(expect_record, actual); - } - - #[test] - fn test_convert_batch_with_override_sequence() { - let metadata = build_test_region_metadata(); - let write_format = - PrimaryKeyWriteFormat::new(metadata).with_override_sequence(Some(415411)); - - let num_rows = 4; - let batch = new_batch(b"test", 1, 2, num_rows); - let columns: Vec = vec![ - Arc::new(Int64Array::from(vec![2; num_rows])), // field1 - Arc::new(Int64Array::from(vec![3; num_rows])), // field0 - Arc::new(TimestampMillisecondArray::from(vec![1, 2, 3, 4])), // ts - build_test_pk_array(&[(b"test".to_vec(), num_rows)]), // primary key - Arc::new(UInt64Array::from(vec![415411; num_rows])), // sequence - Arc::new(UInt8Array::from(vec![TEST_OP_TYPE; num_rows])), // op type - ]; - let expect_record = RecordBatch::try_new(build_test_arrow_schema(), columns).unwrap(); - - let actual = write_format.convert_batch(&batch).unwrap(); - assert_eq!(expect_record, actual); - } - #[test] fn test_projection_indices() { let metadata = build_test_region_metadata(); @@ -1867,4 +1790,100 @@ mod tests { let result = format.convert_batch(record_batch.clone(), None).unwrap(); assert_eq!(record_batch, result); } + + #[test] + fn test_convert_flat_batch() { + let metadata = build_test_region_metadata(); + let write_format = PrimaryKeyWriteFormat::new(metadata); + + let num_rows = 4; + // Build a flat record batch: tag0, tag1, field1, field0, ts, __primary_key, __sequence, __op_type + let flat_columns: Vec = input_columns_for_flat_batch(num_rows); + let flat_batch = RecordBatch::try_new(build_test_flat_sst_schema(), flat_columns).unwrap(); + + // num_fields = 2 (field1, field0) + let result = write_format.convert_flat_batch(&flat_batch, 2).unwrap(); + + // Expected: tag columns stripped, only field1, field0, ts, __primary_key, __sequence, __op_type + let expected_columns: Vec = vec![ + Arc::new(Int64Array::from(vec![2; num_rows])), // field1 + Arc::new(Int64Array::from(vec![3; num_rows])), // field0 + Arc::new(TimestampMillisecondArray::from(vec![1, 2, 3, 4])), // ts + build_test_pk_array(&[(b"test".to_vec(), num_rows)]), // __primary_key + Arc::new(UInt64Array::from(vec![TEST_SEQUENCE; num_rows])), // __sequence + Arc::new(UInt8Array::from(vec![TEST_OP_TYPE; num_rows])), // __op_type + ]; + let expected = RecordBatch::try_new(build_test_arrow_schema(), expected_columns).unwrap(); + + assert_eq!(expected, result); + } + + #[test] + fn test_convert_flat_batch_with_override_sequence() { + let metadata = build_test_region_metadata(); + let write_format = PrimaryKeyWriteFormat::new(metadata).with_override_sequence(Some(999)); + + let num_rows = 4; + let flat_columns: Vec = input_columns_for_flat_batch(num_rows); + let flat_batch = RecordBatch::try_new(build_test_flat_sst_schema(), flat_columns).unwrap(); + + let result = write_format.convert_flat_batch(&flat_batch, 2).unwrap(); + + let expected_columns: Vec = vec![ + Arc::new(Int64Array::from(vec![2; num_rows])), // field1 + Arc::new(Int64Array::from(vec![3; num_rows])), // field0 + Arc::new(TimestampMillisecondArray::from(vec![1, 2, 3, 4])), // ts + build_test_pk_array(&[(b"test".to_vec(), num_rows)]), // __primary_key + Arc::new(UInt64Array::from(vec![999; num_rows])), // overridden __sequence + Arc::new(UInt8Array::from(vec![TEST_OP_TYPE; num_rows])), // __op_type + ]; + let expected = RecordBatch::try_new(build_test_arrow_schema(), expected_columns).unwrap(); + + assert_eq!(expected, result); + } + + #[test] + fn test_convert_flat_batch_no_tags() { + // Test with a region that has no primary key columns (no tags to strip). + let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "field0", + ConcreteDataType::int64_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 2, + }); + let metadata = Arc::new(builder.build().unwrap()); + let write_format = PrimaryKeyWriteFormat::new(metadata); + + let num_rows = 3; + // No tag columns, so flat batch is: field0, ts, __primary_key, __sequence, __op_type + let sst_schema = write_format.arrow_schema().clone(); + let columns: Vec = vec![ + Arc::new(Int64Array::from(vec![10; num_rows])), // field0 + Arc::new(TimestampMillisecondArray::from(vec![1, 2, 3])), // ts + build_test_pk_array(&[(b"".to_vec(), num_rows)]), // __primary_key + Arc::new(UInt64Array::from(vec![TEST_SEQUENCE; num_rows])), // __sequence + Arc::new(UInt8Array::from(vec![TEST_OP_TYPE; num_rows])), // __op_type + ]; + let flat_batch = RecordBatch::try_new(sst_schema.clone(), columns.clone()).unwrap(); + + // num_fields = 1, num_tag_columns = 5 - 1 - 4 = 0, so nothing is stripped + let result = write_format.convert_flat_batch(&flat_batch, 1).unwrap(); + let expected = RecordBatch::try_new(sst_schema, columns).unwrap(); + + assert_eq!(expected, result); + } } diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 5dbe026b238..9561a07df35 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -21,9 +21,8 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use api::v1::SemanticType; -use async_trait::async_trait; use common_recordbatch::filter::SimpleFilterEvaluator; -use common_telemetry::{debug, tracing, warn}; +use common_telemetry::{tracing, warn}; use datafusion_expr::Expr; use datatypes::arrow::array::ArrayRef; use datatypes::arrow::datatypes::Field; @@ -35,29 +34,28 @@ use mito_codec::row_converter::{PrimaryKeyFilter, build_primary_key_codec}; use object_store::ObjectStore; use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, RowSelection}; use parquet::arrow::{FieldLevels, ProjectionMask, parquet_to_arrow_field_levels}; -use parquet::file::metadata::{KeyValue, PageIndexPolicy, ParquetMetaData}; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; use partition::expr::PartitionExpr; -use snafu::{OptionExt, ResultExt}; +use snafu::ResultExt; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataRef}; use store_api::region_request::PathType; use store_api::storage::{ColumnId, FileId}; use table::predicate::Predicate; -use crate::cache::CacheStrategy; use crate::cache::index::result_cache::PredicateKey; +use crate::cache::{CacheStrategy, CachedSstMeta}; #[cfg(feature = "vector_index")] use crate::error::ApplyVectorIndexSnafu; use crate::error::{ - ArrowReaderSnafu, InvalidMetadataSnafu, InvalidParquetSnafu, ReadDataPartSnafu, - ReadParquetSnafu, Result, SerializePartitionExprSnafu, + ArrowReaderSnafu, ReadDataPartSnafu, ReadParquetSnafu, Result, SerializePartitionExprSnafu, }; use crate::metrics::{ PRECISE_FILTER_ROWS_TOTAL, READ_ROW_GROUPS_TOTAL, READ_ROWS_IN_ROW_GROUP_TOTAL, READ_ROWS_TOTAL, READ_STAGE_ELAPSED, }; use crate::read::flat_projection::CompactionProjectionMapper; -use crate::read::prune::{PruneReader, Source}; +use crate::read::prune::FlatPruneReader; use crate::read::{Batch, BatchReader}; use crate::sst::file::FileHandle; use crate::sst::index::bloom_filter::applier::{ @@ -71,6 +69,7 @@ use crate::sst::index::inverted_index::applier::{ }; #[cfg(feature = "vector_index")] use crate::sst::index::vector_index::applier::VectorIndexApplierRef; +use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; use crate::sst::parquet::file_range::{ FileRangeContext, FileRangeContextRef, PartitionFilterContext, PreFilterMode, RangeBase, row_group_contains_delete, @@ -80,7 +79,6 @@ use crate::sst::parquet::metadata::MetadataLoader; use crate::sst::parquet::row_group::{InMemoryRowGroup, ParquetFetchMetrics}; use crate::sst::parquet::row_selection::RowGroupSelection; use crate::sst::parquet::stats::RowGroupPruningStats; -use crate::sst::parquet::{DEFAULT_READ_BATCH_SIZE, PARQUET_METADATA_KEY}; use crate::sst::tag_maybe_to_dictionary_field; const INDEX_TYPE_FULLTEXT: &str = "fulltext"; @@ -303,7 +301,8 @@ impl ParquetReaderBuilder { pub async fn build(&self) -> Result> { let mut metrics = ReaderMetrics::default(); - let Some((context, selection)) = self.build_reader_input(&mut metrics).await? else { + let Some((context, selection)) = self.build_reader_input_inner(&mut metrics, true).await? + else { return Ok(None); }; ParquetReader::new(Arc::new(context), selection) @@ -325,12 +324,14 @@ impl ParquetReaderBuilder { &self, metrics: &mut ReaderMetrics, ) -> Result> { - self.build_reader_input_inner(metrics).await + self.build_reader_input_inner(metrics, self.flat_format) + .await } async fn build_reader_input_inner( &self, metrics: &mut ReaderMetrics, + flat_format: bool, ) -> Result> { let start = Instant::now(); @@ -338,7 +339,7 @@ impl ParquetReaderBuilder { let file_size = self.file_handle.meta_ref().file_size; // Loads parquet metadata of the file. - let (parquet_meta, cache_miss) = self + let (sst_meta, cache_miss) = self .read_parquet_metadata( &file_path, file_size, @@ -346,9 +347,8 @@ impl ParquetReaderBuilder { self.page_index_policy, ) .await?; - // Decodes region metadata. - let key_value_meta = parquet_meta.file_metadata().key_value_metadata(); - let region_meta = Arc::new(Self::get_region_metadata(&file_path, key_value_meta)?); + let parquet_meta = sst_meta.parquet_metadata(); + let region_meta = sst_meta.region_metadata(); let region_partition_expr_str = self .expected_metadata .as_ref() @@ -373,7 +373,7 @@ impl ParquetReaderBuilder { // before compat handling. let compaction_projection_mapper = if self.compaction && !is_same_region_partition - && self.flat_format + && flat_format && region_meta.primary_key_encoding == PrimaryKeyEncoding::Sparse { Some(CompactionProjectionMapper::try_new(®ion_meta)?) @@ -385,7 +385,7 @@ impl ParquetReaderBuilder { ReadFormat::new( region_meta.clone(), Some(column_ids), - self.flat_format, + flat_format, Some(parquet_meta.file_metadata().schema_descr().num_columns()), &file_path, skip_auto_convert, @@ -401,7 +401,7 @@ impl ParquetReaderBuilder { ReadFormat::new( region_meta.clone(), Some(&column_ids), - self.flat_format, + flat_format, Some(parquet_meta.file_metadata().schema_descr().num_columns()), &file_path, skip_auto_convert, @@ -613,42 +613,15 @@ impl ParquetReaderBuilder { })) } - /// Decodes region metadata from key value. - fn get_region_metadata( - file_path: &str, - key_value_meta: Option<&Vec>, - ) -> Result { - let key_values = key_value_meta.context(InvalidParquetSnafu { - file: file_path, - reason: "missing key value meta", - })?; - let meta_value = key_values - .iter() - .find(|kv| kv.key == PARQUET_METADATA_KEY) - .with_context(|| InvalidParquetSnafu { - file: file_path, - reason: format!("key {} not found", PARQUET_METADATA_KEY), - })?; - let json = meta_value - .value - .as_ref() - .with_context(|| InvalidParquetSnafu { - file: file_path, - reason: format!("No value for key {}", PARQUET_METADATA_KEY), - })?; - - RegionMetadata::from_json(json).context(InvalidMetadataSnafu) - } - /// Reads parquet metadata of specific file. - /// Returns (metadata, cache_miss_flag). + /// Returns (fused metadata, cache_miss_flag). async fn read_parquet_metadata( &self, file_path: &str, file_size: u64, cache_metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, - ) -> Result<(Arc, bool)> { + ) -> Result<(Arc, bool)> { let start = Instant::now(); let _t = READ_STAGE_ELAPSED .with_label_values(&["read_parquet_metadata"]) @@ -658,7 +631,7 @@ impl ParquetReaderBuilder { // Tries to get from cache with metrics tracking. if let Some(metadata) = self .cache_strategy - .get_parquet_meta_data(file_id, cache_metrics, page_index_policy) + .get_sst_meta_data(file_id, cache_metrics, page_index_policy) .await { cache_metrics.metadata_load_cost += start.elapsed(); @@ -671,10 +644,10 @@ impl ParquetReaderBuilder { metadata_loader.with_page_index_policy(page_index_policy); let metadata = metadata_loader.load(cache_metrics).await?; - let metadata = Arc::new(metadata); + let metadata = Arc::new(CachedSstMeta::try_new(file_path, metadata)?); // Cache the metadata. self.cache_strategy - .put_parquet_meta_data(file_id, metadata.clone()); + .put_sst_meta_data(file_id, metadata.clone()); cache_metrics.metadata_load_cost += start.elapsed(); Ok((metadata, true)) @@ -1766,24 +1739,6 @@ impl RowGroupReaderBuilder { } } -/// The state of a [ParquetReader]. -enum ReaderState { - /// The reader is reading a row group. - Readable(PruneReader), - /// The reader is exhausted. - Exhausted(ReaderMetrics), -} - -impl ReaderState { - /// Returns the metrics of the reader. - fn metrics(&self) -> ReaderMetrics { - match self { - ReaderState::Readable(reader) => reader.metrics(), - ReaderState::Exhausted(m) => m.clone(), - } - } -} - /// The filter to evaluate or the prune result of the default value. pub(crate) enum MaybeFilter { /// The filter to evaluate. @@ -1894,13 +1849,12 @@ pub struct ParquetReader { /// Row group selection to read. selection: RowGroupSelection, /// Reader of current row group. - reader_state: ReaderState, + reader: Option, /// Metrics for tracking row group fetch operations. fetch_metrics: ParquetFetchMetrics, } -#[async_trait] -impl BatchReader for ParquetReader { +impl ParquetReader { #[tracing::instrument( skip_all, fields( @@ -1908,18 +1862,20 @@ impl BatchReader for ParquetReader { file_id = %self.context.reader_builder().file_handle.file_id() ) )] - async fn next_batch(&mut self) -> Result> { - let ReaderState::Readable(reader) = &mut self.reader_state else { - return Ok(None); - }; + pub async fn next_record_batch(&mut self) -> Result> { + loop { + if let Some(reader) = &mut self.reader { + if let Some(batch) = reader.next_batch()? { + return Ok(Some(batch)); + } + self.reader = None; + continue; + } - // We don't collect the elapsed time if the reader returns an error. - if let Some(batch) = reader.next_batch().await? { - return Ok(Some(batch)); - } + let Some((row_group_idx, row_selection)) = self.selection.pop_first() else { + return Ok(None); + }; - // No more items in current row group, reads next row group. - while let Some((row_group_idx, row_selection)) = self.selection.pop_first() { let parquet_reader = self .context .reader_builder() @@ -1930,54 +1886,14 @@ impl BatchReader for ParquetReader { ) .await?; - // Resets the parquet reader. - // Compute skip_fields for this row group let skip_fields = self.context.should_skip_fields(row_group_idx); - reader.reset_source( - Source::RowGroup(RowGroupReader::new(self.context.clone(), parquet_reader)), + self.reader = Some(FlatPruneReader::new_with_row_group_reader( + self.context.clone(), + FlatRowGroupReader::new(self.context.clone(), parquet_reader), skip_fields, - ); - if let Some(batch) = reader.next_batch().await? { - return Ok(Some(batch)); - } + )); } - - // The reader is exhausted. - self.reader_state = ReaderState::Exhausted(reader.metrics().clone()); - Ok(None) } -} - -impl Drop for ParquetReader { - fn drop(&mut self) { - let metrics = self.reader_state.metrics(); - debug!( - "Read parquet {} {}, range: {:?}, {}/{} row groups, metrics: {:?}", - self.context.reader_builder().file_handle.region_id(), - self.context.reader_builder().file_handle.file_id(), - self.context.reader_builder().file_handle.time_range(), - metrics.filter_metrics.rg_total - - metrics.filter_metrics.rg_inverted_filtered - - metrics.filter_metrics.rg_minmax_filtered - - metrics.filter_metrics.rg_fulltext_filtered - - metrics.filter_metrics.rg_bloom_filtered, - metrics.filter_metrics.rg_total, - metrics - ); - - // Report metrics. - READ_STAGE_ELAPSED - .with_label_values(&["build_parquet_reader"]) - .observe(metrics.build_cost.as_secs_f64()); - READ_STAGE_ELAPSED - .with_label_values(&["scan_row_groups"]) - .observe(metrics.scan_cost.as_secs_f64()); - metrics.observe_rows("parquet_reader"); - metrics.filter_metrics.observe(); - } -} - -impl ParquetReader { /// Creates a new reader. #[tracing::instrument( skip_all, @@ -1990,28 +1906,27 @@ impl ParquetReader { context: FileRangeContextRef, mut selection: RowGroupSelection, ) -> Result { + debug_assert!(context.read_format().as_flat().is_some()); let fetch_metrics = ParquetFetchMetrics::default(); - // No more items in current row group, reads next row group. - let reader_state = if let Some((row_group_idx, row_selection)) = selection.pop_first() { + let reader = if let Some((row_group_idx, row_selection)) = selection.pop_first() { let parquet_reader = context .reader_builder() .build(row_group_idx, Some(row_selection), Some(&fetch_metrics)) .await?; - // Compute skip_fields once for this row group let skip_fields = context.should_skip_fields(row_group_idx); - ReaderState::Readable(PruneReader::new_with_row_group_reader( + Some(FlatPruneReader::new_with_row_group_reader( context.clone(), - RowGroupReader::new(context.clone(), parquet_reader), + FlatRowGroupReader::new(context.clone(), parquet_reader), skip_fields, )) } else { - ReaderState::Exhausted(ReaderMetrics::default()) + None }; Ok(ParquetReader { context, selection, - reader_state, + reader, fetch_metrics, }) } diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index b207f11ef8d..4e75073e269 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -50,7 +50,7 @@ use crate::config::{IndexBuildMode, IndexConfig}; use crate::error::{ InvalidMetadataSnafu, OpenDalSnafu, Result, UnexpectedSnafu, WriteParquetSnafu, }; -use crate::read::{Batch, FlatSource, Source}; +use crate::read::FlatSource; use crate::sst::file::RegionFileId; use crate::sst::index::{IndexOutput, Indexer, IndexerBuilder}; use crate::sst::parquet::flat_format::{FlatWriteFormat, time_index_column_index}; @@ -60,6 +60,35 @@ use crate::sst::{ DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, }; +/// Converts a flat RecordBatch for writing to parquet. +enum FlatBatchConverter { + /// Write as-is in flat format. + Flat(FlatWriteFormat), + /// Convert flat batch to primary-key format by stripping tag columns. + PrimaryKey { + format: PrimaryKeyWriteFormat, + num_fields: usize, + }, +} + +impl FlatBatchConverter { + fn arrow_schema(&self) -> &SchemaRef { + match self { + FlatBatchConverter::Flat(f) => f.arrow_schema(), + FlatBatchConverter::PrimaryKey { format, .. } => format.arrow_schema(), + } + } + + fn convert_batch(&self, batch: &RecordBatch) -> Result { + match self { + FlatBatchConverter::Flat(f) => f.convert_batch(batch), + FlatBatchConverter::PrimaryKey { format, num_fields } => { + format.convert_flat_batch(batch, *num_fields) + } + } + } +} + /// Parquet SST writer. pub struct ParquetWriter<'a, F: WriterFactory, I: IndexerBuilder, P: FilePathProvider> { /// Path provider that creates SST and index file paths according to file id. @@ -240,81 +269,6 @@ where Ok(()) } - /// Iterates source and writes all rows to Parquet file. - /// - /// Returns the [SstInfo] if the SST is written. - pub async fn write_all( - &mut self, - source: Source, - override_sequence: Option, // override the `sequence` field from `Source` - opts: &WriteOptions, - ) -> Result { - let res = self - .write_all_without_cleaning(source, override_sequence, opts) - .await; - if res.is_err() { - // Clean tmp files explicitly on failure. - let file_id = self.current_file; - if let Some(cleaner) = &self.file_cleaner { - cleaner.clean_by_file_id(file_id).await; - } - } - res - } - - async fn write_all_without_cleaning( - &mut self, - mut source: Source, - override_sequence: Option, // override the `sequence` field from `Source` - opts: &WriteOptions, - ) -> Result { - let mut results = smallvec![]; - let write_format = PrimaryKeyWriteFormat::new(self.metadata.clone()) - .with_override_sequence(override_sequence); - let mut stats = SourceStats::default(); - - while let Some(res) = self - .write_next_batch(&mut source, &write_format, opts) - .await - .transpose() - { - match res { - Ok(mut batch) => { - stats.update(&batch); - let start = Instant::now(); - // safety: self.current_indexer must be set when first batch has been written. - match self.index_config.build_mode { - IndexBuildMode::Sync => { - self.current_indexer - .as_mut() - .unwrap() - .update(&mut batch) - .await; - } - IndexBuildMode::Async => {} - } - self.metrics.update_index += start.elapsed(); - if let Some(max_file_size) = opts.max_file_size - && self.bytes_written.load(Ordering::Relaxed) > max_file_size - { - self.finish_current_file(&mut results, &mut stats).await?; - } - } - Err(e) => { - if let Some(indexer) = &mut self.current_indexer { - indexer.abort().await; - } - return Err(e); - } - } - } - - self.finish_current_file(&mut results, &mut stats).await?; - - // object_store.write will make sure all bytes are written or an error is raised. - Ok(results) - } - /// Iterates FlatSource and writes all RecordBatch in flat format to Parquet file. /// /// Returns the [SstInfo] if the SST is written. @@ -324,11 +278,15 @@ where override_sequence: Option, opts: &WriteOptions, ) -> Result { - let res = self - .write_all_flat_without_cleaning(source, override_sequence, opts) - .await; + let converter = FlatBatchConverter::Flat( + FlatWriteFormat::new( + self.metadata.clone(), + &FlatSchemaOptions::from_encoding(self.metadata.primary_key_encoding), + ) + .with_override_sequence(override_sequence), + ); + let res = self.write_all_flat_inner(source, &converter, opts).await; if res.is_err() { - // Clean tmp files explicitly on failure. let file_id = self.current_file; if let Some(cleaner) = &self.file_cleaner { cleaner.clean_by_file_id(file_id).await; @@ -337,36 +295,58 @@ where res } - async fn write_all_flat_without_cleaning( + /// Iterates FlatSource and writes all RecordBatch in primary-key format to Parquet file. + /// + /// Returns the [SstInfo] if the SST is written. + pub async fn write_all_flat_as_primary_key( &mut self, - mut source: FlatSource, + source: FlatSource, override_sequence: Option, opts: &WriteOptions, + ) -> Result { + let num_fields = self.metadata.field_columns().count(); + let converter = FlatBatchConverter::PrimaryKey { + format: PrimaryKeyWriteFormat::new(self.metadata.clone()) + .with_override_sequence(override_sequence), + num_fields, + }; + let res = self.write_all_flat_inner(source, &converter, opts).await; + if res.is_err() { + let file_id = self.current_file; + if let Some(cleaner) = &self.file_cleaner { + cleaner.clean_by_file_id(file_id).await; + } + } + res + } + + async fn write_all_flat_inner( + &mut self, + mut source: FlatSource, + converter: &FlatBatchConverter, + opts: &WriteOptions, ) -> Result { let mut results = smallvec![]; - let flat_format = FlatWriteFormat::new( - self.metadata.clone(), - &FlatSchemaOptions::from_encoding(self.metadata.primary_key_encoding), - ) - .with_override_sequence(override_sequence); let mut stats = SourceStats::default(); while let Some(record_batch) = self - .write_next_flat_batch(&mut source, &flat_format, opts) + .write_next_flat_batch(&mut source, converter, opts) .await .transpose() { match record_batch { Ok(batch) => { stats.update_flat(&batch)?; - let start = Instant::now(); - // safety: self.current_indexer must be set when first batch has been written. - self.current_indexer - .as_mut() - .unwrap() - .update_flat(&batch) - .await; - self.metrics.update_index += start.elapsed(); + if matches!(self.index_config.build_mode, IndexBuildMode::Sync) { + let start = Instant::now(); + // safety: self.current_indexer must be set when first batch has been written. + self.current_indexer + .as_mut() + .unwrap() + .update_flat(&batch) + .await; + self.metrics.update_index += start.elapsed(); + } if let Some(max_file_size) = opts.max_file_size && self.bytes_written.load(Ordering::Relaxed) > max_file_size { @@ -411,34 +391,10 @@ where .set_column_compression(op_type_col, Compression::UNCOMPRESSED) } - async fn write_next_batch( - &mut self, - source: &mut Source, - write_format: &PrimaryKeyWriteFormat, - opts: &WriteOptions, - ) -> Result> { - let start = Instant::now(); - let Some(batch) = source.next_batch().await? else { - return Ok(None); - }; - self.metrics.iter_source += start.elapsed(); - - let arrow_batch = write_format.convert_batch(&batch)?; - - let start = Instant::now(); - self.maybe_init_writer(write_format.arrow_schema(), opts) - .await? - .write(&arrow_batch) - .await - .context(WriteParquetSnafu)?; - self.metrics.write_batch += start.elapsed(); - Ok(Some(batch)) - } - async fn write_next_flat_batch( &mut self, source: &mut FlatSource, - flat_format: &FlatWriteFormat, + converter: &FlatBatchConverter, opts: &WriteOptions, ) -> Result> { let start = Instant::now(); @@ -447,15 +403,16 @@ where }; self.metrics.iter_source += start.elapsed(); - let arrow_batch = flat_format.convert_batch(&record_batch)?; + let arrow_batch = converter.convert_batch(&record_batch)?; let start = Instant::now(); - self.maybe_init_writer(flat_format.arrow_schema(), opts) + self.maybe_init_writer(converter.arrow_schema(), opts) .await? .write(&arrow_batch) .await .context(WriteParquetSnafu)?; self.metrics.write_batch += start.elapsed(); + // Return original flat batch for stats/indexer which use flat layout. Ok(Some(record_batch)) } @@ -515,26 +472,6 @@ struct SourceStats { } impl SourceStats { - fn update(&mut self, batch: &Batch) { - if batch.is_empty() { - return; - } - - self.num_rows += batch.num_rows(); - self.series_estimator.update(batch); - // Safety: batch is not empty. - let (min_in_batch, max_in_batch) = ( - batch.first_timestamp().unwrap(), - batch.last_timestamp().unwrap(), - ); - if let Some(time_range) = &mut self.time_range { - time_range.0 = time_range.0.min(min_in_batch); - time_range.1 = time_range.1.max(max_in_batch); - } else { - self.time_range = Some((min_in_batch, max_in_batch)); - } - } - fn update_flat(&mut self, record_batch: &RecordBatch) -> Result<()> { if record_batch.num_rows() == 0 { return Ok(()); diff --git a/src/mito2/src/test_util/memtable_util.rs b/src/mito2/src/test_util/memtable_util.rs index 7ddac4ee0d1..58ea49fa412 100644 --- a/src/mito2/src/test_util/memtable_util.rs +++ b/src/mito2/src/test_util/memtable_util.rs @@ -83,16 +83,6 @@ impl Memtable for EmptyMemtable { Ok(()) } - #[cfg(any(test, feature = "test"))] - fn iter( - &self, - _projection: Option<&[ColumnId]>, - _filters: Option, - _sequence: Option, - ) -> Result { - Ok(Box::new(std::iter::empty())) - } - fn ranges( &self, _projection: Option<&[ColumnId]>, diff --git a/src/mito2/src/test_util/sst_util.rs b/src/mito2/src/test_util/sst_util.rs index 389d9bf107d..e9515030c05 100644 --- a/src/mito2/src/test_util/sst_util.rs +++ b/src/mito2/src/test_util/sst_util.rs @@ -18,7 +18,11 @@ use std::sync::Arc; use api::v1::{OpType, SemanticType}; use common_time::Timestamp; -use datatypes::arrow::array::{BinaryArray, TimestampMillisecondArray, UInt8Array, UInt64Array}; +use datatypes::arrow::array::{ + ArrayRef, BinaryDictionaryBuilder, RecordBatch, StringDictionaryBuilder, + TimestampMillisecondArray, UInt8Array, UInt64Array, +}; +use datatypes::arrow::datatypes::UInt32Type; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; use datatypes::value::ValueRef; @@ -32,8 +36,9 @@ use store_api::metric_engine_consts::{ use store_api::storage::consts::ReservedColumnId; use store_api::storage::{FileId, RegionId}; -use crate::read::{Batch, BatchBuilder, Source}; +use crate::read::{Batch, FlatSource, Source}; use crate::sst::file::{FileHandle, FileMeta}; +use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; use crate::test_util::{VecBatchReader, new_batch_builder, new_noop_file_purger}; /// Test region id. @@ -246,34 +251,68 @@ pub fn new_batch_by_range(tags: &[&str], start: usize, end: usize) -> Batch { new_batch_with_custom_sequence(tags, start, end, 1000) } -pub fn new_batch_with_binary(tags: &[&str], start: usize, end: usize) -> Batch { +/// Creates a flat format RecordBatch for testing. +/// Similar to `new_batch_by_range` but returns a RecordBatch in flat format. +pub fn new_record_batch_by_range(tags: &[&str], start: usize, end: usize) -> RecordBatch { + new_record_batch_with_custom_sequence(tags, start, end, 1000) +} + +/// Creates a flat format RecordBatch for testing with a custom sequence. +pub fn new_record_batch_with_custom_sequence( + tags: &[&str], + start: usize, + end: usize, + sequence: u64, +) -> RecordBatch { assert!(end >= start); + let metadata = Arc::new(sst_region_metadata()); + let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default()); + + let num_rows = end - start; + let mut columns = Vec::new(); + + // Add primary key columns (tag_0, tag_1) as dictionary arrays + let mut tag_0_builder = StringDictionaryBuilder::::new(); + let mut tag_1_builder = StringDictionaryBuilder::::new(); + + for _ in 0..num_rows { + tag_0_builder.append_value(tags[0]); + tag_1_builder.append_value(tags[1]); + } + + columns.push(Arc::new(tag_0_builder.finish()) as ArrayRef); + columns.push(Arc::new(tag_1_builder.finish()) as ArrayRef); + + // Add field column (field_0) + let field_values: Vec = (start..end).map(|v| v as u64).collect(); + columns.push(Arc::new(UInt64Array::from(field_values))); + + // Add time index column (ts) + let timestamps: Vec = (start..end).map(|v| v as i64).collect(); + columns.push(Arc::new(TimestampMillisecondArray::from(timestamps))); + + // Add encoded primary key column let pk = new_primary_key(tags); - let timestamps: Vec<_> = (start..end).map(|v| v as i64).collect(); - let sequences = vec![1000; end - start]; - let op_types = vec![OpType::Put; end - start]; + let mut pk_builder = BinaryDictionaryBuilder::::new(); + for _ in 0..num_rows { + pk_builder.append(&pk).unwrap(); + } + columns.push(Arc::new(pk_builder.finish())); - let field: Vec<_> = (start..end) - .map(|_v| "some data".as_bytes().to_vec()) - .collect(); + // Add sequence column + columns.push(Arc::new(UInt64Array::from_value(sequence, num_rows))); - let mut builder = BatchBuilder::new(pk); - builder - .timestamps_array(Arc::new(TimestampMillisecondArray::from_iter_values( - timestamps.iter().copied(), - ))) - .unwrap() - .sequences_array(Arc::new(UInt64Array::from_iter_values( - sequences.iter().copied(), - ))) - .unwrap() - .op_types_array(Arc::new(UInt8Array::from_iter_values( - op_types.iter().map(|v| *v as u8), - ))) - .unwrap() - .push_field_array(1, Arc::new(BinaryArray::from_iter_values(field))) - .unwrap(); - builder.build().unwrap() + // Add op_type column + columns.push(Arc::new(UInt8Array::from_value( + OpType::Put as u8, + num_rows, + ))); + RecordBatch::try_new(flat_schema, columns).unwrap() +} + +/// Creates a FlatSource from flat format RecordBatches. +pub fn new_flat_source_from_record_batches(batches: Vec) -> FlatSource { + FlatSource::Iter(Box::new(batches.into_iter().map(Ok))) } /// Creates a new region metadata for testing SSTs with binary datatype. diff --git a/src/servers/dashboard/VERSION b/src/servers/dashboard/VERSION index 03ee1a5314b..87a1cf595a6 100644 --- a/src/servers/dashboard/VERSION +++ b/src/servers/dashboard/VERSION @@ -1 +1 @@ -v0.11.13 +v0.12.0 diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index ca6a77a0775..506a240cac3 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -78,7 +78,7 @@ use crate::metrics_handler::MetricsHandler; use crate::prometheus_handler::PrometheusHandlerRef; use crate::query_handler::sql::ServerSqlQueryHandlerRef; use crate::query_handler::{ - InfluxdbLineProtocolHandlerRef, JaegerQueryHandlerRef, LogQueryHandlerRef, + DashboardHandlerRef, InfluxdbLineProtocolHandlerRef, JaegerQueryHandlerRef, LogQueryHandlerRef, OpenTelemetryProtocolHandlerRef, OpentsdbProtocolHandlerRef, PipelineHandlerRef, PromStoreProtocolHandlerRef, }; @@ -112,8 +112,8 @@ pub mod utils; use result::HttpOutputWriter; pub(crate) use timeout::DynamicTimeoutLayer; +mod client_ip; use crate::prom_remote_write::validation::PromValidationMode; - mod hints; mod read_preference; #[cfg(any(test, feature = "testing"))] @@ -507,6 +507,11 @@ pub struct GreptimeOptionsConfigState { pub greptime_config_options: String, } +#[derive(Clone)] +pub struct DashboardState { + pub handler: DashboardHandlerRef, +} + pub struct HttpServerBuilder { options: HttpOptions, plugins: Plugins, @@ -703,6 +708,16 @@ impl HttpServerBuilder { } } + pub fn with_dashboard_handler(self, handler: DashboardHandlerRef) -> Self { + Self { + router: self.router.nest( + &format!("/{HTTP_API_VERSION}/dashboards"), + HttpServer::route_dashboard(handler), + ), + ..self + } + } + pub fn with_extra_router(self, router: Router) -> Self { Self { router: self.router.merge(router), @@ -868,6 +883,7 @@ impl HttpServer { authorize::check_http_auth, )) .layer(middleware::from_fn(hints::extract_hints)) + .layer(middleware::from_fn(client_ip::log_error_with_client_ip)) .layer(middleware::from_fn( read_preference::extract_read_preference, )), @@ -1169,6 +1185,26 @@ impl HttpServer { ) .with_state(handler) } + + #[cfg(feature = "dashboard")] + fn route_dashboard(handler: DashboardHandlerRef) -> Router { + use crate::http::dashboard::{add_dashboard, delete_dashboard, list_dashboards}; + + Router::new() + .route("/", routing::get(list_dashboards)) + .route("/{dashboard_name}", routing::post(add_dashboard)) + .route("/{dashboard_name}", routing::delete(delete_dashboard)) + .layer( + ServiceBuilder::new() + .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), + ) + .with_state(DashboardState { handler }) + } + + #[cfg(not(feature = "dashboard"))] + fn route_dashboard(handler: DashboardHandlerRef) -> Router { + Router::new().with_state(DashboardState { handler }) + } } pub const HTTP_SERVER: &str = "HTTP_SERVER"; @@ -1212,7 +1248,10 @@ impl Server for HttpServer { error!(e; "Failed to set TCP_NODELAY on incoming connection"); } }); - let serve = axum::serve(listener, app.into_make_service()); + let serve = axum::serve( + listener, + app.into_make_service_with_connect_info::(), + ); // FIXME(yingwen): Support keepalive. // See: diff --git a/src/servers/src/http/client_ip.rs b/src/servers/src/http/client_ip.rs new file mode 100644 index 00000000000..70df554ebb5 --- /dev/null +++ b/src/servers/src/http/client_ip.rs @@ -0,0 +1,109 @@ +// 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. + +use std::net::SocketAddr; + +use axum::body::Body; +use axum::extract::{ConnectInfo, MatchedPath}; +use axum::http::Request; +use axum::middleware::Next; +use axum::response::Response; +use common_telemetry::warn; + +/// Middleware that logs HTTP error responses (4xx/5xx) with client IP address. +/// +/// Extracts client address from [`ConnectInfo`] if available. +pub async fn log_error_with_client_ip(req: Request, next: Next) -> Response { + let request_info = req + .extensions() + .get::>() + .map(|c| c.0) + .map(|addr| { + let method = req.method().clone(); + let uri = req.uri().clone(); + let matched_path = req.extensions().get::().cloned(); + (addr, method, uri, matched_path) + }); + + let response = next.run(req).await; + + if (response.status().is_client_error() || response.status().is_server_error()) + && let Some((addr, method, uri, matched_path)) = request_info + { + warn!( + "HTTP error response {} for {} {} (matched: {}) from client {}", + response.status(), + method, + uri, + matched_path + .as_ref() + .map(|p| p.as_str()) + .unwrap_or(""), + addr + ); + } + + response +} + +#[cfg(test)] +mod tests { + use axum::Router; + use axum::routing::get; + use http::StatusCode; + use tower::ServiceExt; + + use super::*; + + #[tokio::test] + async fn test_middleware_passes_error_response() { + async fn not_found_handler() -> StatusCode { + StatusCode::NOT_FOUND + } + + let app = Router::new() + .route("/not-found", get(not_found_handler)) + .layer(axum::middleware::from_fn(log_error_with_client_ip)); + + let response = app + .oneshot( + Request::builder() + .uri("/not-found") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::NOT_FOUND); + } + + #[tokio::test] + async fn test_middleware_passes_success_response() { + async fn ok_handler() -> StatusCode { + StatusCode::OK + } + + let app = Router::new() + .route("/ok", get(ok_handler)) + .layer(axum::middleware::from_fn(log_error_with_client_ip)); + + let response = app + .oneshot(Request::builder().uri("/ok").body(Body::empty()).unwrap()) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + } +} diff --git a/src/servers/src/http/dashboard.rs b/src/servers/src/http/dashboard.rs index bdb98490f09..ea894ca7d01 100644 --- a/src/servers/src/http/dashboard.rs +++ b/src/servers/src/http/dashboard.rs @@ -12,14 +12,21 @@ // See the License for the specific language governing permissions and // limitations under the License. -use axum::body::Body; +use std::sync::Arc; +use std::time::Instant; + +use axum::body::{Body, Bytes}; +use axum::extract::{Extension, Path, State}; use axum::http::{StatusCode, Uri, header}; use axum::response::Response; -use common_telemetry::debug; +use common_telemetry::{debug, error}; use rust_embed::RustEmbed; -use snafu::ResultExt; +use session::context::{Channel, QueryContext}; +use snafu::{ResultExt, ensure}; -use crate::error::{BuildHttpResponseSnafu, Result}; +use crate::error::{BuildHttpResponseSnafu, InvalidParameterSnafu, Result}; +use crate::http::DashboardState; +use crate::http::result::greptime_manage_resp::{DashboardOutput, GreptimedbManageResponse}; #[derive(RustEmbed)] #[folder = "dashboard/dist/"] @@ -61,3 +68,102 @@ fn get_assets(path: &str) -> Result { } .context(BuildHttpResponseSnafu) } + +#[axum_macros::debug_handler] +pub async fn add_dashboard( + State(state): State, + Path(dashboard_name): Path, + Extension(mut query_ctx): Extension, + payload: Bytes, +) -> Result { + let start = Instant::now(); + let handler = state.handler; + ensure!( + !dashboard_name.is_empty(), + InvalidParameterSnafu { + reason: "dashboard_name is required in path", + } + ); + + let definition = String::from_utf8_lossy(&payload).to_string(); + + query_ctx.set_channel(Channel::HttpSql); + let query_ctx = Arc::new(query_ctx); + + handler + .save(&dashboard_name, &definition, query_ctx) + .await + .map(|_| { + GreptimedbManageResponse::from_dashboard( + dashboard_name, + start.elapsed().as_millis() as u64, + ) + }) + .map_err(|e| { + error!(e; "failed to save dashboard"); + e + }) +} + +#[axum_macros::debug_handler] +pub async fn list_dashboards( + State(state): State, + Extension(mut query_ctx): Extension, +) -> Result { + let start = Instant::now(); + let handler = state.handler; + + query_ctx.set_channel(Channel::HttpSql); + let query_ctx = Arc::new(query_ctx); + + handler + .list(query_ctx) + .await + .map(|dashboards| { + let outputs: Vec = dashboards + .into_iter() + .map(|d| DashboardOutput { + name: d.name, + definition: d.definition, + }) + .collect(); + GreptimedbManageResponse::from_dashboards(outputs, start.elapsed().as_millis() as u64) + }) + .map_err(|e| { + error!(e; "failed to list dashboards"); + e + }) +} + +#[axum_macros::debug_handler] +pub async fn delete_dashboard( + State(state): State, + Extension(mut query_ctx): Extension, + Path(dashboard_name): Path, +) -> Result { + let start = Instant::now(); + let handler = state.handler; + ensure!( + !dashboard_name.is_empty(), + InvalidParameterSnafu { + reason: "dashboard_name is required", + } + ); + + query_ctx.set_channel(Channel::HttpSql); + let query_ctx = Arc::new(query_ctx); + + handler + .delete(&dashboard_name, query_ctx) + .await + .map(|_| { + GreptimedbManageResponse::from_dashboard( + dashboard_name, + start.elapsed().as_millis() as u64, + ) + }) + .map_err(|e| { + error!(e; "failed to delete dashboard"); + e + }) +} diff --git a/src/servers/src/http/result/error_result.rs b/src/servers/src/http/result/error_result.rs index 7b70066b685..9bd6e1a7a39 100644 --- a/src/servers/src/http/result/error_result.rs +++ b/src/servers/src/http/result/error_result.rs @@ -32,17 +32,24 @@ pub struct ErrorResponse { impl ErrorResponse { pub fn from_error(error: impl ErrorExt) -> Self { let code = error.status_code(); - if code.should_log_error() { error!(error; "Failed to handle HTTP request"); } else { debug!("Failed to handle HTTP request, err: {:?}", error); } - - Self::from_error_message(code, error.output_msg()) + ErrorResponse { + code: code as u32, + error: error.output_msg(), + execution_time_ms: 0, + } } pub fn from_error_message(code: StatusCode, msg: String) -> Self { + if code.should_log_error() { + error!("Failed to handle HTTP request: {}", msg); + } else { + debug!("Failed to handle HTTP request: {}", msg); + } ErrorResponse { code: code as u32, error: msg, diff --git a/src/servers/src/http/result/greptime_manage_resp.rs b/src/servers/src/http/result/greptime_manage_resp.rs index 3f7f3c6eeca..2b3a5d455c6 100644 --- a/src/servers/src/http/result/greptime_manage_resp.rs +++ b/src/servers/src/http/result/greptime_manage_resp.rs @@ -62,6 +62,25 @@ impl GreptimedbManageResponse { } } + pub fn from_dashboard(name: String, execution_time_ms: u64) -> Self { + GreptimedbManageResponse { + manage_result: ManageResult::Dashboards { + dashboards: vec![DashboardOutput { + name, + definition: String::new(), + }], + }, + execution_time_ms, + } + } + + pub fn from_dashboards(dashboards: Vec, execution_time_ms: u64) -> Self { + GreptimedbManageResponse { + manage_result: ManageResult::Dashboards { dashboards }, + execution_time_ms, + } + } + pub fn with_execution_time(mut self, execution_time: u64) -> Self { self.execution_time_ms = execution_time; self @@ -77,6 +96,7 @@ impl GreptimedbManageResponse { pub enum ManageResult { Pipelines { pipelines: Vec }, Sql { sql: SqlOutput }, + Dashboards { dashboards: Vec }, } #[derive(Serialize, Deserialize, Debug)] @@ -87,6 +107,13 @@ pub struct PipelineOutput { pipeline: Option, } +#[derive(Serialize, Deserialize, Debug)] +pub struct DashboardOutput { + pub name: String, + #[serde(skip_serializing_if = "String::is_empty")] + pub definition: String, +} + #[derive(Serialize, Deserialize, Debug)] pub struct SqlOutput { pub(crate) sql: String, diff --git a/src/servers/src/query_handler.rs b/src/servers/src/query_handler.rs index 60efe69faad..21c7646560b 100644 --- a/src/servers/src/query_handler.rs +++ b/src/servers/src/query_handler.rs @@ -44,6 +44,12 @@ use pipeline::{GreptimePipelineParams, Pipeline, PipelineInfo, PipelineVersion, use serde_json::Value; use session::context::{QueryContext, QueryContextRef}; +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct DashboardDefinition { + pub name: String, + pub definition: String, +} + use crate::error::Result; use crate::http::jaeger::QueryTraceParams; use crate::influxdb::InfluxdbRequest; @@ -176,6 +182,18 @@ pub trait PipelineHandler { ) -> Result<(String, TimestampNanosecond)>; } +/// Handling dashboard as code CRUD +pub type DashboardHandlerRef = Arc; + +#[async_trait] +pub trait DashboardHandler { + async fn save(&self, name: &str, definition: &str, ctx: QueryContextRef) -> Result<()>; + + async fn list(&self, ctx: QueryContextRef) -> Result>; + + async fn delete(&self, name: &str, ctx: QueryContextRef) -> Result<()>; +} + /// Handle log query requests. #[async_trait] pub trait LogQueryHandler { diff --git a/src/store-api/src/metadata.rs b/src/store-api/src/metadata.rs index d571a5392ff..0c663bccc05 100644 --- a/src/store-api/src/metadata.rs +++ b/src/store-api/src/metadata.rs @@ -18,8 +18,8 @@ use std::any::Any; use std::collections::{HashMap, HashSet}; -use std::fmt; use std::sync::Arc; +use std::{fmt, mem}; use api::v1::SemanticType; use api::v1::column_def::try_as_column_schema; @@ -99,6 +99,12 @@ impl ColumnMetadata { pub fn is_same_datatype(&self, other: &Self) -> bool { self.column_schema.data_type == other.column_schema.data_type } + + /// Returns the estimated memory footprint of this metadata. + pub fn estimated_size(&self) -> usize { + mem::size_of_val(self) - mem::size_of_val(&self.column_schema) + + self.column_schema.estimated_size() + } } #[cfg_attr(doc, aquamarine::aquamarine)] @@ -226,6 +232,25 @@ impl RegionMetadata { serde_json::from_str(s).context(SerdeJsonSnafu) } + /// Returns the estimated memory footprint of this metadata. + pub fn estimated_size(&self) -> usize { + mem::size_of_val(self) + + mem::size_of::() * self.column_metadatas.capacity() + + self + .column_metadatas + .iter() + .map(|column| column.estimated_size() - mem::size_of::()) + .sum::() + + mem::size_of::() * self.primary_key.capacity() + + mem::size_of::<(ColumnId, usize)>() * self.id_to_index.capacity() + + self.schema.estimated_size() + + self + .partition_expr + .as_ref() + .map(|expr| expr.capacity()) + .unwrap_or_default() + } + /// Encode the metadata to a JSON string. pub fn to_json(&self) -> Result { serde_json::to_string(&self).context(SerdeJsonSnafu) diff --git a/tests-fuzz/Cargo.toml b/tests-fuzz/Cargo.toml index a537ca06871..bc687092c00 100644 --- a/tests-fuzz/Cargo.toml +++ b/tests-fuzz/Cargo.toml @@ -100,6 +100,13 @@ test = false bench = false doc = false +[[bin]] +name = "fuzz_repartition_metric_table" +path = "targets/ddl/fuzz_repartition_metric_table.rs" +test = false +bench = false +doc = false + [[bin]] name = "fuzz_alter_table" path = "targets/ddl/fuzz_alter_table.rs" diff --git a/tests-fuzz/README.md b/tests-fuzz/README.md index 6807e19a1c1..cc9d7eb84ee 100644 --- a/tests-fuzz/README.md +++ b/tests-fuzz/README.md @@ -66,3 +66,23 @@ GT_FUZZ_OVERRIDE_SEED=6666 GT_FUZZ_OVERRIDE_ACTIONS=175 cargo fuzz run fuzz_targ ``` For more details, visit [cargo fuzz](https://rust-fuzz.github.io/book/cargo-fuzz/tutorial.html) or run the command `cargo fuzz --help`. + +## Repartition Metric Dump Artifacts + +For `fuzz_repartition_metric_table`, dump artifacts are written under one run directory. + +- Table data snapshots: `.table-data.csv` +- SQL traces per logical table: `.trace.sql` +- Seed metadata: `seed.meta` + +SQL trace behavior: + +- Insert SQL is appended after successful execution with comment fields including + `started_at_ms` and `elapsed_ms`. +- Repartition events are broadcast to all logical table trace files with comment fields including + `action_idx`, `started_at_ms`, `elapsed_ms`, and SQL text. + +Run directory lifecycle: + +- On success, the run directory is cleaned up. +- On failure, the run directory is retained for CI/local diffing. diff --git a/tests-fuzz/src/fake.rs b/tests-fuzz/src/fake.rs index aa92e0293ae..8910a39206c 100644 --- a/tests-fuzz/src/fake.rs +++ b/tests-fuzz/src/fake.rs @@ -65,6 +65,26 @@ where _v: PhantomData, } +pub struct ConstGenerator { + value: V, +} + +impl ConstGenerator { + pub fn new(value: V) -> Self { + Self { value } + } +} + +impl Random for ConstGenerator +where + R: Rng, + V: Clone, +{ + fn choose(&self, _rng: &mut R, amount: usize) -> Vec { + vec![self.value.clone(); amount] + } +} + pub fn random_capitalize_map(rng: &mut R, s: Ident) -> Ident { let mut v = s.value.chars().collect::>(); diff --git a/tests-fuzz/src/generator/create_expr.rs b/tests-fuzz/src/generator/create_expr.rs index fae6a95eda5..261a310db2c 100644 --- a/tests-fuzz/src/generator/create_expr.rs +++ b/tests-fuzz/src/generator/create_expr.rs @@ -193,6 +193,26 @@ fn generate_partition_def( } } +fn generate_metric_partition(partitions: usize) -> Option<(Column, PartitionDef)> { + if partitions <= 1 { + return None; + } + + let partition_column = Column { + name: Ident::new("host"), + column_type: ConcreteDataType::string_datatype(), + options: vec![ColumnOption::PrimaryKey], + }; + let bounds = generate_partition_bounds(&partition_column.column_type, partitions - 1); + let partitions = SimplePartitions::new(partition_column.name.clone(), bounds); + let partition_def = PartitionDef { + columns: vec![partitions.column_name.clone()], + exprs: partitions.generate().unwrap(), + }; + + Some((partition_column, partition_def)) +} + /// Generate a physical table with 2 columns: ts of TimestampType::Millisecond as time index and val of Float64Type. #[derive(Builder)] #[builder(pattern = "owned")] @@ -201,6 +221,8 @@ pub struct CreatePhysicalTableExprGenerator { name_generator: Box>, #[builder(default = "false")] if_not_exists: bool, + #[builder(default = "0")] + partition: usize, #[builder(default, setter(into))] with_clause: HashMap, } @@ -215,25 +237,35 @@ impl Generator for CreatePhysicalTableExpr options.insert(key.clone(), Value::from(value.clone())); } + let mut columns = vec![ + Column { + name: Ident::new("ts"), + column_type: ConcreteDataType::timestamp_millisecond_datatype(), + options: vec![ColumnOption::TimeIndex], + }, + Column { + name: Ident::new("val"), + column_type: ConcreteDataType::float64_datatype(), + options: vec![], + }, + ]; + + let mut partition = None; + let mut primary_keys = vec![]; + if let Some((partition_column, partition_def)) = generate_metric_partition(self.partition) { + columns.push(partition_column); + partition = Some(partition_def); + primary_keys.push(columns.len() - 1); + } + Ok(CreateTableExpr { table_name: self.name_generator.generate(rng), - columns: vec![ - Column { - name: Ident::new("ts"), - column_type: ConcreteDataType::timestamp_millisecond_datatype(), - options: vec![ColumnOption::TimeIndex], - }, - Column { - name: Ident::new("val"), - column_type: ConcreteDataType::float64_datatype(), - options: vec![], - }, - ], + columns, if_not_exists: self.if_not_exists, - partition: None, + partition, engine: "metric".to_string(), options, - primary_keys: vec![], + primary_keys, }) } } @@ -245,6 +277,8 @@ pub struct CreateLogicalTableExprGenerator { physical_table_ctx: TableContextRef, labels: usize, if_not_exists: bool, + #[builder(default = "true")] + include_partition_column: bool, #[builder(default = "Box::new(WordGenerator)")] name_generator: Box>, } @@ -253,11 +287,11 @@ impl Generator for CreateLogicalTableExprG type Error = Error; fn generate(&self, rng: &mut R) -> Result { - // Currently we mock the usage of GreptimeDB as Prometheus' backend, the physical table must have two columns. + // Currently we mock the usage of GreptimeDB as Prometheus' backend, the physical table must have ts and val. ensure!( - self.physical_table_ctx.columns.len() == 2, + self.physical_table_ctx.columns.len() >= 2, error::UnexpectedSnafu { - violated: "The physical table must have two columns" + violated: "The physical table must have at least two columns" } ); @@ -265,9 +299,16 @@ impl Generator for CreateLogicalTableExprG let logical_table_name = self .physical_table_ctx .generate_unique_table_name(rng, self.name_generator.as_ref()); + let mut physical_columns = self.physical_table_ctx.columns.clone(); + if !self.include_partition_column + && let Some(partition_def) = &self.physical_table_ctx.partition + { + physical_columns.retain(|column| !partition_def.columns.contains(&column.name)); + } + let mut logical_table = CreateTableExpr { table_name: logical_table_name, - columns: self.physical_table_ctx.columns.clone(), + columns: physical_columns, if_not_exists: self.if_not_exists, partition: None, engine: "metric".to_string(), @@ -459,6 +500,58 @@ mod tests { })); } + #[test] + fn test_create_physical_table_expr_generator_with_partition() { + let mut rng = rand::rng(); + let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default() + .partition(3) + .if_not_exists(false) + .build() + .unwrap() + .generate(&mut rng) + .unwrap(); + + assert_eq!(physical_table_expr.engine, "metric"); + assert!(physical_table_expr.partition.is_some()); + assert_eq!(physical_table_expr.partition.unwrap().exprs.len(), 3); + } + + #[test] + fn test_create_logical_table_expr_generator_without_partition_column() { + let mut rng = rand::rng(); + let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default() + .partition(3) + .if_not_exists(false) + .build() + .unwrap() + .generate(&mut rng) + .unwrap(); + let partition_columns = physical_table_expr + .partition + .as_ref() + .unwrap() + .columns + .clone(); + let physical_table_ctx = Arc::new(TableContext::from(&physical_table_expr)); + + let logical_table_expr = CreateLogicalTableExprGeneratorBuilder::default() + .physical_table_ctx(physical_table_ctx) + .labels(3) + .include_partition_column(false) + .if_not_exists(false) + .build() + .unwrap() + .generate(&mut rng) + .unwrap(); + + assert!( + logical_table_expr + .columns + .iter() + .all(|column| !partition_columns.contains(&column.name)) + ); + } + #[test] fn test_create_logical_table_expr_generator_deterministic() { let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0); diff --git a/tests-fuzz/src/ir.rs b/tests-fuzz/src/ir.rs index e8c15dcf950..ce1628cd619 100644 --- a/tests-fuzz/src/ir.rs +++ b/tests-fuzz/src/ir.rs @@ -20,6 +20,7 @@ pub(crate) mod insert_expr; pub(crate) mod partition_expr; pub(crate) mod repartition_expr; pub(crate) mod select_expr; +pub(crate) mod string_value; use core::fmt; use std::collections::HashMap; @@ -126,20 +127,7 @@ pub fn generate_partition_bounds(datatype: &ConcreteDataType, bounds: usize) -> ConcreteDataType::Int64(_) => generate_values!(i64, bounds), ConcreteDataType::Float32(_) => generate_values!(f32, bounds), ConcreteDataType::Float64(_) => generate_values!(f64, bounds), - ConcreteDataType::String(_) => { - let base = b'A'; - let range = b'z' - b'A'; - let step = range / (bounds as u8 + 1); - (1..=bounds) - .map(|i| { - Value::from( - char::from(base + step * i as u8) - .escape_default() - .to_string(), - ) - }) - .collect() - } + ConcreteDataType::String(_) => string_value::generate_partition_bounds(bounds), _ => unimplemented!("unsupported type: {datatype}"), } } @@ -157,10 +145,7 @@ pub fn generate_random_value( ConcreteDataType::Int64(_) => Value::from(rng.random::()), ConcreteDataType::Float32(_) => Value::from(rng.random::()), ConcreteDataType::Float64(_) => Value::from(rng.random::()), - ConcreteDataType::String(_) => match random_str { - Some(random) => Value::from(random.generate(rng).value), - None => Value::from(rng.random::().to_string()), - }, + ConcreteDataType::String(_) => string_value::generate_data_string_value(rng, random_str), ConcreteDataType::Date(_) => generate_random_date(rng), _ => unimplemented!("unsupported type: {datatype}"), @@ -341,21 +326,7 @@ pub fn generate_partition_value( } } datatypes::data_type::ConcreteDataType::String(_) => { - let upper = match first { - datatypes::value::Value::String(v) => v.as_utf8(), - _ => "", - }; - if bound_idx == 0 { - if upper <= "A" { - datatypes::value::Value::from("") - } else { - datatypes::value::Value::from("A") - } - } else if bound_idx < bounds.len() { - bounds[bound_idx - 1].clone() - } else { - last.clone() - } + string_value::generate_partition_value(bounds, bound_idx) } _ => unimplemented!("unsupported partition column type: {column_type}"), } diff --git a/tests-fuzz/src/ir/partition_expr.rs b/tests-fuzz/src/ir/partition_expr.rs index c91dd487ae2..908223366ca 100644 --- a/tests-fuzz/src/ir/partition_expr.rs +++ b/tests-fuzz/src/ir/partition_expr.rs @@ -20,7 +20,7 @@ use snafu::ensure; use crate::context::TableContext; use crate::error::{self, Result}; -use crate::ir::{Ident, generate_random_value}; +use crate::ir::{Ident, generate_random_value, string_value}; /// A partitioning scheme that divides a single column into multiple ranges based on provided bounds. /// @@ -245,6 +245,10 @@ pub fn generate_unique_bound( datatype: &ConcreteDataType, bounds: &[Value], ) -> Result { + if matches!(datatype, ConcreteDataType::String(_)) { + return string_value::generate_unique_partition_bound(rng, bounds); + } + for _ in 0..16 { let candidate = generate_random_value(rng, datatype, None); if !bounds.contains(&candidate) { diff --git a/tests-fuzz/src/ir/string_value.rs b/tests-fuzz/src/ir/string_value.rs new file mode 100644 index 00000000000..6a53aa69de7 --- /dev/null +++ b/tests-fuzz/src/ir/string_value.rs @@ -0,0 +1,162 @@ +// 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. + +use datatypes::value::Value; +use rand::Rng; + +use crate::error::{self, Result}; +use crate::generator::Random; +use crate::ir::Ident; + +const READABLE_CHARSET: &[u8] = b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"; + +fn readable_token(index: usize) -> String { + let base = READABLE_CHARSET.len(); + let mut n = index + 1; + let mut buf = Vec::new(); + + while n > 0 { + let rem = (n - 1) % base; + buf.push(READABLE_CHARSET[rem] as char); + n = (n - 1) / base; + } + + buf.iter().rev().collect() +} + +pub fn generate_data_string_value( + rng: &mut R, + random_str: Option<&dyn Random>, +) -> Value { + match random_str { + Some(random) => Value::from(random.generate(rng).value), + None => { + let idx = rng.random_range(0..(READABLE_CHARSET.len() * READABLE_CHARSET.len() * 4)); + Value::from(readable_token(idx)) + } + } +} + +/// Generates ordered readable string bounds for partition expressions. +pub fn generate_partition_bounds(bounds: usize) -> Vec { + let token_space = READABLE_CHARSET.len() * READABLE_CHARSET.len() * 1024; + (1..=bounds) + .map(|i| { + let idx = i * token_space / (bounds + 1); + Value::from(readable_token(idx)) + }) + .collect() +} + +/// Picks a representative string value for the target partition range. +pub fn generate_partition_value(bounds: &[Value], bound_idx: usize) -> Value { + let first = bounds.first().unwrap(); + let last = bounds.last().unwrap(); + let upper = match first { + Value::String(v) => v.as_utf8(), + _ => "", + }; + + if bound_idx == 0 { + if upper <= "0" { + Value::from("") + } else { + Value::from("0") + } + } else if bound_idx < bounds.len() { + bounds[bound_idx - 1].clone() + } else { + last.clone() + } +} + +/// Generates a unique readable bound not present in existing bounds. +pub fn generate_unique_partition_bound(rng: &mut R, bounds: &[Value]) -> Result { + let search_space = READABLE_CHARSET.len() * READABLE_CHARSET.len() * 1024; + let start = rng.random_range(0..search_space); + for offset in 0..search_space { + let idx = start + offset; + let candidate = Value::from(readable_token(idx)); + if !bounds.contains(&candidate) { + return Ok(candidate); + } + } + + error::UnexpectedSnafu { + violated: "unable to generate unique string partition bound".to_string(), + } + .fail() +} + +#[cfg(test)] +mod tests { + use rand::SeedableRng; + use rand_chacha::ChaCha8Rng; + + use super::*; + + #[test] + fn test_readable_token_grows_length() { + assert_eq!("0", readable_token(0)); + assert_eq!("9", readable_token(9)); + assert_eq!("A", readable_token(10)); + assert_eq!("z", readable_token(61)); + assert_eq!("00", readable_token(62)); + } + + #[test] + fn test_generate_partition_bounds_are_readable_and_unique() { + let bounds = generate_partition_bounds(8); + assert_eq!(8, bounds.len()); + + let mut values = bounds + .iter() + .map(|v| match v { + Value::String(s) => s.as_utf8().to_string(), + _ => panic!("expected string value"), + }) + .collect::>(); + let mut dedup = values.clone(); + dedup.sort(); + dedup.dedup(); + assert_eq!(values.len(), dedup.len()); + + for s in values.drain(..) { + assert!(s.chars().all(|c| c.is_ascii_alphanumeric())); + } + } + + #[test] + fn test_generate_partition_value_for_string_bounds() { + let bounds = vec![Value::from("A"), Value::from("M")]; + assert_eq!(Value::from("0"), generate_partition_value(&bounds, 0)); + assert_eq!(Value::from("A"), generate_partition_value(&bounds, 1)); + assert_eq!(Value::from("M"), generate_partition_value(&bounds, 2)); + } + + #[test] + fn test_generate_unique_partition_bound_not_in_existing() { + let mut rng = ChaCha8Rng::seed_from_u64(42); + let bounds = vec![Value::from("0"), Value::from("1"), Value::from("2")]; + let candidate = generate_unique_partition_bound(&mut rng, &bounds).unwrap(); + assert!(!bounds.contains(&candidate)); + match candidate { + Value::String(s) => { + assert!(!s.as_utf8().is_empty()); + assert!(s.as_utf8().chars().all(|c| c.is_ascii_alphanumeric())); + } + _ => panic!("expected string value"), + } + } +} diff --git a/tests-fuzz/src/translator.rs b/tests-fuzz/src/translator.rs index 673b543f2c0..4c5e0bb6a42 100644 --- a/tests-fuzz/src/translator.rs +++ b/tests-fuzz/src/translator.rs @@ -13,6 +13,8 @@ // limitations under the License. mod common; +/// Translator that converts insert expressions into CSV records. +pub mod csv; pub mod mysql; pub mod postgres; diff --git a/tests-fuzz/src/translator/csv.rs b/tests-fuzz/src/translator/csv.rs new file mode 100644 index 00000000000..e95956862c6 --- /dev/null +++ b/tests-fuzz/src/translator/csv.rs @@ -0,0 +1,121 @@ +// 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. + +use crate::error::Error; +use crate::ir::insert_expr::{InsertIntoExpr, RowValue}; +use crate::translator::DslTranslator; + +/// One CSV record converted from an insert row. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CsvRecord { + /// Cell values in column order. + pub values: Vec, +} + +/// CSV records converted from an insert expression. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CsvRecords { + /// Target table name from insert expression. + pub table_name: String, + /// Header values from insert columns. + pub headers: Vec, + /// Converted row records. + pub records: Vec, +} + +/// Translates `InsertIntoExpr` into CSV-writer-ready records. +pub struct InsertExprToCsvRecordsTranslator; + +impl DslTranslator for InsertExprToCsvRecordsTranslator { + type Error = Error; + + fn translate(&self, input: &InsertIntoExpr) -> Result { + let headers = input + .columns + .iter() + .map(|column| column.name.to_string()) + .collect::>(); + let records = input + .values_list + .iter() + .map(|row| CsvRecord { + values: row.iter().map(Self::format_row_value).collect(), + }) + .collect::>(); + + Ok(CsvRecords { + table_name: input.table_name.to_string(), + headers, + records, + }) + } +} + +impl InsertExprToCsvRecordsTranslator { + fn format_row_value(value: &RowValue) -> String { + match value { + RowValue::Value(datatypes::value::Value::Null) => String::new(), + RowValue::Value(v) => v.to_string(), + RowValue::Default => "DEFAULT".to_string(), + } + } +} + +#[cfg(test)] +mod tests { + use datatypes::data_type::ConcreteDataType; + + use super::InsertExprToCsvRecordsTranslator; + use crate::ir::create_expr::ColumnOption; + use crate::ir::insert_expr::{InsertIntoExpr, RowValue}; + use crate::ir::{Column, Ident}; + use crate::translator::DslTranslator; + + #[test] + fn test_translate_insert_expr_to_csv_records() { + let input = InsertIntoExpr { + table_name: Ident::new("metric_a"), + omit_column_list: false, + columns: vec![ + Column { + name: "host".into(), + column_type: ConcreteDataType::string_datatype(), + options: vec![ColumnOption::PrimaryKey], + }, + Column { + name: "value".into(), + column_type: ConcreteDataType::float64_datatype(), + options: vec![], + }, + ], + values_list: vec![ + vec![ + RowValue::Value(datatypes::value::Value::String("web-1".into())), + RowValue::Value(datatypes::value::Value::Int32(15)), + ], + vec![ + RowValue::Value(datatypes::value::Value::Null), + RowValue::Default, + ], + ], + }; + + let output = InsertExprToCsvRecordsTranslator.translate(&input).unwrap(); + assert_eq!(output.table_name, "metric_a"); + assert_eq!(output.headers, vec!["host", "value"]); + assert_eq!(output.records.len(), 2); + assert_eq!(output.records[0].values, vec!["web-1", "15"]); + assert_eq!(output.records[1].values, vec!["", "DEFAULT"]); + } +} diff --git a/tests-fuzz/src/utils.rs b/tests-fuzz/src/utils.rs index 0780f6c93d9..d55abab3c2f 100644 --- a/tests-fuzz/src/utils.rs +++ b/tests-fuzz/src/utils.rs @@ -15,6 +15,8 @@ pub mod cluster_info; pub mod config; pub mod crd; +/// CSV dump writer utilities for fuzz tests. +pub mod csv_dump_writer; pub mod health; pub mod migration; pub mod partition; @@ -22,10 +24,15 @@ pub mod pod_failure; pub mod procedure; #[cfg(feature = "unstable")] pub mod process; +pub mod retry; +/// SQL dump writer utilities for fuzz tests. +pub mod sql_dump_writer; pub mod wait; use std::env; +use std::str::FromStr; +use common_base::readable_size::ReadableSize; use common_telemetry::info; use common_telemetry::tracing::log::LevelFilter; use paste::paste; @@ -126,6 +133,14 @@ pub const GT_FUZZ_INPUT_MAX_COLUMNS: &str = "GT_FUZZ_INPUT_MAX_COLUMNS"; pub const GT_FUZZ_INPUT_MAX_ALTER_ACTIONS: &str = "GT_FUZZ_INPUT_MAX_ALTER_ACTIONS"; pub const GT_FUZZ_INPUT_MAX_INSERT_ACTIONS: &str = "GT_FUZZ_INPUT_MAX_INSERT_ACTIONS"; pub const FUZZ_OVERRIDE_PREFIX: &str = "GT_FUZZ_OVERRIDE_"; +/// Enables CSV dump generation for fuzz runs. +pub const GT_FUZZ_DUMP_TABLE_CSV: &str = "GT_FUZZ_DUMP_TABLE_CSV"; +/// Base directory for CSV dump sessions. +pub const GT_FUZZ_DUMP_DIR: &str = "GT_FUZZ_DUMP_DIR"; +/// Directory suffix used by one CSV dump session. +pub const GT_FUZZ_DUMP_SUFFIX: &str = "GT_FUZZ_DUMP_SUFFIX"; +/// Max in-memory CSV buffer size before auto flush. +pub const GT_FUZZ_DUMP_BUFFER_MAX_BYTES: &str = "GT_FUZZ_DUMP_BUFFER_MAX_BYTES"; /// Reads an override value for a fuzz parameter from env `GT_FUZZ_OVERRIDE_`. pub fn get_fuzz_override(name: &str) -> Option @@ -137,6 +152,33 @@ where env::var(&key).ok().and_then(|v| v.parse().ok()) } +/// Returns CSV dump base directory. +pub fn get_gt_fuzz_dump_dir() -> String { + let _ = dotenv::dotenv(); + env::var(GT_FUZZ_DUMP_DIR).unwrap_or_else(|_| "/tmp/greptime-fuzz-dumps".to_string()) +} + +/// Returns CSV dump directory suffix. +pub fn get_gt_fuzz_dump_suffix() -> String { + let _ = dotenv::dotenv(); + env::var(GT_FUZZ_DUMP_SUFFIX).unwrap_or_else(|_| ".repartition-metric-csv".to_string()) +} + +/// Returns max CSV in-memory buffer size. +pub fn get_gt_fuzz_dump_buffer_max_bytes() -> usize { + let _ = dotenv::dotenv(); + env::var(GT_FUZZ_DUMP_BUFFER_MAX_BYTES) + .ok() + .and_then(|value| { + value.parse::().ok().or_else(|| { + ReadableSize::from_str(&value) + .ok() + .map(|size| size.as_bytes() as usize) + }) + }) + .unwrap_or(8 * 1024 * 1024) +} + macro_rules! make_get_from_env_helper { ($key:expr, $default: expr) => { paste! { diff --git a/tests-fuzz/src/utils/csv_dump_writer.rs b/tests-fuzz/src/utils/csv_dump_writer.rs new file mode 100644 index 00000000000..de16a23c24d --- /dev/null +++ b/tests-fuzz/src/utils/csv_dump_writer.rs @@ -0,0 +1,383 @@ +// 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. + +use std::collections::{HashMap, HashSet}; +use std::fs::{File, OpenOptions, create_dir_all, remove_dir_all}; +use std::io::Write; +use std::path::{Path, PathBuf}; + +use common_telemetry::{info, warn}; +use common_time::util::current_time_millis; +use snafu::ResultExt; + +use crate::error::{self, Result}; +use crate::translator::csv::CsvRecords; +use crate::utils::{ + get_gt_fuzz_dump_buffer_max_bytes, get_gt_fuzz_dump_dir, get_gt_fuzz_dump_suffix, +}; + +/// Metadata for one CSV dump session. +#[derive(Debug, Clone)] +pub struct CsvDumpMetadata { + /// Fuzz target name. + pub target: String, + /// Seed used by current fuzz input. + pub seed: u64, + /// Repartition action count. + pub actions: usize, + /// Initial partition count. + pub partitions: usize, + /// Logical table count. + pub tables: usize, + /// Session start time in unix milliseconds. + pub started_at_unix_ms: i64, +} + +impl CsvDumpMetadata { + /// Builds dump metadata with current timestamp. + pub fn new( + target: impl Into, + seed: u64, + actions: usize, + partitions: usize, + tables: usize, + ) -> Self { + Self { + target: target.into(), + seed, + actions, + partitions, + tables, + started_at_unix_ms: current_time_millis(), + } + } +} + +/// Session writer for staged CSV dump records. +#[derive(Debug)] +pub struct CsvDumpSession { + /// Session metadata. + pub metadata: CsvDumpMetadata, + /// Session directory path. + pub run_dir: PathBuf, + /// Max in-memory buffer size before auto flush. + pub max_buffer_bytes: usize, + records: Vec, + buffered_bytes: usize, + written_tables: HashSet, + full_headers_by_table: HashMap>, +} + +impl CsvDumpSession { + /// Creates session directory and writes seed metadata file. + pub fn new(metadata: CsvDumpMetadata) -> Result { + Self::new_with_buffer_limit(metadata, get_gt_fuzz_dump_buffer_max_bytes()) + } + + /// Creates session with a custom in-memory buffer limit. + pub fn new_with_buffer_limit( + metadata: CsvDumpMetadata, + max_buffer_bytes: usize, + ) -> Result { + let run_dir = build_run_dir(&metadata); + create_dir_all(&run_dir).context(error::CreateFileSnafu { + path: run_dir.to_string_lossy().to_string(), + })?; + write_seed_meta(&run_dir, &metadata)?; + info!( + "Create csv dump session, target: {}, run_dir: {}, max_buffer_bytes: {}", + metadata.target, + run_dir.display(), + max_buffer_bytes + ); + + Ok(Self { + metadata, + run_dir, + max_buffer_bytes, + records: Vec::new(), + buffered_bytes: 0, + written_tables: HashSet::new(), + full_headers_by_table: HashMap::new(), + }) + } + + /// Appends one table CSV records batch with full table headers. + pub fn append(&mut self, records: CsvRecords, full_headers: Vec) -> Result<()> { + self.full_headers_by_table + .entry(records.table_name.clone()) + .or_insert(full_headers); + self.buffered_bytes += estimate_csv_records_size(&records); + self.records.push(records); + if self.buffered_bytes >= self.max_buffer_bytes { + self.flush_buffered_records()?; + } + Ok(()) + } + + /// Flushes all appended batches to CSV files. + pub fn flush_all(&mut self) -> Result<()> { + self.flush_buffered_records() + } + + /// Removes session directory after successful validation. + pub fn cleanup_on_success(&self) -> std::io::Result<()> { + match remove_dir_all(&self.run_dir) { + Ok(_) => { + info!( + "Cleanup csv dump directory on success: {}", + self.run_dir.display() + ); + Ok(()) + } + Err(err) => { + warn!( + "Cleanup csv dump directory failed: {}, error: {:?}", + self.run_dir.display(), + err + ); + Err(err) + } + } + } + + fn flush_buffered_records(&mut self) -> Result<()> { + if self.records.is_empty() { + return Ok(()); + } + for batch in &self.records { + write_batch_csv( + &self.run_dir, + batch, + &mut self.written_tables, + &self.full_headers_by_table, + )?; + } + self.records.clear(); + self.buffered_bytes = 0; + Ok(()) + } +} + +fn write_seed_meta(run_dir: &Path, metadata: &CsvDumpMetadata) -> Result<()> { + let path = run_dir.join("seed.meta"); + let mut file = File::create(&path).context(error::CreateFileSnafu { + path: path.to_string_lossy().to_string(), + })?; + + let content = format!( + "target={}\nseed={}\nactions={}\npartitions={}\ntables={}\nstarted_at_unix_ms={}\n", + metadata.target, + metadata.seed, + metadata.actions, + metadata.partitions, + metadata.tables, + metadata.started_at_unix_ms, + ); + file.write_all(content.as_bytes()) + .context(error::WriteFileSnafu { + path: path.to_string_lossy().to_string(), + }) +} + +fn write_batch_csv( + run_dir: &Path, + batch: &CsvRecords, + written_tables: &mut HashSet, + full_headers_by_table: &HashMap>, +) -> Result<()> { + let output_headers = full_headers_by_table + .get(&batch.table_name) + .cloned() + .unwrap_or_else(|| batch.headers.clone()); + let file_name = format!("{}.table-data.csv", sanitize_file_name(&batch.table_name)); + let path = run_dir.join(file_name); + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(&path) + .context(error::CreateFileSnafu { + path: path.to_string_lossy().to_string(), + })?; + + if written_tables.insert(batch.table_name.clone()) { + file.write_all(join_line(&output_headers).as_bytes()) + .context(error::WriteFileSnafu { + path: path.to_string_lossy().to_string(), + })?; + file.write_all(b"\n").context(error::WriteFileSnafu { + path: path.to_string_lossy().to_string(), + })?; + } + + let header_index = batch + .headers + .iter() + .enumerate() + .map(|(idx, header)| (header.as_str(), idx)) + .collect::>(); + + for record in &batch.records { + let aligned_values = output_headers + .iter() + .map(|header| { + header_index + .get(header.as_str()) + .and_then(|idx| record.values.get(*idx)) + .cloned() + .unwrap_or_default() + }) + .collect::>(); + file.write_all(join_line(&aligned_values).as_bytes()) + .context(error::WriteFileSnafu { + path: path.to_string_lossy().to_string(), + })?; + file.write_all(b"\n").context(error::WriteFileSnafu { + path: path.to_string_lossy().to_string(), + })?; + } + + Ok(()) +} + +fn estimate_csv_records_size(records: &CsvRecords) -> usize { + let headers = records.headers.iter().map(String::len).sum::(); + let rows = records + .records + .iter() + .flat_map(|record| record.values.iter()) + .map(String::len) + .sum::(); + headers + rows +} + +fn join_line(cells: &[String]) -> String { + cells + .iter() + .map(|cell| escape_csv_cell(cell)) + .collect::>() + .join(",") +} + +fn escape_csv_cell(value: &str) -> String { + if value.contains([',', '"', '\n', '\r']) { + format!("\"{}\"", value.replace('"', "\"\"")) + } else { + value.to_string() + } +} + +fn sanitize_file_name(raw: &str) -> String { + raw.chars() + .map(|ch| { + if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' { + ch + } else { + '_' + } + }) + .collect() +} + +fn build_run_dir(metadata: &CsvDumpMetadata) -> PathBuf { + let base = PathBuf::from(get_gt_fuzz_dump_dir()); + let suffix = get_gt_fuzz_dump_suffix(); + let name = format!( + "{}_seed_{}_actions_{}_ts_{}{}", + metadata.target, metadata.seed, metadata.actions, metadata.started_at_unix_ms, suffix + ); + base.join(name) +} + +#[cfg(test)] +mod tests { + use super::{CsvDumpMetadata, CsvDumpSession}; + use crate::translator::csv::{CsvRecord, CsvRecords}; + + #[test] + fn test_create_session_and_flush() { + let mut session = CsvDumpSession::new_with_buffer_limit( + CsvDumpMetadata::new("fuzz_case", 1, 2, 3, 4), + 1024, + ) + .unwrap(); + session + .append( + CsvRecords { + table_name: "metric-a".to_string(), + headers: vec!["host".to_string(), "value".to_string()], + records: vec![CsvRecord { + values: vec!["web-1".to_string(), "10".to_string()], + }], + }, + vec!["host".to_string(), "value".to_string()], + ) + .unwrap(); + session.flush_all().unwrap(); + + assert!(session.run_dir.exists()); + assert!(session.run_dir.join("seed.meta").exists()); + assert!(session.run_dir.join("metric-a.table-data.csv").exists()); + } + + #[test] + fn test_auto_flush_on_buffer_limit() { + let mut session = + CsvDumpSession::new_with_buffer_limit(CsvDumpMetadata::new("fuzz_case", 5, 2, 3, 4), 1) + .unwrap(); + session + .append( + CsvRecords { + table_name: "metric-b".to_string(), + headers: vec!["host".to_string()], + records: vec![CsvRecord { + values: vec!["web-2".to_string()], + }], + }, + vec!["host".to_string()], + ) + .unwrap(); + + assert!(session.run_dir.join("metric-b.table-data.csv").exists()); + assert_eq!(session.buffered_bytes, 0); + } + + #[test] + fn test_flush_with_partial_headers_uses_full_headers() { + let mut session = CsvDumpSession::new_with_buffer_limit( + CsvDumpMetadata::new("fuzz_case", 7, 2, 3, 4), + 1024, + ) + .unwrap(); + session + .append( + CsvRecords { + table_name: "metric-c".to_string(), + headers: vec!["host".to_string(), "value".to_string()], + records: vec![CsvRecord { + values: vec!["web-3".to_string(), "12".to_string()], + }], + }, + vec!["host".to_string(), "idc".to_string(), "value".to_string()], + ) + .unwrap(); + session.flush_all().unwrap(); + + let file = + std::fs::read_to_string(session.run_dir.join("metric-c.table-data.csv")).unwrap(); + let mut lines = file.lines(); + assert_eq!(lines.next().unwrap(), "host,idc,value"); + assert_eq!(lines.next().unwrap(), "web-3,,12"); + } +} diff --git a/tests-fuzz/src/utils/retry.rs b/tests-fuzz/src/utils/retry.rs new file mode 100644 index 00000000000..06d1ede54fa --- /dev/null +++ b/tests-fuzz/src/utils/retry.rs @@ -0,0 +1,49 @@ +// 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. + +use std::future::Future; +use std::time::Duration; + +use common_telemetry::warn; + +pub async fn retry_with_backoff( + mut operation: F, + max_attempts: usize, + init_backoff: Duration, + max_backoff: Duration, +) -> Result +where + F: FnMut() -> Fut, + Fut: Future>, + E: std::fmt::Debug, +{ + let mut backoff = init_backoff; + for attempt in 0..max_attempts { + match operation().await { + Ok(result) => return Ok(result), + Err(err) if attempt + 1 == max_attempts => return Err(err), + Err(err) => { + let current_attempt = attempt + 1; + warn!( + "Retryable operation failed, attempt: {}, max_attempts: {}, backoff: {:?}, error: {:?}", + current_attempt, max_attempts, backoff, err + ); + tokio::time::sleep(backoff).await; + backoff = std::cmp::min(backoff * 2, max_backoff); + } + } + } + + panic!("retry loop should always return") +} diff --git a/tests-fuzz/src/utils/sql_dump_writer.rs b/tests-fuzz/src/utils/sql_dump_writer.rs new file mode 100644 index 00000000000..6f098d95846 --- /dev/null +++ b/tests-fuzz/src/utils/sql_dump_writer.rs @@ -0,0 +1,267 @@ +// 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. + +use std::collections::HashMap; +use std::fs::{OpenOptions, create_dir_all}; +use std::io::Write; +use std::path::PathBuf; + +use snafu::ResultExt; + +use crate::error::{self, Result}; +use crate::utils::get_gt_fuzz_dump_buffer_max_bytes; + +/// Session writer for table-scoped SQL trace files. +#[derive(Debug)] +pub struct SqlDumpSession { + /// Session directory path. + pub run_dir: PathBuf, + /// Max in-memory buffer size before auto flush. + pub max_buffer_bytes: usize, + buffered_bytes: usize, + entries_by_table: HashMap>, +} + +impl SqlDumpSession { + /// Creates SQL dump session with default buffer limit. + pub fn new(run_dir: PathBuf) -> Result { + Self::new_with_buffer_limit(run_dir, get_gt_fuzz_dump_buffer_max_bytes()) + } + + /// Creates SQL dump session with custom buffer limit. + pub fn new_with_buffer_limit(run_dir: PathBuf, max_buffer_bytes: usize) -> Result { + create_dir_all(&run_dir).context(error::CreateFileSnafu { + path: run_dir.to_string_lossy().to_string(), + })?; + + Ok(Self { + run_dir, + max_buffer_bytes, + buffered_bytes: 0, + entries_by_table: HashMap::new(), + }) + } + + /// Appends one SQL statement for a logical table. + pub fn append_sql(&mut self, table: &str, sql: &str, comment: Option<&str>) -> Result<()> { + let entry = format_sql_entry(sql, comment); + self.push_entry(table, entry)?; + Ok(()) + } + + /// Broadcasts one comment event to all table trace files. + pub fn broadcast_event(&mut self, tables: I, event: &str, sql: &str) -> Result<()> + where + I: IntoIterator, + T: AsRef, + { + let entry = format_sql_entry(sql, Some(event)); + for table in tables { + self.push_entry(table.as_ref(), entry.clone())?; + } + Ok(()) + } + + /// Flushes all staged SQL traces to table-scoped files. + pub fn flush_all(&mut self) -> Result<()> { + self.flush_buffered_entries() + } + + fn push_entry(&mut self, table: &str, entry: String) -> Result<()> { + self.buffered_bytes += entry.len(); + self.entries_by_table + .entry(table.to_string()) + .or_default() + .push(entry); + + if self.buffered_bytes >= self.max_buffer_bytes { + self.flush_buffered_entries()?; + } + Ok(()) + } + + fn flush_buffered_entries(&mut self) -> Result<()> { + if self.entries_by_table.is_empty() { + return Ok(()); + } + + for (table, entries) in &self.entries_by_table { + let path = self + .run_dir + .join(format!("{}.trace.sql", sanitize_file_name(table))); + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(&path) + .context(error::CreateFileSnafu { + path: path.to_string_lossy().to_string(), + })?; + + for entry in entries { + file.write_all(entry.as_bytes()) + .context(error::WriteFileSnafu { + path: path.to_string_lossy().to_string(), + })?; + file.write_all(b"\n").context(error::WriteFileSnafu { + path: path.to_string_lossy().to_string(), + })?; + } + } + + self.entries_by_table.clear(); + self.buffered_bytes = 0; + Ok(()) + } +} + +fn format_sql_entry(sql: &str, comment: Option<&str>) -> String { + let normalized_sql = normalize_sql(sql); + if let Some(comment) = comment { + format!("{}\n{normalized_sql}", format_comment(comment)) + } else { + normalized_sql + } +} + +fn format_comment(comment: &str) -> String { + comment + .lines() + .map(|line| format!("-- {line}")) + .collect::>() + .join("\n") +} + +fn normalize_sql(sql: &str) -> String { + let trimmed = sql.trim_end(); + if trimmed.ends_with(';') { + trimmed.to_string() + } else { + format!("{trimmed};") + } +} + +fn sanitize_file_name(raw: &str) -> String { + raw.chars() + .map(|ch| { + if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' { + ch + } else { + '_' + } + }) + .collect() +} + +#[cfg(test)] +mod tests { + use std::time::{SystemTime, UNIX_EPOCH}; + + use super::SqlDumpSession; + + #[test] + fn test_append_sql_writes_table_trace_file() { + let run_dir = std::env::temp_dir().join(format!( + "tests-fuzz-sql-dump-{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() + )); + + let mut session = SqlDumpSession::new_with_buffer_limit(run_dir.clone(), 1024).unwrap(); + session + .append_sql( + "metric-a", + "INSERT INTO t VALUES(1)", + Some("kind=insert elapsed_ms=10"), + ) + .unwrap(); + session.flush_all().unwrap(); + + let content = std::fs::read_to_string(run_dir.join("metric-a.trace.sql")).unwrap(); + assert!(content.contains("-- kind=insert elapsed_ms=10")); + assert!(content.contains("INSERT INTO t VALUES(1);")); + } + + #[test] + fn test_broadcast_event_writes_to_all_tables() { + let run_dir = std::env::temp_dir().join(format!( + "tests-fuzz-sql-broadcast-{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() + )); + + let mut session = SqlDumpSession::new_with_buffer_limit(run_dir.clone(), 1024).unwrap(); + session + .broadcast_event( + ["metric-a", "metric-b"], + "repartition action_idx=3", + "ALTER TABLE t REPARTITION", + ) + .unwrap(); + session.flush_all().unwrap(); + + let content_a = std::fs::read_to_string(run_dir.join("metric-a.trace.sql")).unwrap(); + let content_b = std::fs::read_to_string(run_dir.join("metric-b.trace.sql")).unwrap(); + assert!(content_a.contains("-- repartition action_idx=3")); + assert!(content_a.contains("ALTER TABLE t REPARTITION;")); + assert!(content_b.contains("-- repartition action_idx=3")); + assert!(content_b.contains("ALTER TABLE t REPARTITION;")); + } + + #[test] + fn test_multiline_comment_is_prefixed_per_line() { + let run_dir = std::env::temp_dir().join(format!( + "tests-fuzz-sql-dump-comment-{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() + )); + + let mut session = SqlDumpSession::new_with_buffer_limit(run_dir.clone(), 1024).unwrap(); + session + .append_sql( + "metric-a", + "INSERT INTO t VALUES(1)", + Some("kind=insert\nstarted_at_ms=1 elapsed_ms=2"), + ) + .unwrap(); + session.flush_all().unwrap(); + + let content = std::fs::read_to_string(run_dir.join("metric-a.trace.sql")).unwrap(); + assert!(content.contains("-- kind=insert\n-- started_at_ms=1 elapsed_ms=2")); + } + + #[test] + fn test_auto_flush_on_buffer_limit() { + let run_dir = std::env::temp_dir().join(format!( + "tests-fuzz-sql-dump-limit-{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() + )); + + let mut session = SqlDumpSession::new_with_buffer_limit(run_dir.clone(), 1).unwrap(); + session + .append_sql("metric-a", "INSERT INTO t VALUES(1)", None) + .unwrap(); + + assert!(run_dir.join("metric-a.trace.sql").exists()); + assert_eq!(session.buffered_bytes, 0); + } +} diff --git a/tests-fuzz/targets/ddl/fuzz_repartition_metric_table.rs b/tests-fuzz/targets/ddl/fuzz_repartition_metric_table.rs new file mode 100644 index 00000000000..7932bc7759c --- /dev/null +++ b/tests-fuzz/targets/ddl/fuzz_repartition_metric_table.rs @@ -0,0 +1,684 @@ +// 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. + +#![no_main] + +use std::collections::{BTreeMap, HashMap}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use arbitrary::{Arbitrary, Unstructured}; +use common_telemetry::{info, warn}; +use common_time::Timestamp; +use common_time::util::current_time_millis; +use libfuzzer_sys::fuzz_target; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaChaRng; +use snafu::{ResultExt, ensure}; +use sqlx::{MySql, Pool}; +use tests_fuzz::context::{TableContext, TableContextRef}; +use tests_fuzz::error::{self, Result}; +use tests_fuzz::fake::{ + ConstGenerator, MappedGenerator, WordGenerator, merge_two_word_map_fn, random_capitalize_map, + uppercase_and_keyword_backtick_map, +}; +use tests_fuzz::generator::Generator; +use tests_fuzz::generator::create_expr::{ + CreateLogicalTableExprGeneratorBuilder, CreatePhysicalTableExprGeneratorBuilder, +}; +use tests_fuzz::generator::insert_expr::InsertExprGeneratorBuilder; +use tests_fuzz::generator::repartition_expr::{ + MergePartitionExprGeneratorBuilder, SplitPartitionExprGeneratorBuilder, +}; +use tests_fuzz::ir::{ + CreateTableExpr, Ident, InsertIntoExpr, RepartitionExpr, generate_random_value, + generate_unique_timestamp_for_mysql_with_clock, +}; +use tests_fuzz::translator::DslTranslator; +use tests_fuzz::translator::csv::InsertExprToCsvRecordsTranslator; +use tests_fuzz::translator::mysql::create_expr::CreateTableExprTranslator; +use tests_fuzz::translator::mysql::insert_expr::InsertIntoExprTranslator; +use tests_fuzz::translator::mysql::repartition_expr::RepartitionExprTranslator; +use tests_fuzz::utils::csv_dump_writer::{CsvDumpMetadata, CsvDumpSession}; +use tests_fuzz::utils::retry::retry_with_backoff; +use tests_fuzz::utils::sql_dump_writer::SqlDumpSession; +use tests_fuzz::utils::{ + Connections, get_fuzz_override, get_gt_fuzz_input_max_alter_actions, + get_gt_fuzz_input_max_tables, init_greptime_connections_via_env, +}; +use tests_fuzz::validator::row::count_values; +use tokio::sync::{mpsc, oneshot}; + +const BARRIER_ACK_TIMEOUT_SECS: u64 = 10; +const VALIDATE_QUERY_MAX_ATTEMPTS: usize = 6; +const VALIDATE_QUERY_INIT_BACKOFF: Duration = Duration::from_millis(50); +const VALIDATE_QUERY_MAX_BACKOFF: Duration = Duration::from_millis(800); + +#[derive(Clone)] +struct FuzzContext { + greptime: Pool, +} + +impl FuzzContext { + async fn close(self) { + self.greptime.close().await; + } +} + +#[derive(Clone, Debug)] +struct FuzzInput { + seed: u64, + actions: usize, + partitions: usize, + tables: usize, +} + +fn generate_create_physical_table_expr( + partitions: usize, + rng: &mut R, +) -> Result { + CreatePhysicalTableExprGeneratorBuilder::default() + .name_generator(Box::new(ConstGenerator::new(Ident::new( + "fuzz_repartition_metric_physical", + )))) + .if_not_exists(rng.random_bool(0.5)) + .partition(partitions) + .build() + .unwrap() + .generate(rng) +} + +fn generate_create_logical_table_expr( + physical_table_ctx: TableContextRef, + include_partition_column: bool, + rng: &mut R, +) -> Result { + CreateLogicalTableExprGeneratorBuilder::default() + .name_generator(Box::new(MappedGenerator::new( + WordGenerator, + merge_two_word_map_fn(random_capitalize_map, uppercase_and_keyword_backtick_map), + ))) + .physical_table_ctx(physical_table_ctx) + .labels(rng.random_range(1..=5)) + .if_not_exists(rng.random_bool(0.5)) + .include_partition_column(include_partition_column) + .build() + .unwrap() + .generate(rng) +} + +fn generate_insert_expr( + rows: usize, + rng: &mut R, + table_ctx: TableContextRef, + clock: Arc>, +) -> Result { + let ts_value_generator = generate_unique_timestamp_for_mysql_with_clock(clock); + InsertExprGeneratorBuilder::default() + .omit_column_list(false) + .table_ctx(table_ctx) + .rows(rows) + .value_generator(Box::new(generate_random_value)) + .ts_value_generator(ts_value_generator) + .build() + .unwrap() + .generate(rng) +} + +async fn create_metric_tables( + ctx: &FuzzContext, + rng: &mut R, + partitions: usize, + table_count: usize, +) -> Result<( + TableContextRef, + BTreeMap, + HashMap, + String, +)> { + let create_physical_expr = generate_create_physical_table_expr(partitions, rng)?; + let translator = CreateTableExprTranslator; + let create_physical_sql = translator.translate(&create_physical_expr)?; + let result = sqlx::query(&create_physical_sql) + .execute(&ctx.greptime) + .await + .context(error::ExecuteQuerySnafu { + sql: &create_physical_sql, + })?; + info!("Create physical table: {create_physical_sql}, result: {result:?}"); + let physical_table_ctx = Arc::new(TableContext::from(&create_physical_expr)); + ensure!( + physical_table_ctx.partition.is_some(), + error::AssertSnafu { + reason: "Physical metric table must have partition".to_string() + } + ); + + let mut logical_tables = BTreeMap::new(); + let mut create_logical_sqls = HashMap::new(); + let max_attempts = table_count * 3; + for _ in 0..max_attempts { + if logical_tables.len() >= table_count { + break; + } + + let include_partition_column = rng.random_bool(0.5); + let create_logical_expr = generate_create_logical_table_expr( + physical_table_ctx.clone(), + include_partition_column, + rng, + )?; + if logical_tables.contains_key(&create_logical_expr.table_name) { + continue; + } + + let create_logical_sql = translator.translate(&create_logical_expr)?; + let result = sqlx::query(&create_logical_sql) + .execute(&ctx.greptime) + .await + .context(error::ExecuteQuerySnafu { + sql: &create_logical_sql, + })?; + info!("Create logical table: {create_logical_sql}, result: {result:?}"); + let logical_ctx = Arc::new(TableContext::from(&create_logical_expr)); + create_logical_sqls.insert(logical_ctx.name.to_string(), create_logical_sql); + logical_tables.insert(logical_ctx.name.clone(), logical_ctx); + } + + ensure!( + !logical_tables.is_empty(), + error::AssertSnafu { + reason: "No logical table created".to_string() + } + ); + + Ok(( + physical_table_ctx, + logical_tables, + create_logical_sqls, + create_physical_sql, + )) +} + +async fn execute_insert_with_retry(ctx: &FuzzContext, sql: &str) -> Result<()> { + let mut delay = Duration::from_millis(100); + let mut attempt = 0; + let max_attempts = 10; + loop { + match sqlx::query(sql) + .persistent(false) + .execute(&ctx.greptime) + .await + { + Ok(_) => return Ok(()), + Err(err) => { + tokio::time::sleep(delay).await; + delay = std::cmp::min(delay * 2, Duration::from_secs(1)); + attempt += 1; + warn!("Execute insert with retry: {sql}, attempt: {attempt}, error: {err:?}"); + if attempt >= max_attempts { + return Err(err).context(error::ExecuteQuerySnafu { sql }); + } + } + } + } +} + +struct SharedState { + clock: Arc>, + inserted_rows: HashMap, + csv_dump_session: Option, + sql_dump_session: Option, + running: bool, +} + +enum WriterControl { + Barrier { + epoch: usize, + ack: oneshot::Sender<()>, + }, + Resume { + epoch: usize, + }, + Stop, +} + +fn handle_writer_control(control: WriterControl, paused: &mut bool) -> bool { + match control { + WriterControl::Barrier { epoch, ack } => { + info!("Writer received barrier control, epoch: {epoch}"); + *paused = true; + let _ = ack.send(()); + false + } + WriterControl::Resume { epoch } => { + info!("Writer received resume control, epoch: {epoch}"); + *paused = false; + false + } + WriterControl::Stop => { + info!("Writer received stop control"); + true + } + } +} + +async fn write_loop( + mut rng: R, + ctx: FuzzContext, + logical_tables: BTreeMap, + shared_state: Arc>, + mut control_rx: mpsc::UnboundedReceiver, +) -> Result<()> { + info!("Start write loop"); + let mut paused = false; + loop { + while let Ok(control) = control_rx.try_recv() { + if handle_writer_control(control, &mut paused) { + return Ok(()); + } + } + + if paused { + match control_rx.recv().await { + Some(control) => { + if handle_writer_control(control, &mut paused) { + return Ok(()); + } + } + None => return Ok(()), + } + continue; + } + + let (running, clock) = { + let state = shared_state.lock().unwrap(); + (state.running, state.clock.clone()) + }; + if !running { + break; + } + + for table_ctx in logical_tables.values() { + let rows = rng.random_range(1..=3); + let insert_expr = + generate_insert_expr(rows, &mut rng, table_ctx.clone(), clock.clone())?; + let translator = InsertIntoExprTranslator; + let sql = translator.translate(&insert_expr)?; + let inserted = insert_expr.values_list.len() as u64; + let csv_records = InsertExprToCsvRecordsTranslator.translate(&insert_expr)?; + let table_name = table_ctx.name.to_string(); + let full_headers = table_ctx + .columns + .iter() + .map(|column| column.name.value.clone()) + .collect::>(); + + let started_at_ms = current_time_millis(); + let now = Instant::now(); + execute_insert_with_retry(&ctx, &sql).await?; + let elapsed = now.elapsed(); + info!("Execute insert sql: {sql}, elapsed: {elapsed:?}"); + + let mut state = shared_state.lock().unwrap(); + if let Some(csv_dump_session) = state.csv_dump_session.as_mut() { + csv_dump_session.append(csv_records, full_headers)?; + } + if let Some(sql_dump_session) = state.sql_dump_session.as_mut() { + let comment = format!( + "kind=insert table={} started_at_ms={} elapsed_ms={}", + table_name, + started_at_ms, + elapsed.as_millis() + ); + sql_dump_session.append_sql(&table_name, &sql, Some(&comment))?; + } + *state.inserted_rows.entry(table_name).or_insert(0) += inserted; + } + + tokio::time::sleep(Duration::from_millis(100)).await; + } + info!("Write loop ended"); + + Ok(()) +} + +async fn validate_rows( + ctx: &FuzzContext, + logical_tables: &BTreeMap, + inserted_rows: &HashMap, +) -> Result<()> { + for table_ctx in logical_tables.values() { + let expected = *inserted_rows.get(&table_ctx.name.to_string()).unwrap_or(&0) as usize; + let count_sql = format!("SELECT COUNT(1) AS count FROM {}", table_ctx.name); + let count = retry_with_backoff( + || count_values(&ctx.greptime, &count_sql), + VALIDATE_QUERY_MAX_ATTEMPTS, + VALIDATE_QUERY_INIT_BACKOFF, + VALIDATE_QUERY_MAX_BACKOFF, + ) + .await?; + let distinct_count_sql = format!( + "SELECT COUNT(DISTINCT {}) AS count FROM {}", + table_ctx.timestamp_column().unwrap().name, + table_ctx.name + ); + let distinct_count = retry_with_backoff( + || count_values(&ctx.greptime, &distinct_count_sql), + VALIDATE_QUERY_MAX_ATTEMPTS, + VALIDATE_QUERY_INIT_BACKOFF, + VALIDATE_QUERY_MAX_BACKOFF, + ) + .await?; + info!( + "Validate rows for table: {}, expected: {}, count: {}, distinct_count: {}", + table_ctx.name, expected, count.count as usize, distinct_count.count as usize + ); + assert_eq!(count.count as usize, expected); + + assert_eq!(distinct_count.count as usize, expected); + } + Ok(()) +} + +fn flush_dump_sessions_and_snapshot( + shared_state: &Arc>, +) -> Result> { + let mut state = shared_state.lock().unwrap(); + if let Some(csv_dump_session) = state.csv_dump_session.as_mut() { + csv_dump_session.flush_all()?; + } + if let Some(sql_dump_session) = state.sql_dump_session.as_mut() { + sql_dump_session.flush_all()?; + } + Ok(state.inserted_rows.clone()) +} + +async fn cleanup_tables( + ctx: &FuzzContext, + physical_table_ctx: &TableContextRef, + logical_tables: &BTreeMap, +) -> Result<()> { + for table_ctx in logical_tables.values() { + let drop_logical_sql = format!("DROP TABLE {}", table_ctx.name); + let result = sqlx::query(&drop_logical_sql) + .execute(&ctx.greptime) + .await + .context(error::ExecuteQuerySnafu { + sql: &drop_logical_sql, + })?; + info!("Drop logical table: {drop_logical_sql}, result: {result:?}"); + } + + let drop_physical_sql = format!("DROP TABLE {}", physical_table_ctx.name); + let result = sqlx::query(&drop_physical_sql) + .execute(&ctx.greptime) + .await + .context(error::ExecuteQuerySnafu { + sql: &drop_physical_sql, + })?; + info!("Drop physical table: {drop_physical_sql}, result: {result:?}"); + Ok(()) +} + +fn repartition_operation( + table_ctx: &TableContextRef, + rng: &mut R, +) -> Result { + let split = rng.random_bool(0.5); + if table_ctx.partition.as_ref().unwrap().exprs.len() <= 2 || split { + let expr = SplitPartitionExprGeneratorBuilder::default() + .table_ctx(table_ctx.clone()) + .build() + .unwrap() + .generate(rng)?; + Ok(RepartitionExpr::Split(expr)) + } else { + let expr = MergePartitionExprGeneratorBuilder::default() + .table_ctx(table_ctx.clone()) + .build() + .unwrap() + .generate(rng)?; + Ok(RepartitionExpr::Merge(expr)) + } +} + +impl Arbitrary<'_> for FuzzInput { + fn arbitrary(u: &mut Unstructured<'_>) -> arbitrary::Result { + let seed = get_fuzz_override::("SEED").unwrap_or(u.int_in_range(u64::MIN..=u64::MAX)?); + let mut rng = ChaChaRng::seed_from_u64(seed); + let partitions = + get_fuzz_override::("PARTITIONS").unwrap_or_else(|| rng.random_range(2..8)); + let max_tables = get_gt_fuzz_input_max_tables(); + let tables = get_fuzz_override::("TABLES") + .unwrap_or_else(|| rng.random_range(1..=std::cmp::max(1, max_tables))); + let max_actions = get_gt_fuzz_input_max_alter_actions(); + let actions = get_fuzz_override::("ACTIONS") + .unwrap_or_else(|| rng.random_range(1..max_actions)); + + Ok(FuzzInput { + seed, + actions, + partitions, + tables, + }) + } +} + +async fn execute_repartition_metric_table(ctx: FuzzContext, input: FuzzInput) -> Result<()> { + info!("input: {input:?}"); + let mut rng = ChaChaRng::seed_from_u64(input.seed); + let clock = Arc::new(Mutex::new(Timestamp::current_millis())); + + let (mut physical_table_ctx, logical_tables, create_logical_sqls, create_physical_sql) = + create_metric_tables(&ctx, &mut rng, input.partitions, input.tables).await?; + + let mut inserted_rows = HashMap::with_capacity(logical_tables.len()); + for table_ctx in logical_tables.values() { + inserted_rows.insert(table_ctx.name.to_string(), 0); + } + let csv_dump_session = CsvDumpSession::new(CsvDumpMetadata::new( + "fuzz_repartition_metric_table", + input.seed, + input.actions, + input.partitions, + input.tables, + ))?; + let sql_dump_session = SqlDumpSession::new(csv_dump_session.run_dir.clone())?; + let logical_table_names = logical_tables + .values() + .map(|table_ctx| table_ctx.name.to_string()) + .collect::>(); + + let mut sql_dump_session = sql_dump_session; + sql_dump_session.append_sql( + &physical_table_ctx.name.to_string(), + &create_physical_sql, + Some("kind=create_physical_table"), + )?; + for table_name in &logical_table_names { + if let Some(create_sql) = create_logical_sqls.get(table_name) { + sql_dump_session.append_sql( + table_name, + create_sql, + Some("kind=create_logical_table"), + )?; + } + } + + let shared_state = Arc::new(Mutex::new(SharedState { + clock, + inserted_rows, + csv_dump_session: Some(csv_dump_session), + sql_dump_session: Some(sql_dump_session), + running: true, + })); + let writer_rng = ChaChaRng::seed_from_u64(input.seed ^ 0xA5A5_A5A5_A5A5_A5A5); + let (control_tx, control_rx) = mpsc::unbounded_channel::(); + let writer_task = tokio::spawn(write_loop( + writer_rng, + ctx.clone(), + logical_tables.clone(), + shared_state.clone(), + control_rx, + )); + tokio::time::sleep(Duration::from_millis(100)).await; + + for i in 0..input.actions { + let partition_num = physical_table_ctx.partition.as_ref().unwrap().exprs.len(); + info!( + "partition_num: {partition_num}, action: {}/{}, table: {}, logical table num: {}", + i + 1, + input.actions, + physical_table_ctx.name, + logical_tables.len() + ); + + let repartition_expr = repartition_operation(&physical_table_ctx, &mut rng)?; + let translator = RepartitionExprTranslator; + let sql = translator.translate(&repartition_expr)?; + info!("Repartition sql: {sql}"); + let started_at_ms = current_time_millis(); + let now = Instant::now(); + let result = sqlx::query(&sql) + .execute(&ctx.greptime) + .await + .context(error::ExecuteQuerySnafu { sql: &sql })?; + let elapsed = now.elapsed(); + info!("Repartition result: {result:?}, elapsed: {elapsed:?}"); + + physical_table_ctx = Arc::new( + Arc::unwrap_or_clone(physical_table_ctx) + .repartition(repartition_expr) + .unwrap(), + ); + + let partition_entries = tests_fuzz::validator::partition::fetch_partitions_info_schema( + &ctx.greptime, + "public".into(), + &physical_table_ctx.name, + ) + .await?; + tests_fuzz::validator::partition::assert_partitions( + physical_table_ctx.partition.as_ref().unwrap(), + &partition_entries, + )?; + + { + let mut state = shared_state.lock().unwrap(); + if let Some(sql_dump_session) = state.sql_dump_session.as_mut() { + let repartition_comment = format!( + "kind=repartition table={} action_idx={} started_at_ms={} elapsed_ms={}", + physical_table_ctx.name, + i + 1, + started_at_ms, + elapsed.as_millis() + ); + sql_dump_session.append_sql( + &physical_table_ctx.name.to_string(), + &sql, + Some(&repartition_comment), + )?; + let event = format!( + "repartition action_idx={} started_at_ms={} elapsed_ms={} sql={}", + i + 1, + started_at_ms, + elapsed.as_millis(), + sql + ); + sql_dump_session.broadcast_event(logical_table_names.iter(), &event, &sql)?; + } + } + + let (ack_tx, ack_rx) = oneshot::channel(); + control_tx + .send(WriterControl::Barrier { + epoch: i + 1, + ack: ack_tx, + }) + .expect("barrier control send must succeed"); + tokio::time::timeout(Duration::from_secs(BARRIER_ACK_TIMEOUT_SECS), ack_rx) + .await + .expect("barrier ack timeout") + .expect("barrier ack dropped"); + + let inserted_rows_snapshot = flush_dump_sessions_and_snapshot(&shared_state)?; + info!("validate rows, epoch: {}", i + 1); + validate_rows(&ctx, &logical_tables, &inserted_rows_snapshot).await?; + + control_tx + .send(WriterControl::Resume { epoch: i + 1 }) + .expect("resume control send must succeed"); + } + + let _ = control_tx.send(WriterControl::Stop); + shared_state.lock().unwrap().running = false; + writer_task.await.unwrap().unwrap(); + let inserted_rows = flush_dump_sessions_and_snapshot(&shared_state)?; + let (mut csv_dump_session, mut sql_dump_session) = { + let mut state = shared_state.lock().unwrap(); + (state.csv_dump_session.take(), state.sql_dump_session.take()) + }; + + let run_result = async { + validate_rows(&ctx, &logical_tables, &inserted_rows).await?; + cleanup_tables(&ctx, &physical_table_ctx, &logical_tables).await?; + Ok(()) + } + .await; + + if let Some(csv_dump_session) = csv_dump_session.take() { + match &run_result { + Ok(_) => { + if let Err(err) = csv_dump_session.cleanup_on_success() { + warn!( + "Cleanup csv dump directory failed, path: {}, error: {:?}", + csv_dump_session.run_dir.display(), + err + ); + } + } + Err(_) => { + warn!( + "Keep csv dump directory for failure analysis, path: {}", + csv_dump_session.run_dir.display() + ); + } + } + } + if let Some(sql_dump_session) = sql_dump_session.take() + && run_result.is_err() + { + warn!( + "Keep sql dump directory for failure analysis, path: {}", + sql_dump_session.run_dir.display() + ); + } + + ctx.close().await; + run_result +} + +fuzz_target!(|input: FuzzInput| { + common_telemetry::init_default_ut_logging(); + common_runtime::block_on_global(async { + let Connections { mysql } = init_greptime_connections_via_env().await; + let ctx = FuzzContext { + greptime: mysql.expect("mysql connection init must be succeed"), + }; + execute_repartition_metric_table(ctx, input) + .await + .unwrap_or_else(|err| panic!("fuzz test must be succeed: {err:?}")); + }) +}); diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 0c6b965fd34..ec35205a553 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -5,7 +5,7 @@ edition.workspace = true license.workspace = true [features] -dashboard = [] +dashboard = ["servers/dashboard"] vector_index = [] [lints] diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index fd0d1ef3c4b..2bf6e812c78 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -534,6 +534,7 @@ pub async fn setup_test_http_app_with_frontend_and_custom_options( .with_influxdb_handler(instance.fe_instance().clone()) .with_otlp_handler(instance.fe_instance().clone(), true) .with_jaeger_handler(instance.fe_instance().clone()) + .with_dashboard_handler(instance.fe_instance().clone()) .with_greptime_config_options(instance.opts.to_toml().unwrap()); if let Some(user_provider) = user_provider { diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 68fa2a228d0..c259d3ff241 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -106,6 +106,7 @@ macro_rules! http_tests { test_config_api, test_dynamic_tracer_toggle, test_dashboard_path, + test_dashboard_api, test_prometheus_remote_write, test_prometheus_remote_special_labels, test_prometheus_remote_schema_labels, @@ -1720,6 +1721,121 @@ pub async fn test_dashboard_path(store_type: StorageType) { #[cfg(not(feature = "dashboard"))] pub async fn test_dashboard_path(_: StorageType) {} +#[cfg(feature = "dashboard")] +pub async fn test_dashboard_api(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = setup_test_http_app_with_frontend(store_type, "dashboard_api").await; + let client = TestClient::new(app).await; + + // 1. List dashboards - should be empty initially + let res = client.get("/v1/dashboards").send().await; + assert_eq!(res.status(), StatusCode::OK); + let body: Value = res.json().await; + let dashboards = body.get("dashboards").unwrap().as_array().unwrap(); + assert!(dashboards.is_empty()); + + // 2. Save a dashboard + let dashboard_definition = r#"{"title": "My Dashboard", "panels": []}"#; + let res = client + .post("/v1/dashboards/test_dashboard") + .body(dashboard_definition) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body: Value = res.json().await; + let dashboards = body.get("dashboards").unwrap().as_array().unwrap(); + assert_eq!(dashboards.len(), 1); + assert_eq!(dashboards[0].get("name").unwrap(), "test_dashboard"); + + // 3. Save another dashboard + let res = client + .post("/v1/dashboards/another_dashboard") + .body(r#"{"title": "Another Dashboard"}"#) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + // 4. List dashboards - should have 2 + let res = client.get("/v1/dashboards").send().await; + assert_eq!(res.status(), StatusCode::OK); + let body: Value = res.json().await; + let dashboards = body.get("dashboards").unwrap().as_array().unwrap(); + assert_eq!(dashboards.len(), 2); + + let names: Vec<&str> = dashboards + .iter() + .map(|d| d.get("name").unwrap().as_str().unwrap()) + .collect(); + assert!(names.contains(&"test_dashboard")); + assert!(names.contains(&"another_dashboard")); + + // 5. Update a dashboard by posting again with new definition + let updated_definition = r#"{"title": "Updated Dashboard", "panels": [{"id": 1}]}"#; + let res = client + .post("/v1/dashboards/test_dashboard") + .body(updated_definition) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + let body: Value = res.json().await; + let dashboards = body.get("dashboards").unwrap().as_array().unwrap(); + assert_eq!(dashboards.len(), 1); + assert_eq!(dashboards[0].get("name").unwrap(), "test_dashboard"); + + // Verify the definition was updated by listing again + let res = client.get("/v1/dashboards").send().await; + assert_eq!(res.status(), StatusCode::OK); + let body: Value = res.json().await; + let dashboards = body.get("dashboards").unwrap().as_array().unwrap(); + assert_eq!(dashboards.len(), 2); + + // Find test_dashboard and verify it has updated definition + let test_db = dashboards + .iter() + .find(|d| d.get("name").unwrap() == "test_dashboard") + .unwrap(); + assert_eq!( + test_db.get("definition").unwrap(), + r#"{"title": "Updated Dashboard", "panels": [{"id": 1}]}"# + ); + + // 6. Delete one dashboard + let res = client.delete("/v1/dashboards/test_dashboard").send().await; + assert_eq!(res.status(), StatusCode::OK); + let body: Value = res.json().await; + let dashboards = body.get("dashboards").unwrap().as_array().unwrap(); + assert_eq!(dashboards.len(), 1); + assert_eq!(dashboards[0].get("name").unwrap(), "test_dashboard"); + + // 7. List dashboards - should have 1 + let res = client.get("/v1/dashboards").send().await; + assert_eq!(res.status(), StatusCode::OK); + let body: Value = res.json().await; + let dashboards = body.get("dashboards").unwrap().as_array().unwrap(); + assert_eq!(dashboards.len(), 1); + assert_eq!(dashboards[0].get("name").unwrap(), "another_dashboard"); + + // 8. Delete the remaining dashboard + let res = client + .delete("/v1/dashboards/another_dashboard") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + // 9. List dashboards - should be empty + let res = client.get("/v1/dashboards").send().await; + assert_eq!(res.status(), StatusCode::OK); + let body: Value = res.json().await; + let dashboards = body.get("dashboards").unwrap().as_array().unwrap(); + assert!(dashboards.is_empty()); + + guard.remove_all().await; +} + +#[cfg(not(feature = "dashboard"))] +pub async fn test_dashboard_api(_: StorageType) {} + pub async fn test_prometheus_remote_write(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = diff --git a/tests/cases/standalone/common/types/json/json.result b/tests/cases/standalone/common/types/json/json.result index 8c4755f4ae4..8fad9632b12 100644 --- a/tests/cases/standalone/common/types/json/json.result +++ b/tests/cases/standalone/common/types/json/json.result @@ -37,22 +37,23 @@ INSERT INTO jsons VALUES('[null]', 0), } ] } -}}', 11); +}}', 11), +('{"a":"abc\u2028tom"}', 12); -Affected Rows: 12 +Affected Rows: 13 -INSERT INTO jsons VALUES(parse_json('[null]'), 12), -(parse_json('[true]'), 13), -(parse_json('[false]'), 14), -(parse_json('[0]'), 15), -(parse_json('["foo"]'), 16), -(parse_json('[]'), 17), -(parse_json('{}'), 18), -(parse_json('[0,1]'), 19), -(parse_json('{"foo":"bar"}'), 20), -(parse_json('{"a":null,"foo":"bar"}'), 21), -(parse_json('[-1]'), 22), -(parse_json('[-2147483648]'), 23), +INSERT INTO jsons VALUES(parse_json('[null]'), 1000), +(parse_json('[true]'), 1001), +(parse_json('[false]'), 1002), +(parse_json('[0]'), 1003), +(parse_json('["foo"]'), 1004), +(parse_json('[]'), 1005), +(parse_json('{}'), 1006), +(parse_json('[0,1]'), 1007), +(parse_json('{"foo":"bar"}'), 1008), +(parse_json('{"a":null,"foo":"bar"}'), 1009), +(parse_json('[-1]'), 1010), +(parse_json('[-2147483648]'), 1011), (parse_json('{"entities": { "description": { "urls": [ @@ -76,9 +77,10 @@ INSERT INTO jsons VALUES(parse_json('[null]'), 12), } ] } - }}'), 24); + }}'), 1012), +(parse_json('{"a":"abc\u2028tom"}'), 1013); -Affected Rows: 13 +Affected Rows: 14 SELECT json_to_string(j), t FROM jsons; @@ -97,25 +99,27 @@ SELECT json_to_string(j), t FROM jsons; | {"a":null,"foo":"bar"} | 1970-01-01T00:00:00.009 | | [-1] | 1970-01-01T00:00:00.010 | | {"entities":{"description":{"urls":[{"display_url":"pixiv.net/member.php?id=…","expanded_url":"http://www.pixiv.net/member.php?id=4776","indices":[58,80],"url":"http://t.co/QMLJeFmfMT"},{"display_url":"ask.fm/KATANA77","expanded_url":"http://ask.fm/KATANA77","indices":[95,117],"url":"http://t.co/LU8T7vmU3h"}]}}} | 1970-01-01T00:00:00.011 | -| [null] | 1970-01-01T00:00:00.012 | -| [true] | 1970-01-01T00:00:00.013 | -| [false] | 1970-01-01T00:00:00.014 | -| [0] | 1970-01-01T00:00:00.015 | -| ["foo"] | 1970-01-01T00:00:00.016 | -| [] | 1970-01-01T00:00:00.017 | -| {} | 1970-01-01T00:00:00.018 | -| [0,1] | 1970-01-01T00:00:00.019 | -| {"foo":"bar"} | 1970-01-01T00:00:00.020 | -| {"a":null,"foo":"bar"} | 1970-01-01T00:00:00.021 | -| [-1] | 1970-01-01T00:00:00.022 | -| [-2147483648] | 1970-01-01T00:00:00.023 | -| {"entities":{"description":{"urls":[{"display_url":"pixiv.net/member.php?id=…","expanded_url":"http://www.pixiv.net/member.php?id=4776","indices":[58,80],"url":"http://t.co/QMLJeFmfMT"},{"display_url":"ask.fm/KATANA77","expanded_url":"http://ask.fm/KATANA77","indices":[95,117],"url":"http://t.co/LU8T7vmU3h"}]}}} | 1970-01-01T00:00:00.024 | +| {"a":"abc\u2028tom"} | 1970-01-01T00:00:00.012 | +| [null] | 1970-01-01T00:00:01 | +| [true] | 1970-01-01T00:00:01.001 | +| [false] | 1970-01-01T00:00:01.002 | +| [0] | 1970-01-01T00:00:01.003 | +| ["foo"] | 1970-01-01T00:00:01.004 | +| [] | 1970-01-01T00:00:01.005 | +| {} | 1970-01-01T00:00:01.006 | +| [0,1] | 1970-01-01T00:00:01.007 | +| {"foo":"bar"} | 1970-01-01T00:00:01.008 | +| {"a":null,"foo":"bar"} | 1970-01-01T00:00:01.009 | +| [-1] | 1970-01-01T00:00:01.010 | +| [-2147483648] | 1970-01-01T00:00:01.011 | +| {"entities":{"description":{"urls":[{"display_url":"pixiv.net/member.php?id=…","expanded_url":"http://www.pixiv.net/member.php?id=4776","indices":[58,80],"url":"http://t.co/QMLJeFmfMT"},{"display_url":"ask.fm/KATANA77","expanded_url":"http://ask.fm/KATANA77","indices":[95,117],"url":"http://t.co/LU8T7vmU3h"}]}}} | 1970-01-01T00:00:01.012 | +| {"a":"abc\u2028tom"} | 1970-01-01T00:00:01.013 | +---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------+ --Insert invalid json strings-- DELETE FROM jsons; -Affected Rows: 25 +Affected Rows: 27 INSERT INTO jsons VALUES(parse_json('{"a":1, "b":2, "c":3'), 4); diff --git a/tests/cases/standalone/common/types/json/json.sql b/tests/cases/standalone/common/types/json/json.sql index 868edc59e85..5a521ee1c6d 100644 --- a/tests/cases/standalone/common/types/json/json.sql +++ b/tests/cases/standalone/common/types/json/json.sql @@ -35,20 +35,21 @@ INSERT INTO jsons VALUES('[null]', 0), } ] } -}}', 11); +}}', 11), +('{"a":"abc\u2028tom"}', 12); -INSERT INTO jsons VALUES(parse_json('[null]'), 12), -(parse_json('[true]'), 13), -(parse_json('[false]'), 14), -(parse_json('[0]'), 15), -(parse_json('["foo"]'), 16), -(parse_json('[]'), 17), -(parse_json('{}'), 18), -(parse_json('[0,1]'), 19), -(parse_json('{"foo":"bar"}'), 20), -(parse_json('{"a":null,"foo":"bar"}'), 21), -(parse_json('[-1]'), 22), -(parse_json('[-2147483648]'), 23), +INSERT INTO jsons VALUES(parse_json('[null]'), 1000), +(parse_json('[true]'), 1001), +(parse_json('[false]'), 1002), +(parse_json('[0]'), 1003), +(parse_json('["foo"]'), 1004), +(parse_json('[]'), 1005), +(parse_json('{}'), 1006), +(parse_json('[0,1]'), 1007), +(parse_json('{"foo":"bar"}'), 1008), +(parse_json('{"a":null,"foo":"bar"}'), 1009), +(parse_json('[-1]'), 1010), +(parse_json('[-2147483648]'), 1011), (parse_json('{"entities": { "description": { "urls": [ @@ -72,7 +73,8 @@ INSERT INTO jsons VALUES(parse_json('[null]'), 12), } ] } - }}'), 24); + }}'), 1012), +(parse_json('{"a":"abc\u2028tom"}'), 1013); SELECT json_to_string(j), t FROM jsons; diff --git a/tests/conf/datanode-test.toml.template b/tests/conf/datanode-test.toml.template index 4cb0423c72f..3ec8a2f6956 100644 --- a/tests/conf/datanode-test.toml.template +++ b/tests/conf/datanode-test.toml.template @@ -28,7 +28,7 @@ type = 'File' data_home = '{data_home}' [meta_client_options] -metasrv_addrs = ['{metasrv_addr}'] +metasrv_addrs = ['{addrs.metasrv_addr}'] timeout_millis = 3000 connect_timeout_millis = 5000 tcp_nodelay = false diff --git a/tests/conf/frontend-test.toml.template b/tests/conf/frontend-test.toml.template index de4ce86adca..25d44ff6e4f 100644 --- a/tests/conf/frontend-test.toml.template +++ b/tests/conf/frontend-test.toml.template @@ -1,3 +1,3 @@ [grpc] -bind_addr = "{grpc_addr}" -server_addr = "{grpc_addr}" +bind_addr = "{addrs.grpc_addr}" +server_addr = "{addrs.grpc_addr}" diff --git a/tests/conf/standalone-test.toml.template b/tests/conf/standalone-test.toml.template index 509eac7ca67..50c014e9910 100644 --- a/tests/conf/standalone-test.toml.template +++ b/tests/conf/standalone-test.toml.template @@ -26,12 +26,12 @@ type = 'File' data_home = '{data_home}' [grpc] -bind_addr = '{grpc_addr}' +bind_addr = '{addrs.grpc_addr}' runtime_size = 8 [mysql] enable = true -addr = "{mysql_addr}" +addr = "{addrs.mysql_addr}" runtime_size = 2 prepared_stmt_cache_size= 10000 @@ -40,7 +40,7 @@ mode = "disable" [postgres] enable = true -addr = "{postgres_addr}" +addr = "{addrs.postgres_addr}" runtime_size = 2 [procedure] diff --git a/tests/runner/src/server_mode.rs b/tests/runner/src/server_mode.rs index 172baf32ff6..1f7cb72bf4f 100644 --- a/tests/runner/src/server_mode.rs +++ b/tests/runner/src/server_mode.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::path::Path; use std::sync::{Mutex, OnceLock}; @@ -96,15 +96,7 @@ struct ConfigContext { use_etcd: bool, store_addrs: String, instance_id: usize, - // for following addrs, leave it empty if not needed - // required for datanode - metasrv_addr: String, - // for frontend and standalone - grpc_addr: String, - // for standalone - mysql_addr: String, - // for standalone - postgres_addr: String, + addrs: HashMap, // enable flat format for storage engine enable_flat_format: bool, } @@ -275,40 +267,26 @@ impl ServerMode { let procedure_dir = data_home.join("procedure").display().to_string(); // Get the required addresses based on server mode - let (metasrv_addr, grpc_addr, mysql_addr, postgres_addr) = match self { + let addrs: HashMap = match self { ServerMode::Standalone { rpc_bind_addr, mysql_addr, postgres_addr, - .. - } => ( - String::new(), - rpc_bind_addr.clone(), - mysql_addr.clone(), - postgres_addr.clone(), - ), - ServerMode::Frontend { - rpc_bind_addr, - mysql_addr, - postgres_addr, - .. - } => ( - String::new(), - rpc_bind_addr.clone(), - mysql_addr.clone(), - postgres_addr.clone(), - ), - ServerMode::Datanode { - rpc_bind_addr, - metasrv_addr, - .. - } => ( - metasrv_addr.clone(), - rpc_bind_addr.clone(), - String::new(), - String::new(), - ), - _ => (String::new(), String::new(), String::new(), String::new()), + http_addr, + } => [ + ("http_addr".to_string(), http_addr.clone()), + ("grpc_addr".to_string(), rpc_bind_addr.clone()), + ("mysql_addr".to_string(), mysql_addr.clone()), + ("postgres_addr".to_string(), postgres_addr.clone()), + ] + .into(), + ServerMode::Frontend { rpc_bind_addr, .. } => { + [("grpc_addr".to_string(), rpc_bind_addr.clone())].into() + } + ServerMode::Datanode { metasrv_addr, .. } => { + [("metasrv_addr".to_string(), metasrv_addr.clone())].into() + } + _ => HashMap::new(), }; let ctx = ConfigContext { @@ -326,10 +304,7 @@ impl ServerMode { .collect::>() .join(","), instance_id: id, - metasrv_addr, - grpc_addr, - mysql_addr, - postgres_addr, + addrs, enable_flat_format: db_ctx.store_config().enable_flat_format, };