Add UDF's for sketch support (#6349)

* add UDF's for sketch support

* substrait test

* edge case with empty groups

* address df-review

* match on datadog sketches

* Fix sketch quantile clippy warning

* cleanup dead code in tests

* Use conventional test helper modules

Author: alexanderbianchi

quickwit/Cargo.lock

@@ -15,17 +15,22 @@
 /// Returns whether the given index ID corresponds to a metrics index.
 ///
 /// Metrics indexes use the Parquet/DataFusion pipeline instead of the Tantivy pipeline.
-/// An index is considered a metrics index if it starts with "otel-metrics" or "metrics-".
+/// An index is considered a metrics index if it uses one of the BYOC or OSS
+/// parquet metrics prefixes.
 pub fn is_metrics_index(index_id: &str) -> bool {
-    index_id.starts_with("otel-metrics") || index_id.starts_with("metrics-")
+    ["datadog-metrics", "metrics-", "otel-metrics"]
+        .iter()
+        .any(|prefix| index_id.starts_with(prefix))
 }
 
 /// Returns whether the given index ID corresponds to a sketches index.
 ///
 /// Sketches indexes use the Parquet/DataFusion pipeline with sketch-specific
 /// processors and writers.
 pub fn is_sketches_index(index_id: &str) -> bool {
-    index_id.starts_with("sketches-")
+    ["datadog-sketches", "sketches-"]
+        .iter()
+        .any(|prefix| index_id.starts_with(prefix))
 }
 
 /// Returns whether the given index ID uses the Parquet/DataFusion pipeline.
@@ -44,6 +49,10 @@ mod tests {
         assert!(is_metrics_index("otel-metrics"));
         assert!(is_metrics_index("otel-metrics-custom"));
 
+        // BYOC metrics indexes
+        assert!(is_metrics_index("datadog-metrics"));
+        assert!(is_metrics_index("datadog-metrics-v2"));
+
         // Generic metrics indexes
         assert!(is_metrics_index("metrics-default"));
         assert!(is_metrics_index("metrics-"));
@@ -60,13 +69,18 @@ mod tests {
 
     #[test]
     fn test_is_sketches_index() {
+        assert!(is_sketches_index("datadog-sketches"));
+        assert!(is_sketches_index("datadog-sketches-v2"));
         assert!(is_sketches_index("sketches-default"));
+        assert!(!is_sketches_index("datadog-metrics"));
         assert!(!is_sketches_index("otel-metrics"));
         assert!(!is_sketches_index("my-index"));
     }
 
     #[test]
     fn test_is_parquet_pipeline_index() {
+        assert!(is_parquet_pipeline_index("datadog-metrics"));
+        assert!(is_parquet_pipeline_index("datadog-sketches"));
         assert!(is_parquet_pipeline_index("otel-metrics"));
         assert!(is_parquet_pipeline_index("sketches-default"));
         assert!(!is_parquet_pipeline_index("otel-logs-v0_7"));

quickwit/quickwit-common/src/metrics_specific.rs

@@ -43,6 +43,7 @@ object_store = "0.13"
 bytesize = { workspace = true }
 datafusion = "53"
 datafusion-substrait = "53"
+parquet = { workspace = true }
 prost = { workspace = true }
 serde_json = { workspace = true }
 tempfile = { workspace = true }

quickwit/quickwit-datafusion/Cargo.toml

@@ -24,6 +24,18 @@
 //!     service      VARCHAR,
 //!     env          VARCHAR
 //! ) STORED AS metrics LOCATION 'my-metrics';
+//!
+//! CREATE EXTERNAL TABLE "my-sketches" (
+//!     metric_name    VARCHAR NOT NULL,
+//!     timestamp_secs BIGINT UNSIGNED NOT NULL,
+//!     count          BIGINT UNSIGNED NOT NULL,
+//!     sum            DOUBLE  NOT NULL,
+//!     min            DOUBLE  NOT NULL,
+//!     max            DOUBLE  NOT NULL,
+//!     flags          INT UNSIGNED NOT NULL,
+//!     keys           ARRAY<SMALLINT> NOT NULL,
+//!     counts         ARRAY<BIGINT UNSIGNED> NOT NULL
+//! ) STORED AS sketches LOCATION 'my-sketches';
 //! ```
 
 use std::sync::Arc;
@@ -34,22 +46,32 @@ use datafusion::arrow;
 use datafusion::catalog::{Session, TableProviderFactory};
 use datafusion::error::{DataFusionError, Result as DFResult};
 use datafusion::logical_expr::CreateExternalTable;
+use quickwit_parquet_engine::split::ParquetSplitKind;
 
 use super::index_resolver::MetricsIndexResolver;
 use super::table_provider::MetricsTableProvider;
 
 /// The file type string used in `STORED AS metrics`.
 pub const METRICS_FILE_TYPE: &str = "metrics";
+/// The file type string used in `STORED AS sketches`.
+pub const SKETCHES_FILE_TYPE: &str = "sketches";
 
 /// Creates `MetricsTableProvider` instances from `CREATE EXTERNAL TABLE` DDL.
 #[derive(Debug)]
 pub struct MetricsTableProviderFactory {
     index_resolver: Arc<dyn MetricsIndexResolver>,
+    split_kind: ParquetSplitKind,
 }
 
 impl MetricsTableProviderFactory {
-    pub fn new(index_resolver: Arc<dyn MetricsIndexResolver>) -> Self {
-        Self { index_resolver }
+    pub fn new(
+        index_resolver: Arc<dyn MetricsIndexResolver>,
+        split_kind: ParquetSplitKind,
+    ) -> Self {
+        Self {
+            index_resolver,
+            split_kind,
+        }
     }
 }
 
@@ -66,7 +88,10 @@ impl TableProviderFactory for MetricsTableProviderFactory {
             cmd.location.clone()
         };
 
-        let (split_provider, index_uri) = self.index_resolver.resolve(&index_name).await?;
+        let (split_provider, index_uri) = self
+            .index_resolver
+            .resolve(&index_name, self.split_kind)
+            .await?;
 
         let arrow_schema: SchemaRef = Arc::new(cmd.schema.as_arrow().clone());
 

quickwit/quickwit-datafusion/src/sources/metrics/factory.rs

@@ -26,6 +26,7 @@ use async_trait::async_trait;
 use datafusion::error::Result as DFResult;
 use quickwit_common::uri::Uri;
 use quickwit_metastore::{IndexMetadataResponseExt, ListIndexesMetadataResponseExt};
+use quickwit_parquet_engine::split::ParquetSplitKind;
 use quickwit_proto::metastore::{
     IndexMetadataRequest, ListIndexesMetadataRequest, MetastoreService, MetastoreServiceClient,
 };
@@ -40,7 +41,11 @@ pub trait MetricsIndexResolver: Send + Sync + std::fmt::Debug {
     /// Returns the split provider and storage URI for `index_name`. The
     /// `ObjectStore` for that URI is built on demand by the registry the
     /// first time DataFusion reads from it.
-    async fn resolve(&self, index_name: &str) -> DFResult<(Arc<dyn MetricsSplitProvider>, Uri)>;
+    async fn resolve(
+        &self,
+        index_name: &str,
+        split_kind: ParquetSplitKind,
+    ) -> DFResult<(Arc<dyn MetricsSplitProvider>, Uri)>;
 
     async fn list_index_names(&self) -> DFResult<Vec<String>>;
 }
@@ -71,8 +76,12 @@ impl std::fmt::Debug for MetastoreIndexResolver {
 
 #[async_trait]
 impl MetricsIndexResolver for MetastoreIndexResolver {
-    async fn resolve(&self, index_name: &str) -> DFResult<(Arc<dyn MetricsSplitProvider>, Uri)> {
-        debug!(index_name, "resolving metrics index");
+    async fn resolve(
+        &self,
+        index_name: &str,
+        split_kind: ParquetSplitKind,
+    ) -> DFResult<(Arc<dyn MetricsSplitProvider>, Uri)> {
+        debug!(index_name, ?split_kind, "resolving parquet index");
 
         let response = self
             .metastore
@@ -93,6 +102,7 @@ impl MetricsIndexResolver for MetastoreIndexResolver {
         let split_provider: Arc<dyn MetricsSplitProvider> = Arc::new(MetastoreSplitProvider::new(
             self.metastore.clone(),
             index_uid,
+            split_kind,
         ));
 
         Ok((split_provider, index_uri))

quickwit/quickwit-datafusion/src/sources/metrics/index_resolver.rs

@@ -21,9 +21,9 @@ use datafusion::error::Result as DFResult;
 use quickwit_metastore::{
     ListParquetSplitsQuery, ListParquetSplitsRequestExt, ListParquetSplitsResponseExt,
 };
-use quickwit_parquet_engine::split::ParquetSplitMetadata;
+use quickwit_parquet_engine::split::{ParquetSplitKind, ParquetSplitMetadata};
 use quickwit_proto::metastore::{
-    ListMetricsSplitsRequest, MetastoreService, MetastoreServiceClient,
+    ListMetricsSplitsRequest, ListSketchSplitsRequest, MetastoreService, MetastoreServiceClient,
 };
 use quickwit_proto::types::IndexUid;
 use tracing::{debug, instrument};
@@ -36,13 +36,19 @@ use super::table_provider::MetricsSplitProvider;
 pub struct MetastoreSplitProvider {
     metastore: MetastoreServiceClient,
     index_uid: IndexUid,
+    split_kind: ParquetSplitKind,
 }
 
 impl MetastoreSplitProvider {
-    pub fn new(metastore: MetastoreServiceClient, index_uid: IndexUid) -> Self {
+    pub fn new(
+        metastore: MetastoreServiceClient,
+        index_uid: IndexUid,
+        split_kind: ParquetSplitKind,
+    ) -> Self {
         Self {
             metastore,
             index_uid,
+            split_kind,
         }
     }
 }
@@ -56,26 +62,45 @@ impl MetricsSplitProvider for MetastoreSplitProvider {
             metric_names = ?query.metric_names,
             time_range_start = ?query.time_range_start,
             time_range_end = ?query.time_range_end,
+            split_kind = ?self.split_kind,
             num_splits,
         )
     )]
     async fn list_splits(&self, query: &MetricsSplitQuery) -> DFResult<Vec<ParquetSplitMetadata>> {
         let metastore_query = to_metastore_query(&self.index_uid, query);
 
-        let request =
-            ListMetricsSplitsRequest::try_from_query(self.index_uid.clone(), &metastore_query)
+        let records = match self.split_kind {
+            ParquetSplitKind::Metrics => {
+                let request = ListMetricsSplitsRequest::try_from_query(
+                    self.index_uid.clone(),
+                    &metastore_query,
+                )
                 .map_err(|err| datafusion::error::DataFusionError::External(Box::new(err)))?;
 
-        let response = self
-            .metastore
-            .clone()
-            .list_metrics_splits(request)
-            .await
-            .map_err(|err| datafusion::error::DataFusionError::External(Box::new(err)))?;
+                self.metastore
+                    .clone()
+                    .list_metrics_splits(request)
+                    .await
+                    .map_err(|err| datafusion::error::DataFusionError::External(Box::new(err)))?
+                    .deserialize_splits()
+                    .map_err(|err| datafusion::error::DataFusionError::External(Box::new(err)))?
+            }
+            ParquetSplitKind::Sketches => {
+                let request = ListSketchSplitsRequest::try_from_query(
+                    self.index_uid.clone(),
+                    &metastore_query,
+                )
+                .map_err(|err| datafusion::error::DataFusionError::External(Box::new(err)))?;
 
-        let records = response
-            .deserialize_splits()
-            .map_err(|err| datafusion::error::DataFusionError::External(Box::new(err)))?;
+                self.metastore
+                    .clone()
+                    .list_sketch_splits(request)
+                    .await
+                    .map_err(|err| datafusion::error::DataFusionError::External(Box::new(err)))?
+                    .deserialize_splits()
+                    .map_err(|err| datafusion::error::DataFusionError::External(Box::new(err)))?
+            }
+        };
 
         // The metastore guarantees only Published splits are returned because
         // `to_metastore_query` sets `split_states = vec![Published]`. No