feat(cudf): Serve KvikIO reads through AsyncDataCache

CMake/resolve_dependency_modules/cudf.cmake

@@ -40,12 +40,12 @@ set(
 set(VELOX_rmm_SOURCE_URL "https://github.com/rapidsai/rmm/archive/${VELOX_rmm_COMMIT}.tar.gz")
 velox_resolve_dependency_url(rmm)
 
-# kvikio commit 257d08b from 2026-10-01 (main branch)
+# kvikio commit 4114a21 from 2026-10-07 (main branch)
 set(VELOX_kvikio_VERSION 26.12)
-set(VELOX_kvikio_COMMIT 257d08b41ef4e566cba6f5f64cb365e681961a3b)
+set(VELOX_kvikio_COMMIT 4114a218184956f4bb652f2028980269d709a7cd)
 set(
   VELOX_kvikio_BUILD_SHA256_CHECKSUM
-  15117616e0f7228e2324e0e64c21f3fa9305ffbe6b92d842f967e9b2ceb47bfc
+  161e986cc0e253c32cdc196c18bc2cdeeae51347b672aee7616a572691b6998a
 )
 set(
   VELOX_kvikio_SOURCE_URL

velox/common/caching/AsyncDataCache.cpp

@@ -251,7 +251,8 @@ std::unique_ptr<AsyncDataCacheEntry> CacheShard::getFreeEntryLocked() {
 std::optional<CachePin> CacheShard::lookupLocked(
     RawFileCacheKey key,
     uint64_t size,
-    folly::SemiFuture<bool>* wait) {
+    folly::SemiFuture<bool>* wait,
+    CacheEntrySizePolicy sizePolicy) {
   ++eventCounter_;
   auto it = entryMap_.find(key);
   if (it == entryMap_.end()) {
@@ -266,7 +267,8 @@ std::optional<CachePin> CacheShard::lookupLocked(
     return CachePin{};
   }
   // size=0 (from find()) always passes since entry size is non-negative.
-  if (foundEntry->size() < size) {
+  if (sizePolicy == CacheEntrySizePolicy::kRequireAtLeast &&
+      foundEntry->size() < size) {
     // This can happen if different load quanta apply to access via different
     // connectors. This is not an error but still worth logging.
     VELOX_CACHE_LOG_EVERY_MS(WARNING, 1'000)
@@ -296,11 +298,12 @@ CachePin CacheShard::findOrCreate(
     RawFileCacheKey key,
     uint64_t size,
     bool contiguous,
-    folly::SemiFuture<bool>* wait) {
+    folly::SemiFuture<bool>* wait,
+    CacheEntrySizePolicy sizePolicy) {
   AsyncDataCacheEntry* entryToInit = nullptr;
   {
     std::lock_guard<std::mutex> l(mutex_);
-    auto result = lookupLocked(key, size, wait);
+    auto result = lookupLocked(key, size, wait, sizePolicy);
     if (result.has_value()) {
       return std::move(result.value());
     }
@@ -874,9 +877,10 @@ CachePin AsyncDataCache::findOrCreate(
     RawFileCacheKey key,
     uint64_t size,
     bool contiguous,
-    folly::SemiFuture<bool>* wait) {
+    folly::SemiFuture<bool>* wait,
+    CacheEntrySizePolicy sizePolicy) {
   const int shard = std::hash<RawFileCacheKey>()(key) & shardMask_;
-  return shards_[shard]->findOrCreate(key, size, contiguous, wait);
+  return shards_[shard]->findOrCreate(key, size, contiguous, wait, sizePolicy);
 }
 
 std::optional<CachePin> AsyncDataCache::find(

velox/common/caching/AsyncDataCache.h

@@ -120,6 +120,13 @@ struct RawFileCacheKey {
     return offset == other.offset && fileNum == other.fileNum;
   }
 };
+
+/// Controls whether an entry that begins at the requested key but is shorter
+/// than the requested size is stale or is a reusable prefix.
+enum class CacheEntrySizePolicy : uint8_t {
+  kRequireAtLeast,
+  kAllowSmaller,
+};
 } // namespace facebook::velox::cache
 
 namespace std {
@@ -625,12 +632,15 @@ class CacheShard {
       : cache_(cache), maxWriteRatio_(maxWriteRatio) {}
 
   /// See AsyncDataCache::findOrCreate. If 'contiguous' is true, the
-  /// entry's data is allocated as a single contiguous region.
+  /// entry's data is allocated as a single contiguous region. 'sizePolicy'
+  /// decides how an existing entry at 'key' that is shorter than 'size' is
+  /// treated. See AsyncDataCache::findOrCreate for details.
   CachePin findOrCreate(
       RawFileCacheKey key,
       uint64_t size,
       bool contiguous = false,
-      folly::SemiFuture<bool>* readyFuture = nullptr);
+      folly::SemiFuture<bool>* readyFuture = nullptr,
+      CacheEntrySizePolicy sizePolicy = CacheEntrySizePolicy::kRequireAtLeast);
 
   /// Finds a cache entry for 'key'. Returns a shared-mode pin if the entry
   /// exists and is not exclusive. Returns an empty pin (inside optional) if
@@ -731,7 +741,8 @@ class CacheShard {
   std::optional<CachePin> lookupLocked(
       RawFileCacheKey key,
       uint64_t size,
-      folly::SemiFuture<bool>* waitFuture);
+      folly::SemiFuture<bool>* waitFuture,
+      CacheEntrySizePolicy sizePolicy = CacheEntrySizePolicy::kRequireAtLeast);
 
   void tryAddFreeEntry(std::unique_ptr<AsyncDataCacheEntry>&& entry);
 
@@ -895,11 +906,22 @@ class AsyncDataCache : public memory::Cache {
   /// When the future is realized, the caller may retry findOrCreate().
   /// runtime error with code kNoCacheSpace if there is no space to
   /// create the new entry after evicting any unpinned content.
+  ///
+  /// 'sizePolicy' decides how an existing shared entry at 'key' that is
+  /// shorter than 'size' is treated. With kRequireAtLeast (the default) the
+  /// entry is stale: it is evicted and a new exclusive entry of 'size' is
+  /// created for the caller to fill. With kAllowSmaller the shorter entry is
+  /// returned as a valid prefix and counted as a hit. The cache does not
+  /// track the uncovered suffix. The caller must read the returned entry's
+  /// size(), and request [key.offset + size(), key.offset + 'size') under a
+  /// separate key, which may in turn hit, miss, or return another shorter
+  /// prefix.
   CachePin findOrCreate(
       RawFileCacheKey key,
       uint64_t size,
       bool contiguous = false,
-      folly::SemiFuture<bool>* waitFuture = nullptr);
+      folly::SemiFuture<bool>* waitFuture = nullptr,
+      CacheEntrySizePolicy sizePolicy = CacheEntrySizePolicy::kRequireAtLeast);
 
   /// Finds a cache entry for 'key'. Returns a shared-mode pin if the entry
   /// exists and is not exclusive. Returns an empty pin (inside optional) if
@@ -978,9 +1000,11 @@ class AsyncDataCache : public memory::Cache {
       const std::vector<RawFileCacheKey>& keys,
       const SizeFunc& sizeFunc,
       const ProcessPin& processPin,
-      bool contiguous = false) {
+      bool contiguous = false,
+      CacheEntrySizePolicy sizePolicy = CacheEntrySizePolicy::kRequireAtLeast) {
     for (size_t i = 0; i < keys.size(); ++i) {
-      auto pin = findOrCreate(keys[i], sizeFunc(i), contiguous);
+      auto pin =
+          findOrCreate(keys[i], sizeFunc(i), contiguous, nullptr, sizePolicy);
       if (pin.empty() || pin.checkedEntry()->isShared()) {
         continue;
       }

velox/common/caching/tests/AsyncDataCacheTest.cpp

@@ -1291,17 +1291,40 @@ TEST_P(AsyncDataCacheTest, staleEntry) {
   ASSERT_EQ(stats.numEntries, 1);
   ASSERT_EQ(stats.numHit, 1);
 
+  // A caller that can assemble one read from several consecutive cache entries
+  // may consume this entry as a valid prefix instead of invalidating and
+  // reloading it.
+  const auto beforePrefix = cache_->refreshStats();
+  auto prefixPin = cache_->findOrCreate(
+      key,
+      2 * size,
+      /*contiguous=*/false,
+      &wait,
+      CacheEntrySizePolicy::kAllowSmaller);
+  ASSERT_FALSE(prefixPin.empty());
+  ASSERT_TRUE(wait.isReady());
+  ASSERT_FALSE(prefixPin.entry()->isExclusive());
+  ASSERT_EQ(prefixPin.entry()->size(), size);
+  prefixPin.clear();
+  stats = cache_->refreshStats();
+  ASSERT_EQ(stats.numEntries, 1);
+  auto delta = stats - beforePrefix;
+  ASSERT_EQ(delta.numStales, 0);
+  ASSERT_EQ(delta.numHit, 1);
+
   // Stale cache access with large cache size.
+  const auto beforeStale = cache_->refreshStats();
   auto stalePin =
       cache_->findOrCreate(key, 2 * size, /*contiguous=*/false, &wait);
   ASSERT_FALSE(stalePin.empty());
   ASSERT_TRUE(wait.isReady());
   ASSERT_TRUE(stalePin.entry()->isExclusive());
   stalePin.entry()->setExclusiveToShared();
   stats = cache_->refreshStats();
-  ASSERT_EQ(stats.numStales, 1);
   ASSERT_EQ(stats.numEntries, 1);
-  ASSERT_EQ(stats.numHit, 1);
+  delta = stats - beforeStale;
+  ASSERT_EQ(delta.numStales, 1);
+  ASSERT_EQ(delta.numHit, 0);
 }
 
 TEST_P(AsyncDataCacheTest, shrinkCache) {

velox/experimental/cudf/CudfConfig.h

@@ -65,6 +65,8 @@ struct CudfConfig {
   static constexpr const char* kUcxExchangeLogLevel{"cudf.exchange_log_level"};
   static constexpr const char* kUcxPartitionedOutputBatchRows{
       "cudf.partitioned_output_batch_rows"};
+  static constexpr const char* kCudfKvikioCacheCopyThreads{
+      "cudf.kvikio_cache_copy_threads"};
   /// Query session configs for the cuDF Operators.
   static constexpr const char* kCudfTopNBatchSize{"cudf.topk_batch_size"};
 
@@ -121,6 +123,10 @@ struct CudfConfig {
   /// disable accumulation.
   int64_t partitionedOutputBatchRows{10'000};
 
+  /// Threads that copy KvikIO data from the async data cache to the GPU. One
+  /// pool for the whole process, created by the first cached GPU read.
+  int32_t kvikioCacheCopyThreads{16};
+
   /// Memory resource for cuDF.
   /// Possible values are (cuda, pool, async, arena, managed, managed_pool).
   std::string memoryResource{"async"};

velox/experimental/cudf/connectors/hive/CMakeLists.txt

@@ -15,6 +15,7 @@
 add_library(
   velox_cudf_hive_connector
   BufferedInputDataSource.cpp
+  CachingDataSource.cpp
   CudfHiveConfig.cpp
   CudfHiveConnector.cpp
   CudfHiveConnectorSplit.cpp