mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2025-12-26 08:00:01 +00:00
Compare commits
46 Commits
v0.10.0-ni
...
v0.9.5
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c02b5dae93 | ||
|
|
081c6d9e74 | ||
|
|
ca6e02980e | ||
|
|
74bdba4613 | ||
|
|
2e0e82ddc8 | ||
|
|
e0c4157ad8 | ||
|
|
613e07afb4 | ||
|
|
0ce93f0b88 | ||
|
|
c231eee7c1 | ||
|
|
176f2df5b3 | ||
|
|
4622412dfe | ||
|
|
59ec90299b | ||
|
|
16b8cdc3d5 | ||
|
|
3197b8b535 | ||
|
|
972c2441af | ||
|
|
bb8b54b5d3 | ||
|
|
b5233e500b | ||
|
|
b61a388d04 | ||
|
|
06e565d25a | ||
|
|
3b2ce31a19 | ||
|
|
a889ea88ca | ||
|
|
2f2b4b306c | ||
|
|
856c0280f5 | ||
|
|
aaa9b32908 | ||
|
|
4bb1f4f184 | ||
|
|
0f907ef99e | ||
|
|
a61c0bd1d8 | ||
|
|
7dd0e3ab37 | ||
|
|
d168bde226 | ||
|
|
4b34f610aa | ||
|
|
695ff1e037 | ||
|
|
288fdc3145 | ||
|
|
a8ed3db0aa | ||
|
|
0dd11f53f5 | ||
|
|
19918928c5 | ||
|
|
5f0a83b2b1 | ||
|
|
71a66d15f7 | ||
|
|
2cdd103874 | ||
|
|
4dea4cac47 | ||
|
|
a283e13da7 | ||
|
|
47a3277d12 | ||
|
|
caf5f2c7a5 | ||
|
|
c1e8084af6 | ||
|
|
6e776d5f98 | ||
|
|
e39a9e6feb | ||
|
|
77af4fd981 |
9
.github/workflows/develop.yml
vendored
9
.github/workflows/develop.yml
vendored
@@ -442,6 +442,13 @@ jobs:
|
||||
minio: true
|
||||
kafka: true
|
||||
values: "with-remote-wal.yaml"
|
||||
include:
|
||||
- target: "fuzz_migrate_mito_regions"
|
||||
mode:
|
||||
name: "Local WAL"
|
||||
minio: true
|
||||
kafka: false
|
||||
values: "with-minio.yaml"
|
||||
steps:
|
||||
- name: Remove unused software
|
||||
run: |
|
||||
@@ -530,7 +537,7 @@ jobs:
|
||||
with:
|
||||
image-registry: localhost:5001
|
||||
values-filename: ${{ matrix.mode.values }}
|
||||
enable-region-failover: true
|
||||
enable-region-failover: ${{ matrix.mode.kafka }}
|
||||
- name: Port forward (mysql)
|
||||
run: |
|
||||
kubectl port-forward service/my-greptimedb-frontend 4002:4002 -n my-greptimedb&
|
||||
|
||||
1013
Cargo.lock
generated
1013
Cargo.lock
generated
File diff suppressed because it is too large
Load Diff
@@ -20,6 +20,7 @@ members = [
|
||||
"src/common/mem-prof",
|
||||
"src/common/meta",
|
||||
"src/common/plugins",
|
||||
"src/common/pprof",
|
||||
"src/common/procedure",
|
||||
"src/common/procedure-test",
|
||||
"src/common/query",
|
||||
@@ -64,7 +65,7 @@ members = [
|
||||
resolver = "2"
|
||||
|
||||
[workspace.package]
|
||||
version = "0.9.3"
|
||||
version = "0.9.5"
|
||||
edition = "2021"
|
||||
license = "Apache-2.0"
|
||||
|
||||
@@ -120,7 +121,7 @@ etcd-client = { version = "0.13" }
|
||||
fst = "0.4.7"
|
||||
futures = "0.3"
|
||||
futures-util = "0.3"
|
||||
greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "36334744c7020734dcb4a6b8d24d52ae7ed53fe1" }
|
||||
greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "b4d301184eb0d01fd4d1042fcc7c5dfb54f3c1e3" }
|
||||
humantime = "2.1"
|
||||
humantime-serde = "1.1"
|
||||
itertools = "0.10"
|
||||
@@ -137,12 +138,13 @@ opentelemetry-proto = { version = "0.5", features = [
|
||||
"metrics",
|
||||
"trace",
|
||||
"with-serde",
|
||||
"logs",
|
||||
] }
|
||||
parquet = { version = "51.0.0", default-features = false, features = ["arrow", "async", "object_store"] }
|
||||
paste = "1.0"
|
||||
pin-project = "1.0"
|
||||
prometheus = { version = "0.13.3", features = ["process"] }
|
||||
promql-parser = { version = "0.4" }
|
||||
promql-parser = { version = "0.4.1" }
|
||||
prost = "0.12"
|
||||
raft-engine = { version = "0.4.1", default-features = false }
|
||||
rand = "0.8"
|
||||
@@ -208,6 +210,7 @@ common-macro = { path = "src/common/macro" }
|
||||
common-mem-prof = { path = "src/common/mem-prof" }
|
||||
common-meta = { path = "src/common/meta" }
|
||||
common-plugins = { path = "src/common/plugins" }
|
||||
common-pprof = { path = "src/common/pprof" }
|
||||
common-procedure = { path = "src/common/procedure" }
|
||||
common-procedure-test = { path = "src/common/procedure-test" }
|
||||
common-query = { path = "src/common/query" }
|
||||
|
||||
@@ -83,6 +83,7 @@
|
||||
| `wal.backoff_max` | String | `10s` | The maximum backoff delay.<br/>**It's only used when the provider is `kafka`**. |
|
||||
| `wal.backoff_base` | Integer | `2` | The exponential backoff rate, i.e. next backoff = base * current backoff.<br/>**It's only used when the provider is `kafka`**. |
|
||||
| `wal.backoff_deadline` | String | `5mins` | The deadline of retries.<br/>**It's only used when the provider is `kafka`**. |
|
||||
| `wal.overwrite_entry_start_id` | Bool | `false` | Ignore missing entries during read WAL.<br/>**It's only used when the provider is `kafka`**.<br/><br/>This option ensures that when Kafka messages are deleted, the system <br/>can still successfully replay memtable data without throwing an <br/>out-of-range error. <br/>However, enabling this option might lead to unexpected data loss, <br/>as the system will skip over missing entries instead of treating <br/>them as critical errors. |
|
||||
| `metadata_store` | -- | -- | Metadata storage options. |
|
||||
| `metadata_store.file_size` | String | `256MB` | Kv file size in bytes. |
|
||||
| `metadata_store.purge_threshold` | String | `4GB` | Kv purge threshold. |
|
||||
@@ -161,8 +162,13 @@
|
||||
| `logging.otlp_endpoint` | String | `http://localhost:4317` | The OTLP tracing endpoint. |
|
||||
| `logging.append_stdout` | Bool | `true` | Whether to append logs to stdout. |
|
||||
| `logging.log_format` | String | `text` | The log format. Can be `text`/`json`. |
|
||||
| `logging.max_log_files` | Integer | `720` | The maximum amount of log files. |
|
||||
| `logging.tracing_sample_ratio` | -- | -- | The percentage of tracing will be sampled and exported.<br/>Valid range `[0, 1]`, 1 means all traces are sampled, 0 means all traces are not sampled, the default value is 1.<br/>ratio > 1 are treated as 1. Fractions < 0 are treated as 0 |
|
||||
| `logging.tracing_sample_ratio.default_ratio` | Float | `1.0` | -- |
|
||||
| `logging.slow_query` | -- | -- | The slow query log options. |
|
||||
| `logging.slow_query.enable` | Bool | `false` | Whether to enable slow query log. |
|
||||
| `logging.slow_query.threshold` | String | Unset | The threshold of slow query. |
|
||||
| `logging.slow_query.sample_ratio` | Float | Unset | The sampling ratio of slow query log. The value should be in the range of (0, 1]. |
|
||||
| `export_metrics` | -- | -- | The datanode can export its metrics and send to Prometheus compatible service (e.g. send to `greptimedb` itself) from remote-write API.<br/>This is only used for `greptimedb` to export its own metrics internally. It's different from prometheus scrape. |
|
||||
| `export_metrics.enable` | Bool | `false` | whether enable export metrics. |
|
||||
| `export_metrics.write_interval` | String | `30s` | The interval of export metrics. |
|
||||
@@ -247,8 +253,13 @@
|
||||
| `logging.otlp_endpoint` | String | `http://localhost:4317` | The OTLP tracing endpoint. |
|
||||
| `logging.append_stdout` | Bool | `true` | Whether to append logs to stdout. |
|
||||
| `logging.log_format` | String | `text` | The log format. Can be `text`/`json`. |
|
||||
| `logging.max_log_files` | Integer | `720` | The maximum amount of log files. |
|
||||
| `logging.tracing_sample_ratio` | -- | -- | The percentage of tracing will be sampled and exported.<br/>Valid range `[0, 1]`, 1 means all traces are sampled, 0 means all traces are not sampled, the default value is 1.<br/>ratio > 1 are treated as 1. Fractions < 0 are treated as 0 |
|
||||
| `logging.tracing_sample_ratio.default_ratio` | Float | `1.0` | -- |
|
||||
| `logging.slow_query` | -- | -- | The slow query log options. |
|
||||
| `logging.slow_query.enable` | Bool | `false` | Whether to enable slow query log. |
|
||||
| `logging.slow_query.threshold` | String | Unset | The threshold of slow query. |
|
||||
| `logging.slow_query.sample_ratio` | Float | Unset | The sampling ratio of slow query log. The value should be in the range of (0, 1]. |
|
||||
| `export_metrics` | -- | -- | The datanode can export its metrics and send to Prometheus compatible service (e.g. send to `greptimedb` itself) from remote-write API.<br/>This is only used for `greptimedb` to export its own metrics internally. It's different from prometheus scrape. |
|
||||
| `export_metrics.enable` | Bool | `false` | whether enable export metrics. |
|
||||
| `export_metrics.write_interval` | String | `30s` | The interval of export metrics. |
|
||||
@@ -312,8 +323,13 @@
|
||||
| `logging.otlp_endpoint` | String | `http://localhost:4317` | The OTLP tracing endpoint. |
|
||||
| `logging.append_stdout` | Bool | `true` | Whether to append logs to stdout. |
|
||||
| `logging.log_format` | String | `text` | The log format. Can be `text`/`json`. |
|
||||
| `logging.max_log_files` | Integer | `720` | The maximum amount of log files. |
|
||||
| `logging.tracing_sample_ratio` | -- | -- | The percentage of tracing will be sampled and exported.<br/>Valid range `[0, 1]`, 1 means all traces are sampled, 0 means all traces are not sampled, the default value is 1.<br/>ratio > 1 are treated as 1. Fractions < 0 are treated as 0 |
|
||||
| `logging.tracing_sample_ratio.default_ratio` | Float | `1.0` | -- |
|
||||
| `logging.slow_query` | -- | -- | The slow query log options. |
|
||||
| `logging.slow_query.enable` | Bool | `false` | Whether to enable slow query log. |
|
||||
| `logging.slow_query.threshold` | String | Unset | The threshold of slow query. |
|
||||
| `logging.slow_query.sample_ratio` | Float | Unset | The sampling ratio of slow query log. The value should be in the range of (0, 1]. |
|
||||
| `export_metrics` | -- | -- | The datanode can export its metrics and send to Prometheus compatible service (e.g. send to `greptimedb` itself) from remote-write API.<br/>This is only used for `greptimedb` to export its own metrics internally. It's different from prometheus scrape. |
|
||||
| `export_metrics.enable` | Bool | `false` | whether enable export metrics. |
|
||||
| `export_metrics.write_interval` | String | `30s` | The interval of export metrics. |
|
||||
@@ -394,6 +410,7 @@
|
||||
| `wal.backoff_deadline` | String | `5mins` | The deadline of retries.<br/>**It's only used when the provider is `kafka`**. |
|
||||
| `wal.create_index` | Bool | `true` | Whether to enable WAL index creation.<br/>**It's only used when the provider is `kafka`**. |
|
||||
| `wal.dump_index_interval` | String | `60s` | The interval for dumping WAL indexes.<br/>**It's only used when the provider is `kafka`**. |
|
||||
| `wal.overwrite_entry_start_id` | Bool | `false` | Ignore missing entries during read WAL.<br/>**It's only used when the provider is `kafka`**.<br/><br/>This option ensures that when Kafka messages are deleted, the system <br/>can still successfully replay memtable data without throwing an <br/>out-of-range error. <br/>However, enabling this option might lead to unexpected data loss, <br/>as the system will skip over missing entries instead of treating <br/>them as critical errors. |
|
||||
| `storage` | -- | -- | The data storage options. |
|
||||
| `storage.data_home` | String | `/tmp/greptimedb/` | The working home directory. |
|
||||
| `storage.type` | String | `File` | The storage type used to store the data.<br/>- `File`: the data is stored in the local file system.<br/>- `S3`: the data is stored in the S3 object storage.<br/>- `Gcs`: the data is stored in the Google Cloud Storage.<br/>- `Azblob`: the data is stored in the Azure Blob Storage.<br/>- `Oss`: the data is stored in the Aliyun OSS. |
|
||||
@@ -464,8 +481,13 @@
|
||||
| `logging.otlp_endpoint` | String | `http://localhost:4317` | The OTLP tracing endpoint. |
|
||||
| `logging.append_stdout` | Bool | `true` | Whether to append logs to stdout. |
|
||||
| `logging.log_format` | String | `text` | The log format. Can be `text`/`json`. |
|
||||
| `logging.max_log_files` | Integer | `720` | The maximum amount of log files. |
|
||||
| `logging.tracing_sample_ratio` | -- | -- | The percentage of tracing will be sampled and exported.<br/>Valid range `[0, 1]`, 1 means all traces are sampled, 0 means all traces are not sampled, the default value is 1.<br/>ratio > 1 are treated as 1. Fractions < 0 are treated as 0 |
|
||||
| `logging.tracing_sample_ratio.default_ratio` | Float | `1.0` | -- |
|
||||
| `logging.slow_query` | -- | -- | The slow query log options. |
|
||||
| `logging.slow_query.enable` | Bool | `false` | Whether to enable slow query log. |
|
||||
| `logging.slow_query.threshold` | String | Unset | The threshold of slow query. |
|
||||
| `logging.slow_query.sample_ratio` | Float | Unset | The sampling ratio of slow query log. The value should be in the range of (0, 1]. |
|
||||
| `export_metrics` | -- | -- | The datanode can export its metrics and send to Prometheus compatible service (e.g. send to `greptimedb` itself) from remote-write API.<br/>This is only used for `greptimedb` to export its own metrics internally. It's different from prometheus scrape. |
|
||||
| `export_metrics.enable` | Bool | `false` | whether enable export metrics. |
|
||||
| `export_metrics.write_interval` | String | `30s` | The interval of export metrics. |
|
||||
@@ -510,7 +532,12 @@
|
||||
| `logging.otlp_endpoint` | String | `http://localhost:4317` | The OTLP tracing endpoint. |
|
||||
| `logging.append_stdout` | Bool | `true` | Whether to append logs to stdout. |
|
||||
| `logging.log_format` | String | `text` | The log format. Can be `text`/`json`. |
|
||||
| `logging.max_log_files` | Integer | `720` | The maximum amount of log files. |
|
||||
| `logging.tracing_sample_ratio` | -- | -- | The percentage of tracing will be sampled and exported.<br/>Valid range `[0, 1]`, 1 means all traces are sampled, 0 means all traces are not sampled, the default value is 1.<br/>ratio > 1 are treated as 1. Fractions < 0 are treated as 0 |
|
||||
| `logging.tracing_sample_ratio.default_ratio` | Float | `1.0` | -- |
|
||||
| `logging.slow_query` | -- | -- | The slow query log options. |
|
||||
| `logging.slow_query.enable` | Bool | `false` | Whether to enable slow query log. |
|
||||
| `logging.slow_query.threshold` | String | Unset | The threshold of slow query. |
|
||||
| `logging.slow_query.sample_ratio` | Float | Unset | The sampling ratio of slow query log. The value should be in the range of (0, 1]. |
|
||||
| `tracing` | -- | -- | The tracing options. Only effect when compiled with `tokio-console` feature. |
|
||||
| `tracing.tokio_console_addr` | String | Unset | The tokio console address. |
|
||||
|
||||
@@ -213,6 +213,17 @@ create_index = true
|
||||
## **It's only used when the provider is `kafka`**.
|
||||
dump_index_interval = "60s"
|
||||
|
||||
## Ignore missing entries during read WAL.
|
||||
## **It's only used when the provider is `kafka`**.
|
||||
##
|
||||
## This option ensures that when Kafka messages are deleted, the system
|
||||
## can still successfully replay memtable data without throwing an
|
||||
## out-of-range error.
|
||||
## However, enabling this option might lead to unexpected data loss,
|
||||
## as the system will skip over missing entries instead of treating
|
||||
## them as critical errors.
|
||||
overwrite_entry_start_id = false
|
||||
|
||||
# The Kafka SASL configuration.
|
||||
# **It's only used when the provider is `kafka`**.
|
||||
# Available SASL mechanisms:
|
||||
@@ -580,12 +591,28 @@ append_stdout = true
|
||||
## The log format. Can be `text`/`json`.
|
||||
log_format = "text"
|
||||
|
||||
## The maximum amount of log files.
|
||||
max_log_files = 720
|
||||
|
||||
## The percentage of tracing will be sampled and exported.
|
||||
## Valid range `[0, 1]`, 1 means all traces are sampled, 0 means all traces are not sampled, the default value is 1.
|
||||
## ratio > 1 are treated as 1. Fractions < 0 are treated as 0
|
||||
[logging.tracing_sample_ratio]
|
||||
default_ratio = 1.0
|
||||
|
||||
## The slow query log options.
|
||||
[logging.slow_query]
|
||||
## Whether to enable slow query log.
|
||||
enable = false
|
||||
|
||||
## The threshold of slow query.
|
||||
## @toml2docs:none-default
|
||||
threshold = "10s"
|
||||
|
||||
## The sampling ratio of slow query log. The value should be in the range of (0, 1].
|
||||
## @toml2docs:none-default
|
||||
sample_ratio = 1.0
|
||||
|
||||
## The datanode can export its metrics and send to Prometheus compatible service (e.g. send to `greptimedb` itself) from remote-write API.
|
||||
## This is only used for `greptimedb` to export its own metrics internally. It's different from prometheus scrape.
|
||||
[export_metrics]
|
||||
|
||||
@@ -78,12 +78,28 @@ append_stdout = true
|
||||
## The log format. Can be `text`/`json`.
|
||||
log_format = "text"
|
||||
|
||||
## The maximum amount of log files.
|
||||
max_log_files = 720
|
||||
|
||||
## The percentage of tracing will be sampled and exported.
|
||||
## Valid range `[0, 1]`, 1 means all traces are sampled, 0 means all traces are not sampled, the default value is 1.
|
||||
## ratio > 1 are treated as 1. Fractions < 0 are treated as 0
|
||||
[logging.tracing_sample_ratio]
|
||||
default_ratio = 1.0
|
||||
|
||||
## The slow query log options.
|
||||
[logging.slow_query]
|
||||
## Whether to enable slow query log.
|
||||
enable = false
|
||||
|
||||
## The threshold of slow query.
|
||||
## @toml2docs:none-default
|
||||
threshold = "10s"
|
||||
|
||||
## The sampling ratio of slow query log. The value should be in the range of (0, 1].
|
||||
## @toml2docs:none-default
|
||||
sample_ratio = 1.0
|
||||
|
||||
## The tracing options. Only effect when compiled with `tokio-console` feature.
|
||||
[tracing]
|
||||
## The tokio console address.
|
||||
|
||||
@@ -185,12 +185,28 @@ append_stdout = true
|
||||
## The log format. Can be `text`/`json`.
|
||||
log_format = "text"
|
||||
|
||||
## The maximum amount of log files.
|
||||
max_log_files = 720
|
||||
|
||||
## The percentage of tracing will be sampled and exported.
|
||||
## Valid range `[0, 1]`, 1 means all traces are sampled, 0 means all traces are not sampled, the default value is 1.
|
||||
## ratio > 1 are treated as 1. Fractions < 0 are treated as 0
|
||||
[logging.tracing_sample_ratio]
|
||||
default_ratio = 1.0
|
||||
|
||||
## The slow query log options.
|
||||
[logging.slow_query]
|
||||
## Whether to enable slow query log.
|
||||
enable = false
|
||||
|
||||
## The threshold of slow query.
|
||||
## @toml2docs:none-default
|
||||
threshold = "10s"
|
||||
|
||||
## The sampling ratio of slow query log. The value should be in the range of (0, 1].
|
||||
## @toml2docs:none-default
|
||||
sample_ratio = 1.0
|
||||
|
||||
## The datanode can export its metrics and send to Prometheus compatible service (e.g. send to `greptimedb` itself) from remote-write API.
|
||||
## This is only used for `greptimedb` to export its own metrics internally. It's different from prometheus scrape.
|
||||
[export_metrics]
|
||||
|
||||
@@ -172,12 +172,28 @@ append_stdout = true
|
||||
## The log format. Can be `text`/`json`.
|
||||
log_format = "text"
|
||||
|
||||
## The maximum amount of log files.
|
||||
max_log_files = 720
|
||||
|
||||
## The percentage of tracing will be sampled and exported.
|
||||
## Valid range `[0, 1]`, 1 means all traces are sampled, 0 means all traces are not sampled, the default value is 1.
|
||||
## ratio > 1 are treated as 1. Fractions < 0 are treated as 0
|
||||
[logging.tracing_sample_ratio]
|
||||
default_ratio = 1.0
|
||||
|
||||
## The slow query log options.
|
||||
[logging.slow_query]
|
||||
## Whether to enable slow query log.
|
||||
enable = false
|
||||
|
||||
## The threshold of slow query.
|
||||
## @toml2docs:none-default
|
||||
threshold = "10s"
|
||||
|
||||
## The sampling ratio of slow query log. The value should be in the range of (0, 1].
|
||||
## @toml2docs:none-default
|
||||
sample_ratio = 1.0
|
||||
|
||||
## The datanode can export its metrics and send to Prometheus compatible service (e.g. send to `greptimedb` itself) from remote-write API.
|
||||
## This is only used for `greptimedb` to export its own metrics internally. It's different from prometheus scrape.
|
||||
[export_metrics]
|
||||
|
||||
@@ -237,6 +237,17 @@ backoff_base = 2
|
||||
## **It's only used when the provider is `kafka`**.
|
||||
backoff_deadline = "5mins"
|
||||
|
||||
## Ignore missing entries during read WAL.
|
||||
## **It's only used when the provider is `kafka`**.
|
||||
##
|
||||
## This option ensures that when Kafka messages are deleted, the system
|
||||
## can still successfully replay memtable data without throwing an
|
||||
## out-of-range error.
|
||||
## However, enabling this option might lead to unexpected data loss,
|
||||
## as the system will skip over missing entries instead of treating
|
||||
## them as critical errors.
|
||||
overwrite_entry_start_id = false
|
||||
|
||||
# The Kafka SASL configuration.
|
||||
# **It's only used when the provider is `kafka`**.
|
||||
# Available SASL mechanisms:
|
||||
@@ -624,12 +635,28 @@ append_stdout = true
|
||||
## The log format. Can be `text`/`json`.
|
||||
log_format = "text"
|
||||
|
||||
## The maximum amount of log files.
|
||||
max_log_files = 720
|
||||
|
||||
## The percentage of tracing will be sampled and exported.
|
||||
## Valid range `[0, 1]`, 1 means all traces are sampled, 0 means all traces are not sampled, the default value is 1.
|
||||
## ratio > 1 are treated as 1. Fractions < 0 are treated as 0
|
||||
[logging.tracing_sample_ratio]
|
||||
default_ratio = 1.0
|
||||
|
||||
## The slow query log options.
|
||||
[logging.slow_query]
|
||||
## Whether to enable slow query log.
|
||||
enable = false
|
||||
|
||||
## The threshold of slow query.
|
||||
## @toml2docs:none-default
|
||||
threshold = "10s"
|
||||
|
||||
## The sampling ratio of slow query log. The value should be in the range of (0, 1].
|
||||
## @toml2docs:none-default
|
||||
sample_ratio = 1.0
|
||||
|
||||
## The datanode can export its metrics and send to Prometheus compatible service (e.g. send to `greptimedb` itself) from remote-write API.
|
||||
## This is only used for `greptimedb` to export its own metrics internally. It's different from prometheus scrape.
|
||||
[export_metrics]
|
||||
|
||||
@@ -48,4 +48,4 @@ Please refer to [SQL query](./query.sql) for GreptimeDB and Clickhouse, and [que
|
||||
|
||||
## Addition
|
||||
- You can tune GreptimeDB's configuration to get better performance.
|
||||
- You can setup GreptimeDB to use S3 as storage, see [here](https://docs.greptime.com/user-guide/operations/configuration/#storage-options).
|
||||
- You can setup GreptimeDB to use S3 as storage, see [here](https://docs.greptime.com/user-guide/deployments/configuration#storage-options).
|
||||
|
||||
@@ -1,15 +1,9 @@
|
||||
# Profiling CPU
|
||||
|
||||
## Build GreptimeDB with `pprof` feature
|
||||
|
||||
```bash
|
||||
cargo build --features=pprof
|
||||
```
|
||||
|
||||
## HTTP API
|
||||
Sample at 99 Hertz, for 5 seconds, output report in [protobuf format](https://github.com/google/pprof/blob/master/proto/profile.proto).
|
||||
```bash
|
||||
curl -s '0:4000/v1/prof/cpu' > /tmp/pprof.out
|
||||
curl -s '0:4000/debug/prof/cpu' > /tmp/pprof.out
|
||||
```
|
||||
|
||||
Then you can use `pprof` command with the protobuf file.
|
||||
@@ -19,10 +13,10 @@ go tool pprof -top /tmp/pprof.out
|
||||
|
||||
Sample at 99 Hertz, for 60 seconds, output report in flamegraph format.
|
||||
```bash
|
||||
curl -s '0:4000/v1/prof/cpu?seconds=60&output=flamegraph' > /tmp/pprof.svg
|
||||
curl -s '0:4000/debug/prof/cpu?seconds=60&output=flamegraph' > /tmp/pprof.svg
|
||||
```
|
||||
|
||||
Sample at 49 Hertz, for 10 seconds, output report in text format.
|
||||
```bash
|
||||
curl -s '0:4000/v1/prof/cpu?seconds=10&frequency=49&output=text' > /tmp/pprof.txt
|
||||
curl -s '0:4000/debug/prof/cpu?seconds=10&frequency=49&output=text' > /tmp/pprof.txt
|
||||
```
|
||||
@@ -12,16 +12,10 @@ brew install jemalloc
|
||||
sudo apt install libjemalloc-dev
|
||||
```
|
||||
|
||||
### [flamegraph](https://github.com/brendangregg/FlameGraph)
|
||||
### [flamegraph](https://github.com/brendangregg/FlameGraph)
|
||||
|
||||
```bash
|
||||
curl https://raw.githubusercontent.com/brendangregg/FlameGraph/master/flamegraph.pl > ./flamegraph.pl
|
||||
```
|
||||
|
||||
### Build GreptimeDB with `mem-prof` feature.
|
||||
|
||||
```bash
|
||||
cargo build --features=mem-prof
|
||||
curl https://raw.githubusercontent.com/brendangregg/FlameGraph/master/flamegraph.pl > ./flamegraph.pl
|
||||
```
|
||||
|
||||
## Profiling
|
||||
@@ -35,7 +29,7 @@ MALLOC_CONF=prof:true,lg_prof_interval:28 ./target/debug/greptime standalone sta
|
||||
Dump memory profiling data through HTTP API:
|
||||
|
||||
```bash
|
||||
curl localhost:4000/v1/prof/mem > greptime.hprof
|
||||
curl localhost:4000/debug/prof/mem > greptime.hprof
|
||||
```
|
||||
|
||||
You can periodically dump profiling data and compare them to find the delta memory usage.
|
||||
@@ -45,6 +39,9 @@ You can periodically dump profiling data and compare them to find the delta memo
|
||||
To create flamegraph according to dumped profiling data:
|
||||
|
||||
```bash
|
||||
jeprof --svg <path_to_greptimedb_binary> --base=<baseline_prof> <profile_data> > output.svg
|
||||
```
|
||||
sudo apt install -y libjemalloc-dev
|
||||
|
||||
jeprof <path_to_greptime_binary> <profile_data> --collapse | ./flamegraph.pl > mem-prof.svg
|
||||
|
||||
jeprof <path_to_greptime_binary> --base <baseline_prof> <profile_data> --collapse | ./flamegraph.pl > output.svg
|
||||
```
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 25 KiB After Width: | Height: | Size: 36 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 18 KiB After Width: | Height: | Size: 25 KiB |
@@ -17,10 +17,11 @@ use std::sync::Arc;
|
||||
use common_base::BitVec;
|
||||
use common_decimal::decimal128::{DECIMAL128_DEFAULT_SCALE, DECIMAL128_MAX_PRECISION};
|
||||
use common_decimal::Decimal128;
|
||||
use common_time::interval::IntervalUnit;
|
||||
use common_time::time::Time;
|
||||
use common_time::timestamp::TimeUnit;
|
||||
use common_time::{Date, DateTime, Interval, Timestamp};
|
||||
use common_time::{
|
||||
Date, DateTime, IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth, Timestamp,
|
||||
};
|
||||
use datatypes::prelude::{ConcreteDataType, ValueRef};
|
||||
use datatypes::scalars::ScalarVector;
|
||||
use datatypes::types::{
|
||||
@@ -456,13 +457,11 @@ pub fn push_vals(column: &mut Column, origin_count: usize, vector: VectorRef) {
|
||||
TimeUnit::Microsecond => values.time_microsecond_values.push(val.value()),
|
||||
TimeUnit::Nanosecond => values.time_nanosecond_values.push(val.value()),
|
||||
},
|
||||
Value::Interval(val) => match val.unit() {
|
||||
IntervalUnit::YearMonth => values.interval_year_month_values.push(val.to_i32()),
|
||||
IntervalUnit::DayTime => values.interval_day_time_values.push(val.to_i64()),
|
||||
IntervalUnit::MonthDayNano => values
|
||||
.interval_month_day_nano_values
|
||||
.push(convert_i128_to_interval(val.to_i128())),
|
||||
},
|
||||
Value::IntervalYearMonth(val) => values.interval_year_month_values.push(val.to_i32()),
|
||||
Value::IntervalDayTime(val) => values.interval_day_time_values.push(val.to_i64()),
|
||||
Value::IntervalMonthDayNano(val) => values
|
||||
.interval_month_day_nano_values
|
||||
.push(convert_month_day_nano_to_pb(val)),
|
||||
Value::Decimal128(val) => values.decimal128_values.push(convert_to_pb_decimal128(val)),
|
||||
Value::List(_) | Value::Duration(_) => unreachable!(),
|
||||
});
|
||||
@@ -507,14 +506,12 @@ fn ddl_request_type(request: &DdlRequest) -> &'static str {
|
||||
}
|
||||
}
|
||||
|
||||
/// Converts an i128 value to google protobuf type [IntervalMonthDayNano].
|
||||
pub fn convert_i128_to_interval(v: i128) -> v1::IntervalMonthDayNano {
|
||||
let interval = Interval::from_i128(v);
|
||||
let (months, days, nanoseconds) = interval.to_month_day_nano();
|
||||
/// Converts an interval to google protobuf type [IntervalMonthDayNano].
|
||||
pub fn convert_month_day_nano_to_pb(v: IntervalMonthDayNano) -> v1::IntervalMonthDayNano {
|
||||
v1::IntervalMonthDayNano {
|
||||
months,
|
||||
days,
|
||||
nanoseconds,
|
||||
months: v.months,
|
||||
days: v.days,
|
||||
nanoseconds: v.nanoseconds,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -562,11 +559,15 @@ pub fn pb_value_to_value_ref<'a>(
|
||||
ValueData::TimeMillisecondValue(t) => ValueRef::Time(Time::new_millisecond(*t)),
|
||||
ValueData::TimeMicrosecondValue(t) => ValueRef::Time(Time::new_microsecond(*t)),
|
||||
ValueData::TimeNanosecondValue(t) => ValueRef::Time(Time::new_nanosecond(*t)),
|
||||
ValueData::IntervalYearMonthValue(v) => ValueRef::Interval(Interval::from_i32(*v)),
|
||||
ValueData::IntervalDayTimeValue(v) => ValueRef::Interval(Interval::from_i64(*v)),
|
||||
ValueData::IntervalYearMonthValue(v) => {
|
||||
ValueRef::IntervalYearMonth(IntervalYearMonth::from_i32(*v))
|
||||
}
|
||||
ValueData::IntervalDayTimeValue(v) => {
|
||||
ValueRef::IntervalDayTime(IntervalDayTime::from_i64(*v))
|
||||
}
|
||||
ValueData::IntervalMonthDayNanoValue(v) => {
|
||||
let interval = Interval::from_month_day_nano(v.months, v.days, v.nanoseconds);
|
||||
ValueRef::Interval(interval)
|
||||
let interval = IntervalMonthDayNano::new(v.months, v.days, v.nanoseconds);
|
||||
ValueRef::IntervalMonthDayNano(interval)
|
||||
}
|
||||
ValueData::Decimal128Value(v) => {
|
||||
// get precision and scale from datatype_extension
|
||||
@@ -657,7 +658,7 @@ pub fn pb_values_to_vector_ref(data_type: &ConcreteDataType, values: Values) ->
|
||||
IntervalType::MonthDayNano(_) => {
|
||||
Arc::new(IntervalMonthDayNanoVector::from_iter_values(
|
||||
values.interval_month_day_nano_values.iter().map(|x| {
|
||||
Interval::from_month_day_nano(x.months, x.days, x.nanoseconds).to_i128()
|
||||
IntervalMonthDayNano::new(x.months, x.days, x.nanoseconds).to_i128()
|
||||
}),
|
||||
))
|
||||
}
|
||||
@@ -802,18 +803,18 @@ pub fn pb_values_to_values(data_type: &ConcreteDataType, values: Values) -> Vec<
|
||||
ConcreteDataType::Interval(IntervalType::YearMonth(_)) => values
|
||||
.interval_year_month_values
|
||||
.into_iter()
|
||||
.map(|v| Value::Interval(Interval::from_i32(v)))
|
||||
.map(|v| Value::IntervalYearMonth(IntervalYearMonth::from_i32(v)))
|
||||
.collect(),
|
||||
ConcreteDataType::Interval(IntervalType::DayTime(_)) => values
|
||||
.interval_day_time_values
|
||||
.into_iter()
|
||||
.map(|v| Value::Interval(Interval::from_i64(v)))
|
||||
.map(|v| Value::IntervalDayTime(IntervalDayTime::from_i64(v)))
|
||||
.collect(),
|
||||
ConcreteDataType::Interval(IntervalType::MonthDayNano(_)) => values
|
||||
.interval_month_day_nano_values
|
||||
.into_iter()
|
||||
.map(|v| {
|
||||
Value::Interval(Interval::from_month_day_nano(
|
||||
Value::IntervalMonthDayNano(IntervalMonthDayNano::new(
|
||||
v.months,
|
||||
v.days,
|
||||
v.nanoseconds,
|
||||
@@ -941,18 +942,16 @@ pub fn to_proto_value(value: Value) -> Option<v1::Value> {
|
||||
value_data: Some(ValueData::TimeNanosecondValue(v.value())),
|
||||
},
|
||||
},
|
||||
Value::Interval(v) => match v.unit() {
|
||||
IntervalUnit::YearMonth => v1::Value {
|
||||
value_data: Some(ValueData::IntervalYearMonthValue(v.to_i32())),
|
||||
},
|
||||
IntervalUnit::DayTime => v1::Value {
|
||||
value_data: Some(ValueData::IntervalDayTimeValue(v.to_i64())),
|
||||
},
|
||||
IntervalUnit::MonthDayNano => v1::Value {
|
||||
value_data: Some(ValueData::IntervalMonthDayNanoValue(
|
||||
convert_i128_to_interval(v.to_i128()),
|
||||
)),
|
||||
},
|
||||
Value::IntervalYearMonth(v) => v1::Value {
|
||||
value_data: Some(ValueData::IntervalYearMonthValue(v.to_i32())),
|
||||
},
|
||||
Value::IntervalDayTime(v) => v1::Value {
|
||||
value_data: Some(ValueData::IntervalDayTimeValue(v.to_i64())),
|
||||
},
|
||||
Value::IntervalMonthDayNano(v) => v1::Value {
|
||||
value_data: Some(ValueData::IntervalMonthDayNanoValue(
|
||||
convert_month_day_nano_to_pb(v),
|
||||
)),
|
||||
},
|
||||
Value::Decimal128(v) => v1::Value {
|
||||
value_data: Some(ValueData::Decimal128Value(convert_to_pb_decimal128(v))),
|
||||
@@ -1044,13 +1043,11 @@ pub fn value_to_grpc_value(value: Value) -> GrpcValue {
|
||||
TimeUnit::Microsecond => ValueData::TimeMicrosecondValue(v.value()),
|
||||
TimeUnit::Nanosecond => ValueData::TimeNanosecondValue(v.value()),
|
||||
}),
|
||||
Value::Interval(v) => Some(match v.unit() {
|
||||
IntervalUnit::YearMonth => ValueData::IntervalYearMonthValue(v.to_i32()),
|
||||
IntervalUnit::DayTime => ValueData::IntervalDayTimeValue(v.to_i64()),
|
||||
IntervalUnit::MonthDayNano => {
|
||||
ValueData::IntervalMonthDayNanoValue(convert_i128_to_interval(v.to_i128()))
|
||||
}
|
||||
}),
|
||||
Value::IntervalYearMonth(v) => Some(ValueData::IntervalYearMonthValue(v.to_i32())),
|
||||
Value::IntervalDayTime(v) => Some(ValueData::IntervalDayTimeValue(v.to_i64())),
|
||||
Value::IntervalMonthDayNano(v) => Some(ValueData::IntervalMonthDayNanoValue(
|
||||
convert_month_day_nano_to_pb(v),
|
||||
)),
|
||||
Value::Decimal128(v) => Some(ValueData::Decimal128Value(convert_to_pb_decimal128(v))),
|
||||
Value::List(_) | Value::Duration(_) => unreachable!(),
|
||||
},
|
||||
@@ -1061,6 +1058,7 @@ pub fn value_to_grpc_value(value: Value) -> GrpcValue {
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use common_time::interval::IntervalUnit;
|
||||
use datatypes::types::{
|
||||
Int32Type, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalYearMonthType,
|
||||
TimeMillisecondType, TimeSecondType, TimestampMillisecondType, TimestampSecondType,
|
||||
@@ -1506,11 +1504,11 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_convert_i128_to_interval() {
|
||||
let i128_val = 3000;
|
||||
let interval = convert_i128_to_interval(i128_val);
|
||||
let i128_val = 3;
|
||||
let interval = convert_month_day_nano_to_pb(IntervalMonthDayNano::from_i128(i128_val));
|
||||
assert_eq!(interval.months, 0);
|
||||
assert_eq!(interval.days, 0);
|
||||
assert_eq!(interval.nanoseconds, 3000);
|
||||
assert_eq!(interval.nanoseconds, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1590,9 +1588,9 @@ mod tests {
|
||||
},
|
||||
);
|
||||
let expect = vec![
|
||||
Value::Interval(Interval::from_year_month(1_i32)),
|
||||
Value::Interval(Interval::from_year_month(2_i32)),
|
||||
Value::Interval(Interval::from_year_month(3_i32)),
|
||||
Value::IntervalYearMonth(IntervalYearMonth::new(1_i32)),
|
||||
Value::IntervalYearMonth(IntervalYearMonth::new(2_i32)),
|
||||
Value::IntervalYearMonth(IntervalYearMonth::new(3_i32)),
|
||||
];
|
||||
assert_eq!(expect, actual);
|
||||
|
||||
@@ -1605,9 +1603,9 @@ mod tests {
|
||||
},
|
||||
);
|
||||
let expect = vec![
|
||||
Value::Interval(Interval::from_i64(1_i64)),
|
||||
Value::Interval(Interval::from_i64(2_i64)),
|
||||
Value::Interval(Interval::from_i64(3_i64)),
|
||||
Value::IntervalDayTime(IntervalDayTime::from_i64(1_i64)),
|
||||
Value::IntervalDayTime(IntervalDayTime::from_i64(2_i64)),
|
||||
Value::IntervalDayTime(IntervalDayTime::from_i64(3_i64)),
|
||||
];
|
||||
assert_eq!(expect, actual);
|
||||
|
||||
@@ -1636,9 +1634,9 @@ mod tests {
|
||||
},
|
||||
);
|
||||
let expect = vec![
|
||||
Value::Interval(Interval::from_month_day_nano(1, 2, 3)),
|
||||
Value::Interval(Interval::from_month_day_nano(5, 6, 7)),
|
||||
Value::Interval(Interval::from_month_day_nano(9, 10, 11)),
|
||||
Value::IntervalMonthDayNano(IntervalMonthDayNano::new(1, 2, 3)),
|
||||
Value::IntervalMonthDayNano(IntervalMonthDayNano::new(5, 6, 7)),
|
||||
Value::IntervalMonthDayNano(IntervalMonthDayNano::new(9, 10, 11)),
|
||||
];
|
||||
assert_eq!(expect, actual);
|
||||
}
|
||||
|
||||
@@ -89,9 +89,8 @@ pub enum Error {
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("Failed to get procedure client in {mode} mode"))]
|
||||
GetProcedureClient {
|
||||
mode: String,
|
||||
#[snafu(display("Failed to get information extension client"))]
|
||||
GetInformationExtension {
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
},
|
||||
@@ -301,7 +300,7 @@ impl ErrorExt for Error {
|
||||
| Error::CacheNotFound { .. }
|
||||
| Error::CastManager { .. }
|
||||
| Error::Json { .. }
|
||||
| Error::GetProcedureClient { .. }
|
||||
| Error::GetInformationExtension { .. }
|
||||
| Error::ProcedureIdNotFound { .. } => StatusCode::Unexpected,
|
||||
|
||||
Error::ViewPlanColumnsChanged { .. } => StatusCode::InvalidArguments,
|
||||
|
||||
@@ -21,7 +21,6 @@ use common_catalog::consts::{
|
||||
DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, INFORMATION_SCHEMA_NAME, NUMBERS_TABLE_ID,
|
||||
PG_CATALOG_NAME,
|
||||
};
|
||||
use common_config::Mode;
|
||||
use common_error::ext::BoxedError;
|
||||
use common_meta::cache::{LayeredCacheRegistryRef, ViewInfoCacheRef};
|
||||
use common_meta::key::catalog_name::CatalogNameKey;
|
||||
@@ -34,7 +33,6 @@ use common_meta::kv_backend::KvBackendRef;
|
||||
use common_procedure::ProcedureManagerRef;
|
||||
use futures_util::stream::BoxStream;
|
||||
use futures_util::{StreamExt, TryStreamExt};
|
||||
use meta_client::client::MetaClient;
|
||||
use moka::sync::Cache;
|
||||
use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef};
|
||||
use session::context::{Channel, QueryContext};
|
||||
@@ -50,7 +48,7 @@ use crate::error::{
|
||||
CacheNotFoundSnafu, GetTableCacheSnafu, InvalidTableInfoInCatalogSnafu, ListCatalogsSnafu,
|
||||
ListSchemasSnafu, ListTablesSnafu, Result, TableMetadataManagerSnafu,
|
||||
};
|
||||
use crate::information_schema::InformationSchemaProvider;
|
||||
use crate::information_schema::{InformationExtensionRef, InformationSchemaProvider};
|
||||
use crate::kvbackend::TableCacheRef;
|
||||
use crate::system_schema::pg_catalog::PGCatalogProvider;
|
||||
use crate::system_schema::SystemSchemaProvider;
|
||||
@@ -63,9 +61,8 @@ use crate::CatalogManager;
|
||||
/// comes from `SystemCatalog`, which is static and read-only.
|
||||
#[derive(Clone)]
|
||||
pub struct KvBackendCatalogManager {
|
||||
mode: Mode,
|
||||
/// Only available in `Distributed` mode.
|
||||
meta_client: Option<Arc<MetaClient>>,
|
||||
/// Provides the extension methods for the `information_schema` tables
|
||||
information_extension: InformationExtensionRef,
|
||||
/// Manages partition rules.
|
||||
partition_manager: PartitionRuleManagerRef,
|
||||
/// Manages table metadata.
|
||||
@@ -82,15 +79,13 @@ const CATALOG_CACHE_MAX_CAPACITY: u64 = 128;
|
||||
|
||||
impl KvBackendCatalogManager {
|
||||
pub fn new(
|
||||
mode: Mode,
|
||||
meta_client: Option<Arc<MetaClient>>,
|
||||
information_extension: InformationExtensionRef,
|
||||
backend: KvBackendRef,
|
||||
cache_registry: LayeredCacheRegistryRef,
|
||||
procedure_manager: Option<ProcedureManagerRef>,
|
||||
) -> Arc<Self> {
|
||||
Arc::new_cyclic(|me| Self {
|
||||
mode,
|
||||
meta_client,
|
||||
information_extension,
|
||||
partition_manager: Arc::new(PartitionRuleManager::new(
|
||||
backend.clone(),
|
||||
cache_registry
|
||||
@@ -118,20 +113,15 @@ impl KvBackendCatalogManager {
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the server running mode.
|
||||
pub fn running_mode(&self) -> &Mode {
|
||||
&self.mode
|
||||
}
|
||||
|
||||
pub fn view_info_cache(&self) -> Result<ViewInfoCacheRef> {
|
||||
self.cache_registry.get().context(CacheNotFoundSnafu {
|
||||
name: "view_info_cache",
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the `[MetaClient]`.
|
||||
pub fn meta_client(&self) -> Option<Arc<MetaClient>> {
|
||||
self.meta_client.clone()
|
||||
/// Returns the [`InformationExtension`].
|
||||
pub fn information_extension(&self) -> InformationExtensionRef {
|
||||
self.information_extension.clone()
|
||||
}
|
||||
|
||||
pub fn partition_manager(&self) -> PartitionRuleManagerRef {
|
||||
|
||||
@@ -32,7 +32,11 @@ use std::collections::HashMap;
|
||||
use std::sync::{Arc, Weak};
|
||||
|
||||
use common_catalog::consts::{self, DEFAULT_CATALOG_NAME, INFORMATION_SCHEMA_NAME};
|
||||
use common_error::ext::ErrorExt;
|
||||
use common_meta::cluster::NodeInfo;
|
||||
use common_meta::datanode::RegionStat;
|
||||
use common_meta::key::flow::FlowMetadataManager;
|
||||
use common_procedure::ProcedureInfo;
|
||||
use common_recordbatch::SendableRecordBatchStream;
|
||||
use datatypes::schema::SchemaRef;
|
||||
use lazy_static::lazy_static;
|
||||
@@ -45,7 +49,7 @@ use views::InformationSchemaViews;
|
||||
|
||||
use self::columns::InformationSchemaColumns;
|
||||
use super::{SystemSchemaProviderInner, SystemTable, SystemTableRef};
|
||||
use crate::error::Result;
|
||||
use crate::error::{Error, Result};
|
||||
use crate::system_schema::information_schema::cluster_info::InformationSchemaClusterInfo;
|
||||
use crate::system_schema::information_schema::flows::InformationSchemaFlows;
|
||||
use crate::system_schema::information_schema::information_memory_table::get_schema_columns;
|
||||
@@ -318,3 +322,39 @@ where
|
||||
InformationTable::to_stream(self, request)
|
||||
}
|
||||
}
|
||||
|
||||
pub type InformationExtensionRef = Arc<dyn InformationExtension<Error = Error> + Send + Sync>;
|
||||
|
||||
/// The `InformationExtension` trait provides the extension methods for the `information_schema` tables.
|
||||
#[async_trait::async_trait]
|
||||
pub trait InformationExtension {
|
||||
type Error: ErrorExt;
|
||||
|
||||
/// Gets the nodes information.
|
||||
async fn nodes(&self) -> std::result::Result<Vec<NodeInfo>, Self::Error>;
|
||||
|
||||
/// Gets the procedures information.
|
||||
async fn procedures(&self) -> std::result::Result<Vec<(String, ProcedureInfo)>, Self::Error>;
|
||||
|
||||
/// Gets the region statistics.
|
||||
async fn region_stats(&self) -> std::result::Result<Vec<RegionStat>, Self::Error>;
|
||||
}
|
||||
|
||||
pub struct NoopInformationExtension;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl InformationExtension for NoopInformationExtension {
|
||||
type Error = Error;
|
||||
|
||||
async fn nodes(&self) -> std::result::Result<Vec<NodeInfo>, Self::Error> {
|
||||
Ok(vec![])
|
||||
}
|
||||
|
||||
async fn procedures(&self) -> std::result::Result<Vec<(String, ProcedureInfo)>, Self::Error> {
|
||||
Ok(vec![])
|
||||
}
|
||||
|
||||
async fn region_stats(&self) -> std::result::Result<Vec<RegionStat>, Self::Error> {
|
||||
Ok(vec![])
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,13 +17,10 @@ use std::time::Duration;
|
||||
|
||||
use arrow_schema::SchemaRef as ArrowSchemaRef;
|
||||
use common_catalog::consts::INFORMATION_SCHEMA_CLUSTER_INFO_TABLE_ID;
|
||||
use common_config::Mode;
|
||||
use common_error::ext::BoxedError;
|
||||
use common_meta::cluster::{ClusterInfo, NodeInfo, NodeStatus};
|
||||
use common_meta::peer::Peer;
|
||||
use common_meta::cluster::NodeInfo;
|
||||
use common_recordbatch::adapter::RecordBatchStreamAdapter;
|
||||
use common_recordbatch::{RecordBatch, SendableRecordBatchStream};
|
||||
use common_telemetry::warn;
|
||||
use common_time::timestamp::Timestamp;
|
||||
use datafusion::execution::TaskContext;
|
||||
use datafusion::physical_plan::stream::RecordBatchStreamAdapter as DfRecordBatchStreamAdapter;
|
||||
@@ -40,7 +37,7 @@ use snafu::ResultExt;
|
||||
use store_api::storage::{ScanRequest, TableId};
|
||||
|
||||
use super::CLUSTER_INFO;
|
||||
use crate::error::{CreateRecordBatchSnafu, InternalSnafu, ListNodesSnafu, Result};
|
||||
use crate::error::{CreateRecordBatchSnafu, InternalSnafu, Result};
|
||||
use crate::system_schema::information_schema::{InformationTable, Predicates};
|
||||
use crate::system_schema::utils;
|
||||
use crate::CatalogManager;
|
||||
@@ -70,7 +67,6 @@ const INIT_CAPACITY: usize = 42;
|
||||
pub(super) struct InformationSchemaClusterInfo {
|
||||
schema: SchemaRef,
|
||||
catalog_manager: Weak<dyn CatalogManager>,
|
||||
start_time_ms: u64,
|
||||
}
|
||||
|
||||
impl InformationSchemaClusterInfo {
|
||||
@@ -78,7 +74,6 @@ impl InformationSchemaClusterInfo {
|
||||
Self {
|
||||
schema: Self::schema(),
|
||||
catalog_manager,
|
||||
start_time_ms: common_time::util::current_time_millis() as u64,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -100,11 +95,7 @@ impl InformationSchemaClusterInfo {
|
||||
}
|
||||
|
||||
fn builder(&self) -> InformationSchemaClusterInfoBuilder {
|
||||
InformationSchemaClusterInfoBuilder::new(
|
||||
self.schema.clone(),
|
||||
self.catalog_manager.clone(),
|
||||
self.start_time_ms,
|
||||
)
|
||||
InformationSchemaClusterInfoBuilder::new(self.schema.clone(), self.catalog_manager.clone())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -144,7 +135,6 @@ impl InformationTable for InformationSchemaClusterInfo {
|
||||
|
||||
struct InformationSchemaClusterInfoBuilder {
|
||||
schema: SchemaRef,
|
||||
start_time_ms: u64,
|
||||
catalog_manager: Weak<dyn CatalogManager>,
|
||||
|
||||
peer_ids: Int64VectorBuilder,
|
||||
@@ -158,11 +148,7 @@ struct InformationSchemaClusterInfoBuilder {
|
||||
}
|
||||
|
||||
impl InformationSchemaClusterInfoBuilder {
|
||||
fn new(
|
||||
schema: SchemaRef,
|
||||
catalog_manager: Weak<dyn CatalogManager>,
|
||||
start_time_ms: u64,
|
||||
) -> Self {
|
||||
fn new(schema: SchemaRef, catalog_manager: Weak<dyn CatalogManager>) -> Self {
|
||||
Self {
|
||||
schema,
|
||||
catalog_manager,
|
||||
@@ -174,56 +160,17 @@ impl InformationSchemaClusterInfoBuilder {
|
||||
start_times: TimestampMillisecondVectorBuilder::with_capacity(INIT_CAPACITY),
|
||||
uptimes: StringVectorBuilder::with_capacity(INIT_CAPACITY),
|
||||
active_times: StringVectorBuilder::with_capacity(INIT_CAPACITY),
|
||||
start_time_ms,
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct the `information_schema.cluster_info` virtual table
|
||||
async fn make_cluster_info(&mut self, request: Option<ScanRequest>) -> Result<RecordBatch> {
|
||||
let predicates = Predicates::from_scan_request(&request);
|
||||
let mode = utils::running_mode(&self.catalog_manager)?.unwrap_or(Mode::Standalone);
|
||||
|
||||
match mode {
|
||||
Mode::Standalone => {
|
||||
let build_info = common_version::build_info();
|
||||
|
||||
self.add_node_info(
|
||||
&predicates,
|
||||
NodeInfo {
|
||||
// For the standalone:
|
||||
// - id always 0
|
||||
// - empty string for peer_addr
|
||||
peer: Peer {
|
||||
id: 0,
|
||||
addr: "".to_string(),
|
||||
},
|
||||
last_activity_ts: -1,
|
||||
status: NodeStatus::Standalone,
|
||||
version: build_info.version.to_string(),
|
||||
git_commit: build_info.commit_short.to_string(),
|
||||
// Use `self.start_time_ms` instead.
|
||||
// It's not precise but enough.
|
||||
start_time_ms: self.start_time_ms,
|
||||
},
|
||||
);
|
||||
}
|
||||
Mode::Distributed => {
|
||||
if let Some(meta_client) = utils::meta_client(&self.catalog_manager)? {
|
||||
let node_infos = meta_client
|
||||
.list_nodes(None)
|
||||
.await
|
||||
.map_err(BoxedError::new)
|
||||
.context(ListNodesSnafu)?;
|
||||
|
||||
for node_info in node_infos {
|
||||
self.add_node_info(&predicates, node_info);
|
||||
}
|
||||
} else {
|
||||
warn!("Could not find meta client in distributed mode.");
|
||||
}
|
||||
}
|
||||
let information_extension = utils::information_extension(&self.catalog_manager)?;
|
||||
let node_infos = information_extension.nodes().await?;
|
||||
for node_info in node_infos {
|
||||
self.add_node_info(&predicates, node_info);
|
||||
}
|
||||
|
||||
self.finish()
|
||||
}
|
||||
|
||||
|
||||
@@ -14,14 +14,10 @@
|
||||
|
||||
use std::sync::{Arc, Weak};
|
||||
|
||||
use api::v1::meta::{ProcedureMeta, ProcedureStatus};
|
||||
use arrow_schema::SchemaRef as ArrowSchemaRef;
|
||||
use common_catalog::consts::INFORMATION_SCHEMA_PROCEDURE_INFO_TABLE_ID;
|
||||
use common_config::Mode;
|
||||
use common_error::ext::BoxedError;
|
||||
use common_meta::ddl::{ExecutorContext, ProcedureExecutor};
|
||||
use common_meta::rpc::procedure;
|
||||
use common_procedure::{ProcedureInfo, ProcedureState};
|
||||
use common_procedure::ProcedureInfo;
|
||||
use common_recordbatch::adapter::RecordBatchStreamAdapter;
|
||||
use common_recordbatch::{RecordBatch, SendableRecordBatchStream};
|
||||
use common_time::timestamp::Timestamp;
|
||||
@@ -38,10 +34,7 @@ use snafu::ResultExt;
|
||||
use store_api::storage::{ScanRequest, TableId};
|
||||
|
||||
use super::PROCEDURE_INFO;
|
||||
use crate::error::{
|
||||
ConvertProtoDataSnafu, CreateRecordBatchSnafu, GetProcedureClientSnafu, InternalSnafu,
|
||||
ListProceduresSnafu, ProcedureIdNotFoundSnafu, Result,
|
||||
};
|
||||
use crate::error::{CreateRecordBatchSnafu, InternalSnafu, Result};
|
||||
use crate::system_schema::information_schema::{InformationTable, Predicates};
|
||||
use crate::system_schema::utils;
|
||||
use crate::CatalogManager;
|
||||
@@ -167,45 +160,11 @@ impl InformationSchemaProcedureInfoBuilder {
|
||||
/// Construct the `information_schema.procedure_info` virtual table
|
||||
async fn make_procedure_info(&mut self, request: Option<ScanRequest>) -> Result<RecordBatch> {
|
||||
let predicates = Predicates::from_scan_request(&request);
|
||||
let mode = utils::running_mode(&self.catalog_manager)?.unwrap_or(Mode::Standalone);
|
||||
match mode {
|
||||
Mode::Standalone => {
|
||||
if let Some(procedure_manager) = utils::procedure_manager(&self.catalog_manager)? {
|
||||
let procedures = procedure_manager
|
||||
.list_procedures()
|
||||
.await
|
||||
.map_err(BoxedError::new)
|
||||
.context(ListProceduresSnafu)?;
|
||||
for procedure in procedures {
|
||||
self.add_procedure(
|
||||
&predicates,
|
||||
procedure.state.as_str_name().to_string(),
|
||||
procedure,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
return GetProcedureClientSnafu { mode: "standalone" }.fail();
|
||||
}
|
||||
}
|
||||
Mode::Distributed => {
|
||||
if let Some(meta_client) = utils::meta_client(&self.catalog_manager)? {
|
||||
let procedures = meta_client
|
||||
.list_procedures(&ExecutorContext::default())
|
||||
.await
|
||||
.map_err(BoxedError::new)
|
||||
.context(ListProceduresSnafu)?;
|
||||
for procedure in procedures.procedures {
|
||||
self.add_procedure_info(&predicates, procedure)?;
|
||||
}
|
||||
} else {
|
||||
return GetProcedureClientSnafu {
|
||||
mode: "distributed",
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let information_extension = utils::information_extension(&self.catalog_manager)?;
|
||||
let procedures = information_extension.procedures().await?;
|
||||
for (status, procedure_info) in procedures {
|
||||
self.add_procedure(&predicates, status, procedure_info);
|
||||
}
|
||||
self.finish()
|
||||
}
|
||||
|
||||
@@ -247,34 +206,6 @@ impl InformationSchemaProcedureInfoBuilder {
|
||||
self.lock_keys.push(Some(&lock_keys));
|
||||
}
|
||||
|
||||
fn add_procedure_info(
|
||||
&mut self,
|
||||
predicates: &Predicates,
|
||||
procedure: ProcedureMeta,
|
||||
) -> Result<()> {
|
||||
let pid = match procedure.id {
|
||||
Some(pid) => pid,
|
||||
None => return ProcedureIdNotFoundSnafu {}.fail(),
|
||||
};
|
||||
let pid = procedure::pb_pid_to_pid(&pid)
|
||||
.map_err(BoxedError::new)
|
||||
.context(ConvertProtoDataSnafu)?;
|
||||
let status = ProcedureStatus::try_from(procedure.status)
|
||||
.map(|v| v.as_str_name())
|
||||
.unwrap_or("Unknown")
|
||||
.to_string();
|
||||
let procedure_info = ProcedureInfo {
|
||||
id: pid,
|
||||
type_name: procedure.type_name,
|
||||
start_time_ms: procedure.start_time_ms,
|
||||
end_time_ms: procedure.end_time_ms,
|
||||
state: ProcedureState::Running,
|
||||
lock_keys: procedure.lock_keys,
|
||||
};
|
||||
self.add_procedure(predicates, status, procedure_info);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<RecordBatch> {
|
||||
let columns: Vec<VectorRef> = vec![
|
||||
Arc::new(self.procedure_ids.finish()),
|
||||
|
||||
@@ -224,8 +224,8 @@ impl InformationSchemaRegionPeersBuilder {
|
||||
let region_id = RegionId::new(table_id, route.region.id.region_number()).as_u64();
|
||||
let peer_id = route.leader_peer.clone().map(|p| p.id);
|
||||
let peer_addr = route.leader_peer.clone().map(|p| p.addr);
|
||||
let status = if let Some(status) = route.leader_status {
|
||||
Some(status.as_ref().to_string())
|
||||
let state = if let Some(state) = route.leader_state {
|
||||
Some(state.as_ref().to_string())
|
||||
} else {
|
||||
// Alive by default
|
||||
Some("ALIVE".to_string())
|
||||
@@ -242,7 +242,7 @@ impl InformationSchemaRegionPeersBuilder {
|
||||
self.peer_ids.push(peer_id);
|
||||
self.peer_addrs.push(peer_addr.as_deref());
|
||||
self.is_leaders.push(Some("Yes"));
|
||||
self.statuses.push(status.as_deref());
|
||||
self.statuses.push(state.as_deref());
|
||||
self.down_seconds
|
||||
.push(route.leader_down_millis().map(|m| m / 1000));
|
||||
}
|
||||
|
||||
@@ -16,13 +16,10 @@ use std::sync::{Arc, Weak};
|
||||
|
||||
use arrow_schema::SchemaRef as ArrowSchemaRef;
|
||||
use common_catalog::consts::INFORMATION_SCHEMA_REGION_STATISTICS_TABLE_ID;
|
||||
use common_config::Mode;
|
||||
use common_error::ext::BoxedError;
|
||||
use common_meta::cluster::ClusterInfo;
|
||||
use common_meta::datanode::RegionStat;
|
||||
use common_recordbatch::adapter::RecordBatchStreamAdapter;
|
||||
use common_recordbatch::{DfSendableRecordBatchStream, RecordBatch, SendableRecordBatchStream};
|
||||
use common_telemetry::tracing::warn;
|
||||
use datafusion::execution::TaskContext;
|
||||
use datafusion::physical_plan::stream::RecordBatchStreamAdapter as DfRecordBatchStreamAdapter;
|
||||
use datafusion::physical_plan::streaming::PartitionStream as DfPartitionStream;
|
||||
@@ -34,7 +31,7 @@ use snafu::ResultExt;
|
||||
use store_api::storage::{ScanRequest, TableId};
|
||||
|
||||
use super::{InformationTable, REGION_STATISTICS};
|
||||
use crate::error::{CreateRecordBatchSnafu, InternalSnafu, ListRegionStatsSnafu, Result};
|
||||
use crate::error::{CreateRecordBatchSnafu, InternalSnafu, Result};
|
||||
use crate::information_schema::Predicates;
|
||||
use crate::system_schema::utils;
|
||||
use crate::CatalogManager;
|
||||
@@ -167,28 +164,11 @@ impl InformationSchemaRegionStatisticsBuilder {
|
||||
request: Option<ScanRequest>,
|
||||
) -> Result<RecordBatch> {
|
||||
let predicates = Predicates::from_scan_request(&request);
|
||||
let mode = utils::running_mode(&self.catalog_manager)?.unwrap_or(Mode::Standalone);
|
||||
|
||||
match mode {
|
||||
Mode::Standalone => {
|
||||
// TODO(weny): implement it
|
||||
}
|
||||
Mode::Distributed => {
|
||||
if let Some(meta_client) = utils::meta_client(&self.catalog_manager)? {
|
||||
let region_stats = meta_client
|
||||
.list_region_stats()
|
||||
.await
|
||||
.map_err(BoxedError::new)
|
||||
.context(ListRegionStatsSnafu)?;
|
||||
for region_stat in region_stats {
|
||||
self.add_region_statistic(&predicates, region_stat);
|
||||
}
|
||||
} else {
|
||||
warn!("Meta client is not available");
|
||||
}
|
||||
}
|
||||
let information_extension = utils::information_extension(&self.catalog_manager)?;
|
||||
let region_stats = information_extension.region_stats().await?;
|
||||
for region_stat in region_stats {
|
||||
self.add_region_statistic(&predicates, region_stat);
|
||||
}
|
||||
|
||||
self.finish()
|
||||
}
|
||||
|
||||
|
||||
@@ -12,48 +12,33 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
pub mod tables;
|
||||
use std::sync::Weak;
|
||||
|
||||
use std::sync::{Arc, Weak};
|
||||
|
||||
use common_config::Mode;
|
||||
use common_meta::key::TableMetadataManagerRef;
|
||||
use common_procedure::ProcedureManagerRef;
|
||||
use meta_client::client::MetaClient;
|
||||
use snafu::OptionExt;
|
||||
|
||||
use crate::error::{Result, UpgradeWeakCatalogManagerRefSnafu};
|
||||
use crate::error::{GetInformationExtensionSnafu, Result, UpgradeWeakCatalogManagerRefSnafu};
|
||||
use crate::information_schema::InformationExtensionRef;
|
||||
use crate::kvbackend::KvBackendCatalogManager;
|
||||
use crate::CatalogManager;
|
||||
|
||||
/// Try to get the server running mode from `[CatalogManager]` weak reference.
|
||||
pub fn running_mode(catalog_manager: &Weak<dyn CatalogManager>) -> Result<Option<Mode>> {
|
||||
pub mod tables;
|
||||
|
||||
/// Try to get the `[InformationExtension]` from `[CatalogManager]` weak reference.
|
||||
pub fn information_extension(
|
||||
catalog_manager: &Weak<dyn CatalogManager>,
|
||||
) -> Result<InformationExtensionRef> {
|
||||
let catalog_manager = catalog_manager
|
||||
.upgrade()
|
||||
.context(UpgradeWeakCatalogManagerRefSnafu)?;
|
||||
|
||||
Ok(catalog_manager
|
||||
let information_extension = catalog_manager
|
||||
.as_any()
|
||||
.downcast_ref::<KvBackendCatalogManager>()
|
||||
.map(|manager| manager.running_mode())
|
||||
.copied())
|
||||
}
|
||||
.map(|manager| manager.information_extension())
|
||||
.context(GetInformationExtensionSnafu)?;
|
||||
|
||||
/// Try to get the `[MetaClient]` from `[CatalogManager]` weak reference.
|
||||
pub fn meta_client(catalog_manager: &Weak<dyn CatalogManager>) -> Result<Option<Arc<MetaClient>>> {
|
||||
let catalog_manager = catalog_manager
|
||||
.upgrade()
|
||||
.context(UpgradeWeakCatalogManagerRefSnafu)?;
|
||||
|
||||
let meta_client = match catalog_manager
|
||||
.as_any()
|
||||
.downcast_ref::<KvBackendCatalogManager>()
|
||||
{
|
||||
None => None,
|
||||
Some(manager) => manager.meta_client(),
|
||||
};
|
||||
|
||||
Ok(meta_client)
|
||||
Ok(information_extension)
|
||||
}
|
||||
|
||||
/// Try to get the `[TableMetadataManagerRef]` from `[CatalogManager]` weak reference.
|
||||
@@ -69,17 +54,3 @@ pub fn table_meta_manager(
|
||||
.downcast_ref::<KvBackendCatalogManager>()
|
||||
.map(|manager| manager.table_metadata_manager_ref().clone()))
|
||||
}
|
||||
|
||||
/// Try to get the `[ProcedureManagerRef]` from `[CatalogManager]` weak reference.
|
||||
pub fn procedure_manager(
|
||||
catalog_manager: &Weak<dyn CatalogManager>,
|
||||
) -> Result<Option<ProcedureManagerRef>> {
|
||||
let catalog_manager = catalog_manager
|
||||
.upgrade()
|
||||
.context(UpgradeWeakCatalogManagerRefSnafu)?;
|
||||
|
||||
Ok(catalog_manager
|
||||
.as_any()
|
||||
.downcast_ref::<KvBackendCatalogManager>()
|
||||
.and_then(|manager| manager.procedure_manager()))
|
||||
}
|
||||
|
||||
@@ -259,7 +259,6 @@ mod tests {
|
||||
|
||||
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
|
||||
use cache::{build_fundamental_cache_registry, with_default_composite_cache_registry};
|
||||
use common_config::Mode;
|
||||
use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder};
|
||||
use common_meta::key::TableMetadataManager;
|
||||
use common_meta::kv_backend::memory::MemoryKvBackend;
|
||||
@@ -269,6 +268,8 @@ mod tests {
|
||||
use datafusion::logical_expr::builder::LogicalTableSource;
|
||||
use datafusion::logical_expr::{col, lit, LogicalPlan, LogicalPlanBuilder};
|
||||
|
||||
use crate::information_schema::NoopInformationExtension;
|
||||
|
||||
struct MockDecoder;
|
||||
impl MockDecoder {
|
||||
pub fn arc() -> Arc<Self> {
|
||||
@@ -323,8 +324,7 @@ mod tests {
|
||||
);
|
||||
|
||||
let catalog_manager = KvBackendCatalogManager::new(
|
||||
Mode::Standalone,
|
||||
None,
|
||||
Arc::new(NoopInformationExtension),
|
||||
backend.clone(),
|
||||
layered_cache_registry,
|
||||
None,
|
||||
|
||||
@@ -10,7 +10,7 @@ name = "greptime"
|
||||
path = "src/bin/greptime.rs"
|
||||
|
||||
[features]
|
||||
default = ["python"]
|
||||
default = ["python", "servers/pprof", "servers/mem-prof"]
|
||||
tokio-console = ["common-telemetry/tokio-console"]
|
||||
python = ["frontend/python"]
|
||||
|
||||
|
||||
@@ -158,7 +158,7 @@ fn create_region_routes(regions: Vec<RegionNumber>) -> Vec<RegionRoute> {
|
||||
addr: String::new(),
|
||||
}),
|
||||
follower_peers: vec![],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -46,12 +46,12 @@ use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
|
||||
use crate::cli::cmd::ReplCommand;
|
||||
use crate::cli::helper::RustylineHelper;
|
||||
use crate::cli::AttachCommand;
|
||||
use crate::error;
|
||||
use crate::error::{
|
||||
CollectRecordBatchesSnafu, ParseSqlSnafu, PlanStatementSnafu, PrettyPrintRecordBatchesSnafu,
|
||||
ReadlineSnafu, ReplCreationSnafu, RequestDatabaseSnafu, Result, StartMetaClientSnafu,
|
||||
SubstraitEncodeLogicalPlanSnafu,
|
||||
};
|
||||
use crate::{error, DistributedInformationExtension};
|
||||
|
||||
/// Captures the state of the repl, gathers commands and executes them one by one
|
||||
pub struct Repl {
|
||||
@@ -275,9 +275,9 @@ async fn create_query_engine(meta_addr: &str) -> Result<DatafusionQueryEngine> {
|
||||
.build(),
|
||||
);
|
||||
|
||||
let information_extension = Arc::new(DistributedInformationExtension::new(meta_client.clone()));
|
||||
let catalog_manager = KvBackendCatalogManager::new(
|
||||
Mode::Distributed,
|
||||
Some(meta_client.clone()),
|
||||
information_extension,
|
||||
cached_meta_backend.clone(),
|
||||
layered_cache_registry,
|
||||
None,
|
||||
|
||||
@@ -41,7 +41,7 @@ use crate::error::{
|
||||
MissingConfigSnafu, Result, ShutdownFlownodeSnafu, StartFlownodeSnafu,
|
||||
};
|
||||
use crate::options::{GlobalOptions, GreptimeOptions};
|
||||
use crate::{log_versions, App};
|
||||
use crate::{log_versions, App, DistributedInformationExtension};
|
||||
|
||||
pub const APP_NAME: &str = "greptime-flownode";
|
||||
|
||||
@@ -269,9 +269,10 @@ impl StartCommand {
|
||||
.build(),
|
||||
);
|
||||
|
||||
let information_extension =
|
||||
Arc::new(DistributedInformationExtension::new(meta_client.clone()));
|
||||
let catalog_manager = KvBackendCatalogManager::new(
|
||||
opts.mode,
|
||||
Some(meta_client.clone()),
|
||||
information_extension,
|
||||
cached_meta_backend.clone(),
|
||||
layered_cache_registry.clone(),
|
||||
None,
|
||||
|
||||
@@ -36,8 +36,8 @@ use frontend::instance::builder::FrontendBuilder;
|
||||
use frontend::instance::{FrontendInstance, Instance as FeInstance};
|
||||
use frontend::server::Services;
|
||||
use meta_client::{MetaClientOptions, MetaClientType};
|
||||
use query::stats::StatementStatistics;
|
||||
use servers::tls::{TlsMode, TlsOption};
|
||||
use servers::Mode;
|
||||
use snafu::{OptionExt, ResultExt};
|
||||
use tracing_appender::non_blocking::WorkerGuard;
|
||||
|
||||
@@ -46,7 +46,7 @@ use crate::error::{
|
||||
Result, StartFrontendSnafu,
|
||||
};
|
||||
use crate::options::{GlobalOptions, GreptimeOptions};
|
||||
use crate::{log_versions, App};
|
||||
use crate::{log_versions, App, DistributedInformationExtension};
|
||||
|
||||
type FrontendOptions = GreptimeOptions<frontend::frontend::FrontendOptions>;
|
||||
|
||||
@@ -315,9 +315,10 @@ impl StartCommand {
|
||||
.build(),
|
||||
);
|
||||
|
||||
let information_extension =
|
||||
Arc::new(DistributedInformationExtension::new(meta_client.clone()));
|
||||
let catalog_manager = KvBackendCatalogManager::new(
|
||||
Mode::Distributed,
|
||||
Some(meta_client.clone()),
|
||||
information_extension,
|
||||
cached_meta_backend.clone(),
|
||||
layered_cache_registry.clone(),
|
||||
None,
|
||||
@@ -352,6 +353,7 @@ impl StartCommand {
|
||||
catalog_manager,
|
||||
Arc::new(client),
|
||||
meta_client,
|
||||
StatementStatistics::new(opts.logging.slow_query.clone()),
|
||||
)
|
||||
.with_plugin(plugins.clone())
|
||||
.with_local_cache_invalidator(layered_cache_registry)
|
||||
|
||||
@@ -15,7 +15,17 @@
|
||||
#![feature(assert_matches, let_chains)]
|
||||
|
||||
use async_trait::async_trait;
|
||||
use catalog::information_schema::InformationExtension;
|
||||
use client::api::v1::meta::ProcedureStatus;
|
||||
use common_error::ext::BoxedError;
|
||||
use common_meta::cluster::{ClusterInfo, NodeInfo};
|
||||
use common_meta::datanode::RegionStat;
|
||||
use common_meta::ddl::{ExecutorContext, ProcedureExecutor};
|
||||
use common_meta::rpc::procedure;
|
||||
use common_procedure::{ProcedureInfo, ProcedureState};
|
||||
use common_telemetry::{error, info};
|
||||
use meta_client::MetaClientRef;
|
||||
use snafu::ResultExt;
|
||||
|
||||
use crate::error::Result;
|
||||
|
||||
@@ -94,3 +104,69 @@ fn log_env_flags() {
|
||||
info!("argument: {}", argument);
|
||||
}
|
||||
}
|
||||
|
||||
pub struct DistributedInformationExtension {
|
||||
meta_client: MetaClientRef,
|
||||
}
|
||||
|
||||
impl DistributedInformationExtension {
|
||||
pub fn new(meta_client: MetaClientRef) -> Self {
|
||||
Self { meta_client }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl InformationExtension for DistributedInformationExtension {
|
||||
type Error = catalog::error::Error;
|
||||
|
||||
async fn nodes(&self) -> std::result::Result<Vec<NodeInfo>, Self::Error> {
|
||||
self.meta_client
|
||||
.list_nodes(None)
|
||||
.await
|
||||
.map_err(BoxedError::new)
|
||||
.context(catalog::error::ListNodesSnafu)
|
||||
}
|
||||
|
||||
async fn procedures(&self) -> std::result::Result<Vec<(String, ProcedureInfo)>, Self::Error> {
|
||||
let procedures = self
|
||||
.meta_client
|
||||
.list_procedures(&ExecutorContext::default())
|
||||
.await
|
||||
.map_err(BoxedError::new)
|
||||
.context(catalog::error::ListProceduresSnafu)?
|
||||
.procedures;
|
||||
let mut result = Vec::with_capacity(procedures.len());
|
||||
for procedure in procedures {
|
||||
let pid = match procedure.id {
|
||||
Some(pid) => pid,
|
||||
None => return catalog::error::ProcedureIdNotFoundSnafu {}.fail(),
|
||||
};
|
||||
let pid = procedure::pb_pid_to_pid(&pid)
|
||||
.map_err(BoxedError::new)
|
||||
.context(catalog::error::ConvertProtoDataSnafu)?;
|
||||
let status = ProcedureStatus::try_from(procedure.status)
|
||||
.map(|v| v.as_str_name())
|
||||
.unwrap_or("Unknown")
|
||||
.to_string();
|
||||
let procedure_info = ProcedureInfo {
|
||||
id: pid,
|
||||
type_name: procedure.type_name,
|
||||
start_time_ms: procedure.start_time_ms,
|
||||
end_time_ms: procedure.end_time_ms,
|
||||
state: ProcedureState::Running,
|
||||
lock_keys: procedure.lock_keys,
|
||||
};
|
||||
result.push((status, procedure_info));
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn region_stats(&self) -> std::result::Result<Vec<RegionStat>, Self::Error> {
|
||||
self.meta_client
|
||||
.list_region_stats()
|
||||
.await
|
||||
.map_err(BoxedError::new)
|
||||
.context(catalog::error::ListRegionStatsSnafu)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -48,6 +48,10 @@ impl Instance {
|
||||
_guard: guard,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_inner(&self) -> &MetasrvInstance {
|
||||
&self.instance
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -86,6 +90,14 @@ impl Command {
|
||||
pub fn load_options(&self, global_options: &GlobalOptions) -> Result<MetasrvOptions> {
|
||||
self.subcmd.load_options(global_options)
|
||||
}
|
||||
|
||||
pub fn config_file(&self) -> &Option<String> {
|
||||
self.subcmd.config_file()
|
||||
}
|
||||
|
||||
pub fn env_prefix(&self) -> &String {
|
||||
self.subcmd.env_prefix()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Parser)]
|
||||
@@ -105,6 +117,18 @@ impl SubCommand {
|
||||
SubCommand::Start(cmd) => cmd.load_options(global_options),
|
||||
}
|
||||
}
|
||||
|
||||
fn config_file(&self) -> &Option<String> {
|
||||
match self {
|
||||
SubCommand::Start(cmd) => &cmd.config_file,
|
||||
}
|
||||
}
|
||||
|
||||
fn env_prefix(&self) -> &String {
|
||||
match self {
|
||||
SubCommand::Start(cmd) => &cmd.env_prefix,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Parser)]
|
||||
|
||||
@@ -17,14 +17,18 @@ use std::{fs, path};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use cache::{build_fundamental_cache_registry, with_default_composite_cache_registry};
|
||||
use catalog::information_schema::InformationExtension;
|
||||
use catalog::kvbackend::KvBackendCatalogManager;
|
||||
use clap::Parser;
|
||||
use client::api::v1::meta::RegionRole;
|
||||
use common_base::Plugins;
|
||||
use common_catalog::consts::{MIN_USER_FLOW_ID, MIN_USER_TABLE_ID};
|
||||
use common_config::{metadata_store_dir, Configurable, KvBackendConfig};
|
||||
use common_error::ext::BoxedError;
|
||||
use common_meta::cache::LayeredCacheRegistryBuilder;
|
||||
use common_meta::cache_invalidator::CacheInvalidatorRef;
|
||||
use common_meta::cluster::{NodeInfo, NodeStatus};
|
||||
use common_meta::datanode::RegionStat;
|
||||
use common_meta::ddl::flow_meta::{FlowMetadataAllocator, FlowMetadataAllocatorRef};
|
||||
use common_meta::ddl::table_meta::{TableMetadataAllocator, TableMetadataAllocatorRef};
|
||||
use common_meta::ddl::{DdlContext, NoopRegionFailureDetectorControl, ProcedureExecutorRef};
|
||||
@@ -33,10 +37,11 @@ use common_meta::key::flow::{FlowMetadataManager, FlowMetadataManagerRef};
|
||||
use common_meta::key::{TableMetadataManager, TableMetadataManagerRef};
|
||||
use common_meta::kv_backend::KvBackendRef;
|
||||
use common_meta::node_manager::NodeManagerRef;
|
||||
use common_meta::peer::Peer;
|
||||
use common_meta::region_keeper::MemoryRegionKeeper;
|
||||
use common_meta::sequence::SequenceBuilder;
|
||||
use common_meta::wal_options_allocator::{WalOptionsAllocator, WalOptionsAllocatorRef};
|
||||
use common_procedure::ProcedureManagerRef;
|
||||
use common_procedure::{ProcedureInfo, ProcedureManagerRef};
|
||||
use common_telemetry::info;
|
||||
use common_telemetry::logging::{LoggingOptions, TracingOptions};
|
||||
use common_time::timezone::set_default_timezone;
|
||||
@@ -44,6 +49,7 @@ use common_version::{short_version, version};
|
||||
use common_wal::config::DatanodeWalConfig;
|
||||
use datanode::config::{DatanodeOptions, ProcedureConfig, RegionEngineConfig, StorageConfig};
|
||||
use datanode::datanode::{Datanode, DatanodeBuilder};
|
||||
use datanode::region_server::RegionServer;
|
||||
use file_engine::config::EngineConfig as FileEngineConfig;
|
||||
use flow::{FlowWorkerManager, FlownodeBuilder, FrontendInvoker};
|
||||
use frontend::frontend::FrontendOptions;
|
||||
@@ -55,6 +61,7 @@ use frontend::service_config::{
|
||||
};
|
||||
use meta_srv::metasrv::{FLOW_ID_SEQ, TABLE_ID_SEQ};
|
||||
use mito2::config::MitoConfig;
|
||||
use query::stats::StatementStatistics;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use servers::export_metrics::ExportMetricsOption;
|
||||
use servers::grpc::GrpcOptions;
|
||||
@@ -477,9 +484,18 @@ impl StartCommand {
|
||||
.build(),
|
||||
);
|
||||
|
||||
let datanode = DatanodeBuilder::new(dn_opts, plugins.clone())
|
||||
.with_kv_backend(kv_backend.clone())
|
||||
.build()
|
||||
.await
|
||||
.context(StartDatanodeSnafu)?;
|
||||
|
||||
let information_extension = Arc::new(StandaloneInformationExtension::new(
|
||||
datanode.region_server(),
|
||||
procedure_manager.clone(),
|
||||
));
|
||||
let catalog_manager = KvBackendCatalogManager::new(
|
||||
dn_opts.mode,
|
||||
None,
|
||||
information_extension,
|
||||
kv_backend.clone(),
|
||||
layered_cache_registry.clone(),
|
||||
Some(procedure_manager.clone()),
|
||||
@@ -488,12 +504,6 @@ impl StartCommand {
|
||||
let table_metadata_manager =
|
||||
Self::create_table_metadata_manager(kv_backend.clone()).await?;
|
||||
|
||||
let datanode = DatanodeBuilder::new(dn_opts, plugins.clone())
|
||||
.with_kv_backend(kv_backend.clone())
|
||||
.build()
|
||||
.await
|
||||
.context(StartDatanodeSnafu)?;
|
||||
|
||||
let flow_metadata_manager = Arc::new(FlowMetadataManager::new(kv_backend.clone()));
|
||||
let flow_builder = FlownodeBuilder::new(
|
||||
Default::default(),
|
||||
@@ -557,6 +567,7 @@ impl StartCommand {
|
||||
catalog_manager.clone(),
|
||||
node_manager.clone(),
|
||||
ddl_task_executor.clone(),
|
||||
StatementStatistics::new(opts.logging.slow_query.clone()),
|
||||
)
|
||||
.with_plugin(plugins.clone())
|
||||
.try_build()
|
||||
@@ -642,6 +653,91 @@ impl StartCommand {
|
||||
}
|
||||
}
|
||||
|
||||
pub struct StandaloneInformationExtension {
|
||||
region_server: RegionServer,
|
||||
procedure_manager: ProcedureManagerRef,
|
||||
start_time_ms: u64,
|
||||
}
|
||||
|
||||
impl StandaloneInformationExtension {
|
||||
pub fn new(region_server: RegionServer, procedure_manager: ProcedureManagerRef) -> Self {
|
||||
Self {
|
||||
region_server,
|
||||
procedure_manager,
|
||||
start_time_ms: common_time::util::current_time_millis() as u64,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl InformationExtension for StandaloneInformationExtension {
|
||||
type Error = catalog::error::Error;
|
||||
|
||||
async fn nodes(&self) -> std::result::Result<Vec<NodeInfo>, Self::Error> {
|
||||
let build_info = common_version::build_info();
|
||||
let node_info = NodeInfo {
|
||||
// For the standalone:
|
||||
// - id always 0
|
||||
// - empty string for peer_addr
|
||||
peer: Peer {
|
||||
id: 0,
|
||||
addr: "".to_string(),
|
||||
},
|
||||
last_activity_ts: -1,
|
||||
status: NodeStatus::Standalone,
|
||||
version: build_info.version.to_string(),
|
||||
git_commit: build_info.commit_short.to_string(),
|
||||
// Use `self.start_time_ms` instead.
|
||||
// It's not precise but enough.
|
||||
start_time_ms: self.start_time_ms,
|
||||
};
|
||||
Ok(vec![node_info])
|
||||
}
|
||||
|
||||
async fn procedures(&self) -> std::result::Result<Vec<(String, ProcedureInfo)>, Self::Error> {
|
||||
self.procedure_manager
|
||||
.list_procedures()
|
||||
.await
|
||||
.map_err(BoxedError::new)
|
||||
.map(|procedures| {
|
||||
procedures
|
||||
.into_iter()
|
||||
.map(|procedure| {
|
||||
let status = procedure.state.as_str_name().to_string();
|
||||
(status, procedure)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.context(catalog::error::ListProceduresSnafu)
|
||||
}
|
||||
|
||||
async fn region_stats(&self) -> std::result::Result<Vec<RegionStat>, Self::Error> {
|
||||
let stats = self
|
||||
.region_server
|
||||
.reportable_regions()
|
||||
.into_iter()
|
||||
.map(|stat| {
|
||||
let region_stat = self
|
||||
.region_server
|
||||
.region_statistic(stat.region_id)
|
||||
.unwrap_or_default();
|
||||
RegionStat {
|
||||
id: stat.region_id,
|
||||
rcus: 0,
|
||||
wcus: 0,
|
||||
approximate_bytes: region_stat.estimated_disk_size() as i64,
|
||||
engine: stat.engine,
|
||||
role: RegionRole::from(stat.role).into(),
|
||||
memtable_size: region_stat.memtable_size,
|
||||
manifest_size: region_stat.manifest_size,
|
||||
sst_size: region_stat.sst_size,
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
Ok(stats)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::default::Default;
|
||||
|
||||
@@ -38,6 +38,18 @@ impl Plugins {
|
||||
self.read().get::<T>().cloned()
|
||||
}
|
||||
|
||||
pub fn get_or_insert<T, F>(&self, f: F) -> T
|
||||
where
|
||||
T: 'static + Send + Sync + Clone,
|
||||
F: FnOnce() -> T,
|
||||
{
|
||||
let mut binding = self.write();
|
||||
if !binding.contains::<T>() {
|
||||
binding.insert(f());
|
||||
}
|
||||
binding.get::<T>().cloned().unwrap()
|
||||
}
|
||||
|
||||
pub fn map_mut<T: 'static + Send + Sync, F, R>(&self, mapper: F) -> R
|
||||
where
|
||||
F: FnOnce(Option<&mut T>) -> R,
|
||||
|
||||
@@ -9,7 +9,7 @@ workspace = true
|
||||
|
||||
[features]
|
||||
default = ["geo"]
|
||||
geo = ["geohash", "h3o"]
|
||||
geo = ["geohash", "h3o", "s2"]
|
||||
|
||||
[dependencies]
|
||||
api.workspace = true
|
||||
@@ -35,6 +35,7 @@ num = "0.4"
|
||||
num-traits = "0.2"
|
||||
once_cell.workspace = true
|
||||
paste = "1.0"
|
||||
s2 = { version = "0.0.12", optional = true }
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
session.workspace = true
|
||||
|
||||
@@ -91,5 +91,12 @@ impl AggregateFunctions {
|
||||
register_aggr_func!("argmin", 1, ArgminAccumulatorCreator);
|
||||
register_aggr_func!("scipystatsnormcdf", 2, ScipyStatsNormCdfAccumulatorCreator);
|
||||
register_aggr_func!("scipystatsnormpdf", 2, ScipyStatsNormPdfAccumulatorCreator);
|
||||
|
||||
#[cfg(feature = "geo")]
|
||||
register_aggr_func!(
|
||||
"json_encode_path",
|
||||
3,
|
||||
super::geo::encoding::JsonPathEncodeFunctionCreator
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,7 +16,10 @@ use std::cmp::Ordering;
|
||||
use std::sync::Arc;
|
||||
|
||||
use common_macro::{as_aggr_func_creator, AggrFuncTypeStore};
|
||||
use common_query::error::{BadAccumulatorImplSnafu, CreateAccumulatorSnafu, Result};
|
||||
use common_query::error::{
|
||||
BadAccumulatorImplSnafu, CreateAccumulatorSnafu, InvalidInputStateSnafu, Result,
|
||||
};
|
||||
use common_query::logical_plan::accumulator::AggrFuncTypeStore;
|
||||
use common_query::logical_plan::{Accumulator, AggregateFunctionCreator};
|
||||
use common_query::prelude::*;
|
||||
use datatypes::prelude::*;
|
||||
|
||||
@@ -16,7 +16,10 @@ use std::cmp::Ordering;
|
||||
use std::sync::Arc;
|
||||
|
||||
use common_macro::{as_aggr_func_creator, AggrFuncTypeStore};
|
||||
use common_query::error::{BadAccumulatorImplSnafu, CreateAccumulatorSnafu, Result};
|
||||
use common_query::error::{
|
||||
BadAccumulatorImplSnafu, CreateAccumulatorSnafu, InvalidInputStateSnafu, Result,
|
||||
};
|
||||
use common_query::logical_plan::accumulator::AggrFuncTypeStore;
|
||||
use common_query::logical_plan::{Accumulator, AggregateFunctionCreator};
|
||||
use common_query::prelude::*;
|
||||
use datatypes::prelude::*;
|
||||
|
||||
@@ -17,8 +17,10 @@ use std::sync::Arc;
|
||||
|
||||
use common_macro::{as_aggr_func_creator, AggrFuncTypeStore};
|
||||
use common_query::error::{
|
||||
CreateAccumulatorSnafu, DowncastVectorSnafu, FromScalarValueSnafu, Result,
|
||||
CreateAccumulatorSnafu, DowncastVectorSnafu, FromScalarValueSnafu, InvalidInputStateSnafu,
|
||||
Result,
|
||||
};
|
||||
use common_query::logical_plan::accumulator::AggrFuncTypeStore;
|
||||
use common_query::logical_plan::{Accumulator, AggregateFunctionCreator};
|
||||
use common_query::prelude::*;
|
||||
use datatypes::prelude::*;
|
||||
|
||||
@@ -17,8 +17,10 @@ use std::sync::Arc;
|
||||
|
||||
use common_macro::{as_aggr_func_creator, AggrFuncTypeStore};
|
||||
use common_query::error::{
|
||||
BadAccumulatorImplSnafu, CreateAccumulatorSnafu, DowncastVectorSnafu, Result,
|
||||
BadAccumulatorImplSnafu, CreateAccumulatorSnafu, DowncastVectorSnafu, InvalidInputStateSnafu,
|
||||
Result,
|
||||
};
|
||||
use common_query::logical_plan::accumulator::AggrFuncTypeStore;
|
||||
use common_query::logical_plan::{Accumulator, AggregateFunctionCreator};
|
||||
use common_query::prelude::*;
|
||||
use datatypes::prelude::*;
|
||||
|
||||
@@ -18,8 +18,9 @@ use std::sync::Arc;
|
||||
use common_macro::{as_aggr_func_creator, AggrFuncTypeStore};
|
||||
use common_query::error::{
|
||||
self, BadAccumulatorImplSnafu, CreateAccumulatorSnafu, DowncastVectorSnafu,
|
||||
FromScalarValueSnafu, InvalidInputColSnafu, Result,
|
||||
FromScalarValueSnafu, InvalidInputColSnafu, InvalidInputStateSnafu, Result,
|
||||
};
|
||||
use common_query::logical_plan::accumulator::AggrFuncTypeStore;
|
||||
use common_query::logical_plan::{Accumulator, AggregateFunctionCreator};
|
||||
use common_query::prelude::*;
|
||||
use datatypes::prelude::*;
|
||||
|
||||
@@ -17,8 +17,10 @@ use std::sync::Arc;
|
||||
use common_macro::{as_aggr_func_creator, AggrFuncTypeStore};
|
||||
use common_query::error::{
|
||||
self, BadAccumulatorImplSnafu, CreateAccumulatorSnafu, DowncastVectorSnafu,
|
||||
FromScalarValueSnafu, GenerateFunctionSnafu, InvalidInputColSnafu, Result,
|
||||
FromScalarValueSnafu, GenerateFunctionSnafu, InvalidInputColSnafu, InvalidInputStateSnafu,
|
||||
Result,
|
||||
};
|
||||
use common_query::logical_plan::accumulator::AggrFuncTypeStore;
|
||||
use common_query::logical_plan::{Accumulator, AggregateFunctionCreator};
|
||||
use common_query::prelude::*;
|
||||
use datatypes::prelude::*;
|
||||
|
||||
@@ -17,8 +17,10 @@ use std::sync::Arc;
|
||||
use common_macro::{as_aggr_func_creator, AggrFuncTypeStore};
|
||||
use common_query::error::{
|
||||
self, BadAccumulatorImplSnafu, CreateAccumulatorSnafu, DowncastVectorSnafu,
|
||||
FromScalarValueSnafu, GenerateFunctionSnafu, InvalidInputColSnafu, Result,
|
||||
FromScalarValueSnafu, GenerateFunctionSnafu, InvalidInputColSnafu, InvalidInputStateSnafu,
|
||||
Result,
|
||||
};
|
||||
use common_query::logical_plan::accumulator::AggrFuncTypeStore;
|
||||
use common_query::logical_plan::{Accumulator, AggregateFunctionCreator};
|
||||
use common_query::prelude::*;
|
||||
use datatypes::prelude::*;
|
||||
|
||||
@@ -14,18 +14,19 @@
|
||||
|
||||
use std::fmt;
|
||||
|
||||
use common_query::error::{InvalidFuncArgsSnafu, Result, UnsupportedInputDataTypeSnafu};
|
||||
use common_query::error::{ArrowComputeSnafu, IntoVectorSnafu, InvalidFuncArgsSnafu, Result};
|
||||
use common_query::prelude::Signature;
|
||||
use datatypes::data_type::DataType;
|
||||
use datatypes::arrow::compute::kernels::numeric;
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::value::ValueRef;
|
||||
use datatypes::vectors::VectorRef;
|
||||
use snafu::ensure;
|
||||
use datatypes::vectors::{Helper, VectorRef};
|
||||
use snafu::{ensure, ResultExt};
|
||||
|
||||
use crate::function::{Function, FunctionContext};
|
||||
use crate::helper;
|
||||
|
||||
/// A function adds an interval value to Timestamp, Date or DateTime, and return the result.
|
||||
/// A function adds an interval value to Timestamp, Date, and return the result.
|
||||
/// The implementation of datetime type is based on Date64 which is incorrect so this function
|
||||
/// doesn't support the datetime type.
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct DateAddFunction;
|
||||
|
||||
@@ -44,7 +45,6 @@ impl Function for DateAddFunction {
|
||||
helper::one_of_sigs2(
|
||||
vec![
|
||||
ConcreteDataType::date_datatype(),
|
||||
ConcreteDataType::datetime_datatype(),
|
||||
ConcreteDataType::timestamp_second_datatype(),
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
ConcreteDataType::timestamp_microsecond_datatype(),
|
||||
@@ -69,64 +69,14 @@ impl Function for DateAddFunction {
|
||||
}
|
||||
);
|
||||
|
||||
let left = &columns[0];
|
||||
let right = &columns[1];
|
||||
let left = columns[0].to_arrow_array();
|
||||
let right = columns[1].to_arrow_array();
|
||||
|
||||
let size = left.len();
|
||||
let left_datatype = columns[0].data_type();
|
||||
match left_datatype {
|
||||
ConcreteDataType::Timestamp(_) => {
|
||||
let mut result = left_datatype.create_mutable_vector(size);
|
||||
for i in 0..size {
|
||||
let ts = left.get(i).as_timestamp();
|
||||
let interval = right.get(i).as_interval();
|
||||
|
||||
let new_ts = match (ts, interval) {
|
||||
(Some(ts), Some(interval)) => ts.add_interval(interval),
|
||||
_ => ts,
|
||||
};
|
||||
|
||||
result.push_value_ref(ValueRef::from(new_ts));
|
||||
}
|
||||
|
||||
Ok(result.to_vector())
|
||||
}
|
||||
ConcreteDataType::Date(_) => {
|
||||
let mut result = left_datatype.create_mutable_vector(size);
|
||||
for i in 0..size {
|
||||
let date = left.get(i).as_date();
|
||||
let interval = right.get(i).as_interval();
|
||||
let new_date = match (date, interval) {
|
||||
(Some(date), Some(interval)) => date.add_interval(interval),
|
||||
_ => date,
|
||||
};
|
||||
|
||||
result.push_value_ref(ValueRef::from(new_date));
|
||||
}
|
||||
|
||||
Ok(result.to_vector())
|
||||
}
|
||||
ConcreteDataType::DateTime(_) => {
|
||||
let mut result = left_datatype.create_mutable_vector(size);
|
||||
for i in 0..size {
|
||||
let datetime = left.get(i).as_datetime();
|
||||
let interval = right.get(i).as_interval();
|
||||
let new_datetime = match (datetime, interval) {
|
||||
(Some(datetime), Some(interval)) => datetime.add_interval(interval),
|
||||
_ => datetime,
|
||||
};
|
||||
|
||||
result.push_value_ref(ValueRef::from(new_datetime));
|
||||
}
|
||||
|
||||
Ok(result.to_vector())
|
||||
}
|
||||
_ => UnsupportedInputDataTypeSnafu {
|
||||
function: NAME,
|
||||
datatypes: columns.iter().map(|c| c.data_type()).collect::<Vec<_>>(),
|
||||
}
|
||||
.fail(),
|
||||
}
|
||||
let result = numeric::add(&left, &right).context(ArrowComputeSnafu)?;
|
||||
let arrow_type = result.data_type().clone();
|
||||
Helper::try_into_vector(result).context(IntoVectorSnafu {
|
||||
data_type: arrow_type,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -144,8 +94,7 @@ mod tests {
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::value::Value;
|
||||
use datatypes::vectors::{
|
||||
DateTimeVector, DateVector, IntervalDayTimeVector, IntervalYearMonthVector,
|
||||
TimestampSecondVector,
|
||||
DateVector, IntervalDayTimeVector, IntervalYearMonthVector, TimestampSecondVector,
|
||||
};
|
||||
|
||||
use super::{DateAddFunction, *};
|
||||
@@ -168,16 +117,15 @@ mod tests {
|
||||
ConcreteDataType::date_datatype(),
|
||||
f.return_type(&[ConcreteDataType::date_datatype()]).unwrap()
|
||||
);
|
||||
assert_eq!(
|
||||
ConcreteDataType::datetime_datatype(),
|
||||
f.return_type(&[ConcreteDataType::datetime_datatype()])
|
||||
.unwrap()
|
||||
);
|
||||
assert!(matches!(f.signature(),
|
||||
assert!(
|
||||
matches!(f.signature(),
|
||||
Signature {
|
||||
type_signature: TypeSignature::OneOf(sigs),
|
||||
volatility: Volatility::Immutable
|
||||
} if sigs.len() == 18));
|
||||
} if sigs.len() == 15),
|
||||
"{:?}",
|
||||
f.signature()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -243,36 +191,4 @@ mod tests {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_datetime_date_add() {
|
||||
let f = DateAddFunction;
|
||||
|
||||
let dates = vec![Some(123), None, Some(42), None];
|
||||
// Intervals in months
|
||||
let intervals = vec![1, 2, 3, 1];
|
||||
let results = [Some(2678400123), None, Some(7776000042), None];
|
||||
|
||||
let date_vector = DateTimeVector::from(dates.clone());
|
||||
let interval_vector = IntervalYearMonthVector::from_vec(intervals);
|
||||
let args: Vec<VectorRef> = vec![Arc::new(date_vector), Arc::new(interval_vector)];
|
||||
let vector = f.eval(FunctionContext::default(), &args).unwrap();
|
||||
|
||||
assert_eq!(4, vector.len());
|
||||
for (i, _t) in dates.iter().enumerate() {
|
||||
let v = vector.get(i);
|
||||
let result = results.get(i).unwrap();
|
||||
|
||||
if result.is_none() {
|
||||
assert_eq!(Value::Null, v);
|
||||
continue;
|
||||
}
|
||||
match v {
|
||||
Value::DateTime(date) => {
|
||||
assert_eq!(date.val(), result.unwrap());
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -14,18 +14,19 @@
|
||||
|
||||
use std::fmt;
|
||||
|
||||
use common_query::error::{InvalidFuncArgsSnafu, Result, UnsupportedInputDataTypeSnafu};
|
||||
use common_query::error::{ArrowComputeSnafu, IntoVectorSnafu, InvalidFuncArgsSnafu, Result};
|
||||
use common_query::prelude::Signature;
|
||||
use datatypes::data_type::DataType;
|
||||
use datatypes::arrow::compute::kernels::numeric;
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::value::ValueRef;
|
||||
use datatypes::vectors::VectorRef;
|
||||
use snafu::ensure;
|
||||
use datatypes::vectors::{Helper, VectorRef};
|
||||
use snafu::{ensure, ResultExt};
|
||||
|
||||
use crate::function::{Function, FunctionContext};
|
||||
use crate::helper;
|
||||
|
||||
/// A function subtracts an interval value to Timestamp, Date or DateTime, and return the result.
|
||||
/// A function subtracts an interval value to Timestamp, Date, and return the result.
|
||||
/// The implementation of datetime type is based on Date64 which is incorrect so this function
|
||||
/// doesn't support the datetime type.
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct DateSubFunction;
|
||||
|
||||
@@ -44,7 +45,6 @@ impl Function for DateSubFunction {
|
||||
helper::one_of_sigs2(
|
||||
vec![
|
||||
ConcreteDataType::date_datatype(),
|
||||
ConcreteDataType::datetime_datatype(),
|
||||
ConcreteDataType::timestamp_second_datatype(),
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
ConcreteDataType::timestamp_microsecond_datatype(),
|
||||
@@ -69,65 +69,14 @@ impl Function for DateSubFunction {
|
||||
}
|
||||
);
|
||||
|
||||
let left = &columns[0];
|
||||
let right = &columns[1];
|
||||
let left = columns[0].to_arrow_array();
|
||||
let right = columns[1].to_arrow_array();
|
||||
|
||||
let size = left.len();
|
||||
let left_datatype = columns[0].data_type();
|
||||
|
||||
match left_datatype {
|
||||
ConcreteDataType::Timestamp(_) => {
|
||||
let mut result = left_datatype.create_mutable_vector(size);
|
||||
for i in 0..size {
|
||||
let ts = left.get(i).as_timestamp();
|
||||
let interval = right.get(i).as_interval();
|
||||
|
||||
let new_ts = match (ts, interval) {
|
||||
(Some(ts), Some(interval)) => ts.sub_interval(interval),
|
||||
_ => ts,
|
||||
};
|
||||
|
||||
result.push_value_ref(ValueRef::from(new_ts));
|
||||
}
|
||||
|
||||
Ok(result.to_vector())
|
||||
}
|
||||
ConcreteDataType::Date(_) => {
|
||||
let mut result = left_datatype.create_mutable_vector(size);
|
||||
for i in 0..size {
|
||||
let date = left.get(i).as_date();
|
||||
let interval = right.get(i).as_interval();
|
||||
let new_date = match (date, interval) {
|
||||
(Some(date), Some(interval)) => date.sub_interval(interval),
|
||||
_ => date,
|
||||
};
|
||||
|
||||
result.push_value_ref(ValueRef::from(new_date));
|
||||
}
|
||||
|
||||
Ok(result.to_vector())
|
||||
}
|
||||
ConcreteDataType::DateTime(_) => {
|
||||
let mut result = left_datatype.create_mutable_vector(size);
|
||||
for i in 0..size {
|
||||
let datetime = left.get(i).as_datetime();
|
||||
let interval = right.get(i).as_interval();
|
||||
let new_datetime = match (datetime, interval) {
|
||||
(Some(datetime), Some(interval)) => datetime.sub_interval(interval),
|
||||
_ => datetime,
|
||||
};
|
||||
|
||||
result.push_value_ref(ValueRef::from(new_datetime));
|
||||
}
|
||||
|
||||
Ok(result.to_vector())
|
||||
}
|
||||
_ => UnsupportedInputDataTypeSnafu {
|
||||
function: NAME,
|
||||
datatypes: columns.iter().map(|c| c.data_type()).collect::<Vec<_>>(),
|
||||
}
|
||||
.fail(),
|
||||
}
|
||||
let result = numeric::sub(&left, &right).context(ArrowComputeSnafu)?;
|
||||
let arrow_type = result.data_type().clone();
|
||||
Helper::try_into_vector(result).context(IntoVectorSnafu {
|
||||
data_type: arrow_type,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -145,8 +94,7 @@ mod tests {
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::value::Value;
|
||||
use datatypes::vectors::{
|
||||
DateTimeVector, DateVector, IntervalDayTimeVector, IntervalYearMonthVector,
|
||||
TimestampSecondVector,
|
||||
DateVector, IntervalDayTimeVector, IntervalYearMonthVector, TimestampSecondVector,
|
||||
};
|
||||
|
||||
use super::{DateSubFunction, *};
|
||||
@@ -174,11 +122,15 @@ mod tests {
|
||||
f.return_type(&[ConcreteDataType::datetime_datatype()])
|
||||
.unwrap()
|
||||
);
|
||||
assert!(matches!(f.signature(),
|
||||
assert!(
|
||||
matches!(f.signature(),
|
||||
Signature {
|
||||
type_signature: TypeSignature::OneOf(sigs),
|
||||
volatility: Volatility::Immutable
|
||||
} if sigs.len() == 18));
|
||||
} if sigs.len() == 15),
|
||||
"{:?}",
|
||||
f.signature()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -250,42 +202,4 @@ mod tests {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_datetime_date_sub() {
|
||||
let f = DateSubFunction;
|
||||
let millis_per_month = 3600 * 24 * 30 * 1000;
|
||||
|
||||
let dates = vec![
|
||||
Some(123 * millis_per_month),
|
||||
None,
|
||||
Some(42 * millis_per_month),
|
||||
None,
|
||||
];
|
||||
// Intervals in months
|
||||
let intervals = vec![1, 2, 3, 1];
|
||||
let results = [Some(316137600000), None, Some(100915200000), None];
|
||||
|
||||
let date_vector = DateTimeVector::from(dates.clone());
|
||||
let interval_vector = IntervalYearMonthVector::from_vec(intervals);
|
||||
let args: Vec<VectorRef> = vec![Arc::new(date_vector), Arc::new(interval_vector)];
|
||||
let vector = f.eval(FunctionContext::default(), &args).unwrap();
|
||||
|
||||
assert_eq!(4, vector.len());
|
||||
for (i, _t) in dates.iter().enumerate() {
|
||||
let v = vector.get(i);
|
||||
let result = results.get(i).unwrap();
|
||||
|
||||
if result.is_none() {
|
||||
assert_eq!(Value::Null, v);
|
||||
continue;
|
||||
}
|
||||
match v {
|
||||
Value::DateTime(date) => {
|
||||
assert_eq!(date.val(), result.unwrap());
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,10 +13,11 @@
|
||||
// limitations under the License.
|
||||
|
||||
use std::sync::Arc;
|
||||
pub(crate) mod encoding;
|
||||
mod geohash;
|
||||
mod h3;
|
||||
|
||||
use geohash::{GeohashFunction, GeohashNeighboursFunction};
|
||||
mod helpers;
|
||||
mod s2;
|
||||
|
||||
use crate::function_registry::FunctionRegistry;
|
||||
|
||||
@@ -25,20 +26,39 @@ pub(crate) struct GeoFunctions;
|
||||
impl GeoFunctions {
|
||||
pub fn register(registry: &FunctionRegistry) {
|
||||
// geohash
|
||||
registry.register(Arc::new(GeohashFunction));
|
||||
registry.register(Arc::new(GeohashNeighboursFunction));
|
||||
// h3 family
|
||||
registry.register(Arc::new(geohash::GeohashFunction));
|
||||
registry.register(Arc::new(geohash::GeohashNeighboursFunction));
|
||||
|
||||
// h3 index
|
||||
registry.register(Arc::new(h3::H3LatLngToCell));
|
||||
registry.register(Arc::new(h3::H3LatLngToCellString));
|
||||
|
||||
// h3 index inspection
|
||||
registry.register(Arc::new(h3::H3CellBase));
|
||||
registry.register(Arc::new(h3::H3CellCenterChild));
|
||||
registry.register(Arc::new(h3::H3CellCenterLat));
|
||||
registry.register(Arc::new(h3::H3CellCenterLng));
|
||||
registry.register(Arc::new(h3::H3CellIsPentagon));
|
||||
registry.register(Arc::new(h3::H3CellParent));
|
||||
registry.register(Arc::new(h3::H3CellResolution));
|
||||
registry.register(Arc::new(h3::H3CellToString));
|
||||
registry.register(Arc::new(h3::H3IsNeighbour));
|
||||
registry.register(Arc::new(h3::H3StringToCell));
|
||||
registry.register(Arc::new(h3::H3CellToString));
|
||||
registry.register(Arc::new(h3::H3CellCenterLatLng));
|
||||
registry.register(Arc::new(h3::H3CellResolution));
|
||||
|
||||
// h3 hierarchical grid
|
||||
registry.register(Arc::new(h3::H3CellCenterChild));
|
||||
registry.register(Arc::new(h3::H3CellParent));
|
||||
registry.register(Arc::new(h3::H3CellToChildren));
|
||||
registry.register(Arc::new(h3::H3CellToChildrenSize));
|
||||
registry.register(Arc::new(h3::H3CellToChildPos));
|
||||
registry.register(Arc::new(h3::H3ChildPosToCell));
|
||||
|
||||
// h3 grid traversal
|
||||
registry.register(Arc::new(h3::H3GridDisk));
|
||||
registry.register(Arc::new(h3::H3GridDiskDistances));
|
||||
registry.register(Arc::new(h3::H3GridDistance));
|
||||
registry.register(Arc::new(h3::H3GridPathCells));
|
||||
|
||||
// s2
|
||||
registry.register(Arc::new(s2::S2LatLngToCell));
|
||||
registry.register(Arc::new(s2::S2CellLevel));
|
||||
registry.register(Arc::new(s2::S2CellToToken));
|
||||
registry.register(Arc::new(s2::S2CellParent));
|
||||
}
|
||||
}
|
||||
|
||||
223
src/common/function/src/scalars/geo/encoding.rs
Normal file
223
src/common/function/src/scalars/geo/encoding.rs
Normal file
@@ -0,0 +1,223 @@
|
||||
// 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::sync::Arc;
|
||||
|
||||
use common_error::ext::{BoxedError, PlainError};
|
||||
use common_error::status_code::StatusCode;
|
||||
use common_macro::{as_aggr_func_creator, AggrFuncTypeStore};
|
||||
use common_query::error::{self, InvalidInputStateSnafu, Result};
|
||||
use common_query::logical_plan::accumulator::AggrFuncTypeStore;
|
||||
use common_query::logical_plan::{Accumulator, AggregateFunctionCreator};
|
||||
use common_query::prelude::AccumulatorCreatorFunction;
|
||||
use common_time::Timestamp;
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::value::{ListValue, Value};
|
||||
use datatypes::vectors::VectorRef;
|
||||
use snafu::{ensure, ResultExt};
|
||||
|
||||
use super::helpers::{ensure_columns_len, ensure_columns_n};
|
||||
|
||||
/// Accumulator of lat, lng, timestamp tuples
|
||||
#[derive(Debug)]
|
||||
pub struct JsonPathAccumulator {
|
||||
timestamp_type: ConcreteDataType,
|
||||
lat: Vec<Option<f64>>,
|
||||
lng: Vec<Option<f64>>,
|
||||
timestamp: Vec<Option<Timestamp>>,
|
||||
}
|
||||
|
||||
impl JsonPathAccumulator {
|
||||
fn new(timestamp_type: ConcreteDataType) -> Self {
|
||||
Self {
|
||||
lat: Vec::default(),
|
||||
lng: Vec::default(),
|
||||
timestamp: Vec::default(),
|
||||
timestamp_type,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Accumulator for JsonPathAccumulator {
|
||||
fn state(&self) -> Result<Vec<Value>> {
|
||||
Ok(vec![
|
||||
Value::List(ListValue::new(
|
||||
self.lat.iter().map(|i| Value::from(*i)).collect(),
|
||||
ConcreteDataType::float64_datatype(),
|
||||
)),
|
||||
Value::List(ListValue::new(
|
||||
self.lng.iter().map(|i| Value::from(*i)).collect(),
|
||||
ConcreteDataType::float64_datatype(),
|
||||
)),
|
||||
Value::List(ListValue::new(
|
||||
self.timestamp.iter().map(|i| Value::from(*i)).collect(),
|
||||
self.timestamp_type.clone(),
|
||||
)),
|
||||
])
|
||||
}
|
||||
|
||||
fn update_batch(&mut self, columns: &[VectorRef]) -> Result<()> {
|
||||
// update batch as in datafusion just provides the accumulator original
|
||||
// input.
|
||||
//
|
||||
// columns is vec of [`lat`, `lng`, `timestamp`]
|
||||
// where
|
||||
// - `lat` is a vector of `Value::Float64` or similar type. Each item in
|
||||
// the vector is a row in given dataset.
|
||||
// - so on so forth for `lng` and `timestamp`
|
||||
ensure_columns_n!(columns, 3);
|
||||
|
||||
let lat = &columns[0];
|
||||
let lng = &columns[1];
|
||||
let ts = &columns[2];
|
||||
|
||||
let size = lat.len();
|
||||
|
||||
for idx in 0..size {
|
||||
self.lat.push(lat.get(idx).as_f64_lossy());
|
||||
self.lng.push(lng.get(idx).as_f64_lossy());
|
||||
self.timestamp.push(ts.get(idx).as_timestamp());
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn merge_batch(&mut self, states: &[VectorRef]) -> Result<()> {
|
||||
// merge batch as in datafusion gives state accumulated from the data
|
||||
// returned from child accumulators' state() call
|
||||
// In our particular implementation, the data structure is like
|
||||
//
|
||||
// states is vec of [`lat`, `lng`, `timestamp`]
|
||||
// where
|
||||
// - `lat` is a vector of `Value::List`. Each item in the list is all
|
||||
// coordinates from a child accumulator.
|
||||
// - so on so forth for `lng` and `timestamp`
|
||||
|
||||
ensure_columns_n!(states, 3);
|
||||
|
||||
let lat_lists = &states[0];
|
||||
let lng_lists = &states[1];
|
||||
let ts_lists = &states[2];
|
||||
|
||||
let len = lat_lists.len();
|
||||
|
||||
for idx in 0..len {
|
||||
if let Some(lat_list) = lat_lists
|
||||
.get(idx)
|
||||
.as_list()
|
||||
.map_err(BoxedError::new)
|
||||
.context(error::ExecuteSnafu)?
|
||||
{
|
||||
for v in lat_list.items() {
|
||||
self.lat.push(v.as_f64_lossy());
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(lng_list) = lng_lists
|
||||
.get(idx)
|
||||
.as_list()
|
||||
.map_err(BoxedError::new)
|
||||
.context(error::ExecuteSnafu)?
|
||||
{
|
||||
for v in lng_list.items() {
|
||||
self.lng.push(v.as_f64_lossy());
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ts_list) = ts_lists
|
||||
.get(idx)
|
||||
.as_list()
|
||||
.map_err(BoxedError::new)
|
||||
.context(error::ExecuteSnafu)?
|
||||
{
|
||||
for v in ts_list.items() {
|
||||
self.timestamp.push(v.as_timestamp());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn evaluate(&self) -> Result<Value> {
|
||||
let mut work_vec: Vec<(&Option<f64>, &Option<f64>, &Option<Timestamp>)> = self
|
||||
.lat
|
||||
.iter()
|
||||
.zip(self.lng.iter())
|
||||
.zip(self.timestamp.iter())
|
||||
.map(|((a, b), c)| (a, b, c))
|
||||
.collect();
|
||||
|
||||
// sort by timestamp, we treat null timestamp as 0
|
||||
work_vec.sort_unstable_by_key(|tuple| tuple.2.unwrap_or_else(|| Timestamp::new_second(0)));
|
||||
|
||||
let result = serde_json::to_string(
|
||||
&work_vec
|
||||
.into_iter()
|
||||
// note that we transform to lng,lat for geojson compatibility
|
||||
.map(|(lat, lng, _)| vec![lng, lat])
|
||||
.collect::<Vec<Vec<&Option<f64>>>>(),
|
||||
)
|
||||
.map_err(|e| {
|
||||
BoxedError::new(PlainError::new(
|
||||
format!("Serialization failure: {}", e),
|
||||
StatusCode::EngineExecuteQuery,
|
||||
))
|
||||
})
|
||||
.context(error::ExecuteSnafu)?;
|
||||
|
||||
Ok(Value::String(result.into()))
|
||||
}
|
||||
}
|
||||
|
||||
/// This function accept rows of lat, lng and timestamp, sort with timestamp and
|
||||
/// encoding them into a geojson-like path.
|
||||
///
|
||||
/// Example:
|
||||
///
|
||||
/// ```sql
|
||||
/// SELECT json_encode_path(lat, lon, timestamp) FROM table [group by ...];
|
||||
/// ```
|
||||
///
|
||||
#[as_aggr_func_creator]
|
||||
#[derive(Debug, Default, AggrFuncTypeStore)]
|
||||
pub struct JsonPathEncodeFunctionCreator {}
|
||||
|
||||
impl AggregateFunctionCreator for JsonPathEncodeFunctionCreator {
|
||||
fn creator(&self) -> AccumulatorCreatorFunction {
|
||||
let creator: AccumulatorCreatorFunction = Arc::new(move |types: &[ConcreteDataType]| {
|
||||
let ts_type = types[2].clone();
|
||||
Ok(Box::new(JsonPathAccumulator::new(ts_type)))
|
||||
});
|
||||
|
||||
creator
|
||||
}
|
||||
|
||||
fn output_type(&self) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::string_datatype())
|
||||
}
|
||||
|
||||
fn state_types(&self) -> Result<Vec<ConcreteDataType>> {
|
||||
let input_types = self.input_types()?;
|
||||
ensure!(input_types.len() == 3, InvalidInputStateSnafu);
|
||||
|
||||
let timestamp_type = input_types[2].clone();
|
||||
|
||||
Ok(vec![
|
||||
ConcreteDataType::list_datatype(ConcreteDataType::float64_datatype()),
|
||||
ConcreteDataType::list_datatype(ConcreteDataType::float64_datatype()),
|
||||
ConcreteDataType::list_datatype(timestamp_type),
|
||||
])
|
||||
}
|
||||
}
|
||||
@@ -16,22 +16,75 @@ use std::str::FromStr;
|
||||
|
||||
use common_error::ext::{BoxedError, PlainError};
|
||||
use common_error::status_code::StatusCode;
|
||||
use common_query::error::{self, InvalidFuncArgsSnafu, Result};
|
||||
use common_query::error::{self, Result};
|
||||
use common_query::prelude::{Signature, TypeSignature};
|
||||
use datafusion::logical_expr::Volatility;
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::scalars::ScalarVectorBuilder;
|
||||
use datatypes::value::Value;
|
||||
use datatypes::scalars::{Scalar, ScalarVectorBuilder};
|
||||
use datatypes::value::{ListValue, Value};
|
||||
use datatypes::vectors::{
|
||||
BooleanVectorBuilder, Float64VectorBuilder, MutableVector, StringVectorBuilder,
|
||||
UInt64VectorBuilder, UInt8VectorBuilder, VectorRef,
|
||||
BooleanVectorBuilder, Int32VectorBuilder, ListVectorBuilder, MutableVector,
|
||||
StringVectorBuilder, UInt64VectorBuilder, UInt8VectorBuilder, VectorRef,
|
||||
};
|
||||
use derive_more::Display;
|
||||
use h3o::{CellIndex, LatLng, Resolution};
|
||||
use snafu::{ensure, ResultExt};
|
||||
use once_cell::sync::Lazy;
|
||||
use snafu::ResultExt;
|
||||
|
||||
use super::helpers::{ensure_and_coerce, ensure_columns_len, ensure_columns_n};
|
||||
use crate::function::{Function, FunctionContext};
|
||||
|
||||
static CELL_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
|
||||
vec![
|
||||
ConcreteDataType::int64_datatype(),
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
]
|
||||
});
|
||||
|
||||
static COORDINATE_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
|
||||
vec![
|
||||
ConcreteDataType::float32_datatype(),
|
||||
ConcreteDataType::float64_datatype(),
|
||||
]
|
||||
});
|
||||
static RESOLUTION_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
|
||||
vec![
|
||||
ConcreteDataType::int8_datatype(),
|
||||
ConcreteDataType::int16_datatype(),
|
||||
ConcreteDataType::int32_datatype(),
|
||||
ConcreteDataType::int64_datatype(),
|
||||
ConcreteDataType::uint8_datatype(),
|
||||
ConcreteDataType::uint16_datatype(),
|
||||
ConcreteDataType::uint32_datatype(),
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
]
|
||||
});
|
||||
static DISTANCE_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
|
||||
vec![
|
||||
ConcreteDataType::int8_datatype(),
|
||||
ConcreteDataType::int16_datatype(),
|
||||
ConcreteDataType::int32_datatype(),
|
||||
ConcreteDataType::int64_datatype(),
|
||||
ConcreteDataType::uint8_datatype(),
|
||||
ConcreteDataType::uint16_datatype(),
|
||||
ConcreteDataType::uint32_datatype(),
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
]
|
||||
});
|
||||
|
||||
static POSITION_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
|
||||
vec![
|
||||
ConcreteDataType::int8_datatype(),
|
||||
ConcreteDataType::int16_datatype(),
|
||||
ConcreteDataType::int32_datatype(),
|
||||
ConcreteDataType::int64_datatype(),
|
||||
ConcreteDataType::uint8_datatype(),
|
||||
ConcreteDataType::uint16_datatype(),
|
||||
ConcreteDataType::uint32_datatype(),
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
]
|
||||
});
|
||||
|
||||
/// Function that returns [h3] encoding cellid for a given geospatial coordinate.
|
||||
///
|
||||
/// [h3]: https://h3geo.org/
|
||||
@@ -50,20 +103,8 @@ impl Function for H3LatLngToCell {
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
let mut signatures = Vec::new();
|
||||
for coord_type in &[
|
||||
ConcreteDataType::float32_datatype(),
|
||||
ConcreteDataType::float64_datatype(),
|
||||
] {
|
||||
for resolution_type in &[
|
||||
ConcreteDataType::int8_datatype(),
|
||||
ConcreteDataType::int16_datatype(),
|
||||
ConcreteDataType::int32_datatype(),
|
||||
ConcreteDataType::int64_datatype(),
|
||||
ConcreteDataType::uint8_datatype(),
|
||||
ConcreteDataType::uint16_datatype(),
|
||||
ConcreteDataType::uint32_datatype(),
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
] {
|
||||
for coord_type in COORDINATE_TYPES.as_slice() {
|
||||
for resolution_type in RESOLUTION_TYPES.as_slice() {
|
||||
signatures.push(TypeSignature::Exact(vec![
|
||||
// latitude
|
||||
coord_type.clone(),
|
||||
@@ -78,15 +119,7 @@ impl Function for H3LatLngToCell {
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure!(
|
||||
columns.len() == 3,
|
||||
InvalidFuncArgsSnafu {
|
||||
err_msg: format!(
|
||||
"The length of the args is not correct, expect 3, provided : {}",
|
||||
columns.len()
|
||||
),
|
||||
}
|
||||
);
|
||||
ensure_columns_n!(columns, 3);
|
||||
|
||||
let lat_vec = &columns[0];
|
||||
let lon_vec = &columns[1];
|
||||
@@ -142,20 +175,8 @@ impl Function for H3LatLngToCellString {
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
let mut signatures = Vec::new();
|
||||
for coord_type in &[
|
||||
ConcreteDataType::float32_datatype(),
|
||||
ConcreteDataType::float64_datatype(),
|
||||
] {
|
||||
for resolution_type in &[
|
||||
ConcreteDataType::int8_datatype(),
|
||||
ConcreteDataType::int16_datatype(),
|
||||
ConcreteDataType::int32_datatype(),
|
||||
ConcreteDataType::int64_datatype(),
|
||||
ConcreteDataType::uint8_datatype(),
|
||||
ConcreteDataType::uint16_datatype(),
|
||||
ConcreteDataType::uint32_datatype(),
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
] {
|
||||
for coord_type in COORDINATE_TYPES.as_slice() {
|
||||
for resolution_type in RESOLUTION_TYPES.as_slice() {
|
||||
signatures.push(TypeSignature::Exact(vec![
|
||||
// latitude
|
||||
coord_type.clone(),
|
||||
@@ -170,15 +191,7 @@ impl Function for H3LatLngToCellString {
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure!(
|
||||
columns.len() == 3,
|
||||
InvalidFuncArgsSnafu {
|
||||
err_msg: format!(
|
||||
"The length of the args is not correct, expect 3, provided : {}",
|
||||
columns.len()
|
||||
),
|
||||
}
|
||||
);
|
||||
ensure_columns_n!(columns, 3);
|
||||
|
||||
let lat_vec = &columns[0];
|
||||
let lon_vec = &columns[1];
|
||||
@@ -234,15 +247,7 @@ impl Function for H3CellToString {
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure!(
|
||||
columns.len() == 1,
|
||||
InvalidFuncArgsSnafu {
|
||||
err_msg: format!(
|
||||
"The length of the args is not correct, expect 1, provided : {}",
|
||||
columns.len()
|
||||
),
|
||||
}
|
||||
);
|
||||
ensure_columns_n!(columns, 1);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let size = cell_vec.len();
|
||||
@@ -280,15 +285,7 @@ impl Function for H3StringToCell {
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure!(
|
||||
columns.len() == 1,
|
||||
InvalidFuncArgsSnafu {
|
||||
err_msg: format!(
|
||||
"The length of the args is not correct, expect 1, provided : {}",
|
||||
columns.len()
|
||||
),
|
||||
}
|
||||
);
|
||||
ensure_columns_n!(columns, 1);
|
||||
|
||||
let string_vec = &columns[0];
|
||||
let size = string_vec.len();
|
||||
@@ -319,18 +316,20 @@ impl Function for H3StringToCell {
|
||||
}
|
||||
}
|
||||
|
||||
/// Function that returns centroid latitude of given cell id
|
||||
/// Function that returns centroid latitude and longitude of given cell id
|
||||
#[derive(Clone, Debug, Default, Display)]
|
||||
#[display("{}", self.name())]
|
||||
pub struct H3CellCenterLat;
|
||||
pub struct H3CellCenterLatLng;
|
||||
|
||||
impl Function for H3CellCenterLat {
|
||||
impl Function for H3CellCenterLatLng {
|
||||
fn name(&self) -> &str {
|
||||
"h3_cell_center_lat"
|
||||
"h3_cell_center_latlng"
|
||||
}
|
||||
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::float64_datatype())
|
||||
Ok(ConcreteDataType::list_datatype(
|
||||
ConcreteDataType::float64_datatype(),
|
||||
))
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
@@ -338,69 +337,26 @@ impl Function for H3CellCenterLat {
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure!(
|
||||
columns.len() == 1,
|
||||
InvalidFuncArgsSnafu {
|
||||
err_msg: format!(
|
||||
"The length of the args is not correct, expect 1, provided : {}",
|
||||
columns.len()
|
||||
),
|
||||
}
|
||||
);
|
||||
ensure_columns_n!(columns, 1);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let size = cell_vec.len();
|
||||
let mut results = Float64VectorBuilder::with_capacity(size);
|
||||
let mut results =
|
||||
ListVectorBuilder::with_type_capacity(ConcreteDataType::float64_datatype(), size);
|
||||
|
||||
for i in 0..size {
|
||||
let cell = cell_from_value(cell_vec.get(i))?;
|
||||
let lat = cell.map(|cell| LatLng::from(cell).lat());
|
||||
let latlng = cell.map(LatLng::from);
|
||||
|
||||
results.push(lat);
|
||||
}
|
||||
|
||||
Ok(results.to_vector())
|
||||
}
|
||||
}
|
||||
|
||||
/// Function that returns centroid longitude of given cell id
|
||||
#[derive(Clone, Debug, Default, Display)]
|
||||
#[display("{}", self.name())]
|
||||
pub struct H3CellCenterLng;
|
||||
|
||||
impl Function for H3CellCenterLng {
|
||||
fn name(&self) -> &str {
|
||||
"h3_cell_center_lng"
|
||||
}
|
||||
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::float64_datatype())
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
signature_of_cell()
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure!(
|
||||
columns.len() == 1,
|
||||
InvalidFuncArgsSnafu {
|
||||
err_msg: format!(
|
||||
"The length of the args is not correct, expect 1, provided : {}",
|
||||
columns.len()
|
||||
),
|
||||
if let Some(latlng) = latlng {
|
||||
let result = ListValue::new(
|
||||
vec![latlng.lat().into(), latlng.lng().into()],
|
||||
ConcreteDataType::float64_datatype(),
|
||||
);
|
||||
results.push(Some(result.as_scalar_ref()));
|
||||
} else {
|
||||
results.push(None);
|
||||
}
|
||||
);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let size = cell_vec.len();
|
||||
let mut results = Float64VectorBuilder::with_capacity(size);
|
||||
|
||||
for i in 0..size {
|
||||
let cell = cell_from_value(cell_vec.get(i))?;
|
||||
let lat = cell.map(|cell| LatLng::from(cell).lng());
|
||||
|
||||
results.push(lat);
|
||||
}
|
||||
|
||||
Ok(results.to_vector())
|
||||
@@ -426,15 +382,7 @@ impl Function for H3CellResolution {
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure!(
|
||||
columns.len() == 1,
|
||||
InvalidFuncArgsSnafu {
|
||||
err_msg: format!(
|
||||
"The length of the args is not correct, expect 1, provided : {}",
|
||||
columns.len()
|
||||
),
|
||||
}
|
||||
);
|
||||
ensure_columns_n!(columns, 1);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let size = cell_vec.len();
|
||||
@@ -470,15 +418,7 @@ impl Function for H3CellBase {
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure!(
|
||||
columns.len() == 1,
|
||||
InvalidFuncArgsSnafu {
|
||||
err_msg: format!(
|
||||
"The length of the args is not correct, expect 1, provided : {}",
|
||||
columns.len()
|
||||
),
|
||||
}
|
||||
);
|
||||
ensure_columns_n!(columns, 1);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let size = cell_vec.len();
|
||||
@@ -514,15 +454,7 @@ impl Function for H3CellIsPentagon {
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure!(
|
||||
columns.len() == 1,
|
||||
InvalidFuncArgsSnafu {
|
||||
err_msg: format!(
|
||||
"The length of the args is not correct, expect 1, provided : {}",
|
||||
columns.len()
|
||||
),
|
||||
}
|
||||
);
|
||||
ensure_columns_n!(columns, 1);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let size = cell_vec.len();
|
||||
@@ -558,15 +490,7 @@ impl Function for H3CellCenterChild {
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure!(
|
||||
columns.len() == 2,
|
||||
InvalidFuncArgsSnafu {
|
||||
err_msg: format!(
|
||||
"The length of the args is not correct, expect 2, provided : {}",
|
||||
columns.len()
|
||||
),
|
||||
}
|
||||
);
|
||||
ensure_columns_n!(columns, 2);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let res_vec = &columns[1];
|
||||
@@ -606,15 +530,7 @@ impl Function for H3CellParent {
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure!(
|
||||
columns.len() == 2,
|
||||
InvalidFuncArgsSnafu {
|
||||
err_msg: format!(
|
||||
"The length of the args is not correct, expect 2, provided : {}",
|
||||
columns.len()
|
||||
),
|
||||
}
|
||||
);
|
||||
ensure_columns_n!(columns, 2);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let res_vec = &columns[1];
|
||||
@@ -633,48 +549,323 @@ impl Function for H3CellParent {
|
||||
}
|
||||
}
|
||||
|
||||
/// Function that checks if two cells are neighbour
|
||||
/// Function that returns children cell list
|
||||
#[derive(Clone, Debug, Default, Display)]
|
||||
#[display("{}", self.name())]
|
||||
pub struct H3IsNeighbour;
|
||||
pub struct H3CellToChildren;
|
||||
|
||||
impl Function for H3IsNeighbour {
|
||||
impl Function for H3CellToChildren {
|
||||
fn name(&self) -> &str {
|
||||
"h3_is_neighbour"
|
||||
"h3_cell_to_children"
|
||||
}
|
||||
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::boolean_datatype())
|
||||
Ok(ConcreteDataType::list_datatype(
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
))
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
signature_of_double_cell()
|
||||
signature_of_cell_and_resolution()
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure!(
|
||||
columns.len() == 2,
|
||||
InvalidFuncArgsSnafu {
|
||||
err_msg: format!(
|
||||
"The length of the args is not correct, expect 2, provided : {}",
|
||||
columns.len()
|
||||
),
|
||||
}
|
||||
);
|
||||
ensure_columns_n!(columns, 2);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let cell2_vec = &columns[1];
|
||||
let res_vec = &columns[1];
|
||||
let size = cell_vec.len();
|
||||
let mut results = BooleanVectorBuilder::with_capacity(size);
|
||||
let mut results =
|
||||
ListVectorBuilder::with_type_capacity(ConcreteDataType::uint64_datatype(), size);
|
||||
|
||||
for i in 0..size {
|
||||
let cell = cell_from_value(cell_vec.get(i))?;
|
||||
let res = value_to_resolution(res_vec.get(i))?;
|
||||
let result = cell.map(|cell| {
|
||||
let children: Vec<Value> = cell
|
||||
.children(res)
|
||||
.map(|child| Value::from(u64::from(child)))
|
||||
.collect();
|
||||
ListValue::new(children, ConcreteDataType::uint64_datatype())
|
||||
});
|
||||
|
||||
if let Some(list_value) = result {
|
||||
results.push(Some(list_value.as_scalar_ref()));
|
||||
} else {
|
||||
results.push(None);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(results.to_vector())
|
||||
}
|
||||
}
|
||||
|
||||
/// Function that returns children cell count
|
||||
#[derive(Clone, Debug, Default, Display)]
|
||||
#[display("{}", self.name())]
|
||||
pub struct H3CellToChildrenSize;
|
||||
|
||||
impl Function for H3CellToChildrenSize {
|
||||
fn name(&self) -> &str {
|
||||
"h3_cell_to_children_size"
|
||||
}
|
||||
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::uint64_datatype())
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
signature_of_cell_and_resolution()
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure_columns_n!(columns, 2);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let res_vec = &columns[1];
|
||||
let size = cell_vec.len();
|
||||
let mut results = UInt64VectorBuilder::with_capacity(size);
|
||||
|
||||
for i in 0..size {
|
||||
let cell = cell_from_value(cell_vec.get(i))?;
|
||||
let res = value_to_resolution(res_vec.get(i))?;
|
||||
let result = cell.map(|cell| cell.children_count(res));
|
||||
results.push(result);
|
||||
}
|
||||
|
||||
Ok(results.to_vector())
|
||||
}
|
||||
}
|
||||
|
||||
/// Function that returns the cell position if its parent at given resolution
|
||||
#[derive(Clone, Debug, Default, Display)]
|
||||
#[display("{}", self.name())]
|
||||
pub struct H3CellToChildPos;
|
||||
|
||||
impl Function for H3CellToChildPos {
|
||||
fn name(&self) -> &str {
|
||||
"h3_cell_to_child_pos"
|
||||
}
|
||||
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::uint64_datatype())
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
signature_of_cell_and_resolution()
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure_columns_n!(columns, 2);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let res_vec = &columns[1];
|
||||
let size = cell_vec.len();
|
||||
let mut results = UInt64VectorBuilder::with_capacity(size);
|
||||
|
||||
for i in 0..size {
|
||||
let cell = cell_from_value(cell_vec.get(i))?;
|
||||
let res = value_to_resolution(res_vec.get(i))?;
|
||||
let result = cell.and_then(|cell| cell.child_position(res));
|
||||
results.push(result);
|
||||
}
|
||||
|
||||
Ok(results.to_vector())
|
||||
}
|
||||
}
|
||||
|
||||
/// Function that returns the cell at given position of the parent at given resolution
|
||||
#[derive(Clone, Debug, Default, Display)]
|
||||
#[display("{}", self.name())]
|
||||
pub struct H3ChildPosToCell;
|
||||
|
||||
impl Function for H3ChildPosToCell {
|
||||
fn name(&self) -> &str {
|
||||
"h3_child_pos_to_cell"
|
||||
}
|
||||
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::uint64_datatype())
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
let mut signatures =
|
||||
Vec::with_capacity(POSITION_TYPES.len() * CELL_TYPES.len() * RESOLUTION_TYPES.len());
|
||||
for position_type in POSITION_TYPES.as_slice() {
|
||||
for cell_type in CELL_TYPES.as_slice() {
|
||||
for resolution_type in RESOLUTION_TYPES.as_slice() {
|
||||
signatures.push(TypeSignature::Exact(vec![
|
||||
position_type.clone(),
|
||||
cell_type.clone(),
|
||||
resolution_type.clone(),
|
||||
]));
|
||||
}
|
||||
}
|
||||
}
|
||||
Signature::one_of(signatures, Volatility::Stable)
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure_columns_n!(columns, 3);
|
||||
|
||||
let pos_vec = &columns[0];
|
||||
let cell_vec = &columns[1];
|
||||
let res_vec = &columns[2];
|
||||
let size = cell_vec.len();
|
||||
let mut results = UInt64VectorBuilder::with_capacity(size);
|
||||
|
||||
for i in 0..size {
|
||||
let cell = cell_from_value(cell_vec.get(i))?;
|
||||
let pos = value_to_position(pos_vec.get(i))?;
|
||||
let res = value_to_resolution(res_vec.get(i))?;
|
||||
let result = cell.and_then(|cell| cell.child_at(pos, res).map(u64::from));
|
||||
results.push(result);
|
||||
}
|
||||
|
||||
Ok(results.to_vector())
|
||||
}
|
||||
}
|
||||
|
||||
/// Function that returns cells with k distances of given cell
|
||||
#[derive(Clone, Debug, Default, Display)]
|
||||
#[display("{}", self.name())]
|
||||
pub struct H3GridDisk;
|
||||
|
||||
impl Function for H3GridDisk {
|
||||
fn name(&self) -> &str {
|
||||
"h3_grid_disk"
|
||||
}
|
||||
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::list_datatype(
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
))
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
signature_of_cell_and_distance()
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure_columns_n!(columns, 2);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let k_vec = &columns[1];
|
||||
let size = cell_vec.len();
|
||||
let mut results =
|
||||
ListVectorBuilder::with_type_capacity(ConcreteDataType::uint64_datatype(), size);
|
||||
|
||||
for i in 0..size {
|
||||
let cell = cell_from_value(cell_vec.get(i))?;
|
||||
let k = value_to_distance(k_vec.get(i))?;
|
||||
|
||||
let result = cell.map(|cell| {
|
||||
let children: Vec<Value> = cell
|
||||
.grid_disk::<Vec<_>>(k)
|
||||
.into_iter()
|
||||
.map(|child| Value::from(u64::from(child)))
|
||||
.collect();
|
||||
ListValue::new(children, ConcreteDataType::uint64_datatype())
|
||||
});
|
||||
|
||||
if let Some(list_value) = result {
|
||||
results.push(Some(list_value.as_scalar_ref()));
|
||||
} else {
|
||||
results.push(None);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(results.to_vector())
|
||||
}
|
||||
}
|
||||
|
||||
/// Function that returns all cells within k distances of given cell
|
||||
#[derive(Clone, Debug, Default, Display)]
|
||||
#[display("{}", self.name())]
|
||||
pub struct H3GridDiskDistances;
|
||||
|
||||
impl Function for H3GridDiskDistances {
|
||||
fn name(&self) -> &str {
|
||||
"h3_grid_disk_distances"
|
||||
}
|
||||
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::list_datatype(
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
))
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
signature_of_cell_and_distance()
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure_columns_n!(columns, 2);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let k_vec = &columns[1];
|
||||
let size = cell_vec.len();
|
||||
let mut results =
|
||||
ListVectorBuilder::with_type_capacity(ConcreteDataType::uint64_datatype(), size);
|
||||
|
||||
for i in 0..size {
|
||||
let cell = cell_from_value(cell_vec.get(i))?;
|
||||
let k = value_to_distance(k_vec.get(i))?;
|
||||
|
||||
let result = cell.map(|cell| {
|
||||
let children: Vec<Value> = cell
|
||||
.grid_disk_distances::<Vec<_>>(k)
|
||||
.into_iter()
|
||||
.map(|(child, _distance)| Value::from(u64::from(child)))
|
||||
.collect();
|
||||
ListValue::new(children, ConcreteDataType::uint64_datatype())
|
||||
});
|
||||
|
||||
if let Some(list_value) = result {
|
||||
results.push(Some(list_value.as_scalar_ref()));
|
||||
} else {
|
||||
results.push(None);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(results.to_vector())
|
||||
}
|
||||
}
|
||||
|
||||
/// Function that returns distance between two cells
|
||||
#[derive(Clone, Debug, Default, Display)]
|
||||
#[display("{}", self.name())]
|
||||
pub struct H3GridDistance;
|
||||
|
||||
impl Function for H3GridDistance {
|
||||
fn name(&self) -> &str {
|
||||
"h3_grid_distance"
|
||||
}
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::int32_datatype())
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
signature_of_double_cells()
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure_columns_n!(columns, 2);
|
||||
|
||||
let cell_this_vec = &columns[0];
|
||||
let cell_that_vec = &columns[1];
|
||||
let size = cell_this_vec.len();
|
||||
|
||||
let mut results = Int32VectorBuilder::with_capacity(size);
|
||||
|
||||
for i in 0..size {
|
||||
let result = match (
|
||||
cell_from_value(cell_vec.get(i))?,
|
||||
cell_from_value(cell2_vec.get(i))?,
|
||||
cell_from_value(cell_this_vec.get(i))?,
|
||||
cell_from_value(cell_that_vec.get(i))?,
|
||||
) {
|
||||
(Some(cell_this), Some(cell_that)) => {
|
||||
let is_neighbour = cell_this
|
||||
.is_neighbor_with(cell_that)
|
||||
let dist = cell_this
|
||||
.grid_distance(cell_that)
|
||||
.map_err(|e| {
|
||||
BoxedError::new(PlainError::new(
|
||||
format!("H3 error: {}", e),
|
||||
@@ -682,7 +873,7 @@ impl Function for H3IsNeighbour {
|
||||
))
|
||||
})
|
||||
.context(error::ExecuteSnafu)?;
|
||||
Some(is_neighbour)
|
||||
Some(dist)
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
@@ -694,6 +885,73 @@ impl Function for H3IsNeighbour {
|
||||
}
|
||||
}
|
||||
|
||||
/// Function that returns path cells between two cells
|
||||
#[derive(Clone, Debug, Default, Display)]
|
||||
#[display("{}", self.name())]
|
||||
pub struct H3GridPathCells;
|
||||
|
||||
impl Function for H3GridPathCells {
|
||||
fn name(&self) -> &str {
|
||||
"h3_grid_path_cells"
|
||||
}
|
||||
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::list_datatype(
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
))
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
signature_of_double_cells()
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure_columns_n!(columns, 2);
|
||||
|
||||
let cell_this_vec = &columns[0];
|
||||
let cell_that_vec = &columns[1];
|
||||
let size = cell_this_vec.len();
|
||||
let mut results =
|
||||
ListVectorBuilder::with_type_capacity(ConcreteDataType::uint64_datatype(), size);
|
||||
|
||||
for i in 0..size {
|
||||
let result = match (
|
||||
cell_from_value(cell_this_vec.get(i))?,
|
||||
cell_from_value(cell_that_vec.get(i))?,
|
||||
) {
|
||||
(Some(cell_this), Some(cell_that)) => {
|
||||
let cells = cell_this
|
||||
.grid_path_cells(cell_that)
|
||||
.and_then(|x| x.collect::<std::result::Result<Vec<CellIndex>, _>>())
|
||||
.map_err(|e| {
|
||||
BoxedError::new(PlainError::new(
|
||||
format!("H3 error: {}", e),
|
||||
StatusCode::EngineExecuteQuery,
|
||||
))
|
||||
})
|
||||
.context(error::ExecuteSnafu)?;
|
||||
Some(ListValue::new(
|
||||
cells
|
||||
.into_iter()
|
||||
.map(|c| Value::from(u64::from(c)))
|
||||
.collect(),
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
))
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
|
||||
if let Some(list_value) = result {
|
||||
results.push(Some(list_value.as_scalar_ref()));
|
||||
} else {
|
||||
results.push(None);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(results.to_vector())
|
||||
}
|
||||
}
|
||||
|
||||
fn value_to_resolution(v: Value) -> Result<Resolution> {
|
||||
let r = match v {
|
||||
Value::Int8(v) => v as u8,
|
||||
@@ -716,26 +974,47 @@ fn value_to_resolution(v: Value) -> Result<Resolution> {
|
||||
.context(error::ExecuteSnafu)
|
||||
}
|
||||
|
||||
fn value_to_position(v: Value) -> Result<u64> {
|
||||
match v {
|
||||
Value::Int8(v) => ensure_and_coerce!(v >= 0, v as u64),
|
||||
Value::Int16(v) => ensure_and_coerce!(v >= 0, v as u64),
|
||||
Value::Int32(v) => ensure_and_coerce!(v >= 0, v as u64),
|
||||
Value::Int64(v) => ensure_and_coerce!(v >= 0, v as u64),
|
||||
Value::UInt8(v) => Ok(v as u64),
|
||||
Value::UInt16(v) => Ok(v as u64),
|
||||
Value::UInt32(v) => Ok(v as u64),
|
||||
Value::UInt64(v) => Ok(v),
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
|
||||
fn value_to_distance(v: Value) -> Result<u32> {
|
||||
match v {
|
||||
Value::Int8(v) => ensure_and_coerce!(v >= 0, v as u32),
|
||||
Value::Int16(v) => ensure_and_coerce!(v >= 0, v as u32),
|
||||
Value::Int32(v) => ensure_and_coerce!(v >= 0, v as u32),
|
||||
Value::Int64(v) => ensure_and_coerce!(v >= 0, v as u32),
|
||||
Value::UInt8(v) => Ok(v as u32),
|
||||
Value::UInt16(v) => Ok(v as u32),
|
||||
Value::UInt32(v) => Ok(v),
|
||||
Value::UInt64(v) => Ok(v as u32),
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
|
||||
fn signature_of_cell() -> Signature {
|
||||
let mut signatures = Vec::new();
|
||||
for cell_type in &[
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
ConcreteDataType::int64_datatype(),
|
||||
] {
|
||||
let mut signatures = Vec::with_capacity(CELL_TYPES.len());
|
||||
for cell_type in CELL_TYPES.as_slice() {
|
||||
signatures.push(TypeSignature::Exact(vec![cell_type.clone()]));
|
||||
}
|
||||
|
||||
Signature::one_of(signatures, Volatility::Stable)
|
||||
}
|
||||
|
||||
fn signature_of_double_cell() -> Signature {
|
||||
let mut signatures = Vec::new();
|
||||
let cell_types = &[
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
ConcreteDataType::int64_datatype(),
|
||||
];
|
||||
for cell_type in cell_types {
|
||||
for cell_type2 in cell_types {
|
||||
fn signature_of_double_cells() -> Signature {
|
||||
let mut signatures = Vec::with_capacity(CELL_TYPES.len() * CELL_TYPES.len());
|
||||
for cell_type in CELL_TYPES.as_slice() {
|
||||
for cell_type2 in CELL_TYPES.as_slice() {
|
||||
signatures.push(TypeSignature::Exact(vec![
|
||||
cell_type.clone(),
|
||||
cell_type2.clone(),
|
||||
@@ -747,21 +1026,9 @@ fn signature_of_double_cell() -> Signature {
|
||||
}
|
||||
|
||||
fn signature_of_cell_and_resolution() -> Signature {
|
||||
let mut signatures = Vec::new();
|
||||
for cell_type in &[
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
ConcreteDataType::int64_datatype(),
|
||||
] {
|
||||
for resolution_type in &[
|
||||
ConcreteDataType::int8_datatype(),
|
||||
ConcreteDataType::int16_datatype(),
|
||||
ConcreteDataType::int32_datatype(),
|
||||
ConcreteDataType::int64_datatype(),
|
||||
ConcreteDataType::uint8_datatype(),
|
||||
ConcreteDataType::uint16_datatype(),
|
||||
ConcreteDataType::uint32_datatype(),
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
] {
|
||||
let mut signatures = Vec::with_capacity(CELL_TYPES.len() * RESOLUTION_TYPES.len());
|
||||
for cell_type in CELL_TYPES.as_slice() {
|
||||
for resolution_type in RESOLUTION_TYPES.as_slice() {
|
||||
signatures.push(TypeSignature::Exact(vec![
|
||||
cell_type.clone(),
|
||||
resolution_type.clone(),
|
||||
@@ -771,6 +1038,19 @@ fn signature_of_cell_and_resolution() -> Signature {
|
||||
Signature::one_of(signatures, Volatility::Stable)
|
||||
}
|
||||
|
||||
fn signature_of_cell_and_distance() -> Signature {
|
||||
let mut signatures = Vec::with_capacity(CELL_TYPES.len() * DISTANCE_TYPES.len());
|
||||
for cell_type in CELL_TYPES.as_slice() {
|
||||
for distance_type in DISTANCE_TYPES.as_slice() {
|
||||
signatures.push(TypeSignature::Exact(vec![
|
||||
cell_type.clone(),
|
||||
distance_type.clone(),
|
||||
]));
|
||||
}
|
||||
}
|
||||
Signature::one_of(signatures, Volatility::Stable)
|
||||
}
|
||||
|
||||
fn cell_from_value(v: Value) -> Result<Option<CellIndex>> {
|
||||
let cell = match v {
|
||||
Value::Int64(v) => Some(
|
||||
|
||||
75
src/common/function/src/scalars/geo/helpers.rs
Normal file
75
src/common/function/src/scalars/geo/helpers.rs
Normal file
@@ -0,0 +1,75 @@
|
||||
// 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.
|
||||
|
||||
macro_rules! ensure_columns_len {
|
||||
($columns:ident) => {
|
||||
snafu::ensure!(
|
||||
$columns.windows(2).all(|c| c[0].len() == c[1].len()),
|
||||
common_query::error::InvalidFuncArgsSnafu {
|
||||
err_msg: "The length of input columns are in different size"
|
||||
}
|
||||
)
|
||||
};
|
||||
($column_a:ident, $column_b:ident, $($column_n:ident),*) => {
|
||||
snafu::ensure!(
|
||||
{
|
||||
let mut result = $column_a.len() == $column_b.len();
|
||||
$(
|
||||
result = result && ($column_a.len() == $column_n.len());
|
||||
)*
|
||||
result
|
||||
}
|
||||
common_query::error::InvalidFuncArgsSnafu {
|
||||
err_msg: "The length of input columns are in different size"
|
||||
}
|
||||
)
|
||||
};
|
||||
}
|
||||
|
||||
pub(super) use ensure_columns_len;
|
||||
|
||||
macro_rules! ensure_columns_n {
|
||||
($columns:ident, $n:literal) => {
|
||||
snafu::ensure!(
|
||||
$columns.len() == $n,
|
||||
common_query::error::InvalidFuncArgsSnafu {
|
||||
err_msg: format!(
|
||||
"The length of arguments is not correct, expect {}, provided : {}",
|
||||
stringify!($n),
|
||||
$columns.len()
|
||||
),
|
||||
}
|
||||
);
|
||||
|
||||
if $n > 1 {
|
||||
ensure_columns_len!($columns);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub(super) use ensure_columns_n;
|
||||
|
||||
macro_rules! ensure_and_coerce {
|
||||
($compare:expr, $coerce:expr) => {{
|
||||
snafu::ensure!(
|
||||
$compare,
|
||||
common_query::error::InvalidFuncArgsSnafu {
|
||||
err_msg: "Argument was outside of acceptable range "
|
||||
}
|
||||
);
|
||||
Ok($coerce)
|
||||
}};
|
||||
}
|
||||
|
||||
pub(super) use ensure_and_coerce;
|
||||
275
src/common/function/src/scalars/geo/s2.rs
Normal file
275
src/common/function/src/scalars/geo/s2.rs
Normal file
@@ -0,0 +1,275 @@
|
||||
// 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 common_query::error::{InvalidFuncArgsSnafu, Result};
|
||||
use common_query::prelude::{Signature, TypeSignature};
|
||||
use datafusion::logical_expr::Volatility;
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::scalars::ScalarVectorBuilder;
|
||||
use datatypes::value::Value;
|
||||
use datatypes::vectors::{MutableVector, StringVectorBuilder, UInt64VectorBuilder, VectorRef};
|
||||
use derive_more::Display;
|
||||
use once_cell::sync::Lazy;
|
||||
use s2::cellid::{CellID, MAX_LEVEL};
|
||||
use s2::latlng::LatLng;
|
||||
use snafu::ensure;
|
||||
|
||||
use crate::function::{Function, FunctionContext};
|
||||
use crate::scalars::geo::helpers::{ensure_and_coerce, ensure_columns_len, ensure_columns_n};
|
||||
|
||||
static CELL_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
|
||||
vec![
|
||||
ConcreteDataType::int64_datatype(),
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
]
|
||||
});
|
||||
|
||||
static COORDINATE_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
|
||||
vec![
|
||||
ConcreteDataType::float32_datatype(),
|
||||
ConcreteDataType::float64_datatype(),
|
||||
]
|
||||
});
|
||||
|
||||
static LEVEL_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
|
||||
vec![
|
||||
ConcreteDataType::int8_datatype(),
|
||||
ConcreteDataType::int16_datatype(),
|
||||
ConcreteDataType::int32_datatype(),
|
||||
ConcreteDataType::int64_datatype(),
|
||||
ConcreteDataType::uint8_datatype(),
|
||||
ConcreteDataType::uint16_datatype(),
|
||||
ConcreteDataType::uint32_datatype(),
|
||||
ConcreteDataType::uint64_datatype(),
|
||||
]
|
||||
});
|
||||
|
||||
/// Function that returns [s2] encoding cellid for a given geospatial coordinate.
|
||||
///
|
||||
/// [s2]: http://s2geometry.io
|
||||
#[derive(Clone, Debug, Default, Display)]
|
||||
#[display("{}", self.name())]
|
||||
pub struct S2LatLngToCell;
|
||||
|
||||
impl Function for S2LatLngToCell {
|
||||
fn name(&self) -> &str {
|
||||
"s2_latlng_to_cell"
|
||||
}
|
||||
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::uint64_datatype())
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
let mut signatures = Vec::with_capacity(COORDINATE_TYPES.len());
|
||||
for coord_type in COORDINATE_TYPES.as_slice() {
|
||||
signatures.push(TypeSignature::Exact(vec![
|
||||
// latitude
|
||||
coord_type.clone(),
|
||||
// longitude
|
||||
coord_type.clone(),
|
||||
]));
|
||||
}
|
||||
Signature::one_of(signatures, Volatility::Stable)
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure_columns_n!(columns, 2);
|
||||
|
||||
let lat_vec = &columns[0];
|
||||
let lon_vec = &columns[1];
|
||||
|
||||
let size = lat_vec.len();
|
||||
let mut results = UInt64VectorBuilder::with_capacity(size);
|
||||
|
||||
for i in 0..size {
|
||||
let lat = lat_vec.get(i).as_f64_lossy();
|
||||
let lon = lon_vec.get(i).as_f64_lossy();
|
||||
|
||||
let result = match (lat, lon) {
|
||||
(Some(lat), Some(lon)) => {
|
||||
let coord = LatLng::from_degrees(lat, lon);
|
||||
ensure!(
|
||||
coord.is_valid(),
|
||||
InvalidFuncArgsSnafu {
|
||||
err_msg: "The input coordinates are invalid",
|
||||
}
|
||||
);
|
||||
let cellid = CellID::from(coord);
|
||||
let encoded: u64 = cellid.0;
|
||||
Some(encoded)
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
|
||||
results.push(result);
|
||||
}
|
||||
|
||||
Ok(results.to_vector())
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the level of current s2 cell
|
||||
#[derive(Clone, Debug, Default, Display)]
|
||||
#[display("{}", self.name())]
|
||||
pub struct S2CellLevel;
|
||||
|
||||
impl Function for S2CellLevel {
|
||||
fn name(&self) -> &str {
|
||||
"s2_cell_level"
|
||||
}
|
||||
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::uint64_datatype())
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
signature_of_cell()
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure_columns_n!(columns, 1);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let size = cell_vec.len();
|
||||
let mut results = UInt64VectorBuilder::with_capacity(size);
|
||||
|
||||
for i in 0..size {
|
||||
let cell = cell_from_value(cell_vec.get(i));
|
||||
let res = cell.map(|cell| cell.level());
|
||||
|
||||
results.push(res);
|
||||
}
|
||||
|
||||
Ok(results.to_vector())
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the string presentation of the cell
|
||||
#[derive(Clone, Debug, Default, Display)]
|
||||
#[display("{}", self.name())]
|
||||
pub struct S2CellToToken;
|
||||
|
||||
impl Function for S2CellToToken {
|
||||
fn name(&self) -> &str {
|
||||
"s2_cell_to_token"
|
||||
}
|
||||
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::string_datatype())
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
signature_of_cell()
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure_columns_n!(columns, 1);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let size = cell_vec.len();
|
||||
let mut results = StringVectorBuilder::with_capacity(size);
|
||||
|
||||
for i in 0..size {
|
||||
let cell = cell_from_value(cell_vec.get(i));
|
||||
let res = cell.map(|cell| cell.to_token());
|
||||
|
||||
results.push(res.as_deref());
|
||||
}
|
||||
|
||||
Ok(results.to_vector())
|
||||
}
|
||||
}
|
||||
|
||||
/// Return parent at given level of current s2 cell
|
||||
#[derive(Clone, Debug, Default, Display)]
|
||||
#[display("{}", self.name())]
|
||||
pub struct S2CellParent;
|
||||
|
||||
impl Function for S2CellParent {
|
||||
fn name(&self) -> &str {
|
||||
"s2_cell_parent"
|
||||
}
|
||||
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::uint64_datatype())
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
signature_of_cell_and_level()
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure_columns_n!(columns, 2);
|
||||
|
||||
let cell_vec = &columns[0];
|
||||
let level_vec = &columns[1];
|
||||
let size = cell_vec.len();
|
||||
let mut results = UInt64VectorBuilder::with_capacity(size);
|
||||
|
||||
for i in 0..size {
|
||||
let cell = cell_from_value(cell_vec.get(i));
|
||||
let level = value_to_level(level_vec.get(i))?;
|
||||
let result = cell.map(|cell| cell.parent(level).0);
|
||||
|
||||
results.push(result);
|
||||
}
|
||||
|
||||
Ok(results.to_vector())
|
||||
}
|
||||
}
|
||||
|
||||
fn signature_of_cell() -> Signature {
|
||||
let mut signatures = Vec::with_capacity(CELL_TYPES.len());
|
||||
for cell_type in CELL_TYPES.as_slice() {
|
||||
signatures.push(TypeSignature::Exact(vec![cell_type.clone()]));
|
||||
}
|
||||
|
||||
Signature::one_of(signatures, Volatility::Stable)
|
||||
}
|
||||
|
||||
fn signature_of_cell_and_level() -> Signature {
|
||||
let mut signatures = Vec::with_capacity(CELL_TYPES.len() * LEVEL_TYPES.len());
|
||||
for cell_type in CELL_TYPES.as_slice() {
|
||||
for level_type in LEVEL_TYPES.as_slice() {
|
||||
signatures.push(TypeSignature::Exact(vec![
|
||||
cell_type.clone(),
|
||||
level_type.clone(),
|
||||
]));
|
||||
}
|
||||
}
|
||||
Signature::one_of(signatures, Volatility::Stable)
|
||||
}
|
||||
|
||||
fn cell_from_value(v: Value) -> Option<CellID> {
|
||||
match v {
|
||||
Value::Int64(v) => Some(CellID(v as u64)),
|
||||
Value::UInt64(v) => Some(CellID(v)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn value_to_level(v: Value) -> Result<u64> {
|
||||
match v {
|
||||
Value::Int8(v) => ensure_and_coerce!(v >= 0 && v <= MAX_LEVEL as i8, v as u64),
|
||||
Value::Int16(v) => ensure_and_coerce!(v >= 0 && v <= MAX_LEVEL as i16, v as u64),
|
||||
Value::Int32(v) => ensure_and_coerce!(v >= 0 && v <= MAX_LEVEL as i32, v as u64),
|
||||
Value::Int64(v) => ensure_and_coerce!(v >= 0 && v <= MAX_LEVEL as i64, v as u64),
|
||||
Value::UInt8(v) => ensure_and_coerce!(v <= MAX_LEVEL as u8, v as u64),
|
||||
Value::UInt16(v) => ensure_and_coerce!(v <= MAX_LEVEL as u16, v as u64),
|
||||
Value::UInt32(v) => ensure_and_coerce!(v <= MAX_LEVEL as u32, v as u64),
|
||||
Value::UInt64(v) => ensure_and_coerce!(v <= MAX_LEVEL, v),
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
@@ -15,6 +15,7 @@
|
||||
use std::sync::Arc;
|
||||
mod json_get;
|
||||
mod json_is;
|
||||
mod json_path_exists;
|
||||
mod json_to_string;
|
||||
mod parse_json;
|
||||
|
||||
@@ -46,5 +47,7 @@ impl JsonFunction {
|
||||
registry.register(Arc::new(JsonIsBool));
|
||||
registry.register(Arc::new(JsonIsArray));
|
||||
registry.register(Arc::new(JsonIsObject));
|
||||
|
||||
registry.register(Arc::new(json_path_exists::JsonPathExistsFunction));
|
||||
}
|
||||
}
|
||||
|
||||
172
src/common/function/src/scalars/json/json_path_exists.rs
Normal file
172
src/common/function/src/scalars/json/json_path_exists.rs
Normal file
@@ -0,0 +1,172 @@
|
||||
// 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::fmt::{self, Display};
|
||||
|
||||
use common_query::error::{InvalidFuncArgsSnafu, Result, UnsupportedInputDataTypeSnafu};
|
||||
use common_query::prelude::Signature;
|
||||
use datafusion::logical_expr::Volatility;
|
||||
use datatypes::data_type::ConcreteDataType;
|
||||
use datatypes::prelude::VectorRef;
|
||||
use datatypes::scalars::ScalarVectorBuilder;
|
||||
use datatypes::vectors::{BooleanVectorBuilder, MutableVector};
|
||||
use snafu::ensure;
|
||||
|
||||
use crate::function::{Function, FunctionContext};
|
||||
|
||||
/// Check if the given JSON data contains the given JSON path.
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct JsonPathExistsFunction;
|
||||
|
||||
const NAME: &str = "json_path_exists";
|
||||
|
||||
impl Function for JsonPathExistsFunction {
|
||||
fn name(&self) -> &str {
|
||||
NAME
|
||||
}
|
||||
|
||||
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
|
||||
Ok(ConcreteDataType::boolean_datatype())
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
Signature::exact(
|
||||
vec![
|
||||
ConcreteDataType::json_datatype(),
|
||||
ConcreteDataType::string_datatype(),
|
||||
],
|
||||
Volatility::Immutable,
|
||||
)
|
||||
}
|
||||
|
||||
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
|
||||
ensure!(
|
||||
columns.len() == 2,
|
||||
InvalidFuncArgsSnafu {
|
||||
err_msg: format!(
|
||||
"The length of the args is not correct, expect exactly two, have: {}",
|
||||
columns.len()
|
||||
),
|
||||
}
|
||||
);
|
||||
let jsons = &columns[0];
|
||||
let paths = &columns[1];
|
||||
|
||||
let size = jsons.len();
|
||||
let datatype = jsons.data_type();
|
||||
let mut results = BooleanVectorBuilder::with_capacity(size);
|
||||
|
||||
match datatype {
|
||||
// JSON data type uses binary vector
|
||||
ConcreteDataType::Binary(_) => {
|
||||
for i in 0..size {
|
||||
let json = jsons.get_ref(i);
|
||||
let path = paths.get_ref(i);
|
||||
|
||||
let json = json.as_binary();
|
||||
let path = path.as_string();
|
||||
let result = match (json, path) {
|
||||
(Ok(Some(json)), Ok(Some(path))) => {
|
||||
let json_path = jsonb::jsonpath::parse_json_path(path.as_bytes());
|
||||
match json_path {
|
||||
Ok(json_path) => jsonb::path_exists(json, json_path).ok(),
|
||||
Err(_) => None,
|
||||
}
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
|
||||
results.push(result);
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
return UnsupportedInputDataTypeSnafu {
|
||||
function: NAME,
|
||||
datatypes: columns.iter().map(|c| c.data_type()).collect::<Vec<_>>(),
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
}
|
||||
|
||||
Ok(results.to_vector())
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for JsonPathExistsFunction {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
write!(f, "JSON_PATH_EXISTS")
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use common_query::prelude::TypeSignature;
|
||||
use datatypes::scalars::ScalarVector;
|
||||
use datatypes::vectors::{BinaryVector, StringVector};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_json_path_exists_function() {
|
||||
let json_path_exists = JsonPathExistsFunction;
|
||||
|
||||
assert_eq!("json_path_exists", json_path_exists.name());
|
||||
assert_eq!(
|
||||
ConcreteDataType::boolean_datatype(),
|
||||
json_path_exists
|
||||
.return_type(&[ConcreteDataType::json_datatype()])
|
||||
.unwrap()
|
||||
);
|
||||
|
||||
assert!(matches!(json_path_exists.signature(),
|
||||
Signature {
|
||||
type_signature: TypeSignature::Exact(valid_types),
|
||||
volatility: Volatility::Immutable
|
||||
} if valid_types == vec![ConcreteDataType::json_datatype(), ConcreteDataType::string_datatype()]
|
||||
));
|
||||
|
||||
let json_strings = [
|
||||
r#"{"a": {"b": 2}, "b": 2, "c": 3}"#,
|
||||
r#"{"a": 4, "b": {"c": 6}, "c": 6}"#,
|
||||
r#"{"a": 7, "b": 8, "c": {"a": 7}}"#,
|
||||
r#"{"a": 7, "b": 8, "c": {"a": 7}}"#,
|
||||
];
|
||||
let paths = vec!["$.a.b.c", "$.b", "$.c.a", ".d"];
|
||||
let results = [false, true, true, false];
|
||||
|
||||
let jsonbs = json_strings
|
||||
.iter()
|
||||
.map(|s| {
|
||||
let value = jsonb::parse_value(s.as_bytes()).unwrap();
|
||||
value.to_vec()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let json_vector = BinaryVector::from_vec(jsonbs);
|
||||
let path_vector = StringVector::from_vec(paths);
|
||||
let args: Vec<VectorRef> = vec![Arc::new(json_vector), Arc::new(path_vector)];
|
||||
let vector = json_path_exists
|
||||
.eval(FunctionContext::default(), &args)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(4, vector.len());
|
||||
for (i, gt) in results.iter().enumerate() {
|
||||
let result = vector.get_ref(i);
|
||||
let result = result.as_boolean().unwrap().unwrap();
|
||||
assert_eq!(*gt, result);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -12,7 +12,7 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use api::helper::{convert_i128_to_interval, convert_to_pb_decimal128};
|
||||
use api::helper::{convert_month_day_nano_to_pb, convert_to_pb_decimal128};
|
||||
use api::v1::column::Values;
|
||||
use common_base::BitVec;
|
||||
use datatypes::types::{IntervalType, TimeType, TimestampType, WrapperType};
|
||||
@@ -211,7 +211,7 @@ pub fn values(arrays: &[VectorRef]) -> Result<Values> {
|
||||
ConcreteDataType::Interval(IntervalType::MonthDayNano(_)),
|
||||
IntervalMonthDayNanoVector,
|
||||
interval_month_day_nano_values,
|
||||
|x| { convert_i128_to_interval(x.into_native()) }
|
||||
|x| { convert_month_day_nano_to_pb(x) }
|
||||
),
|
||||
(
|
||||
ConcreteDataType::Decimal128(_),
|
||||
|
||||
@@ -21,23 +21,19 @@ use syn::{parse_macro_input, DeriveInput, ItemStruct};
|
||||
pub(crate) fn impl_aggr_func_type_store(ast: &DeriveInput) -> TokenStream {
|
||||
let name = &ast.ident;
|
||||
let gen = quote! {
|
||||
use common_query::logical_plan::accumulator::AggrFuncTypeStore;
|
||||
use common_query::error::{InvalidInputStateSnafu, Error as QueryError};
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
|
||||
impl AggrFuncTypeStore for #name {
|
||||
fn input_types(&self) -> std::result::Result<Vec<ConcreteDataType>, QueryError> {
|
||||
impl common_query::logical_plan::accumulator::AggrFuncTypeStore for #name {
|
||||
fn input_types(&self) -> std::result::Result<Vec<datatypes::prelude::ConcreteDataType>, common_query::error::Error> {
|
||||
let input_types = self.input_types.load();
|
||||
snafu::ensure!(input_types.is_some(), InvalidInputStateSnafu);
|
||||
snafu::ensure!(input_types.is_some(), common_query::error::InvalidInputStateSnafu);
|
||||
Ok(input_types.as_ref().unwrap().as_ref().clone())
|
||||
}
|
||||
|
||||
fn set_input_types(&self, input_types: Vec<ConcreteDataType>) -> std::result::Result<(), QueryError> {
|
||||
fn set_input_types(&self, input_types: Vec<datatypes::prelude::ConcreteDataType>) -> std::result::Result<(), common_query::error::Error> {
|
||||
let old = self.input_types.swap(Some(std::sync::Arc::new(input_types.clone())));
|
||||
if let Some(old) = old {
|
||||
snafu::ensure!(old.len() == input_types.len(), InvalidInputStateSnafu);
|
||||
snafu::ensure!(old.len() == input_types.len(), common_query::error::InvalidInputStateSnafu);
|
||||
for (x, y) in old.iter().zip(input_types.iter()) {
|
||||
snafu::ensure!(x == y, InvalidInputStateSnafu);
|
||||
snafu::ensure!(x == y, common_query::error::InvalidInputStateSnafu);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
@@ -51,7 +47,7 @@ pub(crate) fn impl_as_aggr_func_creator(_args: TokenStream, input: TokenStream)
|
||||
let mut item_struct = parse_macro_input!(input as ItemStruct);
|
||||
if let syn::Fields::Named(ref mut fields) = item_struct.fields {
|
||||
let result = syn::Field::parse_named.parse2(quote! {
|
||||
input_types: arc_swap::ArcSwapOption<Vec<ConcreteDataType>>
|
||||
input_types: arc_swap::ArcSwapOption<Vec<datatypes::prelude::ConcreteDataType>>
|
||||
});
|
||||
match result {
|
||||
Ok(field) => fields.named.push(field),
|
||||
|
||||
@@ -24,5 +24,5 @@ struct Foo {}
|
||||
fn test_derive() {
|
||||
let _ = Foo::default();
|
||||
assert_fields!(Foo: input_types);
|
||||
assert_impl_all!(Foo: std::fmt::Debug, Default, AggrFuncTypeStore);
|
||||
assert_impl_all!(Foo: std::fmt::Debug, Default, common_query::logical_plan::accumulator::AggrFuncTypeStore);
|
||||
}
|
||||
|
||||
@@ -187,7 +187,7 @@ mod tests {
|
||||
region: Region::new_test(region_id),
|
||||
leader_peer: Some(Peer::empty(1)),
|
||||
follower_peers: vec![],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
}]),
|
||||
HashMap::new(),
|
||||
|
||||
@@ -107,21 +107,21 @@ async fn test_on_submit_alter_request() {
|
||||
region: Region::new_test(RegionId::new(table_id, 1)),
|
||||
leader_peer: Some(Peer::empty(1)),
|
||||
follower_peers: vec![Peer::empty(5)],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
RegionRoute {
|
||||
region: Region::new_test(RegionId::new(table_id, 2)),
|
||||
leader_peer: Some(Peer::empty(2)),
|
||||
follower_peers: vec![Peer::empty(4)],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
RegionRoute {
|
||||
region: Region::new_test(RegionId::new(table_id, 3)),
|
||||
leader_peer: Some(Peer::empty(3)),
|
||||
follower_peers: vec![],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
]),
|
||||
@@ -193,21 +193,21 @@ async fn test_on_submit_alter_request_with_outdated_request() {
|
||||
region: Region::new_test(RegionId::new(table_id, 1)),
|
||||
leader_peer: Some(Peer::empty(1)),
|
||||
follower_peers: vec![Peer::empty(5)],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
RegionRoute {
|
||||
region: Region::new_test(RegionId::new(table_id, 2)),
|
||||
leader_peer: Some(Peer::empty(2)),
|
||||
follower_peers: vec![Peer::empty(4)],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
RegionRoute {
|
||||
region: Region::new_test(RegionId::new(table_id, 3)),
|
||||
leader_peer: Some(Peer::empty(3)),
|
||||
follower_peers: vec![],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
]),
|
||||
|
||||
@@ -119,21 +119,21 @@ async fn test_on_datanode_drop_regions() {
|
||||
region: Region::new_test(RegionId::new(table_id, 1)),
|
||||
leader_peer: Some(Peer::empty(1)),
|
||||
follower_peers: vec![Peer::empty(5)],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
RegionRoute {
|
||||
region: Region::new_test(RegionId::new(table_id, 2)),
|
||||
leader_peer: Some(Peer::empty(2)),
|
||||
follower_peers: vec![Peer::empty(4)],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
RegionRoute {
|
||||
region: Region::new_test(RegionId::new(table_id, 3)),
|
||||
leader_peer: Some(Peer::empty(3)),
|
||||
follower_peers: vec![],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
]),
|
||||
|
||||
@@ -18,6 +18,7 @@ use common_procedure::error::Error as ProcedureError;
|
||||
use snafu::{ensure, OptionExt, ResultExt};
|
||||
use store_api::metric_engine_consts::LOGICAL_TABLE_METADATA_KEY;
|
||||
use table::metadata::TableId;
|
||||
use table::table_reference::TableReference;
|
||||
|
||||
use crate::ddl::DetectingRegion;
|
||||
use crate::error::{Error, OperateDatanodeSnafu, Result, TableNotFoundSnafu, UnsupportedSnafu};
|
||||
@@ -109,8 +110,8 @@ pub async fn check_and_get_physical_table_id(
|
||||
.table_name_manager()
|
||||
.get(physical_table_name)
|
||||
.await?
|
||||
.context(TableNotFoundSnafu {
|
||||
table_name: physical_table_name.to_string(),
|
||||
.with_context(|| TableNotFoundSnafu {
|
||||
table_name: TableReference::from(physical_table_name).to_string(),
|
||||
})
|
||||
.map(|table| table.table_id())
|
||||
}
|
||||
@@ -123,8 +124,8 @@ pub async fn get_physical_table_id(
|
||||
.table_name_manager()
|
||||
.get(logical_table_name)
|
||||
.await?
|
||||
.context(TableNotFoundSnafu {
|
||||
table_name: logical_table_name.to_string(),
|
||||
.with_context(|| TableNotFoundSnafu {
|
||||
table_name: TableReference::from(logical_table_name).to_string(),
|
||||
})
|
||||
.map(|table| table.table_id())?;
|
||||
|
||||
|
||||
@@ -147,6 +147,20 @@ pub enum Error {
|
||||
source: common_procedure::Error,
|
||||
},
|
||||
|
||||
#[snafu(display("Failed to start procedure manager"))]
|
||||
StartProcedureManager {
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
source: common_procedure::Error,
|
||||
},
|
||||
|
||||
#[snafu(display("Failed to stop procedure manager"))]
|
||||
StopProcedureManager {
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
source: common_procedure::Error,
|
||||
},
|
||||
|
||||
#[snafu(display(
|
||||
"Failed to get procedure output, procedure id: {procedure_id}, error: {err_msg}"
|
||||
))]
|
||||
@@ -715,7 +729,9 @@ impl ErrorExt for Error {
|
||||
|
||||
SubmitProcedure { source, .. }
|
||||
| QueryProcedure { source, .. }
|
||||
| WaitProcedure { source, .. } => source.status_code(),
|
||||
| WaitProcedure { source, .. }
|
||||
| StartProcedureManager { source, .. }
|
||||
| StopProcedureManager { source, .. } => source.status_code(),
|
||||
RegisterProcedureLoader { source, .. } => source.status_code(),
|
||||
External { source, .. } => source.status_code(),
|
||||
OperateDatanode { source, .. } => source.status_code(),
|
||||
|
||||
@@ -137,14 +137,16 @@ pub struct DowngradeRegion {
|
||||
/// `None` stands for don't flush before downgrading the region.
|
||||
#[serde(default)]
|
||||
pub flush_timeout: Option<Duration>,
|
||||
/// Rejects all write requests after flushing.
|
||||
pub reject_write: bool,
|
||||
}
|
||||
|
||||
impl Display for DowngradeRegion {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"DowngradeRegion(region_id={}, flush_timeout={:?})",
|
||||
self.region_id, self.flush_timeout,
|
||||
"DowngradeRegion(region_id={}, flush_timeout={:?}, rejct_write={})",
|
||||
self.region_id, self.flush_timeout, self.reject_write
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -140,7 +140,7 @@ use crate::key::table_route::TableRouteKey;
|
||||
use crate::key::txn_helper::TxnOpGetResponseSet;
|
||||
use crate::kv_backend::txn::{Txn, TxnOp};
|
||||
use crate::kv_backend::KvBackendRef;
|
||||
use crate::rpc::router::{region_distribution, RegionRoute, RegionStatus};
|
||||
use crate::rpc::router::{region_distribution, LeaderState, RegionRoute};
|
||||
use crate::rpc::store::BatchDeleteRequest;
|
||||
use crate::DatanodeId;
|
||||
|
||||
@@ -1126,14 +1126,14 @@ impl TableMetadataManager {
|
||||
next_region_route_status: F,
|
||||
) -> Result<()>
|
||||
where
|
||||
F: Fn(&RegionRoute) -> Option<Option<RegionStatus>>,
|
||||
F: Fn(&RegionRoute) -> Option<Option<LeaderState>>,
|
||||
{
|
||||
let mut new_region_routes = current_table_route_value.region_routes()?.clone();
|
||||
|
||||
let mut updated = 0;
|
||||
for route in &mut new_region_routes {
|
||||
if let Some(status) = next_region_route_status(route) {
|
||||
if route.set_leader_status(status) {
|
||||
if let Some(state) = next_region_route_status(route) {
|
||||
if route.set_leader_state(state) {
|
||||
updated += 1;
|
||||
}
|
||||
}
|
||||
@@ -1280,7 +1280,7 @@ mod tests {
|
||||
use crate::key::{DeserializedValueWithBytes, TableMetadataManager, ViewInfoValue};
|
||||
use crate::kv_backend::memory::MemoryKvBackend;
|
||||
use crate::peer::Peer;
|
||||
use crate::rpc::router::{region_distribution, Region, RegionRoute, RegionStatus};
|
||||
use crate::rpc::router::{region_distribution, LeaderState, Region, RegionRoute};
|
||||
|
||||
#[test]
|
||||
fn test_deserialized_value_with_bytes() {
|
||||
@@ -1324,7 +1324,7 @@ mod tests {
|
||||
},
|
||||
leader_peer: Some(Peer::new(datanode, "a2")),
|
||||
follower_peers: vec![],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
}
|
||||
}
|
||||
@@ -1715,7 +1715,7 @@ mod tests {
|
||||
attrs: BTreeMap::new(),
|
||||
},
|
||||
leader_peer: Some(Peer::new(datanode, "a2")),
|
||||
leader_status: Some(RegionStatus::Downgraded),
|
||||
leader_state: Some(LeaderState::Downgrading),
|
||||
follower_peers: vec![],
|
||||
leader_down_since: Some(current_time_millis()),
|
||||
},
|
||||
@@ -1727,7 +1727,7 @@ mod tests {
|
||||
attrs: BTreeMap::new(),
|
||||
},
|
||||
leader_peer: Some(Peer::new(datanode, "a1")),
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
follower_peers: vec![],
|
||||
leader_down_since: None,
|
||||
},
|
||||
@@ -1750,10 +1750,10 @@ mod tests {
|
||||
|
||||
table_metadata_manager
|
||||
.update_leader_region_status(table_id, ¤t_table_route_value, |region_route| {
|
||||
if region_route.leader_status.is_some() {
|
||||
if region_route.leader_state.is_some() {
|
||||
None
|
||||
} else {
|
||||
Some(Some(RegionStatus::Downgraded))
|
||||
Some(Some(LeaderState::Downgrading))
|
||||
}
|
||||
})
|
||||
.await
|
||||
@@ -1768,8 +1768,8 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
updated_route_value.region_routes().unwrap()[0].leader_status,
|
||||
Some(RegionStatus::Downgraded)
|
||||
updated_route_value.region_routes().unwrap()[0].leader_state,
|
||||
Some(LeaderState::Downgrading)
|
||||
);
|
||||
|
||||
assert!(updated_route_value.region_routes().unwrap()[0]
|
||||
@@ -1777,8 +1777,8 @@ mod tests {
|
||||
.is_some());
|
||||
|
||||
assert_eq!(
|
||||
updated_route_value.region_routes().unwrap()[1].leader_status,
|
||||
Some(RegionStatus::Downgraded)
|
||||
updated_route_value.region_routes().unwrap()[1].leader_state,
|
||||
Some(LeaderState::Downgrading)
|
||||
);
|
||||
assert!(updated_route_value.region_routes().unwrap()[1]
|
||||
.leader_down_since
|
||||
@@ -1943,21 +1943,21 @@ mod tests {
|
||||
region: Region::new_test(RegionId::new(table_id, 1)),
|
||||
leader_peer: Some(Peer::empty(1)),
|
||||
follower_peers: vec![Peer::empty(5)],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
RegionRoute {
|
||||
region: Region::new_test(RegionId::new(table_id, 2)),
|
||||
leader_peer: Some(Peer::empty(2)),
|
||||
follower_peers: vec![Peer::empty(4)],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
RegionRoute {
|
||||
region: Region::new_test(RegionId::new(table_id, 3)),
|
||||
leader_peer: Some(Peer::empty(3)),
|
||||
follower_peers: vec![],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
]),
|
||||
@@ -1996,21 +1996,21 @@ mod tests {
|
||||
region: Region::new_test(RegionId::new(table_id, 1)),
|
||||
leader_peer: Some(Peer::empty(1)),
|
||||
follower_peers: vec![Peer::empty(5)],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
RegionRoute {
|
||||
region: Region::new_test(RegionId::new(table_id, 2)),
|
||||
leader_peer: Some(Peer::empty(2)),
|
||||
follower_peers: vec![Peer::empty(4)],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
RegionRoute {
|
||||
region: Region::new_test(RegionId::new(table_id, 3)),
|
||||
leader_peer: Some(Peer::empty(3)),
|
||||
follower_peers: vec![],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
]),
|
||||
|
||||
@@ -21,6 +21,7 @@ use serde::{Deserialize, Serialize};
|
||||
use snafu::OptionExt;
|
||||
use table::metadata::TableId;
|
||||
use table::table_name::TableName;
|
||||
use table::table_reference::TableReference;
|
||||
|
||||
use super::{MetadataKey, MetadataValue, TABLE_NAME_KEY_PATTERN, TABLE_NAME_KEY_PREFIX};
|
||||
use crate::error::{Error, InvalidMetadataSnafu, Result};
|
||||
@@ -122,6 +123,16 @@ impl From<TableNameKey<'_>> for TableName {
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> From<TableNameKey<'a>> for TableReference<'a> {
|
||||
fn from(value: TableNameKey<'a>) -> Self {
|
||||
Self {
|
||||
catalog: value.catalog,
|
||||
schema: value.schema,
|
||||
table: value.table,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> TryFrom<&'a str> for TableNameKey<'a> {
|
||||
type Error = Error;
|
||||
|
||||
|
||||
@@ -744,6 +744,7 @@ mod tests {
|
||||
use crate::kv_backend::memory::MemoryKvBackend;
|
||||
use crate::kv_backend::{KvBackend, TxnService};
|
||||
use crate::peer::Peer;
|
||||
use crate::rpc::router::Region;
|
||||
use crate::rpc::store::PutRequest;
|
||||
|
||||
#[test]
|
||||
@@ -751,11 +752,43 @@ mod tests {
|
||||
let old_raw_v = r#"{"region_routes":[{"region":{"id":1,"name":"r1","partition":null,"attrs":{}},"leader_peer":{"id":2,"addr":"a2"},"follower_peers":[]},{"region":{"id":1,"name":"r1","partition":null,"attrs":{}},"leader_peer":{"id":2,"addr":"a2"},"follower_peers":[]}],"version":0}"#;
|
||||
let v = TableRouteValue::try_from_raw_value(old_raw_v.as_bytes()).unwrap();
|
||||
|
||||
let new_raw_v = format!("{:?}", v);
|
||||
assert_eq!(
|
||||
new_raw_v,
|
||||
r#"Physical(PhysicalTableRouteValue { region_routes: [RegionRoute { region: Region { id: 1(0, 1), name: "r1", partition: None, attrs: {} }, leader_peer: Some(Peer { id: 2, addr: "a2" }), follower_peers: [], leader_status: None, leader_down_since: None }, RegionRoute { region: Region { id: 1(0, 1), name: "r1", partition: None, attrs: {} }, leader_peer: Some(Peer { id: 2, addr: "a2" }), follower_peers: [], leader_status: None, leader_down_since: None }], version: 0 })"#
|
||||
);
|
||||
let expected_table_route = TableRouteValue::Physical(PhysicalTableRouteValue {
|
||||
region_routes: vec![
|
||||
RegionRoute {
|
||||
region: Region {
|
||||
id: RegionId::new(0, 1),
|
||||
name: "r1".to_string(),
|
||||
partition: None,
|
||||
attrs: Default::default(),
|
||||
},
|
||||
leader_peer: Some(Peer {
|
||||
id: 2,
|
||||
addr: "a2".to_string(),
|
||||
}),
|
||||
follower_peers: vec![],
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
RegionRoute {
|
||||
region: Region {
|
||||
id: RegionId::new(0, 1),
|
||||
name: "r1".to_string(),
|
||||
partition: None,
|
||||
attrs: Default::default(),
|
||||
},
|
||||
leader_peer: Some(Peer {
|
||||
id: 2,
|
||||
addr: "a2".to_string(),
|
||||
}),
|
||||
follower_peers: vec![],
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
},
|
||||
],
|
||||
version: 0,
|
||||
});
|
||||
|
||||
assert_eq!(v, expected_table_route);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
183
src/common/meta/src/leadership_notifier.rs
Normal file
183
src/common/meta/src/leadership_notifier.rs
Normal file
@@ -0,0 +1,183 @@
|
||||
// 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::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use common_telemetry::{error, info};
|
||||
|
||||
use crate::error::Result;
|
||||
|
||||
pub type LeadershipChangeNotifierCustomizerRef = Arc<dyn LeadershipChangeNotifierCustomizer>;
|
||||
|
||||
/// A trait for customizing the leadership change notifier.
|
||||
pub trait LeadershipChangeNotifierCustomizer: Send + Sync {
|
||||
fn customize(&self, notifier: &mut LeadershipChangeNotifier);
|
||||
|
||||
fn add_listener(&self, listener: Arc<dyn LeadershipChangeListener>);
|
||||
}
|
||||
|
||||
/// A trait for handling leadership change events in a distributed system.
|
||||
#[async_trait]
|
||||
pub trait LeadershipChangeListener: Send + Sync {
|
||||
/// Returns the listener name.
|
||||
fn name(&self) -> &str;
|
||||
|
||||
/// Called when the node transitions to the leader role.
|
||||
async fn on_leader_start(&self) -> Result<()>;
|
||||
|
||||
/// Called when the node transitions to the follower role.
|
||||
async fn on_leader_stop(&self) -> Result<()>;
|
||||
}
|
||||
|
||||
/// A notifier for leadership change events.
|
||||
#[derive(Default)]
|
||||
pub struct LeadershipChangeNotifier {
|
||||
listeners: Vec<Arc<dyn LeadershipChangeListener>>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct DefaultLeadershipChangeNotifierCustomizer {
|
||||
listeners: Mutex<Vec<Arc<dyn LeadershipChangeListener>>>,
|
||||
}
|
||||
|
||||
impl DefaultLeadershipChangeNotifierCustomizer {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
listeners: Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl LeadershipChangeNotifierCustomizer for DefaultLeadershipChangeNotifierCustomizer {
|
||||
fn customize(&self, notifier: &mut LeadershipChangeNotifier) {
|
||||
info!("Customizing leadership change notifier");
|
||||
let listeners = self.listeners.lock().unwrap().clone();
|
||||
notifier.listeners.extend(listeners);
|
||||
}
|
||||
|
||||
fn add_listener(&self, listener: Arc<dyn LeadershipChangeListener>) {
|
||||
self.listeners.lock().unwrap().push(listener);
|
||||
}
|
||||
}
|
||||
|
||||
impl LeadershipChangeNotifier {
|
||||
/// Adds a listener to the notifier.
|
||||
pub fn add_listener(&mut self, listener: Arc<dyn LeadershipChangeListener>) {
|
||||
self.listeners.push(listener);
|
||||
}
|
||||
|
||||
/// Notify all listeners that the node has become a leader.
|
||||
pub async fn notify_on_leader_start(&self) {
|
||||
for listener in &self.listeners {
|
||||
if let Err(err) = listener.on_leader_start().await {
|
||||
error!(
|
||||
err;
|
||||
"Failed to notify listener: {}, event 'on_leader_start'",
|
||||
listener.name()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Notify all listeners that the node has become a follower.
|
||||
pub async fn notify_on_leader_stop(&self) {
|
||||
for listener in &self.listeners {
|
||||
if let Err(err) = listener.on_leader_stop().await {
|
||||
error!(
|
||||
err;
|
||||
"Failed to notify listener: {}, event: 'on_follower_start'",
|
||||
listener.name()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::*;
|
||||
|
||||
struct MockListener {
|
||||
name: String,
|
||||
on_leader_start_fn: Option<Box<dyn Fn() -> Result<()> + Send + Sync>>,
|
||||
on_follower_start_fn: Option<Box<dyn Fn() -> Result<()> + Send + Sync>>,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl LeadershipChangeListener for MockListener {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
async fn on_leader_start(&self) -> Result<()> {
|
||||
if let Some(f) = &self.on_leader_start_fn {
|
||||
return f();
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn on_leader_stop(&self) -> Result<()> {
|
||||
if let Some(f) = &self.on_follower_start_fn {
|
||||
return f();
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_leadership_change_notifier() {
|
||||
let mut notifier = LeadershipChangeNotifier::default();
|
||||
let listener1 = Arc::new(MockListener {
|
||||
name: "listener1".to_string(),
|
||||
on_leader_start_fn: None,
|
||||
on_follower_start_fn: None,
|
||||
});
|
||||
let called_on_leader_start = Arc::new(AtomicBool::new(false));
|
||||
let called_on_follower_start = Arc::new(AtomicBool::new(false));
|
||||
let called_on_leader_start_moved = called_on_leader_start.clone();
|
||||
let called_on_follower_start_moved = called_on_follower_start.clone();
|
||||
let listener2 = Arc::new(MockListener {
|
||||
name: "listener2".to_string(),
|
||||
on_leader_start_fn: Some(Box::new(move || {
|
||||
called_on_leader_start_moved.store(true, Ordering::Relaxed);
|
||||
Ok(())
|
||||
})),
|
||||
on_follower_start_fn: Some(Box::new(move || {
|
||||
called_on_follower_start_moved.store(true, Ordering::Relaxed);
|
||||
Ok(())
|
||||
})),
|
||||
});
|
||||
|
||||
notifier.add_listener(listener1);
|
||||
notifier.add_listener(listener2);
|
||||
|
||||
let listener1 = notifier.listeners.first().unwrap();
|
||||
let listener2 = notifier.listeners.get(1).unwrap();
|
||||
|
||||
assert_eq!(listener1.name(), "listener1");
|
||||
assert_eq!(listener2.name(), "listener2");
|
||||
|
||||
notifier.notify_on_leader_start().await;
|
||||
assert!(!called_on_follower_start.load(Ordering::Relaxed));
|
||||
assert!(called_on_leader_start.load(Ordering::Relaxed));
|
||||
|
||||
notifier.notify_on_leader_stop().await;
|
||||
assert!(called_on_follower_start.load(Ordering::Relaxed));
|
||||
assert!(called_on_leader_start.load(Ordering::Relaxed));
|
||||
}
|
||||
}
|
||||
@@ -32,6 +32,7 @@ pub mod heartbeat;
|
||||
pub mod instruction;
|
||||
pub mod key;
|
||||
pub mod kv_backend;
|
||||
pub mod leadership_notifier;
|
||||
pub mod lock_key;
|
||||
pub mod metrics;
|
||||
pub mod node_manager;
|
||||
|
||||
@@ -58,7 +58,7 @@ impl MemoryRegionKeeper {
|
||||
Default::default()
|
||||
}
|
||||
|
||||
/// Returns [OpeningRegionGuard] if Region(`region_id`) on Peer(`datanode_id`) does not exist.
|
||||
/// Returns [OperatingRegionGuard] if Region(`region_id`) on Peer(`datanode_id`) does not exist.
|
||||
pub fn register(
|
||||
&self,
|
||||
datanode_id: DatanodeId,
|
||||
|
||||
@@ -108,16 +108,16 @@ pub fn convert_to_region_peer_map(
|
||||
.collect::<HashMap<_, _>>()
|
||||
}
|
||||
|
||||
/// Returns the HashMap<[RegionNumber], [RegionStatus]>;
|
||||
pub fn convert_to_region_leader_status_map(
|
||||
/// Returns the HashMap<[RegionNumber], [LeaderState]>;
|
||||
pub fn convert_to_region_leader_state_map(
|
||||
region_routes: &[RegionRoute],
|
||||
) -> HashMap<RegionNumber, RegionStatus> {
|
||||
) -> HashMap<RegionNumber, LeaderState> {
|
||||
region_routes
|
||||
.iter()
|
||||
.filter_map(|x| {
|
||||
x.leader_status
|
||||
x.leader_state
|
||||
.as_ref()
|
||||
.map(|status| (x.region.id.region_number(), *status))
|
||||
.map(|state| (x.region.id.region_number(), *state))
|
||||
})
|
||||
.collect::<HashMap<_, _>>()
|
||||
}
|
||||
@@ -205,7 +205,7 @@ impl TableRoute {
|
||||
region,
|
||||
leader_peer,
|
||||
follower_peers,
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
});
|
||||
}
|
||||
@@ -259,9 +259,13 @@ pub struct RegionRoute {
|
||||
pub follower_peers: Vec<Peer>,
|
||||
/// `None` by default.
|
||||
#[builder(setter(into, strip_option), default)]
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub leader_status: Option<RegionStatus>,
|
||||
/// The start time when the leader is in `Downgraded` status.
|
||||
#[serde(
|
||||
default,
|
||||
alias = "leader_status",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub leader_state: Option<LeaderState>,
|
||||
/// The start time when the leader is in `Downgraded` state.
|
||||
#[serde(default)]
|
||||
#[builder(default = "self.default_leader_down_since()")]
|
||||
pub leader_down_since: Option<i64>,
|
||||
@@ -269,76 +273,79 @@ pub struct RegionRoute {
|
||||
|
||||
impl RegionRouteBuilder {
|
||||
fn default_leader_down_since(&self) -> Option<i64> {
|
||||
match self.leader_status {
|
||||
Some(Some(RegionStatus::Downgraded)) => Some(current_time_millis()),
|
||||
match self.leader_state {
|
||||
Some(Some(LeaderState::Downgrading)) => Some(current_time_millis()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The Status of the [Region].
|
||||
/// The State of the [`Region`] Leader.
|
||||
/// TODO(dennis): It's better to add more fine-grained statuses such as `PENDING` etc.
|
||||
#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, AsRefStr)]
|
||||
#[strum(serialize_all = "UPPERCASE")]
|
||||
pub enum RegionStatus {
|
||||
/// The following cases in which the [Region] will be downgraded.
|
||||
pub enum LeaderState {
|
||||
/// The following cases in which the [`Region`] will be downgraded.
|
||||
///
|
||||
/// - The [Region] is unavailable(e.g., Crashed, Network disconnected).
|
||||
/// - The [Region] was planned to migrate to another [Peer].
|
||||
Downgraded,
|
||||
/// - The [`Region`] may be unavailable (e.g., Crashed, Network disconnected).
|
||||
/// - The [`Region`] was planned to migrate to another [`Peer`].
|
||||
#[serde(alias = "Downgraded")]
|
||||
Downgrading,
|
||||
}
|
||||
|
||||
impl RegionRoute {
|
||||
/// Returns true if the Leader [Region] is downgraded.
|
||||
/// Returns true if the Leader [`Region`] is downgraded.
|
||||
///
|
||||
/// The following cases in which the [Region] will be downgraded.
|
||||
/// The following cases in which the [`Region`] will be downgraded.
|
||||
///
|
||||
/// - The [Region] is unavailable(e.g., Crashed, Network disconnected).
|
||||
/// - The [Region] was planned to migrate to another [Peer].
|
||||
/// - The [`Region`] is unavailable(e.g., Crashed, Network disconnected).
|
||||
/// - The [`Region`] was planned to migrate to another [`Peer`].
|
||||
///
|
||||
pub fn is_leader_downgraded(&self) -> bool {
|
||||
matches!(self.leader_status, Some(RegionStatus::Downgraded))
|
||||
pub fn is_leader_downgrading(&self) -> bool {
|
||||
matches!(self.leader_state, Some(LeaderState::Downgrading))
|
||||
}
|
||||
|
||||
/// Marks the Leader [Region] as downgraded.
|
||||
/// Marks the Leader [`Region`] as [`RegionState::Downgrading`].
|
||||
///
|
||||
/// We should downgrade a [Region] before deactivating it:
|
||||
/// We should downgrade a [`Region`] before deactivating it:
|
||||
///
|
||||
/// - During the [Region] Failover Procedure.
|
||||
/// - Migrating a [Region].
|
||||
/// - During the [`Region`] Failover Procedure.
|
||||
/// - Migrating a [`Region`].
|
||||
///
|
||||
/// **Notes:** Meta Server will stop renewing the lease for the downgraded [Region].
|
||||
/// **Notes:** Meta Server will renewing a special lease(`Downgrading`) for the downgrading [`Region`].
|
||||
///
|
||||
/// A downgrading region will reject any write requests, and only allow memetable to be flushed to object storage
|
||||
///
|
||||
pub fn downgrade_leader(&mut self) {
|
||||
self.leader_down_since = Some(current_time_millis());
|
||||
self.leader_status = Some(RegionStatus::Downgraded)
|
||||
self.leader_state = Some(LeaderState::Downgrading)
|
||||
}
|
||||
|
||||
/// Returns how long since the leader is in `Downgraded` status.
|
||||
/// Returns how long since the leader is in `Downgraded` state.
|
||||
pub fn leader_down_millis(&self) -> Option<i64> {
|
||||
self.leader_down_since
|
||||
.map(|start| current_time_millis() - start)
|
||||
}
|
||||
|
||||
/// Sets the leader status.
|
||||
/// Sets the leader state.
|
||||
///
|
||||
/// Returns true if updated.
|
||||
pub fn set_leader_status(&mut self, status: Option<RegionStatus>) -> bool {
|
||||
let updated = self.leader_status != status;
|
||||
pub fn set_leader_state(&mut self, state: Option<LeaderState>) -> bool {
|
||||
let updated = self.leader_state != state;
|
||||
|
||||
match (status, updated) {
|
||||
(Some(RegionStatus::Downgraded), true) => {
|
||||
match (state, updated) {
|
||||
(Some(LeaderState::Downgrading), true) => {
|
||||
self.leader_down_since = Some(current_time_millis());
|
||||
}
|
||||
(Some(RegionStatus::Downgraded), false) => {
|
||||
// Do nothing if leader is still in `Downgraded` status.
|
||||
(Some(LeaderState::Downgrading), false) => {
|
||||
// Do nothing if leader is still in `Downgraded` state.
|
||||
}
|
||||
_ => {
|
||||
self.leader_down_since = None;
|
||||
}
|
||||
}
|
||||
|
||||
self.leader_status = status;
|
||||
self.leader_state = state;
|
||||
updated
|
||||
}
|
||||
}
|
||||
@@ -477,15 +484,15 @@ mod tests {
|
||||
},
|
||||
leader_peer: Some(Peer::new(1, "a1")),
|
||||
follower_peers: vec![Peer::new(2, "a2"), Peer::new(3, "a3")],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
};
|
||||
|
||||
assert!(!region_route.is_leader_downgraded());
|
||||
assert!(!region_route.is_leader_downgrading());
|
||||
|
||||
region_route.downgrade_leader();
|
||||
|
||||
assert!(region_route.is_leader_downgraded());
|
||||
assert!(region_route.is_leader_downgrading());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -499,7 +506,7 @@ mod tests {
|
||||
},
|
||||
leader_peer: Some(Peer::new(1, "a1")),
|
||||
follower_peers: vec![Peer::new(2, "a2"), Peer::new(3, "a3")],
|
||||
leader_status: None,
|
||||
leader_state: None,
|
||||
leader_down_since: None,
|
||||
};
|
||||
|
||||
@@ -510,6 +517,73 @@ mod tests {
|
||||
assert_eq!(decoded, region_route);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_region_route_compatibility() {
|
||||
let region_route = RegionRoute {
|
||||
region: Region {
|
||||
id: 2.into(),
|
||||
name: "r2".to_string(),
|
||||
partition: None,
|
||||
attrs: BTreeMap::new(),
|
||||
},
|
||||
leader_peer: Some(Peer::new(1, "a1")),
|
||||
follower_peers: vec![Peer::new(2, "a2"), Peer::new(3, "a3")],
|
||||
leader_state: Some(LeaderState::Downgrading),
|
||||
leader_down_since: None,
|
||||
};
|
||||
let input = r#"{"region":{"id":2,"name":"r2","partition":null,"attrs":{}},"leader_peer":{"id":1,"addr":"a1"},"follower_peers":[{"id":2,"addr":"a2"},{"id":3,"addr":"a3"}],"leader_state":"Downgraded","leader_down_since":null}"#;
|
||||
let decoded: RegionRoute = serde_json::from_str(input).unwrap();
|
||||
assert_eq!(decoded, region_route);
|
||||
|
||||
let region_route = RegionRoute {
|
||||
region: Region {
|
||||
id: 2.into(),
|
||||
name: "r2".to_string(),
|
||||
partition: None,
|
||||
attrs: BTreeMap::new(),
|
||||
},
|
||||
leader_peer: Some(Peer::new(1, "a1")),
|
||||
follower_peers: vec![Peer::new(2, "a2"), Peer::new(3, "a3")],
|
||||
leader_state: Some(LeaderState::Downgrading),
|
||||
leader_down_since: None,
|
||||
};
|
||||
let input = r#"{"region":{"id":2,"name":"r2","partition":null,"attrs":{}},"leader_peer":{"id":1,"addr":"a1"},"follower_peers":[{"id":2,"addr":"a2"},{"id":3,"addr":"a3"}],"leader_status":"Downgraded","leader_down_since":null}"#;
|
||||
let decoded: RegionRoute = serde_json::from_str(input).unwrap();
|
||||
assert_eq!(decoded, region_route);
|
||||
|
||||
let region_route = RegionRoute {
|
||||
region: Region {
|
||||
id: 2.into(),
|
||||
name: "r2".to_string(),
|
||||
partition: None,
|
||||
attrs: BTreeMap::new(),
|
||||
},
|
||||
leader_peer: Some(Peer::new(1, "a1")),
|
||||
follower_peers: vec![Peer::new(2, "a2"), Peer::new(3, "a3")],
|
||||
leader_state: Some(LeaderState::Downgrading),
|
||||
leader_down_since: None,
|
||||
};
|
||||
let input = r#"{"region":{"id":2,"name":"r2","partition":null,"attrs":{}},"leader_peer":{"id":1,"addr":"a1"},"follower_peers":[{"id":2,"addr":"a2"},{"id":3,"addr":"a3"}],"leader_state":"Downgrading","leader_down_since":null}"#;
|
||||
let decoded: RegionRoute = serde_json::from_str(input).unwrap();
|
||||
assert_eq!(decoded, region_route);
|
||||
|
||||
let region_route = RegionRoute {
|
||||
region: Region {
|
||||
id: 2.into(),
|
||||
name: "r2".to_string(),
|
||||
partition: None,
|
||||
attrs: BTreeMap::new(),
|
||||
},
|
||||
leader_peer: Some(Peer::new(1, "a1")),
|
||||
follower_peers: vec![Peer::new(2, "a2"), Peer::new(3, "a3")],
|
||||
leader_state: Some(LeaderState::Downgrading),
|
||||
leader_down_since: None,
|
||||
};
|
||||
let input = r#"{"region":{"id":2,"name":"r2","partition":null,"attrs":{}},"leader_peer":{"id":1,"addr":"a1"},"follower_peers":[{"id":2,"addr":"a2"},{"id":3,"addr":"a3"}],"leader_status":"Downgrading","leader_down_since":null}"#;
|
||||
let decoded: RegionRoute = serde_json::from_str(input).unwrap();
|
||||
assert_eq!(decoded, region_route);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_de_serialize_partition() {
|
||||
let p = Partition {
|
||||
|
||||
@@ -17,6 +17,7 @@ pub mod kafka;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use common_wal::config::MetasrvWalConfig;
|
||||
use common_wal::options::{KafkaWalOptions, WalOptions, WAL_OPTIONS_KEY};
|
||||
use snafu::ResultExt;
|
||||
@@ -24,6 +25,7 @@ use store_api::storage::{RegionId, RegionNumber};
|
||||
|
||||
use crate::error::{EncodeWalOptionsSnafu, Result};
|
||||
use crate::kv_backend::KvBackendRef;
|
||||
use crate::leadership_notifier::LeadershipChangeListener;
|
||||
use crate::wal_options_allocator::kafka::topic_manager::TopicManager as KafkaTopicManager;
|
||||
|
||||
/// Allocates wal options in region granularity.
|
||||
@@ -94,6 +96,21 @@ impl WalOptionsAllocator {
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LeadershipChangeListener for WalOptionsAllocator {
|
||||
fn name(&self) -> &str {
|
||||
"WalOptionsAllocator"
|
||||
}
|
||||
|
||||
async fn on_leader_start(&self) -> Result<()> {
|
||||
self.start().await
|
||||
}
|
||||
|
||||
async fn on_leader_stop(&self) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Allocates a wal options for each region. The allocated wal options is encoded immediately.
|
||||
pub fn allocate_region_wal_options(
|
||||
regions: Vec<RegionNumber>,
|
||||
|
||||
22
src/common/pprof/Cargo.toml
Normal file
22
src/common/pprof/Cargo.toml
Normal file
@@ -0,0 +1,22 @@
|
||||
[package]
|
||||
name = "common-pprof"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
common-error.workspace = true
|
||||
common-macro.workspace = true
|
||||
prost.workspace = true
|
||||
snafu.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
[target.'cfg(unix)'.dependencies]
|
||||
pprof = { version = "0.13", features = [
|
||||
"flamegraph",
|
||||
"prost-codec",
|
||||
"protobuf",
|
||||
] }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
99
src/common/pprof/src/lib.rs
Normal file
99
src/common/pprof/src/lib.rs
Normal file
@@ -0,0 +1,99 @@
|
||||
// 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.
|
||||
|
||||
#[cfg(unix)]
|
||||
pub mod nix;
|
||||
|
||||
pub mod error {
|
||||
use std::any::Any;
|
||||
|
||||
use common_error::ext::ErrorExt;
|
||||
use common_error::status_code::StatusCode;
|
||||
use common_macro::stack_trace_debug;
|
||||
use snafu::{Location, Snafu};
|
||||
|
||||
#[derive(Snafu)]
|
||||
#[stack_trace_debug]
|
||||
#[snafu(visibility(pub(crate)))]
|
||||
pub enum Error {
|
||||
#[cfg(unix)]
|
||||
#[snafu(display("Pprof error"))]
|
||||
Pprof {
|
||||
#[snafu(source)]
|
||||
error: pprof::Error,
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("Pprof is unsupported on this platform"))]
|
||||
Unsupported {
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
},
|
||||
}
|
||||
|
||||
pub type Result<T> = std::result::Result<T, Error>;
|
||||
|
||||
impl ErrorExt for Error {
|
||||
fn status_code(&self) -> StatusCode {
|
||||
match self {
|
||||
#[cfg(unix)]
|
||||
Error::Pprof { .. } => StatusCode::Unexpected,
|
||||
Error::Unsupported { .. } => StatusCode::Unsupported,
|
||||
}
|
||||
}
|
||||
|
||||
fn as_any(&self) -> &dyn Any {
|
||||
self
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
pub mod dummy {
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::error::{Result, UnsupportedSnafu};
|
||||
|
||||
/// Dummpy CPU profiler utility.
|
||||
#[derive(Debug)]
|
||||
pub struct Profiling {}
|
||||
|
||||
impl Profiling {
|
||||
/// Creates a new profiler.
|
||||
pub fn new(_duration: Duration, _frequency: i32) -> Profiling {
|
||||
Profiling {}
|
||||
}
|
||||
|
||||
/// Profiles and returns a generated text.
|
||||
pub async fn dump_text(&self) -> Result<String> {
|
||||
UnsupportedSnafu {}.fail()
|
||||
}
|
||||
|
||||
/// Profiles and returns a generated flamegraph.
|
||||
pub async fn dump_flamegraph(&self) -> Result<Vec<u8>> {
|
||||
UnsupportedSnafu {}.fail()
|
||||
}
|
||||
|
||||
/// Profiles and returns a generated proto.
|
||||
pub async fn dump_proto(&self) -> Result<Vec<u8>> {
|
||||
UnsupportedSnafu {}.fail()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
pub use dummy::Profiling;
|
||||
#[cfg(unix)]
|
||||
pub use nix::Profiling;
|
||||
78
src/common/pprof/src/nix.rs
Normal file
78
src/common/pprof/src/nix.rs
Normal file
@@ -0,0 +1,78 @@
|
||||
// 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::time::Duration;
|
||||
|
||||
use pprof::protos::Message;
|
||||
use snafu::ResultExt;
|
||||
|
||||
use crate::error::{PprofSnafu, Result};
|
||||
|
||||
/// CPU profiler utility.
|
||||
// Inspired by https://github.com/datafuselabs/databend/blob/67f445e83cd4eceda98f6c1c114858929d564029/src/common/base/src/base/profiling.rs
|
||||
#[derive(Debug)]
|
||||
pub struct Profiling {
|
||||
/// Sample duration.
|
||||
duration: Duration,
|
||||
/// Sample frequency.
|
||||
frequency: i32,
|
||||
}
|
||||
|
||||
impl Profiling {
|
||||
/// Creates a new profiler.
|
||||
pub fn new(duration: Duration, frequency: i32) -> Profiling {
|
||||
Profiling {
|
||||
duration,
|
||||
frequency,
|
||||
}
|
||||
}
|
||||
|
||||
/// Profiles and returns a generated pprof report.
|
||||
pub async fn report(&self) -> Result<pprof::Report> {
|
||||
let guard = pprof::ProfilerGuardBuilder::default()
|
||||
.frequency(self.frequency)
|
||||
.blocklist(&["libc", "libgcc", "pthread", "vdso"])
|
||||
.build()
|
||||
.context(PprofSnafu)?;
|
||||
tokio::time::sleep(self.duration).await;
|
||||
guard.report().build().context(PprofSnafu)
|
||||
}
|
||||
|
||||
/// Profiles and returns a generated text.
|
||||
pub async fn dump_text(&self) -> Result<String> {
|
||||
let report = self.report().await?;
|
||||
let text = format!("{report:?}");
|
||||
Ok(text)
|
||||
}
|
||||
|
||||
/// Profiles and returns a generated flamegraph.
|
||||
pub async fn dump_flamegraph(&self) -> Result<Vec<u8>> {
|
||||
let mut body: Vec<u8> = Vec::new();
|
||||
|
||||
let report = self.report().await?;
|
||||
report.flamegraph(&mut body).context(PprofSnafu)?;
|
||||
|
||||
Ok(body)
|
||||
}
|
||||
|
||||
/// Profiles and returns a generated proto.
|
||||
pub async fn dump_proto(&self) -> Result<Vec<u8>> {
|
||||
let report = self.report().await?;
|
||||
// Generate google’s pprof format report.
|
||||
let profile = report.pprof().context(PprofSnafu)?;
|
||||
let body = profile.encode_to_vec();
|
||||
|
||||
Ok(body)
|
||||
}
|
||||
}
|
||||
@@ -329,6 +329,7 @@ impl ExecutionPlanVisitor for MetricCollector {
|
||||
level: self.current_level,
|
||||
metrics: vec![],
|
||||
});
|
||||
self.current_level += 1;
|
||||
return Ok(true);
|
||||
};
|
||||
|
||||
@@ -365,8 +366,7 @@ impl ExecutionPlanVisitor for MetricCollector {
|
||||
}
|
||||
|
||||
fn post_visit(&mut self, _plan: &dyn ExecutionPlan) -> std::result::Result<bool, Self::Error> {
|
||||
// the last minus will underflow
|
||||
self.current_level = self.current_level.wrapping_sub(1);
|
||||
self.current_level -= 1;
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,6 +17,7 @@ backtrace = "0.3"
|
||||
common-error.workspace = true
|
||||
console-subscriber = { version = "0.1", optional = true }
|
||||
greptime-proto.workspace = true
|
||||
humantime-serde.workspace = true
|
||||
lazy_static.workspace = true
|
||||
once_cell.workspace = true
|
||||
opentelemetry = { version = "0.21.0", default-features = false, features = [
|
||||
|
||||
@@ -15,6 +15,7 @@
|
||||
//! logging stuffs, inspired by databend
|
||||
use std::env;
|
||||
use std::sync::{Arc, Mutex, Once};
|
||||
use std::time::Duration;
|
||||
|
||||
use once_cell::sync::{Lazy, OnceCell};
|
||||
use opentelemetry::{global, KeyValue};
|
||||
@@ -26,7 +27,7 @@ use serde::{Deserialize, Serialize};
|
||||
use tracing_appender::non_blocking::WorkerGuard;
|
||||
use tracing_appender::rolling::{RollingFileAppender, Rotation};
|
||||
use tracing_log::LogTracer;
|
||||
use tracing_subscriber::filter::Targets;
|
||||
use tracing_subscriber::filter::{FilterFn, Targets};
|
||||
use tracing_subscriber::fmt::Layer;
|
||||
use tracing_subscriber::layer::SubscriberExt;
|
||||
use tracing_subscriber::prelude::*;
|
||||
@@ -53,6 +54,9 @@ pub struct LoggingOptions {
|
||||
/// The log format that can be one of "json" or "text". Default is "text".
|
||||
pub log_format: LogFormat,
|
||||
|
||||
/// The maximum number of log files set by default.
|
||||
pub max_log_files: usize,
|
||||
|
||||
/// Whether to append logs to stdout. Default is true.
|
||||
pub append_stdout: bool,
|
||||
|
||||
@@ -64,6 +68,24 @@ pub struct LoggingOptions {
|
||||
|
||||
/// The tracing sample ratio.
|
||||
pub tracing_sample_ratio: Option<TracingSampleOptions>,
|
||||
|
||||
/// The logging options of slow query.
|
||||
pub slow_query: SlowQueryOptions,
|
||||
}
|
||||
|
||||
/// The options of slow query.
|
||||
#[derive(Clone, Debug, Serialize, Deserialize, Default)]
|
||||
#[serde(default)]
|
||||
pub struct SlowQueryOptions {
|
||||
/// Whether to enable slow query log.
|
||||
pub enable: bool,
|
||||
|
||||
/// The threshold of slow queries.
|
||||
#[serde(with = "humantime_serde")]
|
||||
pub threshold: Option<Duration>,
|
||||
|
||||
/// The sample ratio of slow queries.
|
||||
pub sample_ratio: Option<f64>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
@@ -96,6 +118,9 @@ impl Default for LoggingOptions {
|
||||
otlp_endpoint: None,
|
||||
tracing_sample_ratio: None,
|
||||
append_stdout: true,
|
||||
slow_query: SlowQueryOptions::default(),
|
||||
// Rotation hourly, 24 files per day, keeps info log files of 30 days
|
||||
max_log_files: 720,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -186,8 +211,17 @@ pub fn init_global_logging(
|
||||
|
||||
// Configure the file logging layer with rolling policy.
|
||||
let file_logging_layer = if !opts.dir.is_empty() {
|
||||
let rolling_appender =
|
||||
RollingFileAppender::new(Rotation::HOURLY, &opts.dir, "greptimedb");
|
||||
let rolling_appender = RollingFileAppender::builder()
|
||||
.rotation(Rotation::HOURLY)
|
||||
.filename_prefix("greptimedb")
|
||||
.max_log_files(opts.max_log_files)
|
||||
.build(&opts.dir)
|
||||
.unwrap_or_else(|e| {
|
||||
panic!(
|
||||
"initializing rolling file appender at {} failed: {}",
|
||||
&opts.dir, e
|
||||
)
|
||||
});
|
||||
let (writer, guard) = tracing_appender::non_blocking(rolling_appender);
|
||||
guards.push(guard);
|
||||
|
||||
@@ -208,8 +242,17 @@ pub fn init_global_logging(
|
||||
|
||||
// Configure the error file logging layer with rolling policy.
|
||||
let err_file_logging_layer = if !opts.dir.is_empty() {
|
||||
let rolling_appender =
|
||||
RollingFileAppender::new(Rotation::HOURLY, &opts.dir, "greptimedb-err");
|
||||
let rolling_appender = RollingFileAppender::builder()
|
||||
.rotation(Rotation::HOURLY)
|
||||
.filename_prefix("greptimedb-err")
|
||||
.max_log_files(opts.max_log_files)
|
||||
.build(&opts.dir)
|
||||
.unwrap_or_else(|e| {
|
||||
panic!(
|
||||
"initializing rolling file appender at {} failed: {}",
|
||||
&opts.dir, e
|
||||
)
|
||||
});
|
||||
let (writer, guard) = tracing_appender::non_blocking(rolling_appender);
|
||||
guards.push(guard);
|
||||
|
||||
@@ -235,6 +278,51 @@ pub fn init_global_logging(
|
||||
None
|
||||
};
|
||||
|
||||
let slow_query_logging_layer = if !opts.dir.is_empty() && opts.slow_query.enable {
|
||||
let rolling_appender = RollingFileAppender::builder()
|
||||
.rotation(Rotation::HOURLY)
|
||||
.filename_prefix("greptimedb-slow-queries")
|
||||
.max_log_files(opts.max_log_files)
|
||||
.build(&opts.dir)
|
||||
.unwrap_or_else(|e| {
|
||||
panic!(
|
||||
"initializing rolling file appender at {} failed: {}",
|
||||
&opts.dir, e
|
||||
)
|
||||
});
|
||||
let (writer, guard) = tracing_appender::non_blocking(rolling_appender);
|
||||
guards.push(guard);
|
||||
|
||||
// Only logs if the field contains "slow".
|
||||
let slow_query_filter = FilterFn::new(|metadata| {
|
||||
metadata
|
||||
.fields()
|
||||
.iter()
|
||||
.any(|field| field.name().contains("slow"))
|
||||
});
|
||||
|
||||
if opts.log_format == LogFormat::Json {
|
||||
Some(
|
||||
Layer::new()
|
||||
.json()
|
||||
.with_writer(writer)
|
||||
.with_ansi(false)
|
||||
.with_filter(slow_query_filter)
|
||||
.boxed(),
|
||||
)
|
||||
} else {
|
||||
Some(
|
||||
Layer::new()
|
||||
.with_writer(writer)
|
||||
.with_ansi(false)
|
||||
.with_filter(slow_query_filter)
|
||||
.boxed(),
|
||||
)
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// resolve log level settings from:
|
||||
// - options from command line or config files
|
||||
// - environment variable: RUST_LOG
|
||||
@@ -279,6 +367,7 @@ pub fn init_global_logging(
|
||||
.with(stdout_logging_layer)
|
||||
.with(file_logging_layer)
|
||||
.with(err_file_logging_layer)
|
||||
.with(slow_query_logging_layer)
|
||||
};
|
||||
|
||||
// consume the `tracing_opts` to avoid "unused" warnings.
|
||||
@@ -289,7 +378,8 @@ pub fn init_global_logging(
|
||||
.with(dyn_filter)
|
||||
.with(stdout_logging_layer)
|
||||
.with(file_logging_layer)
|
||||
.with(err_file_logging_layer);
|
||||
.with(err_file_logging_layer)
|
||||
.with(slow_query_logging_layer);
|
||||
|
||||
if opts.enable_otlp_tracing {
|
||||
global::set_text_map_propagator(TraceContextPropagator::new());
|
||||
|
||||
@@ -152,6 +152,17 @@ macro_rules! trace {
|
||||
};
|
||||
}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! slow {
|
||||
(target: $target:expr, $($arg:tt)+) => {
|
||||
$crate::log!(target: $target, slow = true, $crate::tracing::Level::INFO, $($arg)+)
|
||||
};
|
||||
|
||||
($($arg:tt)+) => {
|
||||
$crate::log!($crate::tracing::Level::INFO, slow = true, $($arg)+)
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use common_error::mock::MockError;
|
||||
|
||||
@@ -14,13 +14,13 @@
|
||||
|
||||
use std::fmt::{Display, Formatter, Write};
|
||||
|
||||
use chrono::{Datelike, Days, LocalResult, Months, NaiveDate, NaiveTime, TimeZone};
|
||||
use chrono::{Datelike, Days, LocalResult, Months, NaiveDate, NaiveTime, TimeDelta, TimeZone};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use snafu::ResultExt;
|
||||
|
||||
use crate::error::{InvalidDateStrSnafu, ParseDateStrSnafu, Result};
|
||||
use crate::interval::Interval;
|
||||
use crate::interval::{IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth};
|
||||
use crate::timezone::get_timezone;
|
||||
use crate::util::datetime_to_utc;
|
||||
use crate::Timezone;
|
||||
@@ -134,29 +134,64 @@ impl Date {
|
||||
(self.0 as i64) * 24 * 3600
|
||||
}
|
||||
|
||||
/// Adds given Interval to the current date.
|
||||
/// Returns None if the resulting date would be out of range.
|
||||
pub fn add_interval(&self, interval: Interval) -> Option<Date> {
|
||||
// FIXME(yingwen): remove add/sub intervals later
|
||||
/// Adds given [IntervalYearMonth] to the current date.
|
||||
pub fn add_year_month(&self, interval: IntervalYearMonth) -> Option<Date> {
|
||||
let naive_date = self.to_chrono_date()?;
|
||||
|
||||
let (months, days, _) = interval.to_month_day_nano();
|
||||
|
||||
naive_date
|
||||
.checked_add_months(Months::new(months as u32))?
|
||||
.checked_add_days(Days::new(days as u64))
|
||||
.checked_add_months(Months::new(interval.months as u32))
|
||||
.map(Into::into)
|
||||
}
|
||||
|
||||
/// Subtracts given Interval to the current date.
|
||||
/// Returns None if the resulting date would be out of range.
|
||||
pub fn sub_interval(&self, interval: Interval) -> Option<Date> {
|
||||
/// Adds given [IntervalDayTime] to the current date.
|
||||
pub fn add_day_time(&self, interval: IntervalDayTime) -> Option<Date> {
|
||||
let naive_date = self.to_chrono_date()?;
|
||||
|
||||
let (months, days, _) = interval.to_month_day_nano();
|
||||
naive_date
|
||||
.checked_add_days(Days::new(interval.days as u64))?
|
||||
.checked_add_signed(TimeDelta::milliseconds(interval.milliseconds as i64))
|
||||
.map(Into::into)
|
||||
}
|
||||
|
||||
/// Adds given [IntervalMonthDayNano] to the current date.
|
||||
pub fn add_month_day_nano(&self, interval: IntervalMonthDayNano) -> Option<Date> {
|
||||
let naive_date = self.to_chrono_date()?;
|
||||
|
||||
naive_date
|
||||
.checked_sub_months(Months::new(months as u32))?
|
||||
.checked_sub_days(Days::new(days as u64))
|
||||
.checked_add_months(Months::new(interval.months as u32))?
|
||||
.checked_add_days(Days::new(interval.days as u64))?
|
||||
.checked_add_signed(TimeDelta::nanoseconds(interval.nanoseconds))
|
||||
.map(Into::into)
|
||||
}
|
||||
|
||||
/// Subtracts given [IntervalYearMonth] to the current date.
|
||||
pub fn sub_year_month(&self, interval: IntervalYearMonth) -> Option<Date> {
|
||||
let naive_date = self.to_chrono_date()?;
|
||||
|
||||
naive_date
|
||||
.checked_sub_months(Months::new(interval.months as u32))
|
||||
.map(Into::into)
|
||||
}
|
||||
|
||||
/// Subtracts given [IntervalDayTime] to the current date.
|
||||
pub fn sub_day_time(&self, interval: IntervalDayTime) -> Option<Date> {
|
||||
let naive_date = self.to_chrono_date()?;
|
||||
|
||||
naive_date
|
||||
.checked_sub_days(Days::new(interval.days as u64))?
|
||||
.checked_sub_signed(TimeDelta::milliseconds(interval.milliseconds as i64))
|
||||
.map(Into::into)
|
||||
}
|
||||
|
||||
/// Subtracts given [IntervalMonthDayNano] to the current date.
|
||||
pub fn sub_month_day_nano(&self, interval: IntervalMonthDayNano) -> Option<Date> {
|
||||
let naive_date = self.to_chrono_date()?;
|
||||
|
||||
naive_date
|
||||
.checked_sub_months(Months::new(interval.months as u32))?
|
||||
.checked_sub_days(Days::new(interval.days as u64))?
|
||||
.checked_sub_signed(TimeDelta::nanoseconds(interval.nanoseconds))
|
||||
.map(Into::into)
|
||||
}
|
||||
|
||||
@@ -246,12 +281,12 @@ mod tests {
|
||||
fn test_add_sub_interval() {
|
||||
let date = Date::new(1000);
|
||||
|
||||
let interval = Interval::from_year_month(3);
|
||||
let interval = IntervalYearMonth::new(3);
|
||||
|
||||
let new_date = date.add_interval(interval).unwrap();
|
||||
let new_date = date.add_year_month(interval).unwrap();
|
||||
assert_eq!(new_date.val(), 1091);
|
||||
|
||||
assert_eq!(date, new_date.sub_interval(interval).unwrap());
|
||||
assert_eq!(date, new_date.sub_year_month(interval).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -13,16 +13,18 @@
|
||||
// limitations under the License.
|
||||
|
||||
use std::fmt::{Display, Formatter, Write};
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::{Days, LocalResult, Months, NaiveDateTime, TimeZone as ChronoTimeZone, Utc};
|
||||
use chrono::{
|
||||
Days, LocalResult, Months, NaiveDateTime, TimeDelta, TimeZone as ChronoTimeZone, Utc,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use snafu::ResultExt;
|
||||
|
||||
use crate::error::{InvalidDateStrSnafu, Result};
|
||||
use crate::interval::{IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth};
|
||||
use crate::timezone::{get_timezone, Timezone};
|
||||
use crate::util::{datetime_to_utc, format_utc_datetime};
|
||||
use crate::{Date, Interval};
|
||||
use crate::Date;
|
||||
|
||||
const DATETIME_FORMAT: &str = "%F %H:%M:%S%.f";
|
||||
const DATETIME_FORMAT_WITH_TZ: &str = "%F %H:%M:%S%.f%z";
|
||||
@@ -160,32 +162,66 @@ impl DateTime {
|
||||
None => Utc.from_utc_datetime(&v).naive_local(),
|
||||
})
|
||||
}
|
||||
/// Adds given Interval to the current datetime.
|
||||
/// Returns None if the resulting datetime would be out of range.
|
||||
pub fn add_interval(&self, interval: Interval) -> Option<Self> {
|
||||
|
||||
// FIXME(yingwen): remove add/sub intervals later
|
||||
/// Adds given [IntervalYearMonth] to the current datetime.
|
||||
pub fn add_year_month(&self, interval: IntervalYearMonth) -> Option<Self> {
|
||||
let naive_datetime = self.to_chrono_datetime()?;
|
||||
let (months, days, nsecs) = interval.to_month_day_nano();
|
||||
|
||||
let naive_datetime = naive_datetime
|
||||
.checked_add_months(Months::new(months as u32))?
|
||||
.checked_add_days(Days::new(days as u64))?
|
||||
+ Duration::from_nanos(nsecs as u64);
|
||||
|
||||
Some(naive_datetime.into())
|
||||
naive_datetime
|
||||
.checked_add_months(Months::new(interval.months as u32))
|
||||
.map(Into::into)
|
||||
}
|
||||
|
||||
/// Subtracts given Interval to the current datetime.
|
||||
/// Returns None if the resulting datetime would be out of range.
|
||||
pub fn sub_interval(&self, interval: Interval) -> Option<Self> {
|
||||
/// Adds given [IntervalDayTime] to the current datetime.
|
||||
pub fn add_day_time(&self, interval: IntervalDayTime) -> Option<Self> {
|
||||
let naive_datetime = self.to_chrono_datetime()?;
|
||||
let (months, days, nsecs) = interval.to_month_day_nano();
|
||||
|
||||
let naive_datetime = naive_datetime
|
||||
.checked_sub_months(Months::new(months as u32))?
|
||||
.checked_sub_days(Days::new(days as u64))?
|
||||
- Duration::from_nanos(nsecs as u64);
|
||||
naive_datetime
|
||||
.checked_add_days(Days::new(interval.days as u64))?
|
||||
.checked_add_signed(TimeDelta::milliseconds(interval.milliseconds as i64))
|
||||
.map(Into::into)
|
||||
}
|
||||
|
||||
Some(naive_datetime.into())
|
||||
/// Adds given [IntervalMonthDayNano] to the current datetime.
|
||||
pub fn add_month_day_nano(&self, interval: IntervalMonthDayNano) -> Option<Self> {
|
||||
let naive_datetime = self.to_chrono_datetime()?;
|
||||
|
||||
naive_datetime
|
||||
.checked_add_months(Months::new(interval.months as u32))?
|
||||
.checked_add_days(Days::new(interval.days as u64))?
|
||||
.checked_add_signed(TimeDelta::nanoseconds(interval.nanoseconds))
|
||||
.map(Into::into)
|
||||
}
|
||||
|
||||
/// Subtracts given [IntervalYearMonth] to the current datetime.
|
||||
pub fn sub_year_month(&self, interval: IntervalYearMonth) -> Option<Self> {
|
||||
let naive_datetime = self.to_chrono_datetime()?;
|
||||
|
||||
naive_datetime
|
||||
.checked_sub_months(Months::new(interval.months as u32))
|
||||
.map(Into::into)
|
||||
}
|
||||
|
||||
/// Subtracts given [IntervalDayTime] to the current datetime.
|
||||
pub fn sub_day_time(&self, interval: IntervalDayTime) -> Option<Self> {
|
||||
let naive_datetime = self.to_chrono_datetime()?;
|
||||
|
||||
naive_datetime
|
||||
.checked_sub_days(Days::new(interval.days as u64))?
|
||||
.checked_sub_signed(TimeDelta::milliseconds(interval.milliseconds as i64))
|
||||
.map(Into::into)
|
||||
}
|
||||
|
||||
/// Subtracts given [IntervalMonthDayNano] to the current datetime.
|
||||
pub fn sub_month_day_nano(&self, interval: IntervalMonthDayNano) -> Option<Self> {
|
||||
let naive_datetime = self.to_chrono_datetime()?;
|
||||
|
||||
naive_datetime
|
||||
.checked_sub_months(Months::new(interval.months as u32))?
|
||||
.checked_sub_days(Days::new(interval.days as u64))?
|
||||
.checked_sub_signed(TimeDelta::nanoseconds(interval.nanoseconds))
|
||||
.map(Into::into)
|
||||
}
|
||||
|
||||
/// Convert to [common_time::date].
|
||||
@@ -231,12 +267,12 @@ mod tests {
|
||||
fn test_add_sub_interval() {
|
||||
let datetime = DateTime::new(1000);
|
||||
|
||||
let interval = Interval::from_day_time(1, 200);
|
||||
let interval = IntervalDayTime::new(1, 200);
|
||||
|
||||
let new_datetime = datetime.add_interval(interval).unwrap();
|
||||
let new_datetime = datetime.add_day_time(interval).unwrap();
|
||||
assert_eq!(new_datetime.val(), 1000 + 3600 * 24 * 1000 + 200);
|
||||
|
||||
assert_eq!(datetime, new_datetime.sub_interval(interval).unwrap());
|
||||
assert_eq!(datetime, new_datetime.sub_day_time(interval).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -12,18 +12,10 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::cmp::Ordering;
|
||||
use std::default::Default;
|
||||
use std::fmt::{self, Display, Formatter, Write};
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::hash::Hash;
|
||||
|
||||
use arrow::datatypes::IntervalUnit as ArrowIntervalUnit;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use snafu::ResultExt;
|
||||
|
||||
use crate::duration::Duration;
|
||||
use crate::error::{Result, TimestampOverflowSnafu};
|
||||
|
||||
#[derive(
|
||||
Debug, Default, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize,
|
||||
@@ -61,268 +53,269 @@ impl From<ArrowIntervalUnit> for IntervalUnit {
|
||||
}
|
||||
}
|
||||
|
||||
/// Interval Type represents a period of time.
|
||||
/// It is composed of months, days and nanoseconds.
|
||||
/// 3 kinds of interval are supported: year-month, day-time and
|
||||
/// month-day-nano, which will be stored in the following format.
|
||||
/// Interval data format:
|
||||
/// | months | days | nsecs |
|
||||
/// | 4bytes | 4bytes | 8bytes |
|
||||
#[derive(Debug, Clone, Default, Copy, Serialize, Deserialize)]
|
||||
pub struct Interval {
|
||||
months: i32,
|
||||
days: i32,
|
||||
nsecs: i64,
|
||||
unit: IntervalUnit,
|
||||
// The `Value` type requires Serialize, Deserialize.
|
||||
#[derive(
|
||||
Debug, Default, Copy, Clone, Eq, PartialEq, Hash, Ord, PartialOrd, Serialize, Deserialize,
|
||||
)]
|
||||
#[repr(C)]
|
||||
pub struct IntervalYearMonth {
|
||||
/// Number of months
|
||||
pub months: i32,
|
||||
}
|
||||
|
||||
// Nanosecond convert to other time unit
|
||||
pub const NANOS_PER_SEC: i64 = 1_000_000_000;
|
||||
pub const NANOS_PER_MILLI: i64 = 1_000_000;
|
||||
pub const NANOS_PER_MICRO: i64 = 1_000;
|
||||
pub const NANOS_PER_HOUR: i64 = 60 * 60 * NANOS_PER_SEC;
|
||||
pub const NANOS_PER_DAY: i64 = 24 * NANOS_PER_HOUR;
|
||||
pub const NANOS_PER_MONTH: i64 = 30 * NANOS_PER_DAY;
|
||||
|
||||
pub const DAYS_PER_MONTH: i64 = 30;
|
||||
|
||||
impl Interval {
|
||||
/// Creates a new interval from months, days and nanoseconds.
|
||||
/// Precision is nanosecond.
|
||||
pub fn from_month_day_nano(months: i32, days: i32, nsecs: i64) -> Self {
|
||||
Interval {
|
||||
months,
|
||||
days,
|
||||
nsecs,
|
||||
unit: IntervalUnit::MonthDayNano,
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates a new interval from months.
|
||||
pub fn from_year_month(months: i32) -> Self {
|
||||
Interval {
|
||||
months,
|
||||
days: 0,
|
||||
nsecs: 0,
|
||||
unit: IntervalUnit::YearMonth,
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates a new interval from days and milliseconds.
|
||||
pub fn from_day_time(days: i32, millis: i32) -> Self {
|
||||
Interval {
|
||||
months: 0,
|
||||
days,
|
||||
nsecs: (millis as i64) * NANOS_PER_MILLI,
|
||||
unit: IntervalUnit::DayTime,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_duration(&self) -> Result<Duration> {
|
||||
Ok(Duration::new_nanosecond(
|
||||
self.to_nanosecond()
|
||||
.try_into()
|
||||
.context(TimestampOverflowSnafu)?,
|
||||
))
|
||||
}
|
||||
|
||||
/// Return a tuple(months, days, nanoseconds) from the interval.
|
||||
pub fn to_month_day_nano(&self) -> (i32, i32, i64) {
|
||||
(self.months, self.days, self.nsecs)
|
||||
}
|
||||
|
||||
/// Converts the interval to nanoseconds.
|
||||
pub fn to_nanosecond(&self) -> i128 {
|
||||
let days = (self.days as i64) + DAYS_PER_MONTH * (self.months as i64);
|
||||
(self.nsecs as i128) + (NANOS_PER_DAY as i128) * (days as i128)
|
||||
}
|
||||
|
||||
/// Smallest interval value.
|
||||
pub const MIN: Self = Self {
|
||||
months: i32::MIN,
|
||||
days: i32::MIN,
|
||||
nsecs: i64::MIN,
|
||||
unit: IntervalUnit::MonthDayNano,
|
||||
};
|
||||
|
||||
/// Largest interval value.
|
||||
pub const MAX: Self = Self {
|
||||
months: i32::MAX,
|
||||
days: i32::MAX,
|
||||
nsecs: i64::MAX,
|
||||
unit: IntervalUnit::MonthDayNano,
|
||||
};
|
||||
|
||||
/// Returns the justified interval.
|
||||
/// allows you to adjust the interval of 30-day as one month and the interval of 24-hour as one day
|
||||
pub fn justified_interval(&self) -> Self {
|
||||
let mut result = *self;
|
||||
let extra_months_d = self.days as i64 / DAYS_PER_MONTH;
|
||||
let extra_months_nsecs = self.nsecs / NANOS_PER_MONTH;
|
||||
result.days -= (extra_months_d * DAYS_PER_MONTH) as i32;
|
||||
result.nsecs -= extra_months_nsecs * NANOS_PER_MONTH;
|
||||
|
||||
let extra_days = self.nsecs / NANOS_PER_DAY;
|
||||
result.nsecs -= extra_days * NANOS_PER_DAY;
|
||||
|
||||
result.months += extra_months_d as i32 + extra_months_nsecs as i32;
|
||||
result.days += extra_days as i32;
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// Convert Interval to nanoseconds,
|
||||
/// to check whether Interval is positive
|
||||
pub fn is_positive(&self) -> bool {
|
||||
self.to_nanosecond() > 0
|
||||
}
|
||||
|
||||
/// is_zero
|
||||
pub fn is_zero(&self) -> bool {
|
||||
self.months == 0 && self.days == 0 && self.nsecs == 0
|
||||
}
|
||||
|
||||
/// get unit
|
||||
pub fn unit(&self) -> IntervalUnit {
|
||||
self.unit
|
||||
}
|
||||
|
||||
/// Multiple Interval by an integer with overflow check.
|
||||
/// Returns justified Interval, or `None` if overflow occurred.
|
||||
pub fn checked_mul_int<I>(&self, rhs: I) -> Option<Self>
|
||||
where
|
||||
I: TryInto<i32>,
|
||||
{
|
||||
let rhs = rhs.try_into().ok()?;
|
||||
let months = self.months.checked_mul(rhs)?;
|
||||
let days = self.days.checked_mul(rhs)?;
|
||||
let nsecs = self.nsecs.checked_mul(rhs as i64)?;
|
||||
|
||||
Some(
|
||||
Self {
|
||||
months,
|
||||
days,
|
||||
nsecs,
|
||||
unit: self.unit,
|
||||
}
|
||||
.justified_interval(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Convert Interval to ISO 8601 string
|
||||
pub fn to_iso8601_string(self) -> String {
|
||||
IntervalFormat::from(self).to_iso8601_string()
|
||||
}
|
||||
|
||||
/// Convert Interval to postgres verbose string
|
||||
pub fn to_postgres_string(self) -> String {
|
||||
IntervalFormat::from(self).to_postgres_string()
|
||||
}
|
||||
|
||||
/// Convert Interval to sql_standard string
|
||||
pub fn to_sql_standard_string(self) -> String {
|
||||
IntervalFormat::from(self).to_sql_standard_string()
|
||||
}
|
||||
|
||||
/// Interval Type and i128 [IntervalUnit::MonthDayNano] Convert
|
||||
/// v consists of months(i32) | days(i32) | nsecs(i64)
|
||||
pub fn from_i128(v: i128) -> Self {
|
||||
Interval {
|
||||
nsecs: v as i64,
|
||||
days: (v >> 64) as i32,
|
||||
months: (v >> 96) as i32,
|
||||
unit: IntervalUnit::MonthDayNano,
|
||||
}
|
||||
}
|
||||
|
||||
/// `Interval` Type and i64 [IntervalUnit::DayTime] Convert
|
||||
/// v consists of days(i32) | milliseconds(i32)
|
||||
pub fn from_i64(v: i64) -> Self {
|
||||
Interval {
|
||||
nsecs: ((v as i32) as i64) * NANOS_PER_MILLI,
|
||||
days: (v >> 32) as i32,
|
||||
months: 0,
|
||||
unit: IntervalUnit::DayTime,
|
||||
}
|
||||
}
|
||||
|
||||
/// `Interval` Type and i32 [IntervalUnit::YearMonth] Convert
|
||||
/// v consists of months(i32)
|
||||
pub fn from_i32(v: i32) -> Self {
|
||||
Interval {
|
||||
nsecs: 0,
|
||||
days: 0,
|
||||
months: v,
|
||||
unit: IntervalUnit::YearMonth,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_i128(&self) -> i128 {
|
||||
// 128 96 64 0
|
||||
// +-------+-------+-------+-------+-------+-------+-------+-------+
|
||||
// | months | days | nanoseconds |
|
||||
// +-------+-------+-------+-------+-------+-------+-------+-------+
|
||||
let months = (self.months as u128 & u32::MAX as u128) << 96;
|
||||
let days = (self.days as u128 & u32::MAX as u128) << 64;
|
||||
let nsecs = self.nsecs as u128 & u64::MAX as u128;
|
||||
(months | days | nsecs) as i128
|
||||
}
|
||||
|
||||
pub fn to_i64(&self) -> i64 {
|
||||
// 64 32 0
|
||||
// +-------+-------+-------+-------+-------+-------+-------+-------+
|
||||
// | days | milliseconds |
|
||||
// +-------+-------+-------+-------+-------+-------+-------+-------+
|
||||
let days = (self.days as u64 & u32::MAX as u64) << 32;
|
||||
let milliseconds = (self.nsecs / NANOS_PER_MILLI) as u64 & u32::MAX as u64;
|
||||
(days | milliseconds) as i64
|
||||
impl IntervalYearMonth {
|
||||
pub fn new(months: i32) -> Self {
|
||||
Self { months }
|
||||
}
|
||||
|
||||
pub fn to_i32(&self) -> i32 {
|
||||
self.months
|
||||
}
|
||||
|
||||
pub fn from_i32(months: i32) -> Self {
|
||||
Self { months }
|
||||
}
|
||||
|
||||
pub fn negative(&self) -> Self {
|
||||
Self {
|
||||
months: -self.months,
|
||||
days: -self.days,
|
||||
nsecs: -self.nsecs,
|
||||
unit: self.unit,
|
||||
Self::new(-self.months)
|
||||
}
|
||||
|
||||
pub fn to_iso8601_string(&self) -> String {
|
||||
IntervalFormat::from(*self).to_iso8601_string()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<IntervalYearMonth> for IntervalFormat {
|
||||
fn from(interval: IntervalYearMonth) -> Self {
|
||||
IntervalFormat {
|
||||
years: interval.months / 12,
|
||||
months: interval.months % 12,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<i128> for Interval {
|
||||
impl From<i32> for IntervalYearMonth {
|
||||
fn from(v: i32) -> Self {
|
||||
Self::from_i32(v)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<IntervalYearMonth> for i32 {
|
||||
fn from(v: IntervalYearMonth) -> Self {
|
||||
v.to_i32()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<IntervalYearMonth> for serde_json::Value {
|
||||
fn from(v: IntervalYearMonth) -> Self {
|
||||
serde_json::Value::from(v.to_i32())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(
|
||||
Debug, Default, Copy, Clone, Eq, PartialEq, Hash, Ord, PartialOrd, Serialize, Deserialize,
|
||||
)]
|
||||
#[repr(C)]
|
||||
pub struct IntervalDayTime {
|
||||
/// Number of days
|
||||
pub days: i32,
|
||||
/// Number of milliseconds
|
||||
pub milliseconds: i32,
|
||||
}
|
||||
|
||||
impl IntervalDayTime {
|
||||
/// The additive identity i.e. `0`.
|
||||
pub const ZERO: Self = Self::new(0, 0);
|
||||
|
||||
/// The multiplicative inverse, i.e. `-1`.
|
||||
pub const MINUS_ONE: Self = Self::new(-1, -1);
|
||||
|
||||
/// The maximum value that can be represented
|
||||
pub const MAX: Self = Self::new(i32::MAX, i32::MAX);
|
||||
|
||||
/// The minimum value that can be represented
|
||||
pub const MIN: Self = Self::new(i32::MIN, i32::MIN);
|
||||
|
||||
pub const fn new(days: i32, milliseconds: i32) -> Self {
|
||||
Self { days, milliseconds }
|
||||
}
|
||||
|
||||
pub fn to_i64(&self) -> i64 {
|
||||
let d = (self.days as u64 & u32::MAX as u64) << 32;
|
||||
let m = self.milliseconds as u64 & u32::MAX as u64;
|
||||
(d | m) as i64
|
||||
}
|
||||
|
||||
pub fn from_i64(value: i64) -> Self {
|
||||
let days = (value >> 32) as i32;
|
||||
let milliseconds = value as i32;
|
||||
Self { days, milliseconds }
|
||||
}
|
||||
|
||||
pub fn negative(&self) -> Self {
|
||||
Self::new(-self.days, -self.milliseconds)
|
||||
}
|
||||
|
||||
pub fn to_iso8601_string(&self) -> String {
|
||||
IntervalFormat::from(*self).to_iso8601_string()
|
||||
}
|
||||
|
||||
pub fn as_millis(&self) -> i64 {
|
||||
self.days as i64 * MS_PER_DAY + self.milliseconds as i64
|
||||
}
|
||||
}
|
||||
|
||||
impl From<i64> for IntervalDayTime {
|
||||
fn from(v: i64) -> Self {
|
||||
Self::from_i64(v)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<IntervalDayTime> for i64 {
|
||||
fn from(v: IntervalDayTime) -> Self {
|
||||
v.to_i64()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<IntervalDayTime> for serde_json::Value {
|
||||
fn from(v: IntervalDayTime) -> Self {
|
||||
serde_json::Value::from(v.to_i64())
|
||||
}
|
||||
}
|
||||
|
||||
// Millisecond convert to other time unit
|
||||
pub const MS_PER_SEC: i64 = 1_000;
|
||||
pub const MS_PER_MINUTE: i64 = 60 * MS_PER_SEC;
|
||||
pub const MS_PER_HOUR: i64 = 60 * MS_PER_MINUTE;
|
||||
pub const MS_PER_DAY: i64 = 24 * MS_PER_HOUR;
|
||||
pub const NANOS_PER_MILLI: i64 = 1_000_000;
|
||||
|
||||
impl From<IntervalDayTime> for IntervalFormat {
|
||||
fn from(interval: IntervalDayTime) -> Self {
|
||||
IntervalFormat {
|
||||
days: interval.days,
|
||||
hours: interval.milliseconds as i64 / MS_PER_HOUR,
|
||||
minutes: (interval.milliseconds as i64 % MS_PER_HOUR) / MS_PER_MINUTE,
|
||||
seconds: (interval.milliseconds as i64 % MS_PER_MINUTE) / MS_PER_SEC,
|
||||
microseconds: (interval.milliseconds as i64 % MS_PER_SEC) * MS_PER_SEC,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(
|
||||
Debug, Default, Copy, Clone, Eq, PartialEq, Hash, Ord, PartialOrd, Serialize, Deserialize,
|
||||
)]
|
||||
#[repr(C)]
|
||||
pub struct IntervalMonthDayNano {
|
||||
/// Number of months
|
||||
pub months: i32,
|
||||
/// Number of days
|
||||
pub days: i32,
|
||||
/// Number of nanoseconds
|
||||
pub nanoseconds: i64,
|
||||
}
|
||||
|
||||
impl IntervalMonthDayNano {
|
||||
/// The additive identity i.e. `0`.
|
||||
pub const ZERO: Self = Self::new(0, 0, 0);
|
||||
|
||||
/// The multiplicative inverse, i.e. `-1`.
|
||||
pub const MINUS_ONE: Self = Self::new(-1, -1, -1);
|
||||
|
||||
/// The maximum value that can be represented
|
||||
pub const MAX: Self = Self::new(i32::MAX, i32::MAX, i64::MAX);
|
||||
|
||||
/// The minimum value that can be represented
|
||||
pub const MIN: Self = Self::new(i32::MIN, i32::MIN, i64::MIN);
|
||||
|
||||
pub const fn new(months: i32, days: i32, nanoseconds: i64) -> Self {
|
||||
Self {
|
||||
months,
|
||||
days,
|
||||
nanoseconds,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_i128(&self) -> i128 {
|
||||
let m = (self.months as u128 & u32::MAX as u128) << 96;
|
||||
let d = (self.days as u128 & u32::MAX as u128) << 64;
|
||||
let n = self.nanoseconds as u128 & u64::MAX as u128;
|
||||
(m | d | n) as i128
|
||||
}
|
||||
|
||||
pub fn from_i128(value: i128) -> Self {
|
||||
let months = (value >> 96) as i32;
|
||||
let days = (value >> 64) as i32;
|
||||
let nanoseconds = value as i64;
|
||||
Self {
|
||||
months,
|
||||
days,
|
||||
nanoseconds,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn negative(&self) -> Self {
|
||||
Self::new(-self.months, -self.days, -self.nanoseconds)
|
||||
}
|
||||
|
||||
pub fn to_iso8601_string(&self) -> String {
|
||||
IntervalFormat::from(*self).to_iso8601_string()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<i128> for IntervalMonthDayNano {
|
||||
fn from(v: i128) -> Self {
|
||||
Self::from_i128(v)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Interval> for i128 {
|
||||
fn from(v: Interval) -> Self {
|
||||
impl From<IntervalMonthDayNano> for i128 {
|
||||
fn from(v: IntervalMonthDayNano) -> Self {
|
||||
v.to_i128()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Interval> for serde_json::Value {
|
||||
fn from(v: Interval) -> Self {
|
||||
Value::String(v.to_string())
|
||||
impl From<IntervalMonthDayNano> for serde_json::Value {
|
||||
fn from(v: IntervalMonthDayNano) -> Self {
|
||||
serde_json::Value::from(v.to_i128().to_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for Interval {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
|
||||
let mut s = String::new();
|
||||
if self.months != 0 {
|
||||
write!(s, "{} months ", self.months)?;
|
||||
// Nanosecond convert to other time unit
|
||||
pub const NS_PER_SEC: i64 = 1_000_000_000;
|
||||
pub const NS_PER_MINUTE: i64 = 60 * NS_PER_SEC;
|
||||
pub const NS_PER_HOUR: i64 = 60 * NS_PER_MINUTE;
|
||||
pub const NS_PER_DAY: i64 = 24 * NS_PER_HOUR;
|
||||
|
||||
impl From<IntervalMonthDayNano> for IntervalFormat {
|
||||
fn from(interval: IntervalMonthDayNano) -> Self {
|
||||
IntervalFormat {
|
||||
years: interval.months / 12,
|
||||
months: interval.months % 12,
|
||||
days: interval.days,
|
||||
hours: interval.nanoseconds / NS_PER_HOUR,
|
||||
minutes: (interval.nanoseconds % NS_PER_HOUR) / NS_PER_MINUTE,
|
||||
seconds: (interval.nanoseconds % NS_PER_MINUTE) / NS_PER_SEC,
|
||||
microseconds: (interval.nanoseconds % NS_PER_SEC) / 1_000,
|
||||
}
|
||||
if self.days != 0 {
|
||||
write!(s, "{} days ", self.days)?;
|
||||
}
|
||||
if self.nsecs != 0 {
|
||||
write!(s, "{} nsecs", self.nsecs)?;
|
||||
}
|
||||
write!(f, "{}", s.trim())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn interval_year_month_to_month_day_nano(interval: IntervalYearMonth) -> IntervalMonthDayNano {
|
||||
IntervalMonthDayNano {
|
||||
months: interval.months,
|
||||
days: 0,
|
||||
nanoseconds: 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn interval_day_time_to_month_day_nano(interval: IntervalDayTime) -> IntervalMonthDayNano {
|
||||
IntervalMonthDayNano {
|
||||
months: 0,
|
||||
days: interval.days,
|
||||
nanoseconds: interval.milliseconds as i64 * NANOS_PER_MILLI,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -339,31 +332,6 @@ pub struct IntervalFormat {
|
||||
pub microseconds: i64,
|
||||
}
|
||||
|
||||
impl From<Interval> for IntervalFormat {
|
||||
fn from(val: Interval) -> IntervalFormat {
|
||||
let months = val.months;
|
||||
let days = val.days;
|
||||
let microseconds = val.nsecs / NANOS_PER_MICRO;
|
||||
let years = (months - (months % 12)) / 12;
|
||||
let months = months - years * 12;
|
||||
let hours = (microseconds - (microseconds % 3_600_000_000)) / 3_600_000_000;
|
||||
let microseconds = microseconds - hours * 3_600_000_000;
|
||||
let minutes = (microseconds - (microseconds % 60_000_000)) / 60_000_000;
|
||||
let microseconds = microseconds - minutes * 60_000_000;
|
||||
let seconds = (microseconds - (microseconds % 1_000_000)) / 1_000_000;
|
||||
let microseconds = microseconds - seconds * 1_000_000;
|
||||
IntervalFormat {
|
||||
years,
|
||||
months,
|
||||
days,
|
||||
hours,
|
||||
minutes,
|
||||
seconds,
|
||||
microseconds,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl IntervalFormat {
|
||||
/// All the field in the interval is 0
|
||||
pub fn is_zero(&self) -> bool {
|
||||
@@ -540,117 +508,37 @@ fn get_time_part(
|
||||
interval
|
||||
}
|
||||
|
||||
/// IntervalCompare is used to compare two intervals
|
||||
/// It makes interval into nanoseconds style.
|
||||
#[derive(PartialEq, Eq, Hash, PartialOrd, Ord)]
|
||||
struct IntervalCompare(i128);
|
||||
|
||||
impl From<Interval> for IntervalCompare {
|
||||
fn from(interval: Interval) -> Self {
|
||||
Self(interval.to_nanosecond())
|
||||
}
|
||||
}
|
||||
|
||||
impl Ord for Interval {
|
||||
fn cmp(&self, other: &Self) -> Ordering {
|
||||
IntervalCompare::from(*self).cmp(&IntervalCompare::from(*other))
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialOrd for Interval {
|
||||
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
|
||||
Some(self.cmp(other))
|
||||
}
|
||||
}
|
||||
|
||||
impl Eq for Interval {}
|
||||
|
||||
impl PartialEq for Interval {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.cmp(other).is_eq()
|
||||
}
|
||||
}
|
||||
|
||||
impl Hash for Interval {
|
||||
fn hash<H: Hasher>(&self, state: &mut H) {
|
||||
IntervalCompare::from(*self).hash(state)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use super::*;
|
||||
use crate::timestamp::TimeUnit;
|
||||
|
||||
#[test]
|
||||
fn test_from_year_month() {
|
||||
let interval = Interval::from_year_month(1);
|
||||
let interval = IntervalYearMonth::new(1);
|
||||
assert_eq!(interval.months, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_from_date_time() {
|
||||
let interval = Interval::from_day_time(1, 2);
|
||||
let interval = IntervalDayTime::new(1, 2);
|
||||
assert_eq!(interval.days, 1);
|
||||
assert_eq!(interval.nsecs, 2_000_000);
|
||||
assert_eq!(interval.milliseconds, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_to_duration() {
|
||||
let interval = Interval::from_day_time(1, 2);
|
||||
|
||||
let duration = interval.to_duration().unwrap();
|
||||
assert_eq!(86400002000000, duration.value());
|
||||
assert_eq!(TimeUnit::Nanosecond, duration.unit());
|
||||
|
||||
let interval = Interval::from_year_month(12);
|
||||
|
||||
let duration = interval.to_duration().unwrap();
|
||||
assert_eq!(31104000000000000, duration.value());
|
||||
assert_eq!(TimeUnit::Nanosecond, duration.unit());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_interval_is_positive() {
|
||||
let interval = Interval::from_year_month(1);
|
||||
assert!(interval.is_positive());
|
||||
let interval = Interval::from_year_month(-1);
|
||||
assert!(!interval.is_positive());
|
||||
|
||||
let interval = Interval::from_day_time(1, i32::MIN);
|
||||
assert!(!interval.is_positive());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_to_nanosecond() {
|
||||
let interval = Interval::from_year_month(1);
|
||||
assert_eq!(interval.to_nanosecond(), 2592000000000000);
|
||||
let interval = Interval::from_day_time(1, 2);
|
||||
assert_eq!(interval.to_nanosecond(), 86400002000000);
|
||||
|
||||
let max_interval = Interval::from_month_day_nano(i32::MAX, i32::MAX, i64::MAX);
|
||||
assert_eq!(max_interval.to_nanosecond(), 5751829423496836854775807);
|
||||
|
||||
let min_interval = Interval::from_month_day_nano(i32::MIN, i32::MIN, i64::MIN);
|
||||
assert_eq!(min_interval.to_nanosecond(), -5751829426175236854775808);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_interval_is_zero() {
|
||||
let interval = Interval::from_month_day_nano(1, 1, 1);
|
||||
assert!(!interval.is_zero());
|
||||
let interval = Interval::from_month_day_nano(0, 0, 0);
|
||||
assert!(interval.is_zero());
|
||||
fn test_from_month_day_nano() {
|
||||
let interval = IntervalMonthDayNano::new(1, 2, 3);
|
||||
assert_eq!(interval.months, 1);
|
||||
assert_eq!(interval.days, 2);
|
||||
assert_eq!(interval.nanoseconds, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_interval_i128_convert() {
|
||||
let test_interval_eq = |month, day, nano| {
|
||||
let interval = Interval::from_month_day_nano(month, day, nano);
|
||||
let interval = IntervalMonthDayNano::new(month, day, nano);
|
||||
let interval_i128 = interval.to_i128();
|
||||
let interval2 = Interval::from_i128(interval_i128);
|
||||
let interval2 = IntervalMonthDayNano::from_i128(interval_i128);
|
||||
assert_eq!(interval, interval2);
|
||||
};
|
||||
|
||||
@@ -666,11 +554,26 @@ mod tests {
|
||||
test_interval_eq(i32::MAX, i32::MIN, i64::MIN);
|
||||
test_interval_eq(i32::MIN, i32::MAX, i64::MIN);
|
||||
test_interval_eq(i32::MIN, i32::MIN, i64::MIN);
|
||||
|
||||
let interval = IntervalMonthDayNano::from_i128(1);
|
||||
assert_eq!(interval, IntervalMonthDayNano::new(0, 0, 1));
|
||||
assert_eq!(1, IntervalMonthDayNano::new(0, 0, 1).to_i128());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_interval_i64_convert() {
|
||||
let interval = IntervalDayTime::from_i64(1);
|
||||
assert_eq!(interval, IntervalDayTime::new(0, 1));
|
||||
assert_eq!(1, IntervalDayTime::new(0, 1).to_i64());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_interval_format() {
|
||||
let interval = Interval::from_month_day_nano(14, 160, 1000000);
|
||||
let interval = IntervalMonthDayNano {
|
||||
months: 14,
|
||||
days: 160,
|
||||
nanoseconds: 1000000,
|
||||
};
|
||||
let interval_format = IntervalFormat::from(interval);
|
||||
assert_eq!(interval_format.years, 1);
|
||||
assert_eq!(interval_format.months, 2);
|
||||
@@ -681,94 +584,34 @@ mod tests {
|
||||
assert_eq!(interval_format.microseconds, 1000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_interval_hash() {
|
||||
let interval = Interval::from_month_day_nano(1, 31, 1);
|
||||
let interval2 = Interval::from_month_day_nano(2, 1, 1);
|
||||
let mut map = HashMap::new();
|
||||
map.insert(interval, 1);
|
||||
assert_eq!(map.get(&interval2), Some(&1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_interval_mul_int() {
|
||||
let interval = Interval::from_month_day_nano(1, 1, 1);
|
||||
let interval2 = interval.checked_mul_int(2).unwrap();
|
||||
assert_eq!(interval2.months, 2);
|
||||
assert_eq!(interval2.days, 2);
|
||||
assert_eq!(interval2.nsecs, 2);
|
||||
|
||||
// test justified interval
|
||||
let interval = Interval::from_month_day_nano(1, 31, 1);
|
||||
let interval2 = interval.checked_mul_int(2).unwrap();
|
||||
assert_eq!(interval2.months, 4);
|
||||
assert_eq!(interval2.days, 2);
|
||||
assert_eq!(interval2.nsecs, 2);
|
||||
|
||||
// test overflow situation
|
||||
let interval = Interval::from_month_day_nano(i32::MAX, 1, 1);
|
||||
let interval2 = interval.checked_mul_int(2);
|
||||
assert!(interval2.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_display() {
|
||||
let interval = Interval::from_month_day_nano(1, 1, 1);
|
||||
assert_eq!(interval.to_string(), "1 months 1 days 1 nsecs");
|
||||
|
||||
let interval = Interval::from_month_day_nano(14, 31, 10000000000);
|
||||
assert_eq!(interval.to_string(), "14 months 31 days 10000000000 nsecs");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_interval_justified() {
|
||||
let interval = Interval::from_month_day_nano(1, 131, 1).justified_interval();
|
||||
let interval2 = Interval::from_month_day_nano(5, 11, 1);
|
||||
assert_eq!(interval, interval2);
|
||||
|
||||
let interval = Interval::from_month_day_nano(1, 1, NANOS_PER_MONTH + 2 * NANOS_PER_DAY)
|
||||
.justified_interval();
|
||||
let interval2 = Interval::from_month_day_nano(2, 3, 0);
|
||||
assert_eq!(interval, interval2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_serde_json() {
|
||||
let interval = Interval::from_month_day_nano(1, 1, 1);
|
||||
let json = serde_json::to_string(&interval).unwrap();
|
||||
assert_eq!(
|
||||
json,
|
||||
"{\"months\":1,\"days\":1,\"nsecs\":1,\"unit\":\"MonthDayNano\"}"
|
||||
);
|
||||
let interval2: Interval = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(interval, interval2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_to_iso8601_string() {
|
||||
// Test interval zero
|
||||
let interval = Interval::from_month_day_nano(0, 0, 0);
|
||||
let interval = IntervalMonthDayNano::new(0, 0, 0);
|
||||
assert_eq!(interval.to_iso8601_string(), "PT0S");
|
||||
|
||||
let interval = Interval::from_month_day_nano(1, 1, 1);
|
||||
let interval = IntervalMonthDayNano::new(1, 1, 1);
|
||||
assert_eq!(interval.to_iso8601_string(), "P0Y1M1DT0H0M0S");
|
||||
|
||||
let interval = Interval::from_month_day_nano(14, 31, 10000000000);
|
||||
let interval = IntervalMonthDayNano::new(14, 31, 10000000000);
|
||||
assert_eq!(interval.to_iso8601_string(), "P1Y2M31DT0H0M10S");
|
||||
|
||||
let interval = Interval::from_month_day_nano(14, 31, 23210200000000);
|
||||
let interval = IntervalMonthDayNano::new(14, 31, 23210200000000);
|
||||
assert_eq!(interval.to_iso8601_string(), "P1Y2M31DT6H26M50.2S");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_to_postgres_string() {
|
||||
// Test interval zero
|
||||
let interval = Interval::from_month_day_nano(0, 0, 0);
|
||||
assert_eq!(interval.to_postgres_string(), "00:00:00");
|
||||
|
||||
let interval = Interval::from_month_day_nano(23, 100, 23210200000000);
|
||||
let interval = IntervalMonthDayNano::new(0, 0, 0);
|
||||
assert_eq!(
|
||||
interval.to_postgres_string(),
|
||||
IntervalFormat::from(interval).to_postgres_string(),
|
||||
"00:00:00"
|
||||
);
|
||||
|
||||
let interval = IntervalMonthDayNano::new(23, 100, 23210200000000);
|
||||
assert_eq!(
|
||||
IntervalFormat::from(interval).to_postgres_string(),
|
||||
"1 year 11 mons 100 days 06:26:50.200000"
|
||||
);
|
||||
}
|
||||
@@ -776,18 +619,21 @@ mod tests {
|
||||
#[test]
|
||||
fn test_to_sql_standard_string() {
|
||||
// Test zero interval
|
||||
let interval = Interval::from_month_day_nano(0, 0, 0);
|
||||
assert_eq!(interval.to_sql_standard_string(), "0");
|
||||
let interval = IntervalMonthDayNano::new(0, 0, 0);
|
||||
assert_eq!(IntervalFormat::from(interval).to_sql_standard_string(), "0");
|
||||
|
||||
let interval = Interval::from_month_day_nano(23, 100, 23210200000000);
|
||||
let interval = IntervalMonthDayNano::new(23, 100, 23210200000000);
|
||||
assert_eq!(
|
||||
interval.to_sql_standard_string(),
|
||||
IntervalFormat::from(interval).to_sql_standard_string(),
|
||||
"+1-11 +100 +6:26:50.200000"
|
||||
);
|
||||
|
||||
// Test interval without year, month, day
|
||||
let interval = Interval::from_month_day_nano(0, 0, 23210200000000);
|
||||
assert_eq!(interval.to_sql_standard_string(), "6:26:50.200000");
|
||||
let interval = IntervalMonthDayNano::new(0, 0, 23210200000000);
|
||||
assert_eq!(
|
||||
IntervalFormat::from(interval).to_sql_standard_string(),
|
||||
"6:26:50.200000"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -27,7 +27,7 @@ pub mod util;
|
||||
pub use date::Date;
|
||||
pub use datetime::DateTime;
|
||||
pub use duration::Duration;
|
||||
pub use interval::Interval;
|
||||
pub use interval::{IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth};
|
||||
pub use range::RangeMillis;
|
||||
pub use timestamp::Timestamp;
|
||||
pub use timestamp_millis::TimestampMillis;
|
||||
|
||||
@@ -20,16 +20,17 @@ use std::time::Duration;
|
||||
|
||||
use arrow::datatypes::TimeUnit as ArrowTimeUnit;
|
||||
use chrono::{
|
||||
DateTime, Days, LocalResult, Months, NaiveDate, NaiveDateTime, NaiveTime,
|
||||
DateTime, Days, LocalResult, Months, NaiveDate, NaiveDateTime, NaiveTime, TimeDelta,
|
||||
TimeZone as ChronoTimeZone, Utc,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use snafu::{OptionExt, ResultExt};
|
||||
|
||||
use crate::error;
|
||||
use crate::error::{ArithmeticOverflowSnafu, ParseTimestampSnafu, Result, TimestampOverflowSnafu};
|
||||
use crate::interval::{IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth};
|
||||
use crate::timezone::{get_timezone, Timezone};
|
||||
use crate::util::{datetime_to_utc, div_ceil};
|
||||
use crate::{error, Interval};
|
||||
|
||||
/// Timestamp represents the value of units(seconds/milliseconds/microseconds/nanoseconds) elapsed
|
||||
/// since UNIX epoch. The valid value range of [Timestamp] depends on it's unit (all in UTC timezone):
|
||||
@@ -140,40 +141,77 @@ impl Timestamp {
|
||||
})
|
||||
}
|
||||
|
||||
/// Adds given Interval to the current timestamp.
|
||||
/// Returns None if the resulting timestamp would be out of range.
|
||||
pub fn add_interval(&self, interval: Interval) -> Option<Timestamp> {
|
||||
// FIXME(yingwen): remove add/sub intervals later
|
||||
/// Adds given [IntervalYearMonth] to the current timestamp.
|
||||
pub fn add_year_month(&self, interval: IntervalYearMonth) -> Option<Timestamp> {
|
||||
let naive_datetime = self.to_chrono_datetime()?;
|
||||
let (months, days, nsecs) = interval.to_month_day_nano();
|
||||
|
||||
let naive_datetime = naive_datetime
|
||||
.checked_add_months(Months::new(months as u32))?
|
||||
.checked_add_days(Days::new(days as u64))?
|
||||
+ Duration::from_nanos(nsecs as u64);
|
||||
let naive_datetime =
|
||||
naive_datetime.checked_add_months(Months::new(interval.months as u32))?;
|
||||
|
||||
match Timestamp::from_chrono_datetime(naive_datetime) {
|
||||
// Have to convert the new timestamp by the current unit.
|
||||
Some(ts) => ts.convert_to(self.unit),
|
||||
None => None,
|
||||
}
|
||||
// Have to convert the new timestamp by the current unit.
|
||||
Timestamp::from_chrono_datetime(naive_datetime).and_then(|ts| ts.convert_to(self.unit))
|
||||
}
|
||||
|
||||
/// Subtracts given Interval to the current timestamp.
|
||||
/// Returns None if the resulting timestamp would be out of range.
|
||||
pub fn sub_interval(&self, interval: Interval) -> Option<Timestamp> {
|
||||
/// Adds given [IntervalDayTime] to the current timestamp.
|
||||
pub fn add_day_time(&self, interval: IntervalDayTime) -> Option<Timestamp> {
|
||||
let naive_datetime = self.to_chrono_datetime()?;
|
||||
let (months, days, nsecs) = interval.to_month_day_nano();
|
||||
|
||||
let naive_datetime = naive_datetime
|
||||
.checked_sub_months(Months::new(months as u32))?
|
||||
.checked_sub_days(Days::new(days as u64))?
|
||||
- Duration::from_nanos(nsecs as u64);
|
||||
.checked_add_days(Days::new(interval.days as u64))?
|
||||
.checked_add_signed(TimeDelta::milliseconds(interval.milliseconds as i64))?;
|
||||
|
||||
match Timestamp::from_chrono_datetime(naive_datetime) {
|
||||
// Have to convert the new timestamp by the current unit.
|
||||
Some(ts) => ts.convert_to(self.unit),
|
||||
None => None,
|
||||
}
|
||||
// Have to convert the new timestamp by the current unit.
|
||||
Timestamp::from_chrono_datetime(naive_datetime).and_then(|ts| ts.convert_to(self.unit))
|
||||
}
|
||||
|
||||
/// Adds given [IntervalMonthDayNano] to the current timestamp.
|
||||
pub fn add_month_day_nano(&self, interval: IntervalMonthDayNano) -> Option<Timestamp> {
|
||||
let naive_datetime = self.to_chrono_datetime()?;
|
||||
|
||||
let naive_datetime = naive_datetime
|
||||
.checked_add_months(Months::new(interval.months as u32))?
|
||||
.checked_add_days(Days::new(interval.days as u64))?
|
||||
.checked_add_signed(TimeDelta::nanoseconds(interval.nanoseconds))?;
|
||||
|
||||
// Have to convert the new timestamp by the current unit.
|
||||
Timestamp::from_chrono_datetime(naive_datetime).and_then(|ts| ts.convert_to(self.unit))
|
||||
}
|
||||
|
||||
/// Subtracts given [IntervalYearMonth] to the current timestamp.
|
||||
pub fn sub_year_month(&self, interval: IntervalYearMonth) -> Option<Timestamp> {
|
||||
let naive_datetime = self.to_chrono_datetime()?;
|
||||
|
||||
let naive_datetime =
|
||||
naive_datetime.checked_sub_months(Months::new(interval.months as u32))?;
|
||||
|
||||
// Have to convert the new timestamp by the current unit.
|
||||
Timestamp::from_chrono_datetime(naive_datetime).and_then(|ts| ts.convert_to(self.unit))
|
||||
}
|
||||
|
||||
/// Subtracts given [IntervalDayTime] to the current timestamp.
|
||||
pub fn sub_day_time(&self, interval: IntervalDayTime) -> Option<Timestamp> {
|
||||
let naive_datetime = self.to_chrono_datetime()?;
|
||||
|
||||
let naive_datetime = naive_datetime
|
||||
.checked_sub_days(Days::new(interval.days as u64))?
|
||||
.checked_sub_signed(TimeDelta::milliseconds(interval.milliseconds as i64))?;
|
||||
|
||||
// Have to convert the new timestamp by the current unit.
|
||||
Timestamp::from_chrono_datetime(naive_datetime).and_then(|ts| ts.convert_to(self.unit))
|
||||
}
|
||||
|
||||
/// Subtracts given [IntervalMonthDayNano] to the current timestamp.
|
||||
pub fn sub_month_day_nano(&self, interval: IntervalMonthDayNano) -> Option<Timestamp> {
|
||||
let naive_datetime = self.to_chrono_datetime()?;
|
||||
|
||||
let naive_datetime = naive_datetime
|
||||
.checked_sub_months(Months::new(interval.months as u32))?
|
||||
.checked_sub_days(Days::new(interval.days as u64))?
|
||||
.checked_sub_signed(TimeDelta::nanoseconds(interval.nanoseconds))?;
|
||||
|
||||
// Have to convert the new timestamp by the current unit.
|
||||
Timestamp::from_chrono_datetime(naive_datetime).and_then(|ts| ts.convert_to(self.unit))
|
||||
}
|
||||
|
||||
/// Subtracts current timestamp with another timestamp, yielding a duration.
|
||||
@@ -688,13 +726,13 @@ mod tests {
|
||||
fn test_add_sub_interval() {
|
||||
let ts = Timestamp::new(1000, TimeUnit::Millisecond);
|
||||
|
||||
let interval = Interval::from_day_time(1, 200);
|
||||
let interval = IntervalDayTime::new(1, 200);
|
||||
|
||||
let new_ts = ts.add_interval(interval).unwrap();
|
||||
let new_ts = ts.add_day_time(interval).unwrap();
|
||||
assert_eq!(new_ts.unit(), TimeUnit::Millisecond);
|
||||
assert_eq!(new_ts.value(), 1000 + 3600 * 24 * 1000 + 200);
|
||||
|
||||
assert_eq!(ts, new_ts.sub_interval(interval).unwrap());
|
||||
assert_eq!(ts, new_ts.sub_day_time(interval).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -46,6 +46,8 @@ pub struct DatanodeKafkaConfig {
|
||||
pub create_index: bool,
|
||||
#[serde(with = "humantime_serde")]
|
||||
pub dump_index_interval: Duration,
|
||||
/// Ignore missing entries during read WAL.
|
||||
pub overwrite_entry_start_id: bool,
|
||||
}
|
||||
|
||||
impl Default for DatanodeKafkaConfig {
|
||||
@@ -60,6 +62,7 @@ impl Default for DatanodeKafkaConfig {
|
||||
auto_create_topics: true,
|
||||
create_index: true,
|
||||
dump_index_interval: Duration::from_secs(60),
|
||||
overwrite_entry_start_id: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -129,8 +129,10 @@ impl RegionAliveKeeper {
|
||||
let request = RegionRequest::Close(RegionCloseRequest {});
|
||||
if let Err(e) = self.region_server.handle_request(region_id, request).await {
|
||||
if e.status_code() != StatusCode::RegionNotFound {
|
||||
let _ = self.region_server.set_writable(region_id, false);
|
||||
error!(e; "Failed to close staled region {}, set region to readonly.",region_id);
|
||||
let _ = self
|
||||
.region_server
|
||||
.set_region_role(region_id, RegionRole::Follower);
|
||||
error!(e; "Failed to close staled region {}, convert region to follower.", region_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -378,7 +380,7 @@ impl CountdownTask {
|
||||
}
|
||||
},
|
||||
Some(CountdownCommand::Reset((role, deadline))) => {
|
||||
let _ = self.region_server.set_writable(self.region_id, role.writable());
|
||||
let _ = self.region_server.set_region_role(self.region_id, role);
|
||||
trace!(
|
||||
"Reset deadline of region {region_id} to approximately {} seconds later.",
|
||||
(deadline - Instant::now()).as_secs_f32(),
|
||||
@@ -399,8 +401,8 @@ impl CountdownTask {
|
||||
}
|
||||
}
|
||||
() = &mut countdown => {
|
||||
warn!("The region {region_id} lease is expired, set region to readonly.");
|
||||
let _ = self.region_server.set_writable(self.region_id, false);
|
||||
warn!("The region {region_id} lease is expired, convert region to follower.");
|
||||
let _ = self.region_server.set_region_role(self.region_id, RegionRole::Follower);
|
||||
// resets the countdown.
|
||||
let far_future = Instant::now() + Duration::from_secs(86400 * 365 * 30);
|
||||
countdown.as_mut().reset(far_future);
|
||||
@@ -436,7 +438,9 @@ mod test {
|
||||
.handle_request(region_id, RegionRequest::Create(builder.build()))
|
||||
.await
|
||||
.unwrap();
|
||||
region_server.set_writable(region_id, true).unwrap();
|
||||
region_server
|
||||
.set_region_role(region_id, RegionRole::Leader)
|
||||
.unwrap();
|
||||
|
||||
// Register a region before starting.
|
||||
alive_keeper.register_region(region_id).await;
|
||||
|
||||
@@ -47,7 +47,7 @@ use servers::server::ServerHandlers;
|
||||
use servers::Mode;
|
||||
use snafu::{ensure, OptionExt, ResultExt};
|
||||
use store_api::path_utils::{region_dir, WAL_DIR};
|
||||
use store_api::region_engine::RegionEngineRef;
|
||||
use store_api::region_engine::{RegionEngineRef, RegionRole};
|
||||
use store_api::region_request::RegionOpenRequest;
|
||||
use store_api::storage::RegionId;
|
||||
use tokio::fs;
|
||||
@@ -546,9 +546,9 @@ async fn open_all_regions(
|
||||
|
||||
for region_id in open_regions {
|
||||
if open_with_writable {
|
||||
if let Err(e) = region_server.set_writable(region_id, true) {
|
||||
if let Err(e) = region_server.set_region_role(region_id, RegionRole::Leader) {
|
||||
error!(
|
||||
e; "failed to set writable for region {region_id}"
|
||||
e; "failed to convert region {region_id} to leader"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -126,7 +126,9 @@ impl HeartbeatTask {
|
||||
let mut follower_region_lease_count = 0;
|
||||
for lease in &lease.regions {
|
||||
match lease.role() {
|
||||
RegionRole::Leader => leader_region_lease_count += 1,
|
||||
RegionRole::Leader | RegionRole::DowngradingLeader => {
|
||||
leader_region_lease_count += 1
|
||||
}
|
||||
RegionRole::Follower => follower_region_lease_count += 1,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -153,6 +153,7 @@ mod tests {
|
||||
use mito2::engine::MITO_ENGINE_NAME;
|
||||
use mito2::test_util::{CreateRequestBuilder, TestEnv};
|
||||
use store_api::path_utils::region_dir;
|
||||
use store_api::region_engine::RegionRole;
|
||||
use store_api::region_request::{RegionCloseRequest, RegionRequest};
|
||||
use store_api::storage::RegionId;
|
||||
use tokio::sync::mpsc::{self, Receiver};
|
||||
@@ -213,6 +214,7 @@ mod tests {
|
||||
let instruction = Instruction::DowngradeRegion(DowngradeRegion {
|
||||
region_id: RegionId::new(2048, 1),
|
||||
flush_timeout: Some(Duration::from_secs(1)),
|
||||
reject_write: false,
|
||||
});
|
||||
assert!(heartbeat_handler
|
||||
.is_acceptable(&heartbeat_env.create_handler_ctx((meta.clone(), instruction))));
|
||||
@@ -295,7 +297,9 @@ mod tests {
|
||||
}
|
||||
|
||||
assert_matches!(
|
||||
region_server.set_writable(region_id, true).unwrap_err(),
|
||||
region_server
|
||||
.set_region_role(region_id, RegionRole::Leader)
|
||||
.unwrap_err(),
|
||||
error::Error::RegionNotFound { .. }
|
||||
);
|
||||
}
|
||||
@@ -411,6 +415,7 @@ mod tests {
|
||||
let instruction = Instruction::DowngradeRegion(DowngradeRegion {
|
||||
region_id,
|
||||
flush_timeout: Some(Duration::from_secs(1)),
|
||||
reject_write: false,
|
||||
});
|
||||
|
||||
let mut ctx = heartbeat_env.create_handler_ctx((meta, instruction));
|
||||
@@ -433,6 +438,7 @@ mod tests {
|
||||
let instruction = Instruction::DowngradeRegion(DowngradeRegion {
|
||||
region_id: RegionId::new(2048, 1),
|
||||
flush_timeout: Some(Duration::from_secs(1)),
|
||||
reject_write: false,
|
||||
});
|
||||
let mut ctx = heartbeat_env.create_handler_ctx((meta, instruction));
|
||||
let control = heartbeat_handler.handle(&mut ctx).await.unwrap();
|
||||
|
||||
@@ -16,7 +16,7 @@ use common_meta::instruction::{DowngradeRegion, DowngradeRegionReply, Instructio
|
||||
use common_telemetry::tracing::info;
|
||||
use common_telemetry::warn;
|
||||
use futures_util::future::BoxFuture;
|
||||
use store_api::region_engine::SetReadonlyResponse;
|
||||
use store_api::region_engine::{SetRegionRoleStateResponse, SettableRegionRoleState};
|
||||
use store_api::region_request::{RegionFlushRequest, RegionRequest};
|
||||
use store_api::storage::RegionId;
|
||||
|
||||
@@ -24,16 +24,20 @@ use crate::heartbeat::handler::HandlerContext;
|
||||
use crate::heartbeat::task_tracker::WaitResult;
|
||||
|
||||
impl HandlerContext {
|
||||
async fn set_readonly_gracefully(&self, region_id: RegionId) -> InstructionReply {
|
||||
match self.region_server.set_readonly_gracefully(region_id).await {
|
||||
Ok(SetReadonlyResponse::Success { last_entry_id }) => {
|
||||
async fn downgrade_to_follower_gracefully(&self, region_id: RegionId) -> InstructionReply {
|
||||
match self
|
||||
.region_server
|
||||
.set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower)
|
||||
.await
|
||||
{
|
||||
Ok(SetRegionRoleStateResponse::Success { last_entry_id }) => {
|
||||
InstructionReply::DowngradeRegion(DowngradeRegionReply {
|
||||
last_entry_id,
|
||||
exists: true,
|
||||
error: None,
|
||||
})
|
||||
}
|
||||
Ok(SetReadonlyResponse::NotFound) => {
|
||||
Ok(SetRegionRoleStateResponse::NotFound) => {
|
||||
InstructionReply::DowngradeRegion(DowngradeRegionReply {
|
||||
last_entry_id: None,
|
||||
exists: false,
|
||||
@@ -53,10 +57,12 @@ impl HandlerContext {
|
||||
DowngradeRegion {
|
||||
region_id,
|
||||
flush_timeout,
|
||||
reject_write,
|
||||
}: DowngradeRegion,
|
||||
) -> BoxFuture<'static, InstructionReply> {
|
||||
Box::pin(async move {
|
||||
let Some(writable) = self.region_server.is_writable(region_id) else {
|
||||
let Some(writable) = self.region_server.is_region_leader(region_id) else {
|
||||
warn!("Region: {region_id} is not found");
|
||||
return InstructionReply::DowngradeRegion(DowngradeRegionReply {
|
||||
last_entry_id: None,
|
||||
exists: false,
|
||||
@@ -64,61 +70,89 @@ impl HandlerContext {
|
||||
});
|
||||
};
|
||||
|
||||
let region_server_moved = self.region_server.clone();
|
||||
|
||||
// Ignores flush request
|
||||
if !writable {
|
||||
return self.set_readonly_gracefully(region_id).await;
|
||||
return self.downgrade_to_follower_gracefully(region_id).await;
|
||||
}
|
||||
|
||||
let region_server_moved = self.region_server.clone();
|
||||
if let Some(flush_timeout) = flush_timeout {
|
||||
let register_result = self
|
||||
.downgrade_tasks
|
||||
.try_register(
|
||||
// If flush_timeout is not set, directly convert region to follower.
|
||||
let Some(flush_timeout) = flush_timeout else {
|
||||
return self.downgrade_to_follower_gracefully(region_id).await;
|
||||
};
|
||||
|
||||
if reject_write {
|
||||
// Sets region to downgrading, the downgrading region will reject all write requests.
|
||||
match self
|
||||
.region_server
|
||||
.set_region_role_state_gracefully(
|
||||
region_id,
|
||||
Box::pin(async move {
|
||||
info!("Flush region: {region_id} before downgrading region");
|
||||
region_server_moved
|
||||
.handle_request(
|
||||
region_id,
|
||||
RegionRequest::Flush(RegionFlushRequest {
|
||||
row_group_size: None,
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}),
|
||||
SettableRegionRoleState::DowngradingLeader,
|
||||
)
|
||||
.await;
|
||||
|
||||
if register_result.is_busy() {
|
||||
warn!("Another flush task is running for the region: {region_id}");
|
||||
}
|
||||
|
||||
let mut watcher = register_result.into_watcher();
|
||||
let result = self.catchup_tasks.wait(&mut watcher, flush_timeout).await;
|
||||
|
||||
match result {
|
||||
WaitResult::Timeout => {
|
||||
InstructionReply::DowngradeRegion(DowngradeRegionReply {
|
||||
.await
|
||||
{
|
||||
Ok(SetRegionRoleStateResponse::Success { .. }) => {}
|
||||
Ok(SetRegionRoleStateResponse::NotFound) => {
|
||||
warn!("Region: {region_id} is not found");
|
||||
return InstructionReply::DowngradeRegion(DowngradeRegionReply {
|
||||
last_entry_id: None,
|
||||
exists: true,
|
||||
error: Some(format!(
|
||||
"Flush region: {region_id} before downgrading region is timeout"
|
||||
)),
|
||||
})
|
||||
exists: false,
|
||||
error: None,
|
||||
});
|
||||
}
|
||||
WaitResult::Finish(Ok(_)) => self.set_readonly_gracefully(region_id).await,
|
||||
WaitResult::Finish(Err(err)) => {
|
||||
InstructionReply::DowngradeRegion(DowngradeRegionReply {
|
||||
Err(err) => {
|
||||
warn!(err; "Failed to convert region to downgrading leader");
|
||||
return InstructionReply::DowngradeRegion(DowngradeRegionReply {
|
||||
last_entry_id: None,
|
||||
exists: true,
|
||||
error: Some(format!("{err:?}")),
|
||||
})
|
||||
});
|
||||
}
|
||||
}
|
||||
} else {
|
||||
self.set_readonly_gracefully(region_id).await
|
||||
}
|
||||
|
||||
let register_result = self
|
||||
.downgrade_tasks
|
||||
.try_register(
|
||||
region_id,
|
||||
Box::pin(async move {
|
||||
info!("Flush region: {region_id} before converting region to follower");
|
||||
region_server_moved
|
||||
.handle_request(
|
||||
region_id,
|
||||
RegionRequest::Flush(RegionFlushRequest {
|
||||
row_group_size: None,
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
if register_result.is_busy() {
|
||||
warn!("Another flush task is running for the region: {region_id}");
|
||||
}
|
||||
|
||||
let mut watcher = register_result.into_watcher();
|
||||
let result = self.catchup_tasks.wait(&mut watcher, flush_timeout).await;
|
||||
|
||||
match result {
|
||||
WaitResult::Timeout => InstructionReply::DowngradeRegion(DowngradeRegionReply {
|
||||
last_entry_id: None,
|
||||
exists: true,
|
||||
error: Some(format!("Flush region: {region_id} is timeout")),
|
||||
}),
|
||||
WaitResult::Finish(Ok(_)) => self.downgrade_to_follower_gracefully(region_id).await,
|
||||
WaitResult::Finish(Err(err)) => {
|
||||
InstructionReply::DowngradeRegion(DowngradeRegionReply {
|
||||
last_entry_id: None,
|
||||
exists: true,
|
||||
error: Some(format!("{err:?}")),
|
||||
})
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -131,7 +165,7 @@ mod tests {
|
||||
|
||||
use common_meta::instruction::{DowngradeRegion, InstructionReply};
|
||||
use mito2::engine::MITO_ENGINE_NAME;
|
||||
use store_api::region_engine::{RegionRole, SetReadonlyResponse};
|
||||
use store_api::region_engine::{RegionRole, SetRegionRoleStateResponse};
|
||||
use store_api::region_request::RegionRequest;
|
||||
use store_api::storage::RegionId;
|
||||
use tokio::time::Instant;
|
||||
@@ -155,6 +189,7 @@ mod tests {
|
||||
.handle_downgrade_region_instruction(DowngradeRegion {
|
||||
region_id,
|
||||
flush_timeout,
|
||||
reject_write: false,
|
||||
})
|
||||
.await;
|
||||
assert_matches!(reply, InstructionReply::DowngradeRegion(_));
|
||||
@@ -182,8 +217,9 @@ mod tests {
|
||||
|
||||
Ok(0)
|
||||
}));
|
||||
region_engine.handle_set_readonly_gracefully_mock_fn =
|
||||
Some(Box::new(|_| Ok(SetReadonlyResponse::success(Some(1024)))))
|
||||
region_engine.handle_set_readonly_gracefully_mock_fn = Some(Box::new(|_| {
|
||||
Ok(SetRegionRoleStateResponse::success(Some(1024)))
|
||||
}))
|
||||
});
|
||||
mock_region_server.register_test_region(region_id, mock_engine);
|
||||
let handler_context = HandlerContext::new_for_test(mock_region_server);
|
||||
@@ -195,6 +231,7 @@ mod tests {
|
||||
.handle_downgrade_region_instruction(DowngradeRegion {
|
||||
region_id,
|
||||
flush_timeout,
|
||||
reject_write: false,
|
||||
})
|
||||
.await;
|
||||
assert_matches!(reply, InstructionReply::DowngradeRegion(_));
|
||||
@@ -215,8 +252,9 @@ mod tests {
|
||||
MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
|
||||
region_engine.mock_role = Some(Some(RegionRole::Leader));
|
||||
region_engine.handle_request_delay = Some(Duration::from_secs(100));
|
||||
region_engine.handle_set_readonly_gracefully_mock_fn =
|
||||
Some(Box::new(|_| Ok(SetReadonlyResponse::success(Some(1024)))))
|
||||
region_engine.handle_set_readonly_gracefully_mock_fn = Some(Box::new(|_| {
|
||||
Ok(SetRegionRoleStateResponse::success(Some(1024)))
|
||||
}))
|
||||
});
|
||||
mock_region_server.register_test_region(region_id, mock_engine);
|
||||
let handler_context = HandlerContext::new_for_test(mock_region_server);
|
||||
@@ -227,6 +265,7 @@ mod tests {
|
||||
.handle_downgrade_region_instruction(DowngradeRegion {
|
||||
region_id,
|
||||
flush_timeout: Some(flush_timeout),
|
||||
reject_write: false,
|
||||
})
|
||||
.await;
|
||||
assert_matches!(reply, InstructionReply::DowngradeRegion(_));
|
||||
@@ -246,8 +285,9 @@ mod tests {
|
||||
MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
|
||||
region_engine.mock_role = Some(Some(RegionRole::Leader));
|
||||
region_engine.handle_request_delay = Some(Duration::from_millis(300));
|
||||
region_engine.handle_set_readonly_gracefully_mock_fn =
|
||||
Some(Box::new(|_| Ok(SetReadonlyResponse::success(Some(1024)))))
|
||||
region_engine.handle_set_readonly_gracefully_mock_fn = Some(Box::new(|_| {
|
||||
Ok(SetRegionRoleStateResponse::success(Some(1024)))
|
||||
}))
|
||||
});
|
||||
mock_region_server.register_test_region(region_id, mock_engine);
|
||||
let handler_context = HandlerContext::new_for_test(mock_region_server);
|
||||
@@ -263,6 +303,7 @@ mod tests {
|
||||
.handle_downgrade_region_instruction(DowngradeRegion {
|
||||
region_id,
|
||||
flush_timeout,
|
||||
reject_write: false,
|
||||
})
|
||||
.await;
|
||||
assert_matches!(reply, InstructionReply::DowngradeRegion(_));
|
||||
@@ -277,6 +318,7 @@ mod tests {
|
||||
.handle_downgrade_region_instruction(DowngradeRegion {
|
||||
region_id,
|
||||
flush_timeout: Some(Duration::from_millis(500)),
|
||||
reject_write: false,
|
||||
})
|
||||
.await;
|
||||
assert_matches!(reply, InstructionReply::DowngradeRegion(_));
|
||||
@@ -304,8 +346,9 @@ mod tests {
|
||||
}
|
||||
.fail()
|
||||
}));
|
||||
region_engine.handle_set_readonly_gracefully_mock_fn =
|
||||
Some(Box::new(|_| Ok(SetReadonlyResponse::success(Some(1024)))))
|
||||
region_engine.handle_set_readonly_gracefully_mock_fn = Some(Box::new(|_| {
|
||||
Ok(SetRegionRoleStateResponse::success(Some(1024)))
|
||||
}))
|
||||
});
|
||||
mock_region_server.register_test_region(region_id, mock_engine);
|
||||
let handler_context = HandlerContext::new_for_test(mock_region_server);
|
||||
@@ -321,6 +364,7 @@ mod tests {
|
||||
.handle_downgrade_region_instruction(DowngradeRegion {
|
||||
region_id,
|
||||
flush_timeout,
|
||||
reject_write: false,
|
||||
})
|
||||
.await;
|
||||
assert_matches!(reply, InstructionReply::DowngradeRegion(_));
|
||||
@@ -335,6 +379,7 @@ mod tests {
|
||||
.handle_downgrade_region_instruction(DowngradeRegion {
|
||||
region_id,
|
||||
flush_timeout: Some(Duration::from_millis(500)),
|
||||
reject_write: false,
|
||||
})
|
||||
.await;
|
||||
assert_matches!(reply, InstructionReply::DowngradeRegion(_));
|
||||
@@ -356,7 +401,7 @@ mod tests {
|
||||
MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
|
||||
region_engine.mock_role = Some(Some(RegionRole::Leader));
|
||||
region_engine.handle_set_readonly_gracefully_mock_fn =
|
||||
Some(Box::new(|_| Ok(SetReadonlyResponse::NotFound)));
|
||||
Some(Box::new(|_| Ok(SetRegionRoleStateResponse::NotFound)));
|
||||
});
|
||||
mock_region_server.register_test_region(region_id, mock_engine);
|
||||
let handler_context = HandlerContext::new_for_test(mock_region_server);
|
||||
@@ -365,6 +410,7 @@ mod tests {
|
||||
.handle_downgrade_region_instruction(DowngradeRegion {
|
||||
region_id,
|
||||
flush_timeout: None,
|
||||
reject_write: false,
|
||||
})
|
||||
.await;
|
||||
assert_matches!(reply, InstructionReply::DowngradeRegion(_));
|
||||
@@ -396,6 +442,7 @@ mod tests {
|
||||
.handle_downgrade_region_instruction(DowngradeRegion {
|
||||
region_id,
|
||||
flush_timeout: None,
|
||||
reject_write: false,
|
||||
})
|
||||
.await;
|
||||
assert_matches!(reply, InstructionReply::DowngradeRegion(_));
|
||||
|
||||
@@ -31,7 +31,7 @@ impl HandlerContext {
|
||||
}: UpgradeRegion,
|
||||
) -> BoxFuture<'static, InstructionReply> {
|
||||
Box::pin(async move {
|
||||
let Some(writable) = self.region_server.is_writable(region_id) else {
|
||||
let Some(writable) = self.region_server.is_region_leader(region_id) else {
|
||||
return InstructionReply::UpgradeRegion(UpgradeRegionReply {
|
||||
ready: false,
|
||||
exists: false,
|
||||
|
||||
@@ -32,7 +32,7 @@ use common_recordbatch::SendableRecordBatchStream;
|
||||
use common_runtime::Runtime;
|
||||
use common_telemetry::tracing::{self, info_span};
|
||||
use common_telemetry::tracing_context::{FutureExt, TracingContext};
|
||||
use common_telemetry::{error, info, warn};
|
||||
use common_telemetry::{debug, error, info, warn};
|
||||
use dashmap::DashMap;
|
||||
use datafusion::datasource::{provider_as_source, TableProvider};
|
||||
use datafusion::error::Result as DfResult;
|
||||
@@ -54,7 +54,10 @@ use snafu::{ensure, OptionExt, ResultExt};
|
||||
use store_api::metric_engine_consts::{
|
||||
FILE_ENGINE_NAME, LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME,
|
||||
};
|
||||
use store_api::region_engine::{RegionEngineRef, RegionRole, RegionStatistic, SetReadonlyResponse};
|
||||
use store_api::region_engine::{
|
||||
RegionEngineRef, RegionRole, RegionStatistic, SetRegionRoleStateResponse,
|
||||
SettableRegionRoleState,
|
||||
};
|
||||
use store_api::region_request::{
|
||||
AffectedRows, RegionCloseRequest, RegionOpenRequest, RegionRequest,
|
||||
};
|
||||
@@ -274,37 +277,47 @@ impl RegionServer {
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn is_writable(&self, region_id: RegionId) -> Option<bool> {
|
||||
// TODO(weny): Finds a better way.
|
||||
pub fn is_region_leader(&self, region_id: RegionId) -> Option<bool> {
|
||||
self.inner.region_map.get(®ion_id).and_then(|engine| {
|
||||
engine.role(region_id).map(|role| match role {
|
||||
RegionRole::Follower => false,
|
||||
RegionRole::Leader => true,
|
||||
RegionRole::DowngradingLeader => true,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub fn set_writable(&self, region_id: RegionId, writable: bool) -> Result<()> {
|
||||
pub fn set_region_role(&self, region_id: RegionId, role: RegionRole) -> Result<()> {
|
||||
let engine = self
|
||||
.inner
|
||||
.region_map
|
||||
.get(®ion_id)
|
||||
.with_context(|| RegionNotFoundSnafu { region_id })?;
|
||||
engine
|
||||
.set_writable(region_id, writable)
|
||||
.set_region_role(region_id, role)
|
||||
.with_context(|_| HandleRegionRequestSnafu { region_id })
|
||||
}
|
||||
|
||||
pub async fn set_readonly_gracefully(
|
||||
/// Set region role state gracefully.
|
||||
///
|
||||
/// For [SettableRegionRoleState::Follower]:
|
||||
/// After the call returns, the engine ensures that
|
||||
/// no **further** write or flush operations will succeed in this region.
|
||||
///
|
||||
/// For [SettableRegionRoleState::DowngradingLeader]:
|
||||
/// After the call returns, the engine ensures that
|
||||
/// no **further** write operations will succeed in this region.
|
||||
pub async fn set_region_role_state_gracefully(
|
||||
&self,
|
||||
region_id: RegionId,
|
||||
) -> Result<SetReadonlyResponse> {
|
||||
state: SettableRegionRoleState,
|
||||
) -> Result<SetRegionRoleStateResponse> {
|
||||
match self.inner.region_map.get(®ion_id) {
|
||||
Some(engine) => Ok(engine
|
||||
.set_readonly_gracefully(region_id)
|
||||
.set_region_role_state_gracefully(region_id, state)
|
||||
.await
|
||||
.with_context(|_| HandleRegionRequestSnafu { region_id })?),
|
||||
None => Ok(SetReadonlyResponse::NotFound),
|
||||
None => Ok(SetRegionRoleStateResponse::NotFound),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -842,7 +855,7 @@ impl RegionServerInner {
|
||||
info!("Region {region_id} is deregistered from engine {engine_type}");
|
||||
self.region_map
|
||||
.remove(®ion_id)
|
||||
.map(|(id, engine)| engine.set_writable(id, false));
|
||||
.map(|(id, engine)| engine.set_region_role(id, RegionRole::Follower));
|
||||
self.event_listener.on_region_deregistered(region_id);
|
||||
}
|
||||
RegionChange::Catchup => {
|
||||
@@ -880,7 +893,7 @@ impl RegionServerInner {
|
||||
for region in logical_regions {
|
||||
self.region_map
|
||||
.insert(region, RegionEngineWithStatus::Ready(engine.clone()));
|
||||
info!("Logical region {} is registered!", region);
|
||||
debug!("Logical region {} is registered!", region);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -922,17 +935,19 @@ impl RegionServerInner {
|
||||
.iter()
|
||||
.map(|x| (*x.key(), x.value().clone()))
|
||||
.collect::<Vec<_>>();
|
||||
let num_regions = regions.len();
|
||||
|
||||
for (region_id, engine) in regions {
|
||||
let closed = engine
|
||||
.handle_request(region_id, RegionRequest::Close(RegionCloseRequest {}))
|
||||
.await;
|
||||
match closed {
|
||||
Ok(_) => info!("Region {region_id} is closed"),
|
||||
Ok(_) => debug!("Region {region_id} is closed"),
|
||||
Err(e) => warn!("Failed to close region {region_id}, err: {e}"),
|
||||
}
|
||||
}
|
||||
self.region_map.clear();
|
||||
info!("closed {num_regions} regions");
|
||||
|
||||
let engines = self.engines.write().unwrap().drain().collect::<Vec<_>>();
|
||||
for (engine_name, engine) in engines {
|
||||
|
||||
@@ -32,7 +32,8 @@ use query::{QueryEngine, QueryEngineContext};
|
||||
use session::context::QueryContextRef;
|
||||
use store_api::metadata::RegionMetadataRef;
|
||||
use store_api::region_engine::{
|
||||
RegionEngine, RegionRole, RegionScannerRef, RegionStatistic, SetReadonlyResponse,
|
||||
RegionEngine, RegionRole, RegionScannerRef, RegionStatistic, SetRegionRoleStateResponse,
|
||||
SettableRegionRoleState,
|
||||
};
|
||||
use store_api::region_request::{AffectedRows, RegionRequest};
|
||||
use store_api::storage::{RegionId, ScanRequest};
|
||||
@@ -106,7 +107,7 @@ pub type MockRequestHandler =
|
||||
Box<dyn Fn(RegionId, RegionRequest) -> Result<AffectedRows, Error> + Send + Sync>;
|
||||
|
||||
pub type MockSetReadonlyGracefullyHandler =
|
||||
Box<dyn Fn(RegionId) -> Result<SetReadonlyResponse, Error> + Send + Sync>;
|
||||
Box<dyn Fn(RegionId) -> Result<SetRegionRoleStateResponse, Error> + Send + Sync>;
|
||||
|
||||
pub struct MockRegionEngine {
|
||||
sender: Sender<(RegionId, RegionRequest)>,
|
||||
@@ -220,14 +221,15 @@ impl RegionEngine for MockRegionEngine {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn set_writable(&self, _region_id: RegionId, _writable: bool) -> Result<(), BoxedError> {
|
||||
fn set_region_role(&self, _region_id: RegionId, _role: RegionRole) -> Result<(), BoxedError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn set_readonly_gracefully(
|
||||
async fn set_region_role_state_gracefully(
|
||||
&self,
|
||||
region_id: RegionId,
|
||||
) -> Result<SetReadonlyResponse, BoxedError> {
|
||||
_region_role_state: SettableRegionRoleState,
|
||||
) -> Result<SetRegionRoleStateResponse, BoxedError> {
|
||||
if let Some(mock_fn) = &self.handle_set_readonly_gracefully_mock_fn {
|
||||
return mock_fn(region_id).map_err(BoxedError::new);
|
||||
};
|
||||
|
||||
@@ -249,6 +249,15 @@ impl ConcreteDataType {
|
||||
]
|
||||
}
|
||||
|
||||
pub fn timestamps() -> Vec<ConcreteDataType> {
|
||||
vec![
|
||||
ConcreteDataType::timestamp_second_datatype(),
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
ConcreteDataType::timestamp_microsecond_datatype(),
|
||||
ConcreteDataType::timestamp_nanosecond_datatype(),
|
||||
]
|
||||
}
|
||||
|
||||
/// Convert arrow data type to [ConcreteDataType].
|
||||
///
|
||||
/// # Panics
|
||||
|
||||
@@ -12,11 +12,10 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use common_time::interval::Interval;
|
||||
use common_time::{IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth};
|
||||
use paste::paste;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::prelude::{Scalar, Value, ValueRef};
|
||||
use crate::prelude::Scalar;
|
||||
use crate::scalars::ScalarRef;
|
||||
use crate::types::{
|
||||
IntervalDayTimeType, IntervalMonthDayNanoType, IntervalYearMonthType, WrapperType,
|
||||
@@ -26,39 +25,6 @@ use crate::vectors::{IntervalDayTimeVector, IntervalMonthDayNanoVector, Interval
|
||||
macro_rules! define_interval_with_unit {
|
||||
($unit: ident, $native_ty: ty) => {
|
||||
paste! {
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct [<Interval $unit>](pub Interval);
|
||||
|
||||
impl [<Interval $unit>] {
|
||||
pub fn new(val: $native_ty) -> Self {
|
||||
Self(Interval:: [<from_ $native_ty>](val))
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for [<Interval $unit>] {
|
||||
fn default() -> Self {
|
||||
Self::new(0)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<[<Interval $unit>]> for Value {
|
||||
fn from(t: [<Interval $unit>]) -> Value {
|
||||
Value::Interval(t.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<[<Interval $unit>]> for serde_json::Value {
|
||||
fn from(t: [<Interval $unit>]) -> Self {
|
||||
t.0.into()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<[<Interval $unit>]> for ValueRef<'static> {
|
||||
fn from(t: [<Interval $unit>]) -> Self {
|
||||
ValueRef::Interval(t.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl Scalar for [<Interval $unit>] {
|
||||
type VectorType = [<Interval $unit Vector>];
|
||||
type RefType<'a> = [<Interval $unit>];
|
||||
@@ -87,41 +53,11 @@ macro_rules! define_interval_with_unit {
|
||||
type Native = $native_ty;
|
||||
|
||||
fn from_native(value: Self::Native) -> Self {
|
||||
Self::new(value)
|
||||
Self::[<from_ $native_ty>](value)
|
||||
}
|
||||
|
||||
fn into_native(self) -> Self::Native {
|
||||
self.0.[<to_ $native_ty>]()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<$native_ty> for [<Interval $unit>] {
|
||||
fn from(val: $native_ty) -> Self {
|
||||
[<Interval $unit>]::from_native(val as $native_ty)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<[<Interval $unit>]> for $native_ty {
|
||||
fn from(val: [<Interval $unit>]) -> Self {
|
||||
val.0.[<to_ $native_ty>]()
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<Value> for Option<[<Interval $unit>]> {
|
||||
type Error = $crate::error::Error;
|
||||
|
||||
#[inline]
|
||||
fn try_from(from: Value) -> std::result::Result<Self, Self::Error> {
|
||||
match from {
|
||||
Value::Interval(v) if v.unit() == common_time::interval::IntervalUnit::$unit => {
|
||||
Ok(Some([<Interval $unit>](v)))
|
||||
},
|
||||
Value::Null => Ok(None),
|
||||
_ => $crate::error::TryFromValueSnafu {
|
||||
reason: format!("{:?} is not a {}", from, stringify!([<Interval $unit>])),
|
||||
}
|
||||
.fail(),
|
||||
}
|
||||
self.[<to_ $native_ty>]()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -138,17 +74,17 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_interval_scalar() {
|
||||
let interval = IntervalYearMonth::new(1000);
|
||||
let interval = IntervalYearMonth::from(1000);
|
||||
assert_eq!(interval, interval.as_scalar_ref());
|
||||
assert_eq!(interval, interval.to_owned_scalar());
|
||||
assert_eq!(1000, interval.into_native());
|
||||
|
||||
let interval = IntervalDayTime::new(1000);
|
||||
let interval = IntervalDayTime::from(1000);
|
||||
assert_eq!(interval, interval.as_scalar_ref());
|
||||
assert_eq!(interval, interval.to_owned_scalar());
|
||||
assert_eq!(1000, interval.into_native());
|
||||
|
||||
let interval = IntervalMonthDayNano::new(1000);
|
||||
let interval = IntervalMonthDayNano::from(1000);
|
||||
assert_eq!(interval, interval.as_scalar_ref());
|
||||
assert_eq!(interval, interval.to_owned_scalar());
|
||||
assert_eq!(1000, interval.into_native());
|
||||
@@ -156,15 +92,15 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_interval_convert_to_native_type() {
|
||||
let interval = IntervalMonthDayNano::new(1000);
|
||||
let interval = IntervalMonthDayNano::from(1000);
|
||||
let native_value: i128 = interval.into();
|
||||
assert_eq!(native_value, 1000);
|
||||
|
||||
let interval = IntervalDayTime::new(1000);
|
||||
let interval = IntervalDayTime::from(1000);
|
||||
let native_interval: i64 = interval.into();
|
||||
assert_eq!(native_interval, 1000);
|
||||
|
||||
let interval = IntervalYearMonth::new(1000);
|
||||
let interval = IntervalYearMonth::from(1000);
|
||||
let native_interval: i32 = interval.into();
|
||||
assert_eq!(native_interval, 1000);
|
||||
}
|
||||
|
||||
@@ -17,8 +17,9 @@ use arrow::datatypes::{
|
||||
IntervalMonthDayNanoType as ArrowIntervalMonthDayNanoType, IntervalUnit as ArrowIntervalUnit,
|
||||
IntervalYearMonthType as ArrowIntervalYearMonthType,
|
||||
};
|
||||
use common_time::interval::IntervalUnit;
|
||||
use common_time::Interval;
|
||||
use common_time::interval::{
|
||||
IntervalDayTime, IntervalMonthDayNano, IntervalUnit, IntervalYearMonth,
|
||||
};
|
||||
use enum_dispatch::enum_dispatch;
|
||||
use paste::paste;
|
||||
use serde::{Deserialize, Serialize};
|
||||
@@ -26,7 +27,6 @@ use snafu::OptionExt;
|
||||
|
||||
use crate::data_type::ConcreteDataType;
|
||||
use crate::error;
|
||||
use crate::interval::{IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth};
|
||||
use crate::prelude::{
|
||||
DataType, LogicalTypeId, MutableVector, ScalarVectorBuilder, Value, ValueRef, Vector,
|
||||
};
|
||||
@@ -75,7 +75,7 @@ macro_rules! impl_data_type_for_interval {
|
||||
}
|
||||
|
||||
fn default_value(&self) -> Value {
|
||||
Value::Interval(Interval::from_i128(0))
|
||||
Value::[<Interval $unit>]([<Interval $unit>]::default())
|
||||
}
|
||||
|
||||
fn as_arrow_type(&self) -> ArrowDataType {
|
||||
@@ -124,7 +124,7 @@ macro_rules! impl_data_type_for_interval {
|
||||
fn cast_value_ref(value: ValueRef) -> crate::Result<Option<Self::Wrapper>> {
|
||||
match value {
|
||||
ValueRef::Null => Ok(None),
|
||||
ValueRef::Interval(t) => Ok(Some([<Interval $unit>](t))),
|
||||
ValueRef::[<Interval $unit>](t) => Ok(Some(t)),
|
||||
other => error::CastTypeSnafu {
|
||||
msg: format!("Failed to cast value {:?} to {}", other, stringify!([<Interval $unit>])),
|
||||
}
|
||||
|
||||
@@ -16,7 +16,6 @@ use std::cmp::Ordering;
|
||||
use std::fmt;
|
||||
|
||||
use arrow::datatypes::{ArrowNativeType, ArrowPrimitiveType, DataType as ArrowDataType};
|
||||
use common_time::interval::IntervalUnit;
|
||||
use common_time::{Date, DateTime};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use snafu::OptionExt;
|
||||
@@ -30,6 +29,7 @@ use crate::types::{DateTimeType, DateType};
|
||||
use crate::value::{Value, ValueRef};
|
||||
use crate::vectors::{MutableVector, PrimitiveVector, PrimitiveVectorBuilder, Vector};
|
||||
|
||||
// TODO(yingwen): Can we remove `Into<serde_json::Value>`?
|
||||
/// Represents the wrapper type that wraps a native type using the `newtype pattern`,
|
||||
/// such as [Date](`common_time::Date`) is a wrapper type for the underlying native
|
||||
/// type `i32`.
|
||||
@@ -364,11 +364,7 @@ impl DataType for Int64Type {
|
||||
Value::DateTime(v) => Some(Value::Int64(v.val())),
|
||||
Value::Timestamp(v) => Some(Value::Int64(v.value())),
|
||||
Value::Time(v) => Some(Value::Int64(v.value())),
|
||||
Value::Interval(v) => match v.unit() {
|
||||
IntervalUnit::DayTime => Some(Value::Int64(v.to_i64())),
|
||||
IntervalUnit::YearMonth => None,
|
||||
IntervalUnit::MonthDayNano => None,
|
||||
},
|
||||
// We don't allow casting interval type to int.
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
@@ -410,11 +406,7 @@ impl DataType for Int32Type {
|
||||
Value::Float64(v) => num::cast::cast(v).map(Value::Int32),
|
||||
Value::String(v) => v.as_utf8().parse::<i32>().map(Value::Int32).ok(),
|
||||
Value::Date(v) => Some(Value::Int32(v.val())),
|
||||
Value::Interval(v) => match v.unit() {
|
||||
IntervalUnit::YearMonth => Some(Value::Int32(v.to_i32())),
|
||||
IntervalUnit::DayTime => None,
|
||||
IntervalUnit::MonthDayNano => None,
|
||||
},
|
||||
// We don't allow casting interval type to int.
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -78,7 +78,15 @@ impl DataType for StringType {
|
||||
Value::DateTime(v) => Some(Value::String(StringBytes::from(v.to_string()))),
|
||||
Value::Timestamp(v) => Some(Value::String(StringBytes::from(v.to_iso8601_string()))),
|
||||
Value::Time(v) => Some(Value::String(StringBytes::from(v.to_iso8601_string()))),
|
||||
Value::Interval(v) => Some(Value::String(StringBytes::from(v.to_iso8601_string()))),
|
||||
Value::IntervalYearMonth(v) => {
|
||||
Some(Value::String(StringBytes::from(v.to_iso8601_string())))
|
||||
}
|
||||
Value::IntervalDayTime(v) => {
|
||||
Some(Value::String(StringBytes::from(v.to_iso8601_string())))
|
||||
}
|
||||
Value::IntervalMonthDayNano(v) => {
|
||||
Some(Value::String(StringBytes::from(v.to_iso8601_string())))
|
||||
}
|
||||
Value::Duration(v) => Some(Value::String(StringBytes::from(v.to_string()))),
|
||||
Value::Decimal128(v) => Some(Value::String(StringBytes::from(v.to_string()))),
|
||||
|
||||
|
||||
@@ -28,7 +28,7 @@ use common_time::datetime::DateTime;
|
||||
use common_time::interval::IntervalUnit;
|
||||
use common_time::time::Time;
|
||||
use common_time::timestamp::{TimeUnit, Timestamp};
|
||||
use common_time::{Duration, Interval, Timezone};
|
||||
use common_time::{Duration, IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth, Timezone};
|
||||
use datafusion_common::ScalarValue;
|
||||
use greptime_proto::v1::value::ValueData;
|
||||
pub use ordered_float::OrderedFloat;
|
||||
@@ -38,6 +38,7 @@ use snafu::{ensure, ResultExt};
|
||||
|
||||
use crate::error::{self, ConvertArrowArrayToScalarsSnafu, Error, Result, TryFromValueSnafu};
|
||||
use crate::prelude::*;
|
||||
use crate::schema::ColumnSchema;
|
||||
use crate::type_id::LogicalTypeId;
|
||||
use crate::types::{IntervalType, ListType};
|
||||
use crate::vectors::ListVector;
|
||||
@@ -78,7 +79,10 @@ pub enum Value {
|
||||
Timestamp(Timestamp),
|
||||
Time(Time),
|
||||
Duration(Duration),
|
||||
Interval(Interval),
|
||||
// Interval types:
|
||||
IntervalYearMonth(IntervalYearMonth),
|
||||
IntervalDayTime(IntervalDayTime),
|
||||
IntervalMonthDayNano(IntervalMonthDayNano),
|
||||
|
||||
List(ListValue),
|
||||
}
|
||||
@@ -111,7 +115,15 @@ impl Display for Value {
|
||||
Value::DateTime(v) => write!(f, "{v}"),
|
||||
Value::Timestamp(v) => write!(f, "{}", v.to_iso8601_string()),
|
||||
Value::Time(t) => write!(f, "{}", t.to_iso8601_string()),
|
||||
Value::Interval(v) => write!(f, "{}", v.to_iso8601_string()),
|
||||
Value::IntervalYearMonth(v) => {
|
||||
write!(f, "{}", v.to_iso8601_string())
|
||||
}
|
||||
Value::IntervalDayTime(v) => {
|
||||
write!(f, "{}", v.to_iso8601_string())
|
||||
}
|
||||
Value::IntervalMonthDayNano(v) => {
|
||||
write!(f, "{}", v.to_iso8601_string())
|
||||
}
|
||||
Value::Duration(d) => write!(f, "{d}"),
|
||||
Value::List(v) => {
|
||||
let items = v
|
||||
@@ -153,7 +165,15 @@ macro_rules! define_data_type_func {
|
||||
$struct::DateTime(_) => ConcreteDataType::datetime_datatype(),
|
||||
$struct::Time(t) => ConcreteDataType::time_datatype(*t.unit()),
|
||||
$struct::Timestamp(v) => ConcreteDataType::timestamp_datatype(v.unit()),
|
||||
$struct::Interval(v) => ConcreteDataType::interval_datatype(v.unit()),
|
||||
$struct::IntervalYearMonth(_) => {
|
||||
ConcreteDataType::interval_datatype(IntervalUnit::YearMonth)
|
||||
}
|
||||
$struct::IntervalDayTime(_) => {
|
||||
ConcreteDataType::interval_datatype(IntervalUnit::DayTime)
|
||||
}
|
||||
$struct::IntervalMonthDayNano(_) => {
|
||||
ConcreteDataType::interval_datatype(IntervalUnit::MonthDayNano)
|
||||
}
|
||||
$struct::List(list) => ConcreteDataType::list_datatype(list.datatype().clone()),
|
||||
$struct::Duration(d) => ConcreteDataType::duration_datatype(d.unit()),
|
||||
$struct::Decimal128(d) => {
|
||||
@@ -206,7 +226,9 @@ impl Value {
|
||||
Value::List(v) => ValueRef::List(ListValueRef::Ref { val: v }),
|
||||
Value::Timestamp(v) => ValueRef::Timestamp(*v),
|
||||
Value::Time(v) => ValueRef::Time(*v),
|
||||
Value::Interval(v) => ValueRef::Interval(*v),
|
||||
Value::IntervalYearMonth(v) => ValueRef::IntervalYearMonth(*v),
|
||||
Value::IntervalDayTime(v) => ValueRef::IntervalDayTime(*v),
|
||||
Value::IntervalMonthDayNano(v) => ValueRef::IntervalMonthDayNano(*v),
|
||||
Value::Duration(v) => ValueRef::Duration(*v),
|
||||
Value::Decimal128(v) => ValueRef::Decimal128(*v),
|
||||
}
|
||||
@@ -220,14 +242,6 @@ impl Value {
|
||||
}
|
||||
}
|
||||
|
||||
/// Cast Value to Interval. Return None if value is not a valid interval data type.
|
||||
pub fn as_interval(&self) -> Option<Interval> {
|
||||
match self {
|
||||
Value::Interval(i) => Some(*i),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Cast Value to utf8 String. Return None if value is not a valid string data type.
|
||||
pub fn as_string(&self) -> Option<String> {
|
||||
match self {
|
||||
@@ -255,12 +269,35 @@ impl Value {
|
||||
/// Cast Value to [Time]. Return None if value is not a valid time data type.
|
||||
pub fn as_time(&self) -> Option<Time> {
|
||||
match self {
|
||||
Value::Int64(v) => Some(Time::new_millisecond(*v)),
|
||||
Value::Time(t) => Some(*t),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Cast Value to [IntervalYearMonth]. Return None if value is not a valid interval year month data type.
|
||||
pub fn as_interval_year_month(&self) -> Option<IntervalYearMonth> {
|
||||
match self {
|
||||
Value::IntervalYearMonth(v) => Some(*v),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Cast Value to [IntervalDayTime]. Return None if value is not a valid interval day time data type.
|
||||
pub fn as_interval_day_time(&self) -> Option<IntervalDayTime> {
|
||||
match self {
|
||||
Value::IntervalDayTime(v) => Some(*v),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Cast Value to [IntervalMonthDayNano]. Return None if value is not a valid interval month day nano data type.
|
||||
pub fn as_interval_month_day_nano(&self) -> Option<IntervalMonthDayNano> {
|
||||
match self {
|
||||
Value::IntervalMonthDayNano(v) => Some(*v),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Cast Value to u64. Return None if value is not a valid uint64 data type.
|
||||
pub fn as_u64(&self) -> Option<u64> {
|
||||
match self {
|
||||
@@ -321,11 +358,9 @@ impl Value {
|
||||
TimeUnit::Microsecond => LogicalTypeId::TimeMicrosecond,
|
||||
TimeUnit::Nanosecond => LogicalTypeId::TimeNanosecond,
|
||||
},
|
||||
Value::Interval(v) => match v.unit() {
|
||||
IntervalUnit::YearMonth => LogicalTypeId::IntervalYearMonth,
|
||||
IntervalUnit::DayTime => LogicalTypeId::IntervalDayTime,
|
||||
IntervalUnit::MonthDayNano => LogicalTypeId::IntervalMonthDayNano,
|
||||
},
|
||||
Value::IntervalYearMonth(_) => LogicalTypeId::IntervalYearMonth,
|
||||
Value::IntervalDayTime(_) => LogicalTypeId::IntervalDayTime,
|
||||
Value::IntervalMonthDayNano(_) => LogicalTypeId::IntervalMonthDayNano,
|
||||
Value::Duration(d) => match d.unit() {
|
||||
TimeUnit::Second => LogicalTypeId::DurationSecond,
|
||||
TimeUnit::Millisecond => LogicalTypeId::DurationMillisecond,
|
||||
@@ -375,11 +410,9 @@ impl Value {
|
||||
}
|
||||
Value::Timestamp(t) => timestamp_to_scalar_value(t.unit(), Some(t.value())),
|
||||
Value::Time(t) => time_to_scalar_value(*t.unit(), Some(t.value()))?,
|
||||
Value::Interval(v) => match v.unit() {
|
||||
IntervalUnit::YearMonth => ScalarValue::IntervalYearMonth(Some(v.to_i32())),
|
||||
IntervalUnit::DayTime => ScalarValue::IntervalDayTime(Some(v.to_i64())),
|
||||
IntervalUnit::MonthDayNano => ScalarValue::IntervalMonthDayNano(Some(v.to_i128())),
|
||||
},
|
||||
Value::IntervalYearMonth(v) => ScalarValue::IntervalYearMonth(Some(v.to_i32())),
|
||||
Value::IntervalDayTime(v) => ScalarValue::IntervalDayTime(Some(v.to_i64())),
|
||||
Value::IntervalMonthDayNano(v) => ScalarValue::IntervalMonthDayNano(Some(v.to_i128())),
|
||||
Value::Duration(d) => duration_to_scalar_value(d.unit(), Some(d.value())),
|
||||
Value::Decimal128(d) => {
|
||||
let (v, p, s) = d.to_scalar_value();
|
||||
@@ -434,7 +467,9 @@ impl Value {
|
||||
Value::Timestamp(x) => Some(Value::Timestamp(x.negative())),
|
||||
Value::Time(x) => Some(Value::Time(x.negative())),
|
||||
Value::Duration(x) => Some(Value::Duration(x.negative())),
|
||||
Value::Interval(x) => Some(Value::Interval(x.negative())),
|
||||
Value::IntervalYearMonth(x) => Some(Value::IntervalYearMonth(x.negative())),
|
||||
Value::IntervalDayTime(x) => Some(Value::IntervalDayTime(x.negative())),
|
||||
Value::IntervalMonthDayNano(x) => Some(Value::IntervalMonthDayNano(x.negative())),
|
||||
|
||||
Value::Binary(_) | Value::String(_) | Value::Boolean(_) | Value::List(_) => None,
|
||||
}
|
||||
@@ -571,16 +606,6 @@ pub fn scalar_value_to_timestamp(
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert [ScalarValue] to [Interval].
|
||||
pub fn scalar_value_to_interval(scalar: &ScalarValue) -> Option<Interval> {
|
||||
match scalar {
|
||||
ScalarValue::IntervalYearMonth(v) => v.map(Interval::from_i32),
|
||||
ScalarValue::IntervalDayTime(v) => v.map(Interval::from_i64),
|
||||
ScalarValue::IntervalMonthDayNano(v) => v.map(Interval::from_i128),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
macro_rules! impl_ord_for_value_like {
|
||||
($Type: ident, $left: ident, $right: ident) => {
|
||||
if $left.is_null() && !$right.is_null() {
|
||||
@@ -607,7 +632,9 @@ macro_rules! impl_ord_for_value_like {
|
||||
($Type::DateTime(v1), $Type::DateTime(v2)) => v1.cmp(v2),
|
||||
($Type::Timestamp(v1), $Type::Timestamp(v2)) => v1.cmp(v2),
|
||||
($Type::Time(v1), $Type::Time(v2)) => v1.cmp(v2),
|
||||
($Type::Interval(v1), $Type::Interval(v2)) => v1.cmp(v2),
|
||||
($Type::IntervalYearMonth(v1), $Type::IntervalYearMonth(v2)) => v1.cmp(v2),
|
||||
($Type::IntervalDayTime(v1), $Type::IntervalDayTime(v2)) => v1.cmp(v2),
|
||||
($Type::IntervalMonthDayNano(v1), $Type::IntervalMonthDayNano(v2)) => v1.cmp(v2),
|
||||
($Type::Duration(v1), $Type::Duration(v2)) => v1.cmp(v2),
|
||||
($Type::List(v1), $Type::List(v2)) => v1.cmp(v2),
|
||||
_ => panic!(
|
||||
@@ -685,7 +712,9 @@ impl_try_from_value!(Date, Date);
|
||||
impl_try_from_value!(Time, Time);
|
||||
impl_try_from_value!(DateTime, DateTime);
|
||||
impl_try_from_value!(Timestamp, Timestamp);
|
||||
impl_try_from_value!(Interval, Interval);
|
||||
impl_try_from_value!(IntervalYearMonth, IntervalYearMonth);
|
||||
impl_try_from_value!(IntervalDayTime, IntervalDayTime);
|
||||
impl_try_from_value!(IntervalMonthDayNano, IntervalMonthDayNano);
|
||||
impl_try_from_value!(Duration, Duration);
|
||||
impl_try_from_value!(Decimal128, Decimal128);
|
||||
|
||||
@@ -727,7 +756,9 @@ impl_value_from!(Date, Date);
|
||||
impl_value_from!(Time, Time);
|
||||
impl_value_from!(DateTime, DateTime);
|
||||
impl_value_from!(Timestamp, Timestamp);
|
||||
impl_value_from!(Interval, Interval);
|
||||
impl_value_from!(IntervalYearMonth, IntervalYearMonth);
|
||||
impl_value_from!(IntervalDayTime, IntervalDayTime);
|
||||
impl_value_from!(IntervalMonthDayNano, IntervalMonthDayNano);
|
||||
impl_value_from!(Duration, Duration);
|
||||
impl_value_from!(String, String);
|
||||
impl_value_from!(Decimal128, Decimal128);
|
||||
@@ -774,7 +805,9 @@ impl TryFrom<Value> for serde_json::Value {
|
||||
Value::List(v) => serde_json::to_value(v)?,
|
||||
Value::Timestamp(v) => serde_json::to_value(v.value())?,
|
||||
Value::Time(v) => serde_json::to_value(v.value())?,
|
||||
Value::Interval(v) => serde_json::to_value(v.to_i128())?,
|
||||
Value::IntervalYearMonth(v) => serde_json::to_value(v.to_i32())?,
|
||||
Value::IntervalDayTime(v) => serde_json::to_value(v.to_i64())?,
|
||||
Value::IntervalMonthDayNano(v) => serde_json::to_value(v.to_i128())?,
|
||||
Value::Duration(v) => serde_json::to_value(v.value())?,
|
||||
Value::Decimal128(v) => serde_json::to_value(v.to_string())?,
|
||||
};
|
||||
@@ -926,13 +959,13 @@ impl TryFrom<ScalarValue> for Value {
|
||||
.unwrap_or(Value::Null),
|
||||
|
||||
ScalarValue::IntervalYearMonth(t) => t
|
||||
.map(|x| Value::Interval(Interval::from_i32(x)))
|
||||
.map(|x| Value::IntervalYearMonth(IntervalYearMonth::from_i32(x)))
|
||||
.unwrap_or(Value::Null),
|
||||
ScalarValue::IntervalDayTime(t) => t
|
||||
.map(|x| Value::Interval(Interval::from_i64(x)))
|
||||
.map(|x| Value::IntervalDayTime(IntervalDayTime::from_i64(x)))
|
||||
.unwrap_or(Value::Null),
|
||||
ScalarValue::IntervalMonthDayNano(t) => t
|
||||
.map(|x| Value::Interval(Interval::from_i128(x)))
|
||||
.map(|x| Value::IntervalMonthDayNano(IntervalMonthDayNano::from_i128(x)))
|
||||
.unwrap_or(Value::Null),
|
||||
ScalarValue::DurationSecond(d) => d
|
||||
.map(|x| Value::Duration(Duration::new(x, TimeUnit::Second)))
|
||||
@@ -987,7 +1020,9 @@ impl From<ValueRef<'_>> for Value {
|
||||
ValueRef::DateTime(v) => Value::DateTime(v),
|
||||
ValueRef::Timestamp(v) => Value::Timestamp(v),
|
||||
ValueRef::Time(v) => Value::Time(v),
|
||||
ValueRef::Interval(v) => Value::Interval(v),
|
||||
ValueRef::IntervalYearMonth(v) => Value::IntervalYearMonth(v),
|
||||
ValueRef::IntervalDayTime(v) => Value::IntervalDayTime(v),
|
||||
ValueRef::IntervalMonthDayNano(v) => Value::IntervalMonthDayNano(v),
|
||||
ValueRef::Duration(v) => Value::Duration(v),
|
||||
ValueRef::List(v) => v.to_value(),
|
||||
ValueRef::Decimal128(v) => Value::Decimal128(v),
|
||||
@@ -1026,7 +1061,10 @@ pub enum ValueRef<'a> {
|
||||
Timestamp(Timestamp),
|
||||
Time(Time),
|
||||
Duration(Duration),
|
||||
Interval(Interval),
|
||||
// Interval types:
|
||||
IntervalYearMonth(IntervalYearMonth),
|
||||
IntervalDayTime(IntervalDayTime),
|
||||
IntervalMonthDayNano(IntervalMonthDayNano),
|
||||
|
||||
// Compound types:
|
||||
List(ListValueRef<'a>),
|
||||
@@ -1150,9 +1188,19 @@ impl<'a> ValueRef<'a> {
|
||||
impl_as_for_value_ref!(self, Duration)
|
||||
}
|
||||
|
||||
/// Cast itself to [Interval].
|
||||
pub fn as_interval(&self) -> Result<Option<Interval>> {
|
||||
impl_as_for_value_ref!(self, Interval)
|
||||
/// Cast itself to [IntervalYearMonth].
|
||||
pub fn as_interval_year_month(&self) -> Result<Option<IntervalYearMonth>> {
|
||||
impl_as_for_value_ref!(self, IntervalYearMonth)
|
||||
}
|
||||
|
||||
/// Cast itself to [IntervalDayTime].
|
||||
pub fn as_interval_day_time(&self) -> Result<Option<IntervalDayTime>> {
|
||||
impl_as_for_value_ref!(self, IntervalDayTime)
|
||||
}
|
||||
|
||||
/// Cast itself to [IntervalMonthDayNano].
|
||||
pub fn as_interval_month_day_nano(&self) -> Result<Option<IntervalMonthDayNano>> {
|
||||
impl_as_for_value_ref!(self, IntervalMonthDayNano)
|
||||
}
|
||||
|
||||
/// Cast itself to [ListValueRef].
|
||||
@@ -1212,7 +1260,9 @@ impl_value_ref_from!(Date, Date);
|
||||
impl_value_ref_from!(DateTime, DateTime);
|
||||
impl_value_ref_from!(Timestamp, Timestamp);
|
||||
impl_value_ref_from!(Time, Time);
|
||||
impl_value_ref_from!(Interval, Interval);
|
||||
impl_value_ref_from!(IntervalYearMonth, IntervalYearMonth);
|
||||
impl_value_ref_from!(IntervalDayTime, IntervalDayTime);
|
||||
impl_value_ref_from!(IntervalMonthDayNano, IntervalMonthDayNano);
|
||||
impl_value_ref_from!(Duration, Duration);
|
||||
impl_value_ref_from!(Decimal128, Decimal128);
|
||||
|
||||
@@ -1237,37 +1287,52 @@ impl<'a> From<Option<ListValueRef<'a>>> for ValueRef<'a> {
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> TryFrom<ValueRef<'a>> for serde_json::Value {
|
||||
type Error = serde_json::Error;
|
||||
/// transform a [ValueRef] to a [serde_json::Value].
|
||||
/// The json type will be handled specially
|
||||
pub fn transform_value_ref_to_json_value<'a>(
|
||||
value: ValueRef<'a>,
|
||||
schema: &'a ColumnSchema,
|
||||
) -> serde_json::Result<serde_json::Value> {
|
||||
let json_value = match value {
|
||||
ValueRef::Null => serde_json::Value::Null,
|
||||
ValueRef::Boolean(v) => serde_json::Value::Bool(v),
|
||||
ValueRef::UInt8(v) => serde_json::Value::from(v),
|
||||
ValueRef::UInt16(v) => serde_json::Value::from(v),
|
||||
ValueRef::UInt32(v) => serde_json::Value::from(v),
|
||||
ValueRef::UInt64(v) => serde_json::Value::from(v),
|
||||
ValueRef::Int8(v) => serde_json::Value::from(v),
|
||||
ValueRef::Int16(v) => serde_json::Value::from(v),
|
||||
ValueRef::Int32(v) => serde_json::Value::from(v),
|
||||
ValueRef::Int64(v) => serde_json::Value::from(v),
|
||||
ValueRef::Float32(v) => serde_json::Value::from(v.0),
|
||||
ValueRef::Float64(v) => serde_json::Value::from(v.0),
|
||||
ValueRef::String(bytes) => serde_json::Value::String(bytes.to_string()),
|
||||
ValueRef::Binary(bytes) => {
|
||||
if let ConcreteDataType::Json(_) = schema.data_type {
|
||||
match jsonb::from_slice(bytes) {
|
||||
Ok(json) => json.into(),
|
||||
Err(e) => {
|
||||
error!(e; "Failed to parse jsonb");
|
||||
serde_json::Value::Null
|
||||
}
|
||||
}
|
||||
} else {
|
||||
serde_json::to_value(bytes)?
|
||||
}
|
||||
}
|
||||
ValueRef::Date(v) => serde_json::Value::Number(v.val().into()),
|
||||
ValueRef::DateTime(v) => serde_json::Value::Number(v.val().into()),
|
||||
ValueRef::List(v) => serde_json::to_value(v)?,
|
||||
ValueRef::Timestamp(v) => serde_json::to_value(v.value())?,
|
||||
ValueRef::Time(v) => serde_json::to_value(v.value())?,
|
||||
ValueRef::IntervalYearMonth(v) => serde_json::Value::from(v),
|
||||
ValueRef::IntervalDayTime(v) => serde_json::Value::from(v),
|
||||
ValueRef::IntervalMonthDayNano(v) => serde_json::Value::from(v),
|
||||
ValueRef::Duration(v) => serde_json::to_value(v.value())?,
|
||||
ValueRef::Decimal128(v) => serde_json::to_value(v.to_string())?,
|
||||
};
|
||||
|
||||
fn try_from(value: ValueRef<'a>) -> serde_json::Result<serde_json::Value> {
|
||||
let json_value = match value {
|
||||
ValueRef::Null => serde_json::Value::Null,
|
||||
ValueRef::Boolean(v) => serde_json::Value::Bool(v),
|
||||
ValueRef::UInt8(v) => serde_json::Value::from(v),
|
||||
ValueRef::UInt16(v) => serde_json::Value::from(v),
|
||||
ValueRef::UInt32(v) => serde_json::Value::from(v),
|
||||
ValueRef::UInt64(v) => serde_json::Value::from(v),
|
||||
ValueRef::Int8(v) => serde_json::Value::from(v),
|
||||
ValueRef::Int16(v) => serde_json::Value::from(v),
|
||||
ValueRef::Int32(v) => serde_json::Value::from(v),
|
||||
ValueRef::Int64(v) => serde_json::Value::from(v),
|
||||
ValueRef::Float32(v) => serde_json::Value::from(v.0),
|
||||
ValueRef::Float64(v) => serde_json::Value::from(v.0),
|
||||
ValueRef::String(bytes) => serde_json::Value::String(bytes.to_string()),
|
||||
ValueRef::Binary(bytes) => serde_json::to_value(bytes)?,
|
||||
ValueRef::Date(v) => serde_json::Value::Number(v.val().into()),
|
||||
ValueRef::DateTime(v) => serde_json::Value::Number(v.val().into()),
|
||||
ValueRef::List(v) => serde_json::to_value(v)?,
|
||||
ValueRef::Timestamp(v) => serde_json::to_value(v.value())?,
|
||||
ValueRef::Time(v) => serde_json::to_value(v.value())?,
|
||||
ValueRef::Interval(v) => serde_json::to_value(v.to_i128())?,
|
||||
ValueRef::Duration(v) => serde_json::to_value(v.value())?,
|
||||
ValueRef::Decimal128(v) => serde_json::to_value(v.to_string())?,
|
||||
};
|
||||
|
||||
Ok(json_value)
|
||||
}
|
||||
Ok(json_value)
|
||||
}
|
||||
|
||||
/// Reference to a [ListValue].
|
||||
@@ -1359,7 +1424,9 @@ impl<'a> ValueRef<'a> {
|
||||
ValueRef::Timestamp(_) => 16,
|
||||
ValueRef::Time(_) => 16,
|
||||
ValueRef::Duration(_) => 16,
|
||||
ValueRef::Interval(_) => 24,
|
||||
ValueRef::IntervalYearMonth(_) => 4,
|
||||
ValueRef::IntervalDayTime(_) => 8,
|
||||
ValueRef::IntervalMonthDayNano(_) => 16,
|
||||
ValueRef::Decimal128(_) => 32,
|
||||
ValueRef::List(v) => match v {
|
||||
ListValueRef::Indexed { vector, .. } => vector.memory_size() / vector.len(),
|
||||
@@ -1428,7 +1495,9 @@ pub fn column_data_to_json(data: ValueData) -> JsonValue {
|
||||
mod tests {
|
||||
use arrow::datatypes::DataType as ArrowDataType;
|
||||
use common_time::timezone::set_default_timezone;
|
||||
use greptime_proto::v1::{Decimal128 as ProtoDecimal128, IntervalMonthDayNano};
|
||||
use greptime_proto::v1::{
|
||||
Decimal128 as ProtoDecimal128, IntervalMonthDayNano as ProtoIntervalMonthDayNano,
|
||||
};
|
||||
use num_traits::Float;
|
||||
|
||||
use super::*;
|
||||
@@ -1525,11 +1594,13 @@ mod tests {
|
||||
JsonValue::String("interval year [12]".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
column_data_to_json(ValueData::IntervalMonthDayNanoValue(IntervalMonthDayNano {
|
||||
months: 1,
|
||||
days: 2,
|
||||
nanoseconds: 3,
|
||||
})),
|
||||
column_data_to_json(ValueData::IntervalMonthDayNanoValue(
|
||||
ProtoIntervalMonthDayNano {
|
||||
months: 1,
|
||||
days: 2,
|
||||
nanoseconds: 3,
|
||||
}
|
||||
)),
|
||||
JsonValue::String("interval month [1][2][3]".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
@@ -1740,12 +1811,10 @@ mod tests {
|
||||
ScalarValue::IntervalMonthDayNano(None).try_into().unwrap()
|
||||
);
|
||||
assert_eq!(
|
||||
Value::Interval(Interval::from_month_day_nano(1, 1, 1)),
|
||||
ScalarValue::IntervalMonthDayNano(Some(
|
||||
Interval::from_month_day_nano(1, 1, 1).to_i128()
|
||||
))
|
||||
.try_into()
|
||||
.unwrap()
|
||||
Value::IntervalMonthDayNano(IntervalMonthDayNano::new(1, 1, 1)),
|
||||
ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNano::new(1, 1, 1).to_i128()))
|
||||
.try_into()
|
||||
.unwrap()
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
@@ -1975,9 +2044,17 @@ mod tests {
|
||||
&ConcreteDataType::time_nanosecond_datatype(),
|
||||
&Value::Time(Time::new_nanosecond(1)),
|
||||
);
|
||||
check_type_and_value(
|
||||
&ConcreteDataType::interval_year_month_datatype(),
|
||||
&Value::IntervalYearMonth(IntervalYearMonth::new(1)),
|
||||
);
|
||||
check_type_and_value(
|
||||
&ConcreteDataType::interval_day_time_datatype(),
|
||||
&Value::IntervalDayTime(IntervalDayTime::new(1, 2)),
|
||||
);
|
||||
check_type_and_value(
|
||||
&ConcreteDataType::interval_month_day_nano_datatype(),
|
||||
&Value::Interval(Interval::from_month_day_nano(1, 2, 3)),
|
||||
&Value::IntervalMonthDayNano(IntervalMonthDayNano::new(1, 2, 3)),
|
||||
);
|
||||
check_type_and_value(
|
||||
&ConcreteDataType::duration_second_datatype(),
|
||||
@@ -2160,7 +2237,9 @@ mod tests {
|
||||
check_as_value_ref!(Float64, OrderedF64::from(16.0));
|
||||
check_as_value_ref!(Timestamp, Timestamp::new_millisecond(1));
|
||||
check_as_value_ref!(Time, Time::new_millisecond(1));
|
||||
check_as_value_ref!(Interval, Interval::from_month_day_nano(1, 2, 3));
|
||||
check_as_value_ref!(IntervalYearMonth, IntervalYearMonth::new(1));
|
||||
check_as_value_ref!(IntervalDayTime, IntervalDayTime::new(1, 2));
|
||||
check_as_value_ref!(IntervalMonthDayNano, IntervalMonthDayNano::new(1, 2, 3));
|
||||
check_as_value_ref!(Duration, Duration::new_millisecond(1));
|
||||
|
||||
assert_eq!(
|
||||
@@ -2672,9 +2751,11 @@ mod tests {
|
||||
check_value_ref_size_eq(&ValueRef::DateTime(DateTime::new(1)), 8);
|
||||
check_value_ref_size_eq(&ValueRef::Timestamp(Timestamp::new_millisecond(1)), 16);
|
||||
check_value_ref_size_eq(&ValueRef::Time(Time::new_millisecond(1)), 16);
|
||||
check_value_ref_size_eq(&ValueRef::IntervalYearMonth(IntervalYearMonth::new(1)), 4);
|
||||
check_value_ref_size_eq(&ValueRef::IntervalDayTime(IntervalDayTime::new(1, 2)), 8);
|
||||
check_value_ref_size_eq(
|
||||
&ValueRef::Interval(Interval::from_month_day_nano(1, 2, 3)),
|
||||
24,
|
||||
&ValueRef::IntervalMonthDayNano(IntervalMonthDayNano::new(1, 2, 3)),
|
||||
16,
|
||||
);
|
||||
check_value_ref_size_eq(&ValueRef::Duration(Duration::new_millisecond(1)), 16);
|
||||
check_value_ref_size_eq(
|
||||
|
||||
@@ -421,7 +421,7 @@ mod tests {
|
||||
use common_decimal::Decimal128;
|
||||
use common_time::time::Time;
|
||||
use common_time::timestamp::TimeUnit;
|
||||
use common_time::{Date, DateTime, Duration, Interval};
|
||||
use common_time::{Date, DateTime, Duration, IntervalMonthDayNano};
|
||||
|
||||
use super::*;
|
||||
use crate::value::Value;
|
||||
@@ -689,7 +689,10 @@ mod tests {
|
||||
);
|
||||
assert_eq!(3, vector.len());
|
||||
for i in 0..vector.len() {
|
||||
assert_eq!(Value::Interval(Interval::from_i128(2000)), vector.get(i));
|
||||
assert_eq!(
|
||||
Value::IntervalMonthDayNano(IntervalMonthDayNano::from_i128(2000)),
|
||||
vector.get(i)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -26,8 +26,8 @@ use object_store::ObjectStore;
|
||||
use snafu::{ensure, OptionExt};
|
||||
use store_api::metadata::RegionMetadataRef;
|
||||
use store_api::region_engine::{
|
||||
RegionEngine, RegionRole, RegionScannerRef, RegionStatistic, SetReadonlyResponse,
|
||||
SinglePartitionScanner,
|
||||
RegionEngine, RegionRole, RegionScannerRef, RegionStatistic, SetRegionRoleStateResponse,
|
||||
SettableRegionRoleState, SinglePartitionScanner,
|
||||
};
|
||||
use store_api::region_request::{
|
||||
AffectedRows, RegionCloseRequest, RegionCreateRequest, RegionDropRequest, RegionOpenRequest,
|
||||
@@ -113,22 +113,23 @@ impl RegionEngine for FileRegionEngine {
|
||||
None
|
||||
}
|
||||
|
||||
fn set_writable(&self, region_id: RegionId, writable: bool) -> Result<(), BoxedError> {
|
||||
fn set_region_role(&self, region_id: RegionId, role: RegionRole) -> Result<(), BoxedError> {
|
||||
self.inner
|
||||
.set_writable(region_id, writable)
|
||||
.set_region_role(region_id, role)
|
||||
.map_err(BoxedError::new)
|
||||
}
|
||||
|
||||
async fn set_readonly_gracefully(
|
||||
async fn set_region_role_state_gracefully(
|
||||
&self,
|
||||
region_id: RegionId,
|
||||
) -> Result<SetReadonlyResponse, BoxedError> {
|
||||
_region_role_state: SettableRegionRoleState,
|
||||
) -> Result<SetRegionRoleStateResponse, BoxedError> {
|
||||
let exists = self.inner.get_region(region_id).await.is_some();
|
||||
|
||||
if exists {
|
||||
Ok(SetReadonlyResponse::success(None))
|
||||
Ok(SetRegionRoleStateResponse::success(None))
|
||||
} else {
|
||||
Ok(SetReadonlyResponse::NotFound)
|
||||
Ok(SetRegionRoleStateResponse::NotFound)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -189,7 +190,7 @@ impl EngineInner {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn set_writable(&self, _region_id: RegionId, _writable: bool) -> EngineResult<()> {
|
||||
fn set_region_role(&self, _region_id: RegionId, _region_role: RegionRole) -> EngineResult<()> {
|
||||
// TODO(zhongzc): Improve the semantics and implementation of this API.
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -560,7 +560,7 @@ fn reduce_batch_subgraph(
|
||||
.get_mut(i)
|
||||
.context(InternalSnafu{
|
||||
reason: format!(
|
||||
"Output builder should have the same length as the row, expected at most {} but got {}",
|
||||
"Output builder should have the same length as the row, expected at most {} but got {}",
|
||||
column_cnt - 1,
|
||||
i
|
||||
)
|
||||
@@ -1162,7 +1162,9 @@ fn from_val_to_slice_idx(
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
|
||||
use common_time::{DateTime, Interval, Timestamp};
|
||||
use std::time::Duration;
|
||||
|
||||
use common_time::Timestamp;
|
||||
use datatypes::data_type::{ConcreteDataType, ConcreteDataType as CDT};
|
||||
use hydroflow::scheduled::graph::Hydroflow;
|
||||
|
||||
@@ -1214,8 +1216,8 @@ mod test {
|
||||
let expected = TypedPlan {
|
||||
schema: RelationType::new(vec![
|
||||
ColumnType::new(CDT::uint64_datatype(), true), // sum(number)
|
||||
ColumnType::new(CDT::datetime_datatype(), false), // window start
|
||||
ColumnType::new(CDT::datetime_datatype(), false), // window end
|
||||
ColumnType::new(CDT::timestamp_millisecond_datatype(), false), // window start
|
||||
ColumnType::new(CDT::timestamp_millisecond_datatype(), false), // window end
|
||||
])
|
||||
.into_unnamed(),
|
||||
// TODO(discord9): mfp indirectly ref to key columns
|
||||
@@ -1232,7 +1234,10 @@ mod test {
|
||||
.with_types(
|
||||
RelationType::new(vec![
|
||||
ColumnType::new(ConcreteDataType::uint32_datatype(), false),
|
||||
ColumnType::new(ConcreteDataType::datetime_datatype(), false),
|
||||
ColumnType::new(
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
false,
|
||||
),
|
||||
])
|
||||
.into_unnamed(),
|
||||
),
|
||||
@@ -1242,22 +1247,18 @@ mod test {
|
||||
.map(vec![
|
||||
ScalarExpr::Column(1).call_unary(
|
||||
UnaryFunc::TumbleWindowFloor {
|
||||
window_size: Interval::from_month_day_nano(
|
||||
0,
|
||||
0,
|
||||
1_000_000_000,
|
||||
),
|
||||
start_time: Some(DateTime::new(1625097600000)),
|
||||
window_size: Duration::from_nanos(1_000_000_000),
|
||||
start_time: Some(Timestamp::new_millisecond(
|
||||
1625097600000,
|
||||
)),
|
||||
},
|
||||
),
|
||||
ScalarExpr::Column(1).call_unary(
|
||||
UnaryFunc::TumbleWindowCeiling {
|
||||
window_size: Interval::from_month_day_nano(
|
||||
0,
|
||||
0,
|
||||
1_000_000_000,
|
||||
),
|
||||
start_time: Some(DateTime::new(1625097600000)),
|
||||
window_size: Duration::from_nanos(1_000_000_000),
|
||||
start_time: Some(Timestamp::new_millisecond(
|
||||
1625097600000,
|
||||
)),
|
||||
},
|
||||
),
|
||||
])
|
||||
@@ -1278,9 +1279,9 @@ mod test {
|
||||
}
|
||||
.with_types(
|
||||
RelationType::new(vec![
|
||||
ColumnType::new(CDT::datetime_datatype(), false), // window start
|
||||
ColumnType::new(CDT::datetime_datatype(), false), // window end
|
||||
ColumnType::new(CDT::uint64_datatype(), true), //sum(number)
|
||||
ColumnType::new(CDT::timestamp_millisecond_datatype(), false), // window start
|
||||
ColumnType::new(CDT::timestamp_millisecond_datatype(), false), // window end
|
||||
ColumnType::new(CDT::uint64_datatype(), true), //sum(number)
|
||||
])
|
||||
.with_key(vec![1])
|
||||
.with_time_index(Some(0))
|
||||
|
||||
@@ -171,9 +171,13 @@ impl DfScalarFunction {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub struct RawDfScalarFn {
|
||||
/// The raw bytes encoded datafusion scalar function
|
||||
/// The raw bytes encoded datafusion scalar function,
|
||||
/// due to substrait have too many layers of nested struct and `ScalarFunction` 's derive is different
|
||||
/// for simplicity's sake
|
||||
/// so we store bytes instead of `ScalarFunction` here
|
||||
/// but in unit test we will still compare decoded struct(using `f_decoded` field in Debug impl)
|
||||
pub(crate) f: bytes::BytesMut,
|
||||
/// The input schema of the function
|
||||
pub(crate) input_schema: RelationDesc,
|
||||
@@ -181,6 +185,17 @@ pub struct RawDfScalarFn {
|
||||
pub(crate) extensions: FunctionExtensions,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for RawDfScalarFn {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("RawDfScalarFn")
|
||||
.field("f", &self.f)
|
||||
.field("f_decoded", &ScalarFunction::decode(&mut self.f.as_ref()))
|
||||
.field("df_schema", &self.input_schema)
|
||||
.field("extensions", &self.extensions)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl RawDfScalarFn {
|
||||
pub fn from_proto(
|
||||
f: &substrait::substrait_proto_df::proto::expression::ScalarFunction,
|
||||
|
||||
@@ -16,19 +16,18 @@
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use arrow::array::{ArrayRef, BooleanArray};
|
||||
use common_error::ext::BoxedError;
|
||||
use common_time::timestamp::TimeUnit;
|
||||
use common_time::{DateTime, Timestamp};
|
||||
use common_time::Timestamp;
|
||||
use datafusion_expr::Operator;
|
||||
use datatypes::data_type::ConcreteDataType;
|
||||
use datatypes::prelude::DataType;
|
||||
use datatypes::types::cast;
|
||||
use datatypes::value::Value;
|
||||
use datatypes::vectors::{
|
||||
BooleanVector, DateTimeVector, Helper, TimestampMillisecondVector, VectorRef,
|
||||
};
|
||||
use datatypes::vectors::{BooleanVector, Helper, TimestampMillisecondVector, VectorRef};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use smallvec::smallvec;
|
||||
use snafu::{ensure, OptionExt, ResultExt};
|
||||
@@ -52,8 +51,8 @@ pub enum UnmaterializableFunc {
|
||||
CurrentSchema,
|
||||
TumbleWindow {
|
||||
ts: Box<TypedExpr>,
|
||||
window_size: common_time::Interval,
|
||||
start_time: Option<DateTime>,
|
||||
window_size: Duration,
|
||||
start_time: Option<Timestamp>,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -63,7 +62,8 @@ impl UnmaterializableFunc {
|
||||
match self {
|
||||
Self::Now => Signature {
|
||||
input: smallvec![],
|
||||
output: ConcreteDataType::datetime_datatype(),
|
||||
// TODO(yingwen): Maybe return timestamp.
|
||||
output: ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
generic_fn: GenericFn::Now,
|
||||
},
|
||||
Self::CurrentSchema => Signature {
|
||||
@@ -110,12 +110,12 @@ pub enum UnaryFunc {
|
||||
StepTimestamp,
|
||||
Cast(ConcreteDataType),
|
||||
TumbleWindowFloor {
|
||||
window_size: common_time::Interval,
|
||||
start_time: Option<DateTime>,
|
||||
window_size: Duration,
|
||||
start_time: Option<Timestamp>,
|
||||
},
|
||||
TumbleWindowCeiling {
|
||||
window_size: common_time::Interval,
|
||||
start_time: Option<DateTime>,
|
||||
window_size: Duration,
|
||||
start_time: Option<Timestamp>,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -139,8 +139,8 @@ impl UnaryFunc {
|
||||
},
|
||||
},
|
||||
Self::StepTimestamp => Signature {
|
||||
input: smallvec![ConcreteDataType::datetime_datatype()],
|
||||
output: ConcreteDataType::datetime_datatype(),
|
||||
input: smallvec![ConcreteDataType::timestamp_millisecond_datatype()],
|
||||
output: ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
generic_fn: GenericFn::StepTimestamp,
|
||||
},
|
||||
Self::Cast(to) => Signature {
|
||||
@@ -238,19 +238,19 @@ impl UnaryFunc {
|
||||
}
|
||||
}
|
||||
Self::StepTimestamp => {
|
||||
let datetime_array = get_datetime_array(&arg_col)?;
|
||||
let date_array_ref = datetime_array
|
||||
let timestamp_array = get_timestamp_array(&arg_col)?;
|
||||
let timestamp_array_ref = timestamp_array
|
||||
.as_any()
|
||||
.downcast_ref::<arrow::array::Date64Array>()
|
||||
.downcast_ref::<arrow::array::TimestampMillisecondArray>()
|
||||
.context({
|
||||
TypeMismatchSnafu {
|
||||
expected: ConcreteDataType::boolean_datatype(),
|
||||
actual: ConcreteDataType::from_arrow_type(datetime_array.data_type()),
|
||||
actual: ConcreteDataType::from_arrow_type(timestamp_array.data_type()),
|
||||
}
|
||||
})?;
|
||||
|
||||
let ret = arrow::compute::unary(date_array_ref, |arr| arr + 1);
|
||||
let ret = DateTimeVector::from(ret);
|
||||
let ret = arrow::compute::unary(timestamp_array_ref, |arr| arr + 1);
|
||||
let ret = TimestampMillisecondVector::from(ret);
|
||||
Ok(Arc::new(ret))
|
||||
}
|
||||
Self::Cast(to) => {
|
||||
@@ -266,19 +266,19 @@ impl UnaryFunc {
|
||||
window_size,
|
||||
start_time,
|
||||
} => {
|
||||
let datetime_array = get_datetime_array(&arg_col)?;
|
||||
let date_array_ref = datetime_array
|
||||
let timestamp_array = get_timestamp_array(&arg_col)?;
|
||||
let date_array_ref = timestamp_array
|
||||
.as_any()
|
||||
.downcast_ref::<arrow::array::Date64Array>()
|
||||
.downcast_ref::<arrow::array::TimestampMillisecondArray>()
|
||||
.context({
|
||||
TypeMismatchSnafu {
|
||||
expected: ConcreteDataType::boolean_datatype(),
|
||||
actual: ConcreteDataType::from_arrow_type(datetime_array.data_type()),
|
||||
actual: ConcreteDataType::from_arrow_type(timestamp_array.data_type()),
|
||||
}
|
||||
})?;
|
||||
|
||||
let start_time = start_time.map(|t| t.val());
|
||||
let window_size = (window_size.to_nanosecond() / 1_000_000) as repr::Duration; // nanosecond to millisecond
|
||||
let start_time = start_time.map(|t| t.value());
|
||||
let window_size = window_size.as_millis() as repr::Duration;
|
||||
|
||||
let ret = arrow::compute::unary(date_array_ref, |ts| {
|
||||
get_window_start(ts, window_size, start_time)
|
||||
@@ -291,19 +291,19 @@ impl UnaryFunc {
|
||||
window_size,
|
||||
start_time,
|
||||
} => {
|
||||
let datetime_array = get_datetime_array(&arg_col)?;
|
||||
let date_array_ref = datetime_array
|
||||
let timestamp_array = get_timestamp_array(&arg_col)?;
|
||||
let date_array_ref = timestamp_array
|
||||
.as_any()
|
||||
.downcast_ref::<arrow::array::Date64Array>()
|
||||
.downcast_ref::<arrow::array::TimestampMillisecondArray>()
|
||||
.context({
|
||||
TypeMismatchSnafu {
|
||||
expected: ConcreteDataType::boolean_datatype(),
|
||||
actual: ConcreteDataType::from_arrow_type(datetime_array.data_type()),
|
||||
actual: ConcreteDataType::from_arrow_type(timestamp_array.data_type()),
|
||||
}
|
||||
})?;
|
||||
|
||||
let start_time = start_time.map(|t| t.val());
|
||||
let window_size = (window_size.to_nanosecond() / 1_000_000) as repr::Duration; // nanosecond to millisecond
|
||||
let start_time = start_time.map(|t| t.value());
|
||||
let window_size = window_size.as_millis() as repr::Duration;
|
||||
|
||||
let ret = arrow::compute::unary(date_array_ref, |ts| {
|
||||
get_window_start(ts, window_size, start_time) + window_size
|
||||
@@ -330,19 +330,20 @@ impl UnaryFunc {
|
||||
})?;
|
||||
if let Some(window_size) = window_size_untyped.as_string() {
|
||||
// cast as interval
|
||||
cast(
|
||||
let interval = cast(
|
||||
Value::from(window_size),
|
||||
&ConcreteDataType::interval_month_day_nano_datatype(),
|
||||
&ConcreteDataType::interval_day_time_datatype(),
|
||||
)
|
||||
.map_err(BoxedError::new)
|
||||
.context(ExternalSnafu)?
|
||||
.as_interval()
|
||||
.as_interval_day_time()
|
||||
.context(UnexpectedSnafu {
|
||||
reason: "Expect window size arg to be interval after successful cast"
|
||||
.to_string(),
|
||||
})?
|
||||
} else if let Some(interval) = window_size_untyped.as_interval() {
|
||||
interval
|
||||
})?;
|
||||
Duration::from_millis(interval.as_millis() as u64)
|
||||
} else if let Some(interval) = window_size_untyped.as_interval_day_time() {
|
||||
Duration::from_millis(interval.as_millis() as u64)
|
||||
} else {
|
||||
InvalidQuerySnafu {
|
||||
reason: format!(
|
||||
@@ -357,16 +358,19 @@ impl UnaryFunc {
|
||||
let start_time = match args.get(2) {
|
||||
Some(start_time) => {
|
||||
if let Some(value) = start_time.expr.as_literal() {
|
||||
// cast as DateTime
|
||||
let ret = cast(value, &ConcreteDataType::datetime_datatype())
|
||||
.map_err(BoxedError::new)
|
||||
.context(ExternalSnafu)?
|
||||
.as_datetime()
|
||||
.context(UnexpectedSnafu {
|
||||
reason:
|
||||
"Expect start time arg to be datetime after successful cast"
|
||||
.to_string(),
|
||||
})?;
|
||||
// cast as timestamp
|
||||
let ret = cast(
|
||||
value,
|
||||
&ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
)
|
||||
.map_err(BoxedError::new)
|
||||
.context(ExternalSnafu)?
|
||||
.as_timestamp()
|
||||
.context(UnexpectedSnafu {
|
||||
reason:
|
||||
"Expect start time arg to be timestamp after successful cast"
|
||||
.to_string(),
|
||||
})?;
|
||||
Some(ret)
|
||||
} else {
|
||||
UnexpectedSnafu {
|
||||
@@ -446,15 +450,15 @@ impl UnaryFunc {
|
||||
}
|
||||
Self::StepTimestamp => {
|
||||
let ty = arg.data_type();
|
||||
if let Value::DateTime(datetime) = arg {
|
||||
let datetime = DateTime::from(datetime.val() + 1);
|
||||
Ok(Value::from(datetime))
|
||||
if let Value::Timestamp(timestamp) = arg {
|
||||
let timestamp = Timestamp::new_millisecond(timestamp.value() + 1);
|
||||
Ok(Value::from(timestamp))
|
||||
} else if let Ok(v) = value_to_internal_ts(arg) {
|
||||
let datetime = DateTime::from(v + 1);
|
||||
Ok(Value::from(datetime))
|
||||
let timestamp = Timestamp::new_millisecond(v + 1);
|
||||
Ok(Value::from(timestamp))
|
||||
} else {
|
||||
TypeMismatchSnafu {
|
||||
expected: ConcreteDataType::datetime_datatype(),
|
||||
expected: ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
actual: ty,
|
||||
}
|
||||
.fail()?
|
||||
@@ -474,8 +478,8 @@ impl UnaryFunc {
|
||||
start_time,
|
||||
} => {
|
||||
let ts = get_ts_as_millisecond(arg)?;
|
||||
let start_time = start_time.map(|t| t.val());
|
||||
let window_size = (window_size.to_nanosecond() / 1_000_000) as repr::Duration; // nanosecond to millisecond
|
||||
let start_time = start_time.map(|t| t.value());
|
||||
let window_size = window_size.as_millis() as repr::Duration;
|
||||
let window_start = get_window_start(ts, window_size, start_time);
|
||||
|
||||
let ret = Timestamp::new_millisecond(window_start);
|
||||
@@ -486,8 +490,8 @@ impl UnaryFunc {
|
||||
start_time,
|
||||
} => {
|
||||
let ts = get_ts_as_millisecond(arg)?;
|
||||
let start_time = start_time.map(|t| t.val());
|
||||
let window_size = (window_size.to_nanosecond() / 1_000_000) as repr::Duration; // nanosecond to millisecond
|
||||
let start_time = start_time.map(|t| t.value());
|
||||
let window_size = window_size.as_millis() as repr::Duration;
|
||||
let window_start = get_window_start(ts, window_size, start_time);
|
||||
|
||||
let window_end = window_start + window_size;
|
||||
@@ -498,21 +502,22 @@ impl UnaryFunc {
|
||||
}
|
||||
}
|
||||
|
||||
fn get_datetime_array(vector: &VectorRef) -> Result<arrow::array::ArrayRef, EvalError> {
|
||||
fn get_timestamp_array(vector: &VectorRef) -> Result<arrow::array::ArrayRef, EvalError> {
|
||||
let arrow_array = vector.to_arrow_array();
|
||||
let datetime_array =
|
||||
if *arrow_array.data_type() == ConcreteDataType::datetime_datatype().as_arrow_type() {
|
||||
arrow_array
|
||||
} else {
|
||||
arrow::compute::cast(
|
||||
&arrow_array,
|
||||
&ConcreteDataType::datetime_datatype().as_arrow_type(),
|
||||
)
|
||||
.context(ArrowSnafu {
|
||||
context: "Trying to cast to datetime in StepTimestamp",
|
||||
})?
|
||||
};
|
||||
Ok(datetime_array)
|
||||
let timestamp_array = if *arrow_array.data_type()
|
||||
== ConcreteDataType::timestamp_millisecond_datatype().as_arrow_type()
|
||||
{
|
||||
arrow_array
|
||||
} else {
|
||||
arrow::compute::cast(
|
||||
&arrow_array,
|
||||
&ConcreteDataType::timestamp_millisecond_datatype().as_arrow_type(),
|
||||
)
|
||||
.context(ArrowSnafu {
|
||||
context: "Trying to cast to timestamp in StepTimestamp",
|
||||
})?
|
||||
};
|
||||
Ok(timestamp_array)
|
||||
}
|
||||
|
||||
fn get_window_start(
|
||||
@@ -1284,7 +1289,6 @@ where
|
||||
mod test {
|
||||
use std::sync::Arc;
|
||||
|
||||
use common_time::Interval;
|
||||
use datatypes::vectors::Vector;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
@@ -1292,18 +1296,18 @@ mod test {
|
||||
|
||||
#[test]
|
||||
fn test_tumble_batch() {
|
||||
let datetime_vector = DateTimeVector::from_vec(vec![1, 2, 10, 13, 14, 20, 25]);
|
||||
let timestamp_vector = TimestampMillisecondVector::from_vec(vec![1, 2, 10, 13, 14, 20, 25]);
|
||||
let tumble_start = UnaryFunc::TumbleWindowFloor {
|
||||
window_size: Interval::from_day_time(0, 10),
|
||||
window_size: Duration::from_millis(10),
|
||||
start_time: None,
|
||||
};
|
||||
let tumble_end = UnaryFunc::TumbleWindowCeiling {
|
||||
window_size: Interval::from_day_time(0, 10),
|
||||
window_size: Duration::from_millis(10),
|
||||
start_time: None,
|
||||
};
|
||||
|
||||
let len = datetime_vector.len();
|
||||
let batch = Batch::try_new(vec![Arc::new(datetime_vector)], len).unwrap();
|
||||
let len = timestamp_vector.len();
|
||||
let batch = Batch::try_new(vec![Arc::new(timestamp_vector)], len).unwrap();
|
||||
let arg = ScalarExpr::Column(0);
|
||||
|
||||
let start = tumble_start.eval_batch(&batch, &arg).unwrap();
|
||||
@@ -1459,4 +1463,17 @@ mod test {
|
||||
Err(Error::InvalidQuery { .. })
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cast_int() {
|
||||
let interval = cast(
|
||||
Value::from("1 second"),
|
||||
&ConcreteDataType::interval_day_time_datatype(),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
interval,
|
||||
Value::from(common_time::IntervalDayTime::new(0, 1000))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user