diff --git a/cpp/velox/jni/JniHashTable.cc b/cpp/velox/jni/JniHashTable.cc index 7db693e3d00..4bf48d8143a 100644 --- a/cpp/velox/jni/JniHashTable.cc +++ b/cpp/velox/jni/JniHashTable.cc @@ -45,7 +45,11 @@ void JniHashTableContext::finalize(JNIEnv* env) { } } -jlong JniHashTableContext::callJavaGet(const std::string& id) const { +std::optional JniHashTableContext::callJavaGet(const std::string& id) const { + if (vm_ == nullptr) { + return std::nullopt; + } + JNIEnv* env; if (vm_->GetEnv(reinterpret_cast(&env), jniVersion) != JNI_OK) { throw gluten::GlutenException("JNIEnv was not attached to current thread"); @@ -53,6 +57,8 @@ jlong JniHashTableContext::callJavaGet(const std::string& id) const { const jstring s = env->NewStringUTF(id.c_str()); auto result = env->CallStaticLongMethod(jniVeloxBroadcastBuildSideCache_, jniGet_, s); + env->DeleteLocalRef(s); + checkException(env); return result; } @@ -161,7 +167,18 @@ std::shared_ptr nativeHashTableBuild( } long getJoin(const std::string& hashTableId) { - return JniHashTableContext::getInstance().callJavaGet(hashTableId); + const auto handle = JniHashTableContext::getInstance().callJavaGet(hashTableId); + if (!handle.has_value()) { + // The JVM side VeloxBroadcastBuildSideCache holding the pre-built hash tables is unreachable, + // e.g. in the standalone micro benchmark. Report a miss, the same way that cache reports one, + // and let the caller build the hash table from the join's build side input instead. + LOG(WARNING) << "No JVM is attached to this process, cannot look up the pre-built hash table for cache key: " + << hashTableId + << ". A new hash table will be built from the build side input, which is empty unless the stage was " + "dumped with spark.gluten.velox.buildHashTableOncePerExecutor.enabled=false."; + return 0; + } + return handle.value(); } size_t serializedHashTableSize(std::shared_ptr builder) { diff --git a/cpp/velox/jni/JniHashTable.h b/cpp/velox/jni/JniHashTable.h index 48bed2b4f00..531b8c864a8 100644 --- a/cpp/velox/jni/JniHashTable.h +++ b/cpp/velox/jni/JniHashTable.h @@ -18,6 +18,7 @@ #pragma once #include +#include #include "memory/ColumnarBatch.h" #include "memory/VeloxMemoryManager.h" #include "operators/hashjoin/HashTableBuilder.h" @@ -52,7 +53,11 @@ class JniHashTableContext { return hashTableObjStore_.get(); } - jlong callJavaGet(const std::string& id) const; + // Returns the handle registered under the given id by the JVM side + // VeloxBroadcastBuildSideCache, or std::nullopt when that cache cannot be reached at all because + // no JVM is attached to this process. The latter is the case for the standalone micro benchmark + // and the native unit tests, where JNI_OnLoad never runs. + std::optional callJavaGet(const std::string& id) const; private: JniHashTableContext() : hashTableObjStore_(ObjectStore::create()) {} @@ -89,6 +94,8 @@ std::shared_ptr nativeHashTableBuild( std::vector>& batches, std::shared_ptr memoryPool); +// Returns the handle of the pre-built hash table registered for the given id, or 0 if there is +// none. Safe to call from a process with no JVM attached, in which case it always returns 0. long getJoin(const std::string& hashTableId); // Return the exact serialized hash table size for direct buffer allocation. diff --git a/cpp/velox/operators/writer/VeloxColumnarBatchWriter.cc b/cpp/velox/operators/writer/VeloxColumnarBatchWriter.cc index 598199f19b3..d5cad87905b 100644 --- a/cpp/velox/operators/writer/VeloxColumnarBatchWriter.cc +++ b/cpp/velox/operators/writer/VeloxColumnarBatchWriter.cc @@ -57,6 +57,11 @@ arrow::Status VeloxColumnarBatchWriter::write(const std::shared_ptrclose(); return arrow::Status::OK(); } diff --git a/cpp/velox/tests/CMakeLists.txt b/cpp/velox/tests/CMakeLists.txt index 87331de818c..add195804d4 100644 --- a/cpp/velox/tests/CMakeLists.txt +++ b/cpp/velox/tests/CMakeLists.txt @@ -142,6 +142,9 @@ if(ENABLE_S3) endif() add_velox_test(scoped_timer_test SOURCES ScopedTimerTest.cc) add_velox_test(row_based_checksum_test SOURCES RowBasedChecksumTest.cc) +add_velox_test(jni_hash_table_test SOURCES JniHashTableTest.cc) +add_velox_test(velox_whole_stage_dumper_test SOURCES + VeloxWholeStageDumperTest.cc) if(BUILD_EXAMPLES) add_velox_test(my_udf_test SOURCES MyUdfTest.cc) endif() diff --git a/cpp/velox/tests/JniHashTableTest.cc b/cpp/velox/tests/JniHashTableTest.cc new file mode 100644 index 00000000000..aa57339de54 --- /dev/null +++ b/cpp/velox/tests/JniHashTableTest.cc @@ -0,0 +1,50 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "jni/JniHashTable.h" + +#include + +namespace gluten { + +// This test binary, like the micro benchmark, runs without a JVM. JNI_OnLoad is therefore never +// invoked and JniHashTableContext::vm_ stays null. +// +// Looking up the build side of a broadcast hash join used to dereference that null JavaVM +// unconditionally, so converting any substrait plan carrying a non-empty `hashTableId` crashed the +// process. See https://github.com/apache/gluten/issues/12504. +class JniHashTableTest : public ::testing::Test { + protected: + void SetUp() override { + ASSERT_EQ(JniHashTableContext::getInstance().getJavaVM(), nullptr) + << "This test asserts behaviour of a process with no JVM attached"; + } +}; + +TEST_F(JniHashTableTest, getJoinReportsMissWithoutJvm) { + // Must report a miss rather than segfault. SubstraitToVeloxPlanConverter treats 0 as "no + // pre-built table" and falls back to building one from the build side input. + EXPECT_EQ(getJoin("no-such-broadcast-hash-table-id"), 0); +} + +TEST_F(JniHashTableTest, callJavaGetReportsUnreachableCacheWithoutJvm) { + // Keeps "the cache says there is no such table" distinguishable from "the cache cannot be + // reached at all", rather than collapsing both into a handle of 0 at this level. + EXPECT_FALSE(JniHashTableContext::getInstance().callJavaGet("any-id").has_value()); +} + +} // namespace gluten diff --git a/cpp/velox/tests/VeloxWholeStageDumperTest.cc b/cpp/velox/tests/VeloxWholeStageDumperTest.cc new file mode 100644 index 00000000000..3be3aed448f --- /dev/null +++ b/cpp/velox/tests/VeloxWholeStageDumperTest.cc @@ -0,0 +1,117 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "utils/VeloxWholeStageDumper.h" + +#include +#include + +#include +#include + +#include "memory/VeloxColumnarBatch.h" +#include "velox/common/file/FileSystems.h" +#include "velox/exec/tests/utils/TempDirectoryPath.h" +#include "velox/vector/tests/utils/VectorTestBase.h" + +using namespace facebook::velox; + +namespace gluten { +namespace { +class ColumnarBatchArray final : public ColumnarBatchIterator { + public: + explicit ColumnarBatchArray(std::vector> batches) : batches_(std::move(batches)) {} + + std::shared_ptr next() override { + if (cursor_ >= batches_.size()) { + return nullptr; + } + return batches_[cursor_++]; + } + + private: + const std::vector> batches_; + size_t cursor_{0}; +}; +} // namespace + +class VeloxWholeStageDumperTest : public ::testing::Test, public test::VectorTestBase { + protected: + static void SetUpTestCase() { + memory::MemoryManager::testingSetInstance(memory::MemoryManager::Options{}); + filesystems::registerLocalFileSystem(); + } + + std::unique_ptr makeDumper() { + return std::make_unique(taskInfo_, tmpDir_->getPath(), 4096, rootPool_.get()); + } + + std::filesystem::path dataFile(int32_t iteratorIndex) const { + return std::filesystem::path{tmpDir_->getPath()} / + fmt::format("data_{}_{}_{}_{}.parquet", taskInfo_.stageId, taskInfo_.partitionId, taskInfo_.vId, iteratorIndex); + } + + std::shared_ptr newBatch(int32_t numRows) { + auto rowVector = makeRowVector({makeFlatVector(numRows, [](auto row) { return row; })}); + return std::make_shared(std::move(rowVector)); + } + + SparkTaskInfo taskInfo_{.stageId = 1, .partitionId = 2, .taskId = 3, .vId = 4}; + std::shared_ptr tmpDir_{exec::test::TempDirectoryPath::create()}; +}; + +// An input iterator that yields nothing used to leave the parquet writer uninitialized, since the +// writer takes its schema from the first batch, and closing it dereferenced a null writer. The +// build side of a broadcast hash join is exactly such an iterator: its hash table is handed to the +// native operator through a process local cache rather than streamed. +// See https://github.com/apache/gluten/issues/12504. +TEST_F(VeloxWholeStageDumperTest, dumpEmptyInputIterator) { + auto dumper = makeDumper(); + auto input = std::make_shared(std::vector>{}); + + std::shared_ptr dumped; + ASSERT_NO_THROW(dumped = dumper->dumpInputIterator(0, input)); + + // The task keeps running against what the original iterator produced: nothing. + ASSERT_NE(dumped, nullptr); + ASSERT_EQ(dumped->next(), nullptr); + + // No parquet file can be written without a schema, so the gap is recorded explicitly. --data + // binds files to iterator indexes by position, so an unrecorded gap would silently shift every + // later iterator onto the wrong input. + ASSERT_FALSE(std::filesystem::exists(dataFile(0))); + ASSERT_TRUE(std::filesystem::exists(dataFile(0).string() + ".empty")); +} + +TEST_F(VeloxWholeStageDumperTest, dumpInputIterator) { + auto dumper = makeDumper(); + auto input = + std::make_shared(std::vector>{newBatch(10), newBatch(20)}); + + auto dumped = dumper->dumpInputIterator(0, input); + ASSERT_NE(dumped, nullptr); + ASSERT_TRUE(std::filesystem::exists(dataFile(0))); + + // The returned iterator replays the dumped file, so the task sees the same rows. + int32_t numRows = 0; + while (auto batch = dumped->next()) { + numRows += batch->numRows(); + } + ASSERT_EQ(numRows, 30); +} + +} // namespace gluten diff --git a/cpp/velox/utils/VeloxWholeStageDumper.cc b/cpp/velox/utils/VeloxWholeStageDumper.cc index d8a451dac73..875bb453e69 100644 --- a/cpp/velox/utils/VeloxWholeStageDumper.cc +++ b/cpp/velox/utils/VeloxWholeStageDumper.cc @@ -52,6 +52,15 @@ void dumpToStorage(const std::string& saveDir, const std::string& fileName, cons outFile << content; outFile.close(); } + +// Stands in for an input iterator that turned out to be empty, for which no parquet file was +// written. Keeps the running task seeing exactly what the original iterator produced: nothing. +class EmptyColumnarBatchIterator final : public ColumnarBatchIterator { + public: + std::shared_ptr next() override { + return nullptr; + } +}; } // namespace VeloxWholeStageDumper::VeloxWholeStageDumper( @@ -117,11 +126,38 @@ std::shared_ptr VeloxWholeStageDumper::dumpInputIterator( auto writer = std::make_shared( dumpPath, batchSize_, pool_->addAggregateChild(fmt::format("dump_iterator.{}", iteratorIndex))); + bool wroteBatch = false; while (auto cb = inputIterator->next()) { GLUTEN_THROW_NOT_OK(writer->write(cb)); + wroteBatch = true; } GLUTEN_THROW_NOT_OK(writer->close()); + if (!wroteBatch) { + // The writer takes its schema from the first batch, so an iterator that yielded nothing leaves + // no parquet file behind. Record the gap explicitly: the benchmark binds the files passed to + // --data to iterator indexes by position, so a silently absent file shifts every later + // iterator onto the wrong input. The marker keeps the gap visible in the dump directory, which + // is where whoever replays the stage is looking. + // + // The usual cause is the build side of a broadcast hash join, which is handed to the native + // operator through a process local hash table cache rather than streamed. See + // docs/developers/MicroBenchmarks.md. + dumpToStorage( + saveDir_, + fileName + ".empty", + fmt::format( + "Input iterator {} of {} produced no batches, so {} was not written.\n" + "If this is the build side of a broadcast hash join, re-dump with\n" + "spark.gluten.velox.buildHashTableOncePerExecutor.enabled=false to capture it.\n", + iteratorIndex, + taskInfo_.toString(), + fileName)); + LOG(WARNING) << "Input iterator " << iteratorIndex << " of " << taskInfo_ << " produced no batches, so no " + << fileName << " was written. Left a " << fileName << ".empty marker instead."; + return std::make_shared(); + } + // Velox parquet reader requires leaf memory pool. return std::make_shared( dumpPath, batchSize_, pool_->addLeafChild(fmt::format("retrieve_iterator.{}", iteratorIndex))); diff --git a/docs/developers/MicroBenchmarks.md b/docs/developers/MicroBenchmarks.md index ce0996a6cc2..bbb7de9d9bb 100644 --- a/docs/developers/MicroBenchmarks.md +++ b/docs/developers/MicroBenchmarks.md @@ -139,7 +139,8 @@ or 4 types of dumped file: file splits. - Data file(optional): Parquet formatted, file name `data_{stageId}_{partitionId}_{vId}_{iteratorIdx}.parquet`. If the first stage contains one or - more BHJ operators, there can be one or more input data files. The input data files of a first + more BHJ operators, there can be one or more input data files, see + [Broadcast hash join](#broadcast-hash-join). The input data files of a first stage will be loaded as iterators to serve as the inputs for the pipeline: ``` @@ -152,6 +153,38 @@ or 4 types of dumped file: } ``` +### Broadcast hash join + +By default the build side of a BHJ never reaches the native operator as data. Its hash table is +built once per executor (or on the driver and broadcast in serialized form) and handed to the join +through a process-local cache keyed by a hash table id, so the build side input iterator yields no +rows. The standalone benchmark has no such cache, so it would replay the stage with an empty build +side and report timings for a join that produces no output. + +To dump a stage containing a BHJ, add this config to the dump run: + +``` +--conf spark.gluten.velox.buildHashTableOncePerExecutor.enabled=false +``` + +This makes the build side stream to the native operator instead of being cached, so it should be +dumped as its own `data_{stageId}_{partitionId}_{vId}_{iteratorIdx}.parquet` and the benchmark +should replay the join against real data. Note that the replayed stage builds the hash table +itself, so that time is included in the measurement, unlike the production run where the table is +built once per executor. + +> **This recipe has not been verified end to end.** It follows from how the build side relation is +> constructed, but no one has yet dumped a BHJ stage this way and replayed it. Treat it as a +> starting point rather than a known-good workflow, and please report back on +> [#12504](https://github.com/apache/gluten/issues/12504). + +If the stage is dumped without that config, the build side iterator is empty and no parquet file +can be written for it, since the writer takes its schema from the first batch. The dump leaves a +`data_{stageId}_{partitionId}_{vId}_{iteratorIdx}.parquet.empty` marker in its place recording +which iterator was empty. Mind the gap when passing `--data`: files are bound to iterator indexes +by position, so a missing file shifts every later iterator onto the wrong input. The benchmark also +warns that no JVM is attached and that a new hash table will be built from the build side input. + Run benchmark. By default, the result will be printed to stdout. You can use `--noprint-result` to suppress this output.