Use simdjson for binary json construction for improved performance (#…

…10464)
ydb-platform · Oct 21, 2024 · 25311e4 · 25311e4
1 parent ff01ea7
commit 25311e4
Show file tree

Hide file tree

Showing 30 changed files with 557 additions and 317 deletions.
diff --git a/.github/config/muted_ya.txt b/.github/config/muted_ya.txt
@@ -42,6 +42,7 @@ ydb/core/tx/schemeshard/ut_move_reboots TSchemeShardMoveRebootsTest.WithDataAndP
 ydb/core/tx/schemeshard/ut_pq_reboots TPqGroupTestReboots.AlterWithReboots-PQConfigTransactionsAtSchemeShard-false
 ydb/core/tx/schemeshard/ut_restore TImportTests.ShouldSucceedOnManyTables
 ydb/core/tx/schemeshard/ut_split_merge TSchemeShardSplitBySizeTest.Merge1KShards
+ydb/core/tx/tiering/ut ColumnShardTiers.TTLUsage
 ydb/core/tx/tx_proxy/ut_ext_tenant TExtSubDomainTest.CreateTableInsideAndAlterDomainAndTable-AlterDatabaseCreateHiveFirst*
 ydb/core/tx/tx_proxy/ut_storage_tenant TStorageTenantTest.RemoveStoragePoolBeforeDroppingTablet
 ydb/core/util/ut TCircularOperationQueueTest.ShouldShuffle

diff --git a/ydb/core/formats/arrow/accessor/abstract/ya.make b/ydb/core/formats/arrow/accessor/abstract/ya.make
@@ -4,6 +4,7 @@ PEERDIR(
     contrib/libs/apache/arrow
     ydb/library/conclusion
     ydb/services/metadata/abstract
+    ydb/library/actors/core
     ydb/library/formats/arrow/accessor/abstract
     ydb/library/formats/arrow/accessor/common
     ydb/library/formats/arrow/protos

diff --git a/ydb/core/formats/arrow/converter.cpp b/ydb/core/formats/arrow/converter.cpp
@@ -31,8 +31,8 @@ static bool ConvertData(TCell& cell, const NScheme::TTypeInfo& colType, TMemoryP
         }
         case NScheme::NTypeIds::JsonDocument: {
             const auto binaryJson = NBinaryJson::SerializeToBinaryJson(cell.AsBuf());
-            if (!binaryJson.Defined()) {
-                errorMessage = "Invalid JSON for JsonDocument provided";
+            if (binaryJson.IsFail()) {
+                errorMessage = "Invalid JSON for JsonDocument provided: " + binaryJson.GetErrorMessage();
                 return false;
             }
             const auto saved = memPool.AppendString(TStringBuf(binaryJson->Data(), binaryJson->Size()));
@@ -98,8 +98,9 @@ static arrow::Status ConvertColumn(const NScheme::TTypeInfo colType, std::shared
                     }
                 } else {
                     const auto binaryJson = NBinaryJson::SerializeToBinaryJson(valueBuf);
-                    if (!binaryJson.Defined()) {
-                        return arrow::Status::SerializationError("Cannot serialize json: ", valueBuf);
+                    if (binaryJson.IsFail()) {
+                        return arrow::Status::SerializationError(
+                            "Cannot serialize json (", binaryJson.GetErrorMessage(), "): ", valueBuf.SubStr(0, Min(valueBuf.Size(), 1024ul)));
                     }
                     auto appendResult = builder.Append(binaryJson->Data(), binaryJson->Size());
                     if (!appendResult.ok()) {

diff --git a/ydb/core/formats/arrow/ut/ut_arrow.cpp b/ydb/core/formats/arrow/ut/ut_arrow.cpp
@@ -223,7 +223,7 @@ struct TDataRow {
         std::vector<TCell> cells(value.Cells().data(), value.Cells().data() + value.Cells().size());
 
         auto binaryJson = NBinaryJson::SerializeToBinaryJson(TStringBuf(JsonDocument.data(), JsonDocument.size()));
-        UNIT_ASSERT(binaryJson.Defined());
+        UNIT_ASSERT(binaryJson.IsSuccess());
 
         cells[19] = TCell(binaryJson->Data(), binaryJson->Size());
         return TOwnedCellVec(cells);

diff --git a/ydb/core/io_formats/cell_maker/cell_maker.cpp b/ydb/core/io_formats/cell_maker/cell_maker.cpp
@@ -87,8 +87,13 @@ namespace {
             return false;
         }
 
-        result = NBinaryJson::SerializeToBinaryJson(unescaped);
-        return result.Defined();
+        auto serializedJson = NBinaryJson::SerializeToBinaryJson(unescaped);
+        if (serializedJson.IsFail()) {
+            return false;
+        }
+
+        result = serializedJson.DetachResult();
+        return true;
     }
 
     template <>
@@ -395,8 +400,8 @@ bool MakeCell(TCell& cell, const NJson::TJsonValue& value, const NScheme::TTypeI
         case NScheme::NTypeIds::Json:
             return TCellMaker<TString, TStringBuf>::MakeDirect(cell, NFormats::WriteJson(value), pool, err);
         case NScheme::NTypeIds::JsonDocument:
-            if (const auto& result = NBinaryJson::SerializeToBinaryJson(NFormats::WriteJson(value))) {
-                return TCellMaker<TMaybe<NBinaryJson::TBinaryJson>, TStringBuf>::MakeDirect(cell, result, pool, err, &BinaryJsonToStringBuf);
+            if (auto result = NBinaryJson::SerializeToBinaryJson(NFormats::WriteJson(value)); result.IsSuccess()) {
+                return TCellMaker<TMaybe<NBinaryJson::TBinaryJson>, TStringBuf>::MakeDirect(cell, result.DetachResult(), pool, err, &BinaryJsonToStringBuf);
             } else {
                 return false;
             }

diff --git a/ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp b/ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp
@@ -8920,7 +8920,7 @@ Y_UNIT_TEST_SUITE(KqpOlapTypes) {
         testHelper.CreateTable(testTable);
         std::string jsonString = R"({"col1": "val1", "obj": {"obj_col2_int": 16}})";
         auto maybeJsonDoc = NBinaryJson::SerializeToBinaryJson(jsonString);
-        Y_ABORT_UNLESS(maybeJsonDoc.Defined());
+        Y_ABORT_UNLESS(maybeJsonDoc.IsSuccess());
         const std::string jsonBin(maybeJsonDoc->Data(), maybeJsonDoc->Size());
         {
             TTestHelper::TUpdatesBuilder tableInserter(testTable.GetArrowSchema(schema));

diff --git a/ydb/core/tx/columnshard/engines/scheme/defaults/common/ya.make b/ydb/core/tx/columnshard/engines/scheme/defaults/common/ya.make
@@ -9,6 +9,7 @@ PEERDIR(
     contrib/libs/apache/arrow
     ydb/library/conclusion
     ydb/core/scheme_types
+    ydb/library/actors/core
 )
 
 END()
diff --git a/ydb/core/tx/data_events/common/ya.make b/ydb/core/tx/data_events/common/ya.make
@@ -2,6 +2,7 @@ LIBRARY()
 
 PEERDIR(
     ydb/core/protos
+    ydb/library/conclusion
     ydb/library/yql/core/issue/protos
     ydb/public/api/protos
 )

diff --git a/ydb/core/tx/schemeshard/ut_restore/ut_restore.cpp b/ydb/core/tx/schemeshard/ut_restore/ut_restore.cpp
@@ -1196,7 +1196,7 @@ value {
         const TString string = "test string";
         const TString json = R"({"key": "value"})";
         auto binaryJson = NBinaryJson::SerializeToBinaryJson(json);
-        Y_ABORT_UNLESS(binaryJson.Defined());
+        Y_ABORT_UNLESS(binaryJson.IsSuccess());
 
         const std::pair<ui64, ui64> decimal = NYql::NDecimal::MakePair(NYql::NDecimal::FromString("16.17", NScheme::DECIMAL_PRECISION, NScheme::DECIMAL_SCALE));
         const std::pair<ui64, ui64> decimal35 = NYql::NDecimal::MakePair(NYql::NDecimal::FromString("555555555555555.123456789", 35, 10));

diff --git a/ydb/core/ydb_convert/ydb_convert.cpp b/ydb/core/ydb_convert/ydb_convert.cpp
@@ -512,7 +512,7 @@ Y_FORCE_INLINE void ConvertData(NUdf::TDataTypeId typeId, const Ydb::Value& valu
         case NUdf::TDataType<NUdf::TJsonDocument>::Id: {
             CheckTypeId(value.value_case(), Ydb::Value::kTextValue, "JsonDocument");
             const auto binaryJson = NBinaryJson::SerializeToBinaryJson(value.text_value());
-            if (!binaryJson.Defined()) {
+            if (binaryJson.IsFail()) {
                 throw yexception() << "Invalid JsonDocument value";
             }
             res.SetBytes(binaryJson->Data(), binaryJson->Size());
@@ -1238,7 +1238,7 @@ bool CellFromProtoVal(const NScheme::TTypeInfo& type, i32 typmod, const Ydb::Val
         }
     case NScheme::NTypeIds::JsonDocument : {
         const auto binaryJson = NBinaryJson::SerializeToBinaryJson(val.Gettext_value());
-        if (!binaryJson.Defined()) {
+        if (binaryJson.IsFail()) {
             err = "Invalid JSON for JsonDocument provided";
             return false;
         }

diff --git a/ydb/library/binary_json/ut_benchmark/write.cpp b/ydb/library/binary_json/ut_benchmark/write.cpp
@@ -0,0 +1,50 @@
+#include <benchmark/benchmark.h>
+
+#include <util/random/random.h>
+#include <library/cpp/testing/unittest/registar.h>
+#include <library/cpp/json/json_value.h>
+#include <library/cpp/json/json_writer.h>
+
+#include <ydb/library/binary_json/write.h>
+
+// ya test -r -D BENCHMARK_MAKE_LARGE_PART
+#ifndef BENCHMARK_MAKE_LARGE_PART
+#define BENCHMARK_MAKE_LARGE_PART 0
+#endif
+
+using namespace NKikimr::NBinaryJson;
+
+namespace {
+
+static ui64 seed = 0;
+
+NJson::TJsonValue GetTestJson(ui64 depth = 10, ui64 nChildren = 2) {
+    NJson::TJsonValue value;
+    if (depth == 1) {
+        value.SetValue(NUnitTest::RandomString(10, seed++));
+        return value;
+    }
+    for (ui64 i = 0; i < nChildren; ++i) {
+        value.InsertValue(NUnitTest::RandomString(10, seed++), GetTestJson(depth - 1));
+    }
+    return value;
+}
+
+TString GetTestJsonString() {
+    seed = 42;
+    return NJson::WriteJson(GetTestJson(3, 50));
+}
+
+static void BenchWriteSimdJson(benchmark::State& state) {
+  TString value = GetTestJsonString();
+  TStringBuf buf(value);
+  for (auto _ : state) {
+    auto result = SerializeToBinaryJson(buf);
+    benchmark::DoNotOptimize(result);
+    benchmark::ClobberMemory();
+  }
+}
+
+}
+
+BENCHMARK(BenchWriteSimdJson)->MinTime(1);
diff --git a/ydb/library/binary_json/ut_benchmark/ya.make b/ydb/library/binary_json/ut_benchmark/ya.make
@@ -0,0 +1,30 @@
+G_BENCHMARK()
+
+TAG(ya:fat)
+SIZE(LARGE)
+TIMEOUT(600)
+
+IF (BENCHMARK_MAKE_LARGE_PART)
+    CFLAGS(
+        -DBENCHMARK_MAKE_LARGE_PART=1
+    )
+    TIMEOUT(1200)
+ENDIF()
+
+SRCS(
+    write.cpp
+)
+
+PEERDIR(
+    library/cpp/testing/unittest
+    ydb/library/binary_json
+    ydb/library/yql/minikql/dom
+    ydb/library/yql/minikql/invoke_builtins/llvm14
+    ydb/library/yql/public/udf/service/exception_policy
+    ydb/library/yql/core/issue/protos
+    ydb/library/yql/sql/pg_dummy
+)
+
+YQL_LAST_ABI_VERSION()
+
+END()