columnar: support late materialization

contrib/tiflash-columnar-hub/Cargo.toml

@@ -1,5 +1,5 @@
 [workspace]
-members = ["hub-runtime"]
+members = ["gen-proxy-ffi", "hub-runtime"]
 resolver = "2"
 
 [workspace.dependencies]

contrib/tiflash-columnar-hub/Makefile

@@ -26,7 +26,7 @@ ifneq ($(strip $(ENABLE_FEATURES)),)
 FEATURE_ARGS += --features $(ENABLE_FEATURES)
 endif
 
-.PHONY: debug release build clean
+.PHONY: debug release build clean gen_proxy_ffi
 
 build:
 	cargo build --locked --manifest-path hub-runtime/Cargo.toml $(PROFILE_FLAG) $(FEATURE_ARGS)
@@ -41,3 +41,6 @@ release:
 
 clean:
 	cargo clean --manifest-path hub-runtime/Cargo.toml
+
+gen_proxy_ffi:
+	./gen-proxy-ffi.sh

contrib/tiflash-columnar-hub/gen-proxy-ffi.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+hub_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+if [[ ! -f "${hub_dir}/Cargo.toml" ]]; then
+    echo "Cannot find the Hub Cargo workspace." >&2
+    exit 1
+fi
+
+cd "${hub_dir}"
+cargo run --locked --package gen-proxy-ffi

contrib/tiflash-columnar-hub/gen-proxy-ffi/Cargo.toml

@@ -0,0 +1,5 @@
+[package]
+name = "gen-proxy-ffi"
+version = "0.1.0"
+edition = "2021"
+publish = false

contrib/tiflash-columnar-hub/gen-proxy-ffi/src/main.rs

@@ -0,0 +1,74 @@
+use std::{
+    collections::hash_map::DefaultHasher,
+    ffi::OsStr,
+    fs,
+    hash::{Hash, Hasher},
+    path::{Path, PathBuf},
+};
+
+const RAFT_STORE_PROXY_VERSION_PREFIX: &str = "RAFT_STORE_PROXY_VERSION";
+
+fn collect_headers(dir: &Path, headers: &mut Vec<PathBuf>) {
+    for entry in fs::read_dir(dir).expect("Couldn't read FFI directory") {
+        let path = entry.expect("Couldn't read FFI directory entry").path();
+        if path.is_dir() {
+            collect_headers(&path, headers);
+        } else if path.extension() == Some(OsStr::new("h")) {
+            headers.push(path);
+        }
+    }
+}
+
+fn ffi_version(ffi_dir: &Path) -> u64 {
+    let mut headers = Vec::new();
+    collect_headers(ffi_dir, &mut headers);
+    headers.sort();
+
+    let mut hasher = DefaultHasher::new();
+    for header in headers {
+        fs::read_to_string(header)
+            .expect("Couldn't read FFI header")
+            .hash(&mut hasher);
+    }
+    hasher.finish()
+}
+
+fn replace_version(path: &Path, replacement: &str) {
+    let content = fs::read_to_string(path).expect("Couldn't read generated FFI bindings");
+    assert_eq!(content.matches(RAFT_STORE_PROXY_VERSION_PREFIX).count(), 1);
+    let declaration = format!("pub const {}: u64 = ", RAFT_STORE_PROXY_VERSION_PREFIX);
+    let start = content
+        .find(&declaration)
+        .expect("Couldn't find RAFT_STORE_PROXY_VERSION");
+    let end = content[start..]
+        .find(';')
+        .expect("Couldn't find RAFT_STORE_PROXY_VERSION terminator")
+        + start
+        + 1;
+    let mut updated = content;
+    updated.replace_range(start..end, replacement);
+    fs::write(path, updated).expect("Couldn't write generated FFI bindings");
+}
+
+fn main() {
+    let hub_dir = Path::new(env!("CARGO_MANIFEST_DIR")).parent().unwrap();
+    let ffi_dir = hub_dir.join("hub-runtime/ffi/src/RaftStoreProxyFFI");
+    let version = ffi_version(&ffi_dir);
+
+    fs::write(
+        ffi_dir.join("@version"),
+        format!(
+            "#pragma once\n#include <cstdint>\nnamespace DB {{ constexpr uint64_t {} = {}ull; }}",
+            RAFT_STORE_PROXY_VERSION_PREFIX, version
+        ),
+    )
+    .expect("Couldn't write FFI version header");
+    // The Hub bindings include manually maintained declarations, so update only the ABI fingerprint.
+    replace_version(
+        &hub_dir.join("hub-runtime/src/interfaces.rs"),
+        &format!(
+            "pub const {}: u64 = {};",
+            RAFT_STORE_PROXY_VERSION_PREFIX, version
+        ),
+    );
+}

contrib/tiflash-columnar-hub/hub-runtime/ffi/src/RaftStoreProxyFFI/@version

@@ -1,3 +1,3 @@
 #pragma once
 #include <cstdint>
-namespace DB { constexpr uint64_t RAFT_STORE_PROXY_VERSION = 5493270813306750334ull; }
+namespace DB { constexpr uint64_t RAFT_STORE_PROXY_VERSION = 5429784359048998305ull; }

contrib/tiflash-columnar-hub/hub-runtime/ffi/src/RaftStoreProxyFFI/ProxyFFI.h

@@ -225,6 +225,8 @@ struct ColumnarScanStats {
   uint64_t rough_check_unknown_packs;
   uint64_t remote_segments;
   uint64_t total_segments;
+  uint64_t lm_late_packs_loaded;
+  uint64_t lm_late_rows_gathered;
 };
 
 struct SSTReaderInterfaces {
@@ -261,6 +263,32 @@ struct CloudStorageEngineInterfaces {
   ColumnarScanStats (*fn_columnar_scan_stats)(ColumnarReaderPtr);
 };
 
+// Optional extension for the two-phase columnar reader protocol. This is
+// intentionally discovered through an independently exported, weak symbol so
+// that extending it cannot change RaftStoreProxyFFIHelper's ABI.
+struct ColumnarLateMaterializationInterfaces {
+  uint32_t version;
+  uint32_t size;
+  uint64_t (*fn_read_early_block)(ColumnarReaderPtr, uint64_t, BaseBuffView,
+                                  uint64_t *, int64_t *);
+  RustStrWithView (*fn_read_early_column)(ColumnarReaderPtr, uint64_t,
+                                          int64_t);
+  uint64_t (*fn_materialize_selected)(ColumnarReaderPtr, uint64_t, uint8_t,
+                                      BaseBuffView);
+  RustStrWithView (*fn_read_late_column)(ColumnarReaderPtr, uint64_t,
+                                         int64_t);
+  uint8_t (*fn_finish_materialized_block)(ColumnarReaderPtr, uint64_t);
+  uint8_t (*fn_discard_late_materialization_batch)(ColumnarReaderPtr,
+                                                   uint64_t);
+  uint8_t (*fn_is_late_materialization_supported)(ColumnarReaderPtr,
+                                                   BaseBuffView);
+};
+
+#if defined(__GNUC__) || defined(__clang__)
+extern "C" const ColumnarLateMaterializationInterfaces *
+    tiflash_columnar_get_late_materialization_interfaces() __attribute__((weak));
+#endif
+
 enum class MsgPBType : uint32_t {
   ReadIndexResponse = 0,
   ServerInfoResponse,