From 50c6a5f00ea93b41f0acaf7127d89f4d20e21370 Mon Sep 17 00:00:00 2001
From: "R. Tyler Croy" <rtyler@brokenco.de>
Date: Wed, 4 Dec 2024 22:10:09 +0000
Subject: [PATCH] Incorporate the memory improvements of delta-rs 0.22.x

---
 lambdas/query-metrics/Cargo.toml | 7 ++++---
 lambdas/query-metrics/src/cli.rs | 6 ++++--
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/lambdas/query-metrics/Cargo.toml b/lambdas/query-metrics/Cargo.toml
index b62cd1d..69aecc9 100644
--- a/lambdas/query-metrics/Cargo.toml
+++ b/lambdas/query-metrics/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "query-metrics"
-version = "0.4.0"
+version = "0.5.0"
 edition = "2021"
 
 [[bin]]
@@ -18,8 +18,9 @@ aws-sdk-cloudwatch = "1.11.0"
 aws-sdk-config = "1.11.0"
 aws_lambda_events = { version = "0.12.0" }
 base64 = "0.21.7"
-deltalake-core = { version = "0.19.0", features = ["datafusion"] }
-deltalake-aws = { version = "0.1.4" }
+deltalake-core = { version = "0.22.3", features = ["datafusion"] }
+deltalake-aws = { version = "0.5.0" }
+datafusion-functions-json = "0.43.0"
 
 lambda_runtime = "0.8.3"
 serde = { version = "1.0.195", features = ["derive"] }
diff --git a/lambdas/query-metrics/src/cli.rs b/lambdas/query-metrics/src/cli.rs
index 6ff180a..c5a6589 100644
--- a/lambdas/query-metrics/src/cli.rs
+++ b/lambdas/query-metrics/src/cli.rs
@@ -26,7 +26,9 @@ async fn main() -> anyhow::Result<()> {
     for (name, gauges) in conf.gauges.iter() {
         for gauge in gauges.iter() {
             println!("Querying the {name} table");
-            let ctx = SessionContext::new();
+            let mut ctx = SessionContext::new();
+            datafusion_functions_json::register_all(&mut ctx)?;
+
             let table = deltalake_core::open_table(&gauge.url)
                 .await
                 .expect("Failed to register table");
@@ -48,7 +50,7 @@ async fn main() -> anyhow::Result<()> {
                     println!("Counted {count} rows");
                 }
                 config::Measurement::Numeric => {
-                    println!("Need to run dimensional count");
+                    println!("Need to run numeric count");
                     let batches = df.collect().await.expect("Failed to collect batches");
                     let _ = print_batches(&batches);