DNM: stateless upsert, for running ci

Cargo.toml

@@ -46,6 +46,7 @@ members = [
     "src/http-util",
     "src/interchange",
     "src/kafka-util",
+    "src/kv-cache",
     "src/license-keys",
     "src/lowertest",
     "src/lowertest-derive",
@@ -56,7 +57,6 @@ members = [
     "src/metrics",
     "src/mysql-util",
     "src/mz",
-    "src/mz-debug",
     "src/npm",
     "src/orchestrator",
     "src/orchestrator-kubernetes",
@@ -177,7 +177,6 @@ default-members = [
     "src/metrics",
     "src/mysql-util",
     "src/mz",
-    "src/mz-debug",
     "src/npm",
     "src/orchestrator",
     "src/orchestrator-kubernetes",
@@ -316,6 +315,19 @@ postgres-types = { git = "https://github.com/MaterializeInc/rust-postgres" }
 postgres-openssl = { git = "https://github.com/MaterializeInc/rust-postgres" }
 postgres_array = { git = "https://github.com/MaterializeInc/rust-postgres-array" }
 
+parquet = { git = "https://github.com/ptravers/arrow-rs", branch = "hackweek/upsert-in-s3" }
+arrow-data = { git = "https://github.com/ptravers/arrow-rs", branch = "hackweek/upsert-in-s3" }
+arrow-buffer = { git = "https://github.com/ptravers/arrow-rs", branch = "hackweek/upsert-in-s3" }
+arrow-cast = { git = "https://github.com/ptravers/arrow-rs", branch = "hackweek/upsert-in-s3" }
+arrow-ipc = { git = "https://github.com/ptravers/arrow-rs", branch = "hackweek/upsert-in-s3" }
+arrow-select = { git = "https://github.com/ptravers/arrow-rs", branch = "hackweek/upsert-in-s3" }
+arrow-string = { git = "https://github.com/ptravers/arrow-rs", branch = "hackweek/upsert-in-s3" }
+arrow-ord = { git = "https://github.com/ptravers/arrow-rs", branch = "hackweek/upsert-in-s3" }
+arrow-arith = { git = "https://github.com/ptravers/arrow-rs", branch = "hackweek/upsert-in-s3" }
+arrow = { git = "https://github.com/ptravers/arrow-rs", branch = "hackweek/upsert-in-s3" }
+arrow-schema = { git = "https://github.com/ptravers/arrow-rs", branch = "hackweek/upsert-in-s3" }
+arrow-array = { git = "https://github.com/ptravers/arrow-rs", branch = "hackweek/upsert-in-s3" }
+
 # Waiting on https://github.com/MaterializeInc/serde-value/pull/35.
 serde-value = { git = "https://github.com/MaterializeInc/serde-value.git" }
 

WORKSPACE

@@ -502,12 +502,12 @@ crates_repository(
             gen_build_script = False,
             deps = ["@bzip2"],
         )],
-        "lzma-sys": [crate.annotation(
-            additive_build_file = "@//misc/bazel/c_deps:rust-sys/BUILD.lzma-sys.bazel",
-            gen_build_script = False,
-            # Note: This is a target we add from the additive build file above.
-            deps = [":xz"],
-        )],
+        # "lzma-sys": [crate.annotation(
+        #     additive_build_file = "@//misc/bazel/c_deps:rust-sys/BUILD.lzma-sys.bazel",
+        #     gen_build_script = False,
+        #     # Note: This is a target we add from the additive build file above.
+        #     deps = [":xz"],
+        # )],
         "openssl-sys": [crate.annotation(
             build_script_data = [
                 "@openssl//:openssl_lib",
@@ -634,6 +634,7 @@ crates_repository(
         "//:src/http-util/Cargo.toml",
         "//:src/interchange/Cargo.toml",
         "//:src/kafka-util/Cargo.toml",
+        "//:src/kv-cache/Cargo.toml",
         "//:src/license-keys/Cargo.toml",
         "//:src/lowertest-derive/Cargo.toml",
         "//:src/lowertest/Cargo.toml",
@@ -643,7 +644,6 @@ crates_repository(
         "//:src/metrics/Cargo.toml",
         "//:src/mysql-util/Cargo.toml",
         "//:src/mz/Cargo.toml",
-        "//:src/mz-debug/Cargo.toml",
         "//:src/npm/Cargo.toml",
         "//:src/orchestrator-kubernetes/Cargo.toml",
         "//:src/orchestrator-process/Cargo.toml",

misc/bazel/c_deps/rust-sys/BUILD.protobuf-native.bazel

@@ -38,7 +38,7 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_protobuf//src/google/protobuf/compiler:code_generator",
         "@com_google_protobuf//src/google/protobuf/compiler:importer",
-        "@crates_io__cxx-1.0.122//:cxx_cc",
+        "@crates_io__cxx-1.0.158//:cxx_cc",
     ],
 )
 
@@ -82,6 +82,6 @@ cc_library(
     include_prefix = "protobuf-native",
     deps = [
         "@com_google_absl//absl/strings",
-        "@crates_io__cxx-1.0.122//:cxx_cc",
+        "@crates_io__cxx-1.0.158//:cxx_cc",
     ],
 )

src/kv-cache/BUILD.bazel

@@ -0,0 +1,97 @@
+# Code generated by cargo-gazelle DO NOT EDIT
+
+# Copyright Materialize, Inc. and contributors. All rights reserved.
+#
+# Use of this software is governed by the Business Source License
+# included in the LICENSE file at the root of this repository.
+#
+# As of the Change Date specified in that file, in accordance with
+# the Business Source License, use of this software will be governed
+# by the Apache License, Version 2.0.
+
+load("@crates_io//:defs.bzl", "aliases", "all_crate_deps")
+load("@rules_rust//cargo:defs.bzl", "extract_cargo_lints")
+load("@rules_rust//rust:defs.bzl", "rust_doc_test", "rust_library", "rust_test")
+
+package(default_visibility = ["//visibility:public"])
+
+rust_library(
+    name = "mz_kv_cache",
+    srcs = glob(["src/**/*.rs"]),
+    aliases = aliases(
+        normal = True,
+        proc_macro = True,
+    ),
+    compile_data = [],
+    crate_features = [],
+    data = [],
+    lint_config = ":lints",
+    proc_macro_deps = [] + all_crate_deps(proc_macro = True),
+    rustc_env = {},
+    rustc_flags = [],
+    version = "0.1.0",
+    deps = [
+        "//src/persist-client:mz_persist_client",
+        "//src/persist-types:mz_persist_types",
+        "//src/repr:mz_repr",
+        "//src/storage-types:mz_storage_types",
+    ] + all_crate_deps(normal = True),
+)
+
+alias(
+    name = "kv-cache",
+    actual = "mz_kv_cache",
+)
+
+rust_test(
+    name = "mz_kv_cache_lib_tests",
+    size = "medium",
+    aliases = aliases(
+        normal = True,
+        normal_dev = True,
+        proc_macro = True,
+        proc_macro_dev = True,
+    ),
+    compile_data = [],
+    crate = ":mz_kv_cache",
+    crate_features = [],
+    data = [],
+    env = {},
+    lint_config = ":lints",
+    proc_macro_deps = [] + all_crate_deps(
+        proc_macro = True,
+        proc_macro_dev = True,
+    ),
+    rustc_env = {},
+    rustc_flags = [],
+    version = "0.1.0",
+    deps = [
+        "//src/persist-client:mz_persist_client",
+        "//src/persist-types:mz_persist_types",
+        "//src/repr:mz_repr",
+        "//src/storage-types:mz_storage_types",
+    ] + all_crate_deps(
+        normal = True,
+        normal_dev = True,
+    ),
+)
+
+rust_doc_test(
+    name = "mz_kv_cache_doc_test",
+    crate = ":mz_kv_cache",
+    deps = [
+        "//src/persist-client:mz_persist_client",
+        "//src/persist-types:mz_persist_types",
+        "//src/repr:mz_repr",
+        "//src/storage-types:mz_storage_types",
+    ] + all_crate_deps(
+        normal = True,
+        normal_dev = True,
+    ),
+)
+
+extract_cargo_lints(
+    name = "lints",
+    manifest = "Cargo.toml",
+    workspace = "@//:Cargo.toml",
+)

src/kv-cache/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "mz-kv-cache"
+version = "0.1.0"
+edition.workspace = true
+rust-version.workspace = true
+
+[dependencies]
+anyhow = "1.0.98"
+differential-dataflow = "0.15.1"
+mz-persist-types = { path = "../persist-types" }
+mz-persist-client = { path = "../persist-client" }
+mz-repr = { path = "../repr" }
+mz-storage-types = { path = "../storage-types" }
+timely = "0.21.0"
+tracing = "0.1.37"
+
+[lints]
+workspace = true

src/kv-cache/src/lib.rs

@@ -0,0 +1,170 @@
+use std::collections::BTreeSet;
+
+use differential_dataflow::lattice::Lattice;
+use mz_persist_client::lru::Lru;
+use mz_persist_client::read::ReadHandle;
+use mz_persist_types::Codec64;
+use mz_persist_types::bloom_filter::BloomFilter;
+use mz_repr::{DatumVec, Row, TimestampManipulation};
+use mz_storage_types::StorageDiff;
+use mz_storage_types::sources::SourceData;
+use timely::progress::{Antichain, Timestamp};
+
+/// TODO(upsert-in-persist), make this configurable.
+///
+/// 20 MiB
+const LRU_CACHE_SIZE: usize = 20 * 1024 * 1024;
+
+pub struct KeyValueReadHandle<T> {
+    handle: ReadHandle<SourceData, (), T, StorageDiff>,
+    cache: Lru<Row, Row>,
+    cache_upper: T,
+}
+
+impl<T> KeyValueReadHandle<T>
+where
+    T: Timestamp + Lattice + Codec64 + Sync + TimestampManipulation,
+{
+    pub fn new(handle: ReadHandle<SourceData, (), T, StorageDiff>) -> Self {
+        Self {
+            handle,
+            cache: Lru::new(LRU_CACHE_SIZE, move |_, _, _| {}),
+            cache_upper: T::minimum(),
+        }
+    }
+
+    pub async fn get_multi(
+        &mut self,
+        key_columns: &[usize],
+        keys: Vec<Row>,
+        ts: T,
+    ) -> Vec<(Row, Row)> {
+        // Check the cache first.
+        // let mut cached_values = Vec::new();
+        // let mut obtained_keys = BTreeSet::new();
+        // for key in &keys {
+        //     if let Some((_key, cached_val)) = self.cache.get(key) {
+        //         tracing::info!(?key, ?ts, ?self.cache_upper, "found in cache");
+        //         // if self.cache_upper.step_back().unwrap_or(T::minimum()) == ts {
+        //         if self.cache_upper == ts {
+        //             tracing::info!(?key, ?ts, "returning cached value");
+        //             cached_values.push((key.clone(), cached_val.clone()));
+        //             obtained_keys.insert(key.clone());
+        //         }
+        //     }
+        // }
+
+        // // Skip querying for keys that have already been obtained from the cache.
+        // let keys: Vec<_> = keys
+        //     .into_iter()
+        //     .filter(|key| !obtained_keys.contains(key))
+        //     .collect();
+
+        // // If there isn't anything else to obtain then return early!
+        // if keys.is_empty() {
+        //     // return Vec::new();
+        //     return cached_values;
+        // }
+
+        let as_of = Antichain::from_elem(ts);
+        let batch_parts = self.handle.snapshot(as_of).await.expect("OH NO");
+        assert_eq!(key_columns.len(), 1, "support composite keys");
+        let key_col = key_columns[0];
+
+        // We should fetch a RowGroup in a Part if it contains any of our keys.
+        let mut datum_vec_a = DatumVec::new();
+        let mut encode_buffer = Vec::new();
+        let mut should_fetch = |bloom_filter: &BloomFilter| {
+            keys.iter().any(|row| {
+                let datums = datum_vec_a.borrow_with(row);
+                assert_eq!(datums.len(), 1, "composite keys");
+                let key = datums[0];
+                let contains = bloom_filter.contains(key, &mut encode_buffer);
+                // let contains = true;
+                if contains {
+                    tracing::info!("matched bloom filter for key {key}");
+                } else {
+                    tracing::info!("did not match bloom filter for key {key}");
+                }
+                contains
+            })
+        };
+
+        let mut filtered_values = Vec::new();
+        let mut datum_vec_b = DatumVec::new();
+        let mut datum_vec_c = DatumVec::new();
+
+        // TODO(upsert-in-persist)
+        //
+        // There are two more things we can do to make this faster:
+        // 1. Check the statistics for the primary key column on each `Part`
+        //    before even looking at the bloom filters.
+        // 2. Sort the list of `Part`s by their upper in a descending order.
+        //    If a key exists in an upsert source it is guaranteed to have a
+        //    diff of 1 or -1, meaning the latest value for a key is the
+        //    correct value. There is no need to scan for all previous
+        //    instances of the key and then consolidate.
+
+        for part in batch_parts {
+            // Check if this part could possibly match any of the bloom filters.
+            // let could_match = if let Some(mut bloom_filters) = part.pkey_bloom_filters() {
+            //     bloom_filters.any(|filter| should_fetch(filter))
+            // } else {
+            //     // If there are no bloom filters then we could always match.
+            //     true
+            // };
+
+            // If we can't possibly match this part then don't fetch it.
+            // if !could_match {
+            //     continue;
+            // }
+
+            let values = self.handle.fetch_values(&part, &mut should_fetch).await;
+            for ((source_data, _unit_type), _ts, diff) in values {
+                let source_data = source_data.expect("HACK WEEK");
+                let candidate_row = source_data.0.expect("HACK WEEK");
+
+                let maybe_matching_key = {
+                    let candidate_datums = datum_vec_b.borrow_with(&candidate_row);
+                    keys.iter().find(|wanted_row| {
+                        let wanted_datums = datum_vec_c.borrow_with(wanted_row);
+                        assert_eq!(wanted_datums.len(), 1, "composite keys");
+                        wanted_datums[0] == candidate_datums[key_col]
+                    })
+                };
+                if let Some(matching_key) = maybe_matching_key {
+                    tracing::debug!(?matching_key, ?candidate_row, ?diff, "found matching key");
+                    filtered_values.push(((matching_key.clone(), candidate_row), diff));
+                }
+            }
+        }
+
+        differential_dataflow::consolidation::consolidate(&mut filtered_values);
+        if !filtered_values.iter().all(|(_x, diff)| *diff == 1) {
+            tracing::warn!("filtered values: {filtered_values:?}");
+        }
+        // filtered_values.iter().all(|(_x, diff)| *diff == 1);
+
+        filtered_values
+            .into_iter()
+            .map(|(payload, _diff)| payload)
+            // .chain(cached_values)
+            .collect()
+    }
+
+    pub fn apply_changes(&mut self, mut changes: Vec<((Row, Row), T, StorageDiff)>, upper: T) {
+        changes.sort_by(|a, b| (a.1.clone(), a.2).cmp(&(b.1.clone(), b.2)));
+        self.cache_upper = upper;
+
+        for ((key, val), _ts, diff) in changes {
+            if diff == 1 {
+                let weight = val.byte_len();
+                self.cache.insert(key, val, weight);
+            } else if diff == -1 {
+                self.cache.remove(&key);
+            } else {
+                panic!("unexpected diff value {diff}");
+            }
+        }
+    }
+}