From f1ffdee4aa706f191b134ad26152382f5e265175 Mon Sep 17 00:00:00 2001 From: Facebook Community Bot Date: Fri, 19 Apr 2024 12:13:27 -0700 Subject: [PATCH] Re-sync with internal repository (#28) The internal and external repositories are out of sync. This Pull Request attempts to brings them back in sync by patching the GitHub repository. Please carefully review this patch. You must disable ShipIt for your project in order to merge this pull request. DO NOT IMPORT this pull request. Instead, merge it directly on GitHub using the MERGE BUTTON. Re-enable ShipIt after merging. Co-authored-by: Facebook Community Bot <6422482+facebook-github-bot@users.noreply.github.com> --- .clang-format | 87 + .github/workflows/sanity_check.yml | 70 + .gitmodules | 3 + CMake/BuildFlatBuffers.cmake | 449 ++ CMake/FindFlatBuffers.cmake | 61 + CMake/abseil.cmake | 21 + CMakeLists.txt | 115 + CODE_OF_CONDUCT.md | 80 + LICENSE | 201 + Makefile | 77 + dwio/nimble/common/BitEncoder.h | 75 + dwio/nimble/common/Bits.cpp | 183 + dwio/nimble/common/Bits.h | 193 + dwio/nimble/common/Buffer.h | 95 + dwio/nimble/common/CMakeLists.txt | 13 + dwio/nimble/common/Checksum.cpp | 60 + dwio/nimble/common/Checksum.h | 25 + dwio/nimble/common/DefaultMetricsLogger.cpp | 88 + dwio/nimble/common/DefaultMetricsLogger.h | 38 + dwio/nimble/common/EncodingPrimitives.h | 121 + dwio/nimble/common/EncodingType.h | 42 + dwio/nimble/common/Entropy.h | 37 + dwio/nimble/common/Exceptions.h | 470 ++ dwio/nimble/common/FixedBitArray.cpp | 701 +++ dwio/nimble/common/FixedBitArray.h | 125 + dwio/nimble/common/Huffman.cpp | 159 + dwio/nimble/common/Huffman.h | 289 ++ dwio/nimble/common/IndexMap.h | 68 + dwio/nimble/common/MetricsLogger.cpp | 41 + dwio/nimble/common/MetricsLogger.h | 104 + dwio/nimble/common/NimbleCompare.h | 72 + dwio/nimble/common/Rle.h | 42 + dwio/nimble/common/StopWatch.cpp | 47 + dwio/nimble/common/StopWatch.h | 50 + dwio/nimble/common/Types.cpp | 120 + dwio/nimble/common/Types.h | 298 ++ dwio/nimble/common/Varint.cpp | 828 ++++ dwio/nimble/common/Varint.h | 110 + dwio/nimble/common/Vector.h | 286 ++ .../common/benchmarks/VarintBenchmark.cpp | 325 ++ dwio/nimble/common/tests/BitEncoderTests.cpp | 65 + dwio/nimble/common/tests/BitsTests.cpp | 90 + dwio/nimble/common/tests/CMakeLists.txt | 20 + dwio/nimble/common/tests/ExceptionTests.cpp | 242 + .../common/tests/FixedBitArrayTests.cpp | 340 ++ dwio/nimble/common/tests/HuffmanTests.cpp | 85 + dwio/nimble/common/tests/IndexMapTests.cpp | 52 + dwio/nimble/common/tests/NimbleFileWriter.cpp | 52 + dwio/nimble/common/tests/NimbleFileWriter.h | 22 + dwio/nimble/common/tests/StopWatchTests.cpp | 57 + dwio/nimble/common/tests/TestUtils.h | 362 ++ dwio/nimble/common/tests/VarintTests.cpp | 115 + dwio/nimble/common/tests/VectorTests.cpp | 230 + dwio/nimble/encodings/CMakeLists.txt | 15 + dwio/nimble/encodings/Compression.cpp | 75 + dwio/nimble/encodings/Compression.h | 158 + dwio/nimble/encodings/CompressionInternal.h | 36 + dwio/nimble/encodings/CompressionZstd.cpp | 55 + dwio/nimble/encodings/CompressionZstrong.cpp | 214 + dwio/nimble/encodings/ConstantEncoding.h | 115 + dwio/nimble/encodings/DeltaEncoding.h | 331 ++ dwio/nimble/encodings/DictionaryEncoding.h | 175 + dwio/nimble/encodings/Encoding.cpp | 51 + dwio/nimble/encodings/Encoding.h | 286 ++ dwio/nimble/encodings/EncodingFactoryNew.cpp | 368 ++ dwio/nimble/encodings/EncodingFactoryNew.h | 66 + dwio/nimble/encodings/EncodingIdentifier.h | 43 + dwio/nimble/encodings/EncodingLayout.cpp | 119 + dwio/nimble/encodings/EncodingLayout.h | 37 + .../encodings/EncodingLayoutCapture.cpp | 152 + dwio/nimble/encodings/EncodingLayoutCapture.h | 19 + dwio/nimble/encodings/EncodingSelection.h | 245 + .../encodings/EncodingSelectionPolicy.h | 1060 +++++ dwio/nimble/encodings/FixedBitWidthEncoding.h | 216 + .../nimble/encodings/MainlyConstantEncoding.h | 238 + dwio/nimble/encodings/NullableEncoding.h | 259 ++ dwio/nimble/encodings/RleEncoding.cpp | 29 + dwio/nimble/encodings/RleEncoding.h | 231 + dwio/nimble/encodings/SentinelEncoding.h | 402 ++ dwio/nimble/encodings/SparseBoolEncoding.cpp | 114 + dwio/nimble/encodings/SparseBoolEncoding.h | 64 + dwio/nimble/encodings/Statistics.cpp | 277 ++ dwio/nimble/encodings/Statistics.h | 197 + dwio/nimble/encodings/TrivialEncoding.cpp | 231 + dwio/nimble/encodings/TrivialEncoding.h | 204 + dwio/nimble/encodings/VarintEncoding.h | 153 + .../encodings/tests/BucketBenchmarks.cpp | 148 + dwio/nimble/encodings/tests/CMakeLists.txt | 22 + .../encodings/tests/ConstantEncodingTests.cpp | 153 + .../encodings/tests/EncodingLayoutTests.cpp | 378 ++ .../tests/EncodingSelectionTests.cpp | 869 ++++ .../encodings/tests/EncodingTestsNew.cpp | 555 +++ .../tests/MainlyConstantEncodingTests.cpp | 103 + dwio/nimble/encodings/tests/MapBenchmarks.cpp | 463 ++ .../encodings/tests/NullableEncodingTests.cpp | 528 +++ .../encodings/tests/RleEncodingTests.cpp | 152 + .../encodings/tests/SentinelEncodingTests.cpp | 142 + .../encodings/tests/StatisticsTests.cpp | 383 ++ dwio/nimble/encodings/tests/TestGenerator.cpp | 393 ++ dwio/nimble/encodings/tests/TestUtils.h | 226 + dwio/nimble/tablet/CMakeLists.txt | 23 + dwio/nimble/tablet/Compression.cpp | 52 + dwio/nimble/tablet/Compression.h | 21 + dwio/nimble/tablet/Footer.fbs | 44 + dwio/nimble/tablet/Tablet.cpp | 918 ++++ dwio/nimble/tablet/Tablet.h | 391 ++ dwio/nimble/tablet/footer_flatc.sh | 4 + dwio/nimble/tablet/tests/CMakeLists.txt | 14 + dwio/nimble/tablet/tests/TabletTests.cpp | 724 +++ dwio/nimble/tools/CMakeLists.txt | 3 + dwio/nimble/tools/EncodingLayoutTrainer.cpp | 328 ++ dwio/nimble/tools/EncodingLayoutTrainer.h | 30 + dwio/nimble/tools/EncodingSelectionLogger.cpp | 125 + dwio/nimble/tools/EncodingUtilities.cpp | 292 ++ dwio/nimble/tools/EncodingUtilities.h | 42 + dwio/nimble/tools/NimbleDump.cpp | 304 ++ dwio/nimble/tools/NimbleDumpLib.cpp | 687 +++ dwio/nimble/tools/NimbleDumpLib.h | 37 + dwio/nimble/tools/ParallelReader.cpp | 228 + dwio/nimble/tools/ParallelWriter.cpp | 189 + dwio/nimble/velox/BufferGrowthPolicy.cpp | 38 + dwio/nimble/velox/BufferGrowthPolicy.h | 60 + dwio/nimble/velox/CMakeLists.txt | 97 + dwio/nimble/velox/ChunkedStream.cpp | 71 + dwio/nimble/velox/ChunkedStream.h | 55 + dwio/nimble/velox/ChunkedStreamDecoder.cpp | 172 + dwio/nimble/velox/ChunkedStreamDecoder.h | 41 + dwio/nimble/velox/ChunkedStreamWriter.cpp | 45 + dwio/nimble/velox/ChunkedStreamWriter.h | 24 + dwio/nimble/velox/Config.cpp | 139 + dwio/nimble/velox/Config.h | 36 + dwio/nimble/velox/Decoder.h | 26 + dwio/nimble/velox/Deserializer.cpp | 186 + dwio/nimble/velox/Deserializer.h | 28 + dwio/nimble/velox/EncodingLayoutTree.cpp | 177 + dwio/nimble/velox/EncodingLayoutTree.h | 62 + dwio/nimble/velox/FieldReader.cpp | 2457 ++++++++++ dwio/nimble/velox/FieldReader.h | 125 + dwio/nimble/velox/FieldWriter.cpp | 1280 +++++ dwio/nimble/velox/FieldWriter.h | 334 ++ dwio/nimble/velox/FlatMapLayoutPlanner.cpp | 191 + dwio/nimble/velox/FlatMapLayoutPlanner.h | 23 + dwio/nimble/velox/FlushPolicy.cpp | 17 + dwio/nimble/velox/FlushPolicy.h | 65 + dwio/nimble/velox/Metadata.fbs | 14 + dwio/nimble/velox/OrderedRanges.h | 76 + dwio/nimble/velox/Schema.fbs | 49 + dwio/nimble/velox/SchemaBuilder.cpp | 559 +++ dwio/nimble/velox/SchemaBuilder.h | 303 ++ dwio/nimble/velox/SchemaReader.cpp | 521 +++ dwio/nimble/velox/SchemaReader.h | 187 + dwio/nimble/velox/SchemaSerialization.cpp | 213 + dwio/nimble/velox/SchemaSerialization.h | 26 + dwio/nimble/velox/SchemaTypes.cpp | 52 + dwio/nimble/velox/SchemaTypes.h | 102 + dwio/nimble/velox/SchemaUtils.cpp | 154 + dwio/nimble/velox/SchemaUtils.h | 14 + dwio/nimble/velox/Serializer.cpp | 150 + dwio/nimble/velox/Serializer.h | 47 + dwio/nimble/velox/StreamLabels.cpp | 177 + dwio/nimble/velox/StreamLabels.h | 20 + dwio/nimble/velox/TabletSections.h | 10 + dwio/nimble/velox/VeloxReader.cpp | 410 ++ dwio/nimble/velox/VeloxReader.h | 165 + dwio/nimble/velox/VeloxWriter.cpp | 814 ++++ dwio/nimble/velox/VeloxWriter.h | 78 + .../velox/VeloxWriterDefaultMetadataOSS.cpp | 11 + dwio/nimble/velox/VeloxWriterOptions.h | 115 + .../velox/tests/BufferGrowthPolicyTest.cpp | 127 + dwio/nimble/velox/tests/CMakeLists.txt | 38 + .../velox/tests/ChunkedStreamDecoderTests.cpp | 333 ++ .../nimble/velox/tests/ChunkedStreamTests.cpp | 176 + .../velox/tests/EncodingLayoutTreeTests.cpp | 201 + .../velox/tests/FlatMapLayoutPlannerTests.cpp | 416 ++ .../nimble/velox/tests/OrderedRangesTests.cpp | 54 + dwio/nimble/velox/tests/SchemaTests.cpp | 367 ++ dwio/nimble/velox/tests/SchemaUtils.cpp | 182 + dwio/nimble/velox/tests/SchemaUtils.h | 106 + .../nimble/velox/tests/SerializationTests.cpp | 196 + dwio/nimble/velox/tests/TypeTests.cpp | 415 ++ dwio/nimble/velox/tests/VeloxReaderTests.cpp | 4117 +++++++++++++++++ dwio/nimble/velox/tests/VeloxWriterTests.cpp | 1789 +++++++ license.header | 13 + scripts/format-check.py | 284 ++ scripts/git-clang-format | 622 +++ scripts/license-header.py | 259 ++ scripts/util.py | 90 + velox | 1 + 188 files changed, 43009 insertions(+) create mode 100644 .clang-format create mode 100644 .github/workflows/sanity_check.yml create mode 100644 .gitmodules create mode 100644 CMake/BuildFlatBuffers.cmake create mode 100644 CMake/FindFlatBuffers.cmake create mode 100644 CMake/abseil.cmake create mode 100644 CMakeLists.txt create mode 100644 CODE_OF_CONDUCT.md create mode 100644 LICENSE create mode 100644 Makefile create mode 100644 dwio/nimble/common/BitEncoder.h create mode 100644 dwio/nimble/common/Bits.cpp create mode 100644 dwio/nimble/common/Bits.h create mode 100644 dwio/nimble/common/Buffer.h create mode 100644 dwio/nimble/common/CMakeLists.txt create mode 100644 dwio/nimble/common/Checksum.cpp create mode 100644 dwio/nimble/common/Checksum.h create mode 100644 dwio/nimble/common/DefaultMetricsLogger.cpp create mode 100644 dwio/nimble/common/DefaultMetricsLogger.h create mode 100644 dwio/nimble/common/EncodingPrimitives.h create mode 100644 dwio/nimble/common/EncodingType.h create mode 100644 dwio/nimble/common/Entropy.h create mode 100644 dwio/nimble/common/Exceptions.h create mode 100644 dwio/nimble/common/FixedBitArray.cpp create mode 100644 dwio/nimble/common/FixedBitArray.h create mode 100644 dwio/nimble/common/Huffman.cpp create mode 100644 dwio/nimble/common/Huffman.h create mode 100644 dwio/nimble/common/IndexMap.h create mode 100644 dwio/nimble/common/MetricsLogger.cpp create mode 100644 dwio/nimble/common/MetricsLogger.h create mode 100644 dwio/nimble/common/NimbleCompare.h create mode 100644 dwio/nimble/common/Rle.h create mode 100644 dwio/nimble/common/StopWatch.cpp create mode 100644 dwio/nimble/common/StopWatch.h create mode 100644 dwio/nimble/common/Types.cpp create mode 100644 dwio/nimble/common/Types.h create mode 100644 dwio/nimble/common/Varint.cpp create mode 100644 dwio/nimble/common/Varint.h create mode 100644 dwio/nimble/common/Vector.h create mode 100644 dwio/nimble/common/benchmarks/VarintBenchmark.cpp create mode 100644 dwio/nimble/common/tests/BitEncoderTests.cpp create mode 100644 dwio/nimble/common/tests/BitsTests.cpp create mode 100644 dwio/nimble/common/tests/CMakeLists.txt create mode 100644 dwio/nimble/common/tests/ExceptionTests.cpp create mode 100644 dwio/nimble/common/tests/FixedBitArrayTests.cpp create mode 100644 dwio/nimble/common/tests/HuffmanTests.cpp create mode 100644 dwio/nimble/common/tests/IndexMapTests.cpp create mode 100644 dwio/nimble/common/tests/NimbleFileWriter.cpp create mode 100644 dwio/nimble/common/tests/NimbleFileWriter.h create mode 100644 dwio/nimble/common/tests/StopWatchTests.cpp create mode 100644 dwio/nimble/common/tests/TestUtils.h create mode 100644 dwio/nimble/common/tests/VarintTests.cpp create mode 100644 dwio/nimble/common/tests/VectorTests.cpp create mode 100644 dwio/nimble/encodings/CMakeLists.txt create mode 100644 dwio/nimble/encodings/Compression.cpp create mode 100644 dwio/nimble/encodings/Compression.h create mode 100644 dwio/nimble/encodings/CompressionInternal.h create mode 100644 dwio/nimble/encodings/CompressionZstd.cpp create mode 100644 dwio/nimble/encodings/CompressionZstrong.cpp create mode 100644 dwio/nimble/encodings/ConstantEncoding.h create mode 100644 dwio/nimble/encodings/DeltaEncoding.h create mode 100644 dwio/nimble/encodings/DictionaryEncoding.h create mode 100644 dwio/nimble/encodings/Encoding.cpp create mode 100644 dwio/nimble/encodings/Encoding.h create mode 100644 dwio/nimble/encodings/EncodingFactoryNew.cpp create mode 100644 dwio/nimble/encodings/EncodingFactoryNew.h create mode 100644 dwio/nimble/encodings/EncodingIdentifier.h create mode 100644 dwio/nimble/encodings/EncodingLayout.cpp create mode 100644 dwio/nimble/encodings/EncodingLayout.h create mode 100644 dwio/nimble/encodings/EncodingLayoutCapture.cpp create mode 100644 dwio/nimble/encodings/EncodingLayoutCapture.h create mode 100644 dwio/nimble/encodings/EncodingSelection.h create mode 100644 dwio/nimble/encodings/EncodingSelectionPolicy.h create mode 100644 dwio/nimble/encodings/FixedBitWidthEncoding.h create mode 100644 dwio/nimble/encodings/MainlyConstantEncoding.h create mode 100644 dwio/nimble/encodings/NullableEncoding.h create mode 100644 dwio/nimble/encodings/RleEncoding.cpp create mode 100644 dwio/nimble/encodings/RleEncoding.h create mode 100644 dwio/nimble/encodings/SentinelEncoding.h create mode 100644 dwio/nimble/encodings/SparseBoolEncoding.cpp create mode 100644 dwio/nimble/encodings/SparseBoolEncoding.h create mode 100644 dwio/nimble/encodings/Statistics.cpp create mode 100644 dwio/nimble/encodings/Statistics.h create mode 100644 dwio/nimble/encodings/TrivialEncoding.cpp create mode 100644 dwio/nimble/encodings/TrivialEncoding.h create mode 100644 dwio/nimble/encodings/VarintEncoding.h create mode 100644 dwio/nimble/encodings/tests/BucketBenchmarks.cpp create mode 100644 dwio/nimble/encodings/tests/CMakeLists.txt create mode 100644 dwio/nimble/encodings/tests/ConstantEncodingTests.cpp create mode 100644 dwio/nimble/encodings/tests/EncodingLayoutTests.cpp create mode 100644 dwio/nimble/encodings/tests/EncodingSelectionTests.cpp create mode 100644 dwio/nimble/encodings/tests/EncodingTestsNew.cpp create mode 100644 dwio/nimble/encodings/tests/MainlyConstantEncodingTests.cpp create mode 100644 dwio/nimble/encodings/tests/MapBenchmarks.cpp create mode 100644 dwio/nimble/encodings/tests/NullableEncodingTests.cpp create mode 100644 dwio/nimble/encodings/tests/RleEncodingTests.cpp create mode 100644 dwio/nimble/encodings/tests/SentinelEncodingTests.cpp create mode 100644 dwio/nimble/encodings/tests/StatisticsTests.cpp create mode 100644 dwio/nimble/encodings/tests/TestGenerator.cpp create mode 100644 dwio/nimble/encodings/tests/TestUtils.h create mode 100644 dwio/nimble/tablet/CMakeLists.txt create mode 100644 dwio/nimble/tablet/Compression.cpp create mode 100644 dwio/nimble/tablet/Compression.h create mode 100644 dwio/nimble/tablet/Footer.fbs create mode 100644 dwio/nimble/tablet/Tablet.cpp create mode 100644 dwio/nimble/tablet/Tablet.h create mode 100755 dwio/nimble/tablet/footer_flatc.sh create mode 100644 dwio/nimble/tablet/tests/CMakeLists.txt create mode 100644 dwio/nimble/tablet/tests/TabletTests.cpp create mode 100644 dwio/nimble/tools/CMakeLists.txt create mode 100644 dwio/nimble/tools/EncodingLayoutTrainer.cpp create mode 100644 dwio/nimble/tools/EncodingLayoutTrainer.h create mode 100644 dwio/nimble/tools/EncodingSelectionLogger.cpp create mode 100644 dwio/nimble/tools/EncodingUtilities.cpp create mode 100644 dwio/nimble/tools/EncodingUtilities.h create mode 100644 dwio/nimble/tools/NimbleDump.cpp create mode 100644 dwio/nimble/tools/NimbleDumpLib.cpp create mode 100644 dwio/nimble/tools/NimbleDumpLib.h create mode 100644 dwio/nimble/tools/ParallelReader.cpp create mode 100644 dwio/nimble/tools/ParallelWriter.cpp create mode 100644 dwio/nimble/velox/BufferGrowthPolicy.cpp create mode 100644 dwio/nimble/velox/BufferGrowthPolicy.h create mode 100644 dwio/nimble/velox/CMakeLists.txt create mode 100644 dwio/nimble/velox/ChunkedStream.cpp create mode 100644 dwio/nimble/velox/ChunkedStream.h create mode 100644 dwio/nimble/velox/ChunkedStreamDecoder.cpp create mode 100644 dwio/nimble/velox/ChunkedStreamDecoder.h create mode 100644 dwio/nimble/velox/ChunkedStreamWriter.cpp create mode 100644 dwio/nimble/velox/ChunkedStreamWriter.h create mode 100644 dwio/nimble/velox/Config.cpp create mode 100644 dwio/nimble/velox/Config.h create mode 100644 dwio/nimble/velox/Decoder.h create mode 100644 dwio/nimble/velox/Deserializer.cpp create mode 100644 dwio/nimble/velox/Deserializer.h create mode 100644 dwio/nimble/velox/EncodingLayoutTree.cpp create mode 100644 dwio/nimble/velox/EncodingLayoutTree.h create mode 100644 dwio/nimble/velox/FieldReader.cpp create mode 100644 dwio/nimble/velox/FieldReader.h create mode 100644 dwio/nimble/velox/FieldWriter.cpp create mode 100644 dwio/nimble/velox/FieldWriter.h create mode 100644 dwio/nimble/velox/FlatMapLayoutPlanner.cpp create mode 100644 dwio/nimble/velox/FlatMapLayoutPlanner.h create mode 100644 dwio/nimble/velox/FlushPolicy.cpp create mode 100644 dwio/nimble/velox/FlushPolicy.h create mode 100644 dwio/nimble/velox/Metadata.fbs create mode 100644 dwio/nimble/velox/OrderedRanges.h create mode 100644 dwio/nimble/velox/Schema.fbs create mode 100644 dwio/nimble/velox/SchemaBuilder.cpp create mode 100644 dwio/nimble/velox/SchemaBuilder.h create mode 100644 dwio/nimble/velox/SchemaReader.cpp create mode 100644 dwio/nimble/velox/SchemaReader.h create mode 100644 dwio/nimble/velox/SchemaSerialization.cpp create mode 100644 dwio/nimble/velox/SchemaSerialization.h create mode 100644 dwio/nimble/velox/SchemaTypes.cpp create mode 100644 dwio/nimble/velox/SchemaTypes.h create mode 100644 dwio/nimble/velox/SchemaUtils.cpp create mode 100644 dwio/nimble/velox/SchemaUtils.h create mode 100644 dwio/nimble/velox/Serializer.cpp create mode 100644 dwio/nimble/velox/Serializer.h create mode 100644 dwio/nimble/velox/StreamLabels.cpp create mode 100644 dwio/nimble/velox/StreamLabels.h create mode 100644 dwio/nimble/velox/TabletSections.h create mode 100644 dwio/nimble/velox/VeloxReader.cpp create mode 100644 dwio/nimble/velox/VeloxReader.h create mode 100644 dwio/nimble/velox/VeloxWriter.cpp create mode 100644 dwio/nimble/velox/VeloxWriter.h create mode 100644 dwio/nimble/velox/VeloxWriterDefaultMetadataOSS.cpp create mode 100644 dwio/nimble/velox/VeloxWriterOptions.h create mode 100644 dwio/nimble/velox/tests/BufferGrowthPolicyTest.cpp create mode 100644 dwio/nimble/velox/tests/CMakeLists.txt create mode 100644 dwio/nimble/velox/tests/ChunkedStreamDecoderTests.cpp create mode 100644 dwio/nimble/velox/tests/ChunkedStreamTests.cpp create mode 100644 dwio/nimble/velox/tests/EncodingLayoutTreeTests.cpp create mode 100644 dwio/nimble/velox/tests/FlatMapLayoutPlannerTests.cpp create mode 100644 dwio/nimble/velox/tests/OrderedRangesTests.cpp create mode 100644 dwio/nimble/velox/tests/SchemaTests.cpp create mode 100644 dwio/nimble/velox/tests/SchemaUtils.cpp create mode 100644 dwio/nimble/velox/tests/SchemaUtils.h create mode 100644 dwio/nimble/velox/tests/SerializationTests.cpp create mode 100644 dwio/nimble/velox/tests/TypeTests.cpp create mode 100644 dwio/nimble/velox/tests/VeloxReaderTests.cpp create mode 100644 dwio/nimble/velox/tests/VeloxWriterTests.cpp create mode 100644 license.header create mode 100755 scripts/format-check.py create mode 100755 scripts/git-clang-format create mode 100755 scripts/license-header.py create mode 100644 scripts/util.py create mode 160000 velox diff --git a/.clang-format b/.clang-format new file mode 100644 index 0000000..eab4576 --- /dev/null +++ b/.clang-format @@ -0,0 +1,87 @@ +--- +AccessModifierOffset: -1 +AlignAfterOpenBracket: AlwaysBreak +AlignConsecutiveAssignments: false +AlignConsecutiveDeclarations: false +AlignEscapedNewlinesLeft: true +AlignOperands: false +AlignTrailingComments: false +AllowAllParametersOfDeclarationOnNextLine: false +AllowShortBlocksOnASingleLine: false +AllowShortCaseLabelsOnASingleLine: false +AllowShortFunctionsOnASingleLine: Empty +AllowShortIfStatementsOnASingleLine: false +AllowShortLoopsOnASingleLine: false +AlwaysBreakAfterReturnType: None +AlwaysBreakBeforeMultilineStrings: true +AlwaysBreakTemplateDeclarations: true +BinPackArguments: false +BinPackParameters: false +BraceWrapping: + AfterClass: false + AfterControlStatement: false + AfterEnum: false + AfterFunction: false + AfterNamespace: false + AfterObjCDeclaration: false + AfterStruct: false + AfterUnion: false + BeforeCatch: false + BeforeElse: false + IndentBraces: false +BreakBeforeBinaryOperators: None +BreakBeforeBraces: Attach +BreakBeforeTernaryOperators: true +BreakConstructorInitializersBeforeComma: false +BreakAfterJavaFieldAnnotations: false +BreakStringLiterals: false +ColumnLimit: 80 +CommentPragmas: '^ IWYU pragma:' +ConstructorInitializerAllOnOneLineOrOnePerLine: true +ConstructorInitializerIndentWidth: 4 +ContinuationIndentWidth: 4 +Cpp11BracedListStyle: true +DerivePointerAlignment: false +DisableFormat: false +ForEachMacros: [ FOR_EACH, FOR_EACH_R, FOR_EACH_RANGE, ] +IncludeCategories: + - Regex: '^<.*\.h(pp)?>' + Priority: 1 + - Regex: '^<.*' + Priority: 2 + - Regex: '.*' + Priority: 3 +IndentCaseLabels: true +IndentWidth: 2 +IndentWrappedFunctionNames: false +KeepEmptyLinesAtTheStartOfBlocks: false +MacroBlockBegin: '' +MacroBlockEnd: '' +MaxEmptyLinesToKeep: 1 +NamespaceIndentation: None +ObjCBlockIndentWidth: 2 +ObjCSpaceAfterProperty: false +ObjCSpaceBeforeProtocolList: false +PenaltyBreakBeforeFirstCallParameter: 1 +PenaltyBreakComment: 300 +PenaltyBreakFirstLessLess: 120 +PenaltyBreakString: 1000 +PenaltyExcessCharacter: 1000000 +PenaltyReturnTypeOnItsOwnLine: 200 +PointerAlignment: Left +ReflowComments: true +SortIncludes: true +SpaceAfterCStyleCast: false +SpaceBeforeAssignmentOperators: true +SpaceBeforeParens: ControlStatements +SpaceInEmptyParentheses: false +SpacesBeforeTrailingComments: 1 +SpacesInAngles: false +SpacesInContainerLiterals: true +SpacesInCStyleCastParentheses: false +SpacesInParentheses: false +SpacesInSquareBrackets: false +Standard: Cpp11 +TabWidth: 8 +UseTab: Never +... diff --git a/.github/workflows/sanity_check.yml b/.github/workflows/sanity_check.yml new file mode 100644 index 0000000..3bd9860 --- /dev/null +++ b/.github/workflows/sanity_check.yml @@ -0,0 +1,70 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# Licensed 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. +name: Run Sanity Checks + +on: + pull_request: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.repository }}-${{ github.head_ref || github.sha }} + cancel-in-progress: true + +jobs: + sanity-check: + name: ${{ matrix.config.name }} + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + config: + - { name: "Code Format", + command: "format-fix", + message: "Found format issues", + reqs: "regex cmake-format black" + } + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Install Dependencies + run: | + python -m venv check_env + source check_env/bin/activate + pip install ${{ matrix.config.reqs }} + + - name: Check ${{ matrix.config.name }} + run: | + source check_env/bin/activate + make ${{ matrix.config.command }} + + if ! git diff --quiet; then + diff=`git --no-pager diff` + echo "${{ matrix.command.message }} in the following files:" + git --no-pager diff --name-only + echo "Check the Job summary for a copy-pasteable patch." + + echo "> [!IMPORTANT]" >> $GITHUB_STEP_SUMMARY + echo "${{ matrix.config.message }}" >> $GITHUB_STEP_SUMMARY + echo "> Please apply fix using:" >> $GITHUB_STEP_SUMMARY + echo "\`\`\`sh" >> $GITHUB_STEP_SUMMARY + echo "patch -p1 <> $GITHUB_STEP_SUMMARY + echo "$diff" >> $GITHUB_STEP_SUMMARY + echo "EOF" >> $GITHUB_STEP_SUMMARY + echo "\`\`\`" >> $GITHUB_STEP_SUMMARY + exit 1 + fi diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..b558000 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "velox"] + path = velox + url = https://github.com/facebookincubator/velox.git diff --git a/CMake/BuildFlatBuffers.cmake b/CMake/BuildFlatBuffers.cmake new file mode 100644 index 0000000..631e5ad --- /dev/null +++ b/CMake/BuildFlatBuffers.cmake @@ -0,0 +1,449 @@ +# Copyright 2015 Google Inc. All rights reserved. +# +# Licensed 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. + +# General function to create FlatBuffer build rules for the given list of +# schemas. +# +# flatbuffers_schemas: A list of flatbuffer schema files to process. +# +# schema_include_dirs: A list of schema file include directories, which will be +# passed to flatc via the -I parameter. +# +# custom_target_name: The generated files will be added as dependencies for a +# new custom target with this name. You should add that target as a dependency +# for your main target to ensure these files are built. You can also retrieve +# various properties from this target, such as GENERATED_INCLUDES_DIR, +# BINARY_SCHEMAS_DIR, and COPY_TEXT_SCHEMAS_DIR. +# +# additional_dependencies: A list of additional dependencies that you'd like +# all generated files to depend on. Pass in a blank string if you have none. +# +# generated_includes_dir: Where to generate the C++ header files for these +# schemas. The generated includes directory will automatically be added to +# CMake's include directories, and will be where generated header files are +# placed. This parameter is optional; pass in empty string if you don't want to +# generate include files for these schemas. +# +# binary_schemas_dir: If you specify an optional binary schema directory, binary +# schemas will be generated for these schemas as well, and placed into the given +# directory. +# +# copy_text_schemas_dir: If you want all text schemas (including schemas from +# all schema include directories) copied into a directory (for example, if you +# need them within your project to build JSON files), you can specify that +# folder here. All text schemas will be copied to that folder. +# +# IMPORTANT: Make sure you quote all list arguments you pass to this function! +# Otherwise CMake will only pass in the first element. +# Example: build_flatbuffers("${fb_files}" "${include_dirs}" target_name ...) +function(build_flatbuffers flatbuffers_schemas + schema_include_dirs + custom_target_name + additional_dependencies + generated_includes_dir + binary_schemas_dir + copy_text_schemas_dir) + + # Test if including from FindFlatBuffers + if(FLATBUFFERS_FLATC_EXECUTABLE) + set(FLATC_TARGET "") + set(FLATC ${FLATBUFFERS_FLATC_EXECUTABLE}) + elseif(TARGET flatbuffers::flatc) + set(FLATC_TARGET flatbuffers::flatc) + set(FLATC flatbuffers::flatc) + else() + set(FLATC_TARGET flatc) + set(FLATC flatc) + endif() + set(FLATC_SCHEMA_ARGS --gen-mutable) + if(FLATBUFFERS_FLATC_SCHEMA_EXTRA_ARGS) + set(FLATC_SCHEMA_ARGS + ${FLATBUFFERS_FLATC_SCHEMA_EXTRA_ARGS} + ${FLATC_SCHEMA_ARGS} + ) + endif() + + set(working_dir "${CMAKE_CURRENT_SOURCE_DIR}") + + set(schema_glob "*.fbs") + # Generate the include files parameters. + set(include_params "") + set(all_generated_files "") + foreach (include_dir ${schema_include_dirs}) + set(include_params -I ${include_dir} ${include_params}) + if (NOT ${copy_text_schemas_dir} STREQUAL "") + # Copy text schemas from dependent folders. + file(GLOB_RECURSE dependent_schemas ${include_dir}/${schema_glob}) + foreach (dependent_schema ${dependent_schemas}) + file(COPY ${dependent_schema} DESTINATION ${copy_text_schemas_dir}) + endforeach() + endif() + endforeach() + + foreach(schema ${flatbuffers_schemas}) + get_filename_component(filename ${schema} NAME_WE) + # For each schema, do the things we requested. + if (NOT ${generated_includes_dir} STREQUAL "") + set(generated_include ${generated_includes_dir}/${filename}_generated.h) + add_custom_command( + OUTPUT ${generated_include} + COMMAND ${FLATC} ${FLATC_SCHEMA_ARGS} + -o ${generated_includes_dir} + ${include_params} + -c ${schema} + DEPENDS ${FLATC_TARGET} ${schema} ${additional_dependencies} + WORKING_DIRECTORY "${working_dir}") + list(APPEND all_generated_files ${generated_include}) + endif() + + if (NOT ${binary_schemas_dir} STREQUAL "") + set(binary_schema ${binary_schemas_dir}/${filename}.bfbs) + add_custom_command( + OUTPUT ${binary_schema} + COMMAND ${FLATC} -b --schema + -o ${binary_schemas_dir} + ${include_params} + ${schema} + DEPENDS ${FLATC_TARGET} ${schema} ${additional_dependencies} + WORKING_DIRECTORY "${working_dir}") + list(APPEND all_generated_files ${binary_schema}) + endif() + + if (NOT ${copy_text_schemas_dir} STREQUAL "") + file(COPY ${schema} DESTINATION ${copy_text_schemas_dir}) + endif() + endforeach() + + # Create a custom target that depends on all the generated files. + # This is the target that you can depend on to trigger all these + # to be built. + add_custom_target(${custom_target_name} + DEPENDS ${all_generated_files} ${additional_dependencies}) + + # Register the include directory we are using. + if (NOT ${generated_includes_dir} STREQUAL "") + include_directories(${generated_includes_dir}) + set_property(TARGET ${custom_target_name} + PROPERTY GENERATED_INCLUDES_DIR + ${generated_includes_dir}) + endif() + + # Register the binary schemas dir we are using. + if (NOT ${binary_schemas_dir} STREQUAL "") + set_property(TARGET ${custom_target_name} + PROPERTY BINARY_SCHEMAS_DIR + ${binary_schemas_dir}) + endif() + + # Register the text schema copy dir we are using. + if (NOT ${copy_text_schemas_dir} STREQUAL "") + set_property(TARGET ${custom_target_name} + PROPERTY COPY_TEXT_SCHEMAS_DIR + ${copy_text_schemas_dir}) + endif() +endfunction() + +# Creates a target that can be linked against that generates flatbuffer headers. +# +# This function takes a target name and a list of schemas. You can also specify +# other flagc flags using the FLAGS option to change the behavior of the flatc +# tool. +# +# When the target_link_libraries is done within a different directory than +# flatbuffers_generate_headers is called, then the target should also be dependent +# the custom generation target called GENERATE_. +# +# Arguments: +# TARGET: The name of the target to generate. +# SCHEMAS: The list of schema files to generate code for. +# BINARY_SCHEMAS_DIR: Optional. The directory in which to generate binary +# schemas. Binary schemas will only be generated if a path is provided. +# INCLUDE: Optional. Search for includes in the specified paths. (Use this +# instead of "-I " and the FLAGS option so that CMake is aware of +# the directories that need to be searched). +# INCLUDE_PREFIX: Optional. The directory in which to place the generated +# files. Use this instead of the --include-prefix option. +# FLAGS: Optional. A list of any additional flags that you would like to pass +# to flatc. +# +# Example: +# +# flatbuffers_generate_headers( +# TARGET my_generated_headers_target +# INCLUDE_PREFIX ${MY_INCLUDE_PREFIX}" +# SCHEMAS ${MY_SCHEMA_FILES} +# BINARY_SCHEMAS_DIR "${MY_BINARY_SCHEMA_DIRECTORY}" +# FLAGS --gen-object-api) +# +# target_link_libraries(MyExecutableTarget +# PRIVATE my_generated_headers_target +# ) +# +# Optional (only needed within different directory): +# add_dependencies(app GENERATE_my_generated_headers_target) +function(flatbuffers_generate_headers) + # Parse function arguments. + set(options) + set(one_value_args + "TARGET" + "INCLUDE_PREFIX" + "BINARY_SCHEMAS_DIR") + set(multi_value_args + "SCHEMAS" + "INCLUDE" + "FLAGS") + cmake_parse_arguments( + PARSE_ARGV 0 + FLATBUFFERS_GENERATE_HEADERS + "${options}" + "${one_value_args}" + "${multi_value_args}") + + # Test if including from FindFlatBuffers + if(FLATBUFFERS_FLATC_EXECUTABLE) + set(FLATC_TARGET "") + set(FLATC ${FLATBUFFERS_FLATC_EXECUTABLE}) + elseif(TARGET flatbuffers::flatc) + set(FLATC_TARGET flatbuffers::flatc) + set(FLATC flatbuffers::flatc) + else() + set(FLATC_TARGET flatc) + set(FLATC flatc) + endif() + + set(working_dir "${CMAKE_CURRENT_SOURCE_DIR}") + + # Generate the include files parameters. + set(include_params "") + foreach (include_dir ${FLATBUFFERS_GENERATE_HEADERS_INCLUDE}) + set(include_params -I ${include_dir} ${include_params}) + endforeach() + + # Create a directory to place the generated code. + set(generated_target_dir "${CMAKE_CURRENT_BINARY_DIR}/${FLATBUFFERS_GENERATE_HEADERS_TARGET}") + set(generated_include_dir "${generated_target_dir}") + if (NOT ${FLATBUFFERS_GENERATE_HEADERS_INCLUDE_PREFIX} STREQUAL "") + set(generated_include_dir "${generated_include_dir}/${FLATBUFFERS_GENERATE_HEADERS_INCLUDE_PREFIX}") + list(APPEND FLATBUFFERS_GENERATE_HEADERS_FLAGS + "--include-prefix" ${FLATBUFFERS_GENERATE_HEADERS_INCLUDE_PREFIX}) + endif() + + set(generated_custom_commands) + + # Create rules to generate the code for each schema. + foreach(schema ${FLATBUFFERS_GENERATE_HEADERS_SCHEMAS}) + get_filename_component(filename ${schema} NAME_WE) + set(generated_include "${generated_include_dir}/${filename}_generated.h") + + # Generate files for grpc if needed + set(generated_source_file) + if("${FLATBUFFERS_GENERATE_HEADERS_FLAGS}" MATCHES "--grpc") + # Check if schema file contain a rpc_service definition + file(STRINGS ${schema} has_grpc REGEX "rpc_service") + if(has_grpc) + list(APPEND generated_include "${generated_include_dir}/${filename}.grpc.fb.h") + set(generated_source_file "${generated_include_dir}/${filename}.grpc.fb.cc") + endif() + endif() + + add_custom_command( + OUTPUT ${generated_include} ${generated_source_file} + COMMAND ${FLATC} ${FLATC_ARGS} + -o ${generated_include_dir} + ${include_params} + -c ${schema} + ${FLATBUFFERS_GENERATE_HEADERS_FLAGS} + DEPENDS ${FLATC_TARGET} ${schema} + WORKING_DIRECTORY "${working_dir}" + COMMENT "Building ${schema} flatbuffers...") + list(APPEND all_generated_header_files ${generated_include}) + list(APPEND all_generated_source_files ${generated_source_file}) + list(APPEND generated_custom_commands "${generated_include}" "${generated_source_file}") + + # Geneate the binary flatbuffers schemas if instructed to. + if (NOT ${FLATBUFFERS_GENERATE_HEADERS_BINARY_SCHEMAS_DIR} STREQUAL "") + set(binary_schema + "${FLATBUFFERS_GENERATE_HEADERS_BINARY_SCHEMAS_DIR}/${filename}.bfbs") + add_custom_command( + OUTPUT ${binary_schema} + COMMAND ${FLATC} -b --schema + -o ${FLATBUFFERS_GENERATE_HEADERS_BINARY_SCHEMAS_DIR} + ${include_params} + ${schema} + DEPENDS ${FLATC_TARGET} ${schema} + WORKING_DIRECTORY "${working_dir}") + list(APPEND generated_custom_commands "${binary_schema}") + list(APPEND all_generated_binary_files ${binary_schema}) + endif() + endforeach() + + # Create an additional target as add_custom_command scope is only within same directory (CMakeFile.txt) + set(generate_target GENERATE_${FLATBUFFERS_GENERATE_HEADERS_TARGET}) + add_custom_target(${generate_target} ALL + DEPENDS ${generated_custom_commands} + COMMENT "Generating flatbuffer target ${FLATBUFFERS_GENERATE_HEADERS_TARGET}") + + # Set up interface library + add_library(${FLATBUFFERS_GENERATE_HEADERS_TARGET} INTERFACE) + target_sources( + ${FLATBUFFERS_GENERATE_HEADERS_TARGET} + INTERFACE + ${all_generated_header_files} + ${all_generated_binary_files} + ${all_generated_source_files} + ${FLATBUFFERS_GENERATE_HEADERS_SCHEMAS}) + add_dependencies( + ${FLATBUFFERS_GENERATE_HEADERS_TARGET} + ${FLATC} + ${FLATBUFFERS_GENERATE_HEADERS_SCHEMAS}) + target_include_directories( + ${FLATBUFFERS_GENERATE_HEADERS_TARGET} + INTERFACE ${generated_target_dir}) + + # Organize file layout for IDEs. + source_group( + TREE "${generated_target_dir}" + PREFIX "Flatbuffers/Generated/Headers Files" + FILES ${all_generated_header_files}) + source_group( + TREE "${generated_target_dir}" + PREFIX "Flatbuffers/Generated/Source Files" + FILES ${all_generated_source_files}) + source_group( + TREE ${working_dir} + PREFIX "Flatbuffers/Schemas" + FILES ${FLATBUFFERS_GENERATE_HEADERS_SCHEMAS}) + if (NOT ${FLATBUFFERS_GENERATE_HEADERS_BINARY_SCHEMAS_DIR} STREQUAL "") + source_group( + TREE "${FLATBUFFERS_GENERATE_HEADERS_BINARY_SCHEMAS_DIR}" + PREFIX "Flatbuffers/Generated/Binary Schemas" + FILES ${all_generated_binary_files}) + endif() +endfunction() + +# Creates a target that can be linked against that generates flatbuffer binaries +# from json files. +# +# This function takes a target name and a list of schemas and Json files. You +# can also specify other flagc flags and options to change the behavior of the +# flatc compiler. +# +# Adding this target to your executable ensurses that the flatbuffer binaries +# are compiled before your executable is run. +# +# Arguments: +# TARGET: The name of the target to generate. +# JSON_FILES: The list of json files to compile to flatbuffers binaries. +# SCHEMA: The flatbuffers schema of the Json files to be compiled. +# INCLUDE: Optional. Search for includes in the specified paths. (Use this +# instead of "-I " and the FLAGS option so that CMake is aware of +# the directories that need to be searched). +# OUTPUT_DIR: The directly where the generated flatbuffers binaries should be +# placed. +# FLAGS: Optional. A list of any additional flags that you would like to pass +# to flatc. +# +# Example: +# +# flatbuffers_generate_binary_files( +# TARGET my_binary_data +# SCHEMA "${MY_SCHEMA_DIR}/my_example_schema.fbs" +# JSON_FILES ${MY_JSON_FILES} +# OUTPUT_DIR "${MY_BINARY_DATA_DIRECTORY}" +# FLAGS --strict-json) +# +# target_link_libraries(MyExecutableTarget +# PRIVATE my_binary_data +# ) +function(flatbuffers_generate_binary_files) + # Parse function arguments. + set(options) + set(one_value_args + "TARGET" + "SCHEMA" + "OUTPUT_DIR") + set(multi_value_args + "JSON_FILES" + "INCLUDE" + "FLAGS") + cmake_parse_arguments( + PARSE_ARGV 0 + FLATBUFFERS_GENERATE_BINARY_FILES + "${options}" + "${one_value_args}" + "${multi_value_args}") + + # Test if including from FindFlatBuffers + if(FLATBUFFERS_FLATC_EXECUTABLE) + set(FLATC_TARGET "") + set(FLATC ${FLATBUFFERS_FLATC_EXECUTABLE}) + elseif(TARGET flatbuffers::flatc) + set(FLATC_TARGET flatbuffers::flatc) + set(FLATC flatbuffers::flatc) + else() + set(FLATC_TARGET flatc) + set(FLATC flatc) + endif() + + set(working_dir "${CMAKE_CURRENT_SOURCE_DIR}") + + # Generate the include files parameters. + set(include_params "") + foreach (include_dir ${FLATBUFFERS_GENERATE_BINARY_FILES_INCLUDE}) + set(include_params -I ${include_dir} ${include_params}) + endforeach() + + # Create rules to generate the flatbuffers binary for each json file. + foreach(json_file ${FLATBUFFERS_GENERATE_BINARY_FILES_JSON_FILES}) + get_filename_component(filename ${json_file} NAME_WE) + set(generated_binary_file "${FLATBUFFERS_GENERATE_BINARY_FILES_OUTPUT_DIR}/${filename}.bin") + add_custom_command( + OUTPUT ${generated_binary_file} + COMMAND ${FLATC} ${FLATC_ARGS} + -o ${FLATBUFFERS_GENERATE_BINARY_FILES_OUTPUT_DIR} + ${include_params} + -b ${FLATBUFFERS_GENERATE_BINARY_FILES_SCHEMA} ${json_file} + ${FLATBUFFERS_GENERATE_BINARY_FILES_FLAGS} + DEPENDS ${FLATC_TARGET} ${json_file} + WORKING_DIRECTORY "${working_dir}" + COMMENT "Building ${json_file} binary flatbuffers...") + list(APPEND all_generated_binary_files ${generated_binary_file}) + endforeach() + + # Set up interface library + add_library(${FLATBUFFERS_GENERATE_BINARY_FILES_TARGET} INTERFACE) + target_sources( + ${FLATBUFFERS_GENERATE_BINARY_FILES_TARGET} + INTERFACE + ${all_generated_binary_files} + ${FLATBUFFERS_GENERATE_BINARY_FILES_JSON_FILES} + ${FLATBUFFERS_GENERATE_BINARY_FILES_SCHEMA}) + add_dependencies( + ${FLATBUFFERS_GENERATE_BINARY_FILES_TARGET} + ${FLATC}) + + # Organize file layout for IDEs. + source_group( + TREE ${working_dir} + PREFIX "Flatbuffers/JSON Files" + FILES ${FLATBUFFERS_GENERATE_BINARY_FILES_JSON_FILES}) + source_group( + TREE ${working_dir} + PREFIX "Flatbuffers/Schemas" + FILES ${FLATBUFFERS_GENERATE_BINARY_FILES_SCHEMA}) + source_group( + TREE ${FLATBUFFERS_GENERATE_BINARY_FILES_OUTPUT_DIR} + PREFIX "Flatbuffers/Generated/Binary Files" + FILES ${all_generated_binary_files}) +endfunction() diff --git a/CMake/FindFlatBuffers.cmake b/CMake/FindFlatBuffers.cmake new file mode 100644 index 0000000..044cf7c --- /dev/null +++ b/CMake/FindFlatBuffers.cmake @@ -0,0 +1,61 @@ +# Copyright 2014 Stefan.Eilemann@epfl.ch +# Copyright 2014 Google Inc. All rights reserved. +# +# Licensed 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. + +# Find the flatbuffers schema compiler +# +# Output Variables: +# * FLATBUFFERS_FLATC_EXECUTABLE the flatc compiler executable +# * FLATBUFFERS_FOUND +# +# Provides: +# * FLATBUFFERS_GENERATE_C_HEADERS(Name ) creates the C++ headers +# for the given flatbuffer schema files. +# Returns the header files in ${Name}_OUTPUTS + +set(FLATBUFFERS_CMAKE_DIR ${CMAKE_CURRENT_LIST_DIR}) + +find_program(FLATBUFFERS_FLATC_EXECUTABLE NAMES flatc) +find_path(FLATBUFFERS_INCLUDE_DIR NAMES flatbuffers/flatbuffers.h) + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args(FlatBuffers + DEFAULT_MSG FLATBUFFERS_FLATC_EXECUTABLE FLATBUFFERS_INCLUDE_DIR) + +if(FLATBUFFERS_FOUND) + function(FLATBUFFERS_GENERATE_C_HEADERS Name) + set(FLATC_OUTPUTS) + foreach(FILE ${ARGN}) + get_filename_component(FLATC_OUTPUT ${FILE} NAME_WE) + set(FLATC_OUTPUT + "${CMAKE_CURRENT_BINARY_DIR}/${FLATC_OUTPUT}_generated.h") + list(APPEND FLATC_OUTPUTS ${FLATC_OUTPUT}) + + add_custom_command(OUTPUT ${FLATC_OUTPUT} + COMMAND ${FLATBUFFERS_FLATC_EXECUTABLE} + ARGS -c -o "${CMAKE_CURRENT_BINARY_DIR}/" ${FILE} + DEPENDS ${FILE} + COMMENT "Building C++ header for ${FILE}" + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}) + endforeach() + set(${Name}_OUTPUTS ${FLATC_OUTPUTS} PARENT_SCOPE) + endfunction() + + set(FLATBUFFERS_INCLUDE_DIRS ${FLATBUFFERS_INCLUDE_DIR}) + include_directories(${CMAKE_BINARY_DIR}) +else() + set(FLATBUFFERS_INCLUDE_DIR) +endif() + +include("${FLATBUFFERS_CMAKE_DIR}/BuildFlatBuffers.cmake") diff --git a/CMake/abseil.cmake b/CMake/abseil.cmake new file mode 100644 index 0000000..8768f24 --- /dev/null +++ b/CMake/abseil.cmake @@ -0,0 +1,21 @@ +include_guard(GLOBAL) + +# TODO: these variables are named VELOX_* because we are piggy-backing on +# Velox's resolve dependency module for now. We should change and have +# our own in the future. +set(VELOX_ABSEIL_VERSION 20240116.0) +set(VELOX_ABSEIL_BUILD_SHA256_CHECKSUM + "338420448b140f0dfd1a1ea3c3ce71b3bc172071f24f4d9a57d59b45037da440") +set(VELOX_ABSEIL_SOURCE_URL + "https://github.com/abseil/abseil-cpp/archive/refs/tags/${VELOX_ABSEIL_VERSION}.tar.gz") + +resolve_dependency_url(ABSEIL) + +message(STATUS "Building abseil from source") + +FetchContent_Declare( + abseil + URL ${VELOX_ABSEIL_SOURCE_URL} + URL_HASH ${VELOX_ABSEIL_BUILD_SHA256_CHECKSUM}) + +FetchContent_MakeAvailable(abseil) diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..c059817 --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,115 @@ +cmake_minimum_required(VERSION 3.14) + +# Set the project name. +project(Nimble) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED True) + +# Sets new behavior for CMP0135, which controls how timestamps are extracted +# when using ExternalProject_Add(): +# https://cmake.org/cmake/help/latest/policy/CMP0135.html +if(POLICY CMP0135) + cmake_policy(SET CMP0135 NEW) + set(CMAKE_POLICY_DEFAULT_CMP0135 NEW) +endif() + +# Use ThirdPartyToolchain dependencies macros from Velox. +list(PREPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/CMake" + "${PROJECT_SOURCE_DIR}/velox/CMake") +include(ResolveDependency) + +set(VELOX_BUILD_MINIMAL_WITH_DWIO + ON CACHE BOOL + "Velox minimal build with dwio.") +set(VELOX_BUILD_VECTOR_TEST_UTILS + ON CACHE BOOL + "Velox vector test utilities (VectorMaker).") +set(VELOX_DEPENDENCY_SOURCE + AUTO + CACHE + STRING + "Default dependency source: AUTO SYSTEM or BUNDLED." +) + +message(STATUS "Build type: ${CMAKE_BUILD_TYPE}") + +# Ignore known compiler warnings. +check_cxx_compiler_flag("-Wstringop-overread" COMPILER_HAS_W_STRINGOP_OVERREAD) +if(COMPILER_HAS_W_STRINGOP_OVERREAD) + string(APPEND CMAKE_CXX_FLAGS " -Wno-stringop-overread") +endif() + +check_cxx_compiler_flag("-Wdeprecated-declarations" + COMPILER_HAS_W_DEPRECATED_DECLARATIONS) +if(COMPILER_HAS_W_DEPRECATED_DECLARATIONS) + string(APPEND CMAKE_CXX_FLAGS " -Wno-deprecated-declarations") +endif() + +check_cxx_compiler_flag("-Wmaybe-uninitialized" + COMPILER_HAS_W_MAYBE_UNINITIALIZED) +if(COMPILER_HAS_W_MAYBE_UNINITIALIZED) + string(APPEND CMAKE_CXX_FLAGS " -Wno-maybe-uninitialized") +endif() + +check_cxx_compiler_flag("-Wunknown-warning-option" + COMPILER_HAS_W_UNKNOWN_WARNING_OPTION) +if(COMPILER_HAS_W_UNKNOWN_WARNING_OPTION) + string(APPEND CMAKE_CXX_FLAGS " -Wno-unknown-warning-option") +endif() + +check_cxx_compiler_flag("-Wnullability-completeness" + COMPILER_HAS_W_NULLABILITY_COMPLETENESS) +if(COMPILER_HAS_W_NULLABILITY_COMPLETENESS) + string(APPEND CMAKE_CXX_FLAGS " -Wno-nullability-completeness") +endif() + +# Nimble, Velox and folly need to be compiled with the same compiler flags. +execute_process( + COMMAND + bash -c + "( source ${CMAKE_CURRENT_SOURCE_DIR}/velox/scripts/setup-helper-functions.sh && echo -n $(get_cxx_flags $ENV{CPU_TARGET}))" + OUTPUT_VARIABLE SCRIPT_CXX_FLAGS + RESULT_VARIABLE COMMAND_STATUS) + +if(COMMAND_STATUS EQUAL "1") + message(FATAL_ERROR "Unable to determine compiler flags!") +endif() +message("Setting CMAKE_CXX_FLAGS=${SCRIPT_CXX_FLAGS}") +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${SCRIPT_CXX_FLAGS}") + +message("FINAL CMAKE_CXX_FLAGS=${CMAKE_CXX_FLAGS}") + +include(CTest) # include after project() but before add_subdirectory() + +# This doesn't necessarily need to be a dependency (we can check in the +# generated .cpp/.h files), but adding this for convenience for now. +find_package(FlatBuffers REQUIRED) + +set_source(gtest) +resolve_dependency(gtest) + +set_source(glog) +resolve_dependency(glog) + +set_source(gflags) +resolve_dependency(gflags COMPONENTS shared) + +set_source(folly) +resolve_dependency(folly) + +set_source(abseil) +resolve_dependency(abseil) + +# Use xxhash and xsimd from Velox for now. +include_directories(.) +include_directories(SYSTEM velox) +include_directories(SYSTEM velox/velox/external/xxhash) +include_directories(SYSTEM ${CMAKE_BINARY_DIR}/_deps/xsimd-src/include/) + +add_subdirectory(velox) +add_subdirectory(dwio/nimble/common) +add_subdirectory(dwio/nimble/tablet) +add_subdirectory(dwio/nimble/tools) +add_subdirectory(dwio/nimble/encodings) +add_subdirectory(dwio/nimble/velox) diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md new file mode 100644 index 0000000..08b500a --- /dev/null +++ b/CODE_OF_CONDUCT.md @@ -0,0 +1,80 @@ +# Code of Conduct + +## Our Pledge + +In the interest of fostering an open and welcoming environment, we as +contributors and maintainers pledge to make participation in our project and +our community a harassment-free experience for everyone, regardless of age, body +size, disability, ethnicity, sex characteristics, gender identity and expression, +level of experience, education, socio-economic status, nationality, personal +appearance, race, religion, or sexual identity and orientation. + +## Our Standards + +Examples of behavior that contributes to creating a positive environment +include: + +* Using welcoming and inclusive language +* Being respectful of differing viewpoints and experiences +* Gracefully accepting constructive criticism +* Focusing on what is best for the community +* Showing empathy towards other community members + +Examples of unacceptable behavior by participants include: + +* The use of sexualized language or imagery and unwelcome sexual attention or + advances +* Trolling, insulting/derogatory comments, and personal or political attacks +* Public or private harassment +* Publishing others' private information, such as a physical or electronic + address, without explicit permission +* Other conduct which could reasonably be considered inappropriate in a + professional setting + +## Our Responsibilities + +Project maintainers are responsible for clarifying the standards of acceptable +behavior and are expected to take appropriate and fair corrective action in +response to any instances of unacceptable behavior. + +Project maintainers have the right and responsibility to remove, edit, or +reject comments, commits, code, wiki edits, issues, and other contributions +that are not aligned to this Code of Conduct, or to ban temporarily or +permanently any contributor for other behaviors that they deem inappropriate, +threatening, offensive, or harmful. + +## Scope + +This Code of Conduct applies within all project spaces, and it also applies when +an individual is representing the project or its community in public spaces. +Examples of representing a project or community include using an official +project e-mail address, posting via an official social media account, or acting +as an appointed representative at an online or offline event. Representation of +a project may be further defined and clarified by project maintainers. + +This Code of Conduct also applies outside the project spaces when there is a +reasonable belief that an individual's behavior may have a negative impact on +the project or its community. + +## Enforcement + +Instances of abusive, harassing, or otherwise unacceptable behavior may be +reported by contacting the project team at . All +complaints will be reviewed and investigated and will result in a response that +is deemed necessary and appropriate to the circumstances. The project team is +obligated to maintain confidentiality with regard to the reporter of an incident. +Further details of specific enforcement policies may be posted separately. + +Project maintainers who do not follow or enforce the Code of Conduct in good +faith may face temporary or permanent repercussions as determined by other +members of the project's leadership. + +## Attribution + +This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4, +available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html + +[homepage]: https://www.contributor-covenant.org + +For answers to common questions about this code of conduct, see +https://www.contributor-covenant.org/faq diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..b09cd78 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ +Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed 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. diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..9c915c5 --- /dev/null +++ b/Makefile @@ -0,0 +1,77 @@ +.PHONY: all cmake build clean debug release unit submodules + +BUILD_BASE_DIR=_build +BUILD_DIR=release +BUILD_TYPE=Release + +CMAKE_FLAGS := -DCMAKE_BUILD_TYPE=$(BUILD_TYPE) + +# Use Ninja if available. If Ninja is used, pass through parallelism control flags. +USE_NINJA ?= 1 +ifeq ($(USE_NINJA), 1) +ifneq ($(shell which ninja), ) +GENERATOR := -GNinja + +# Ninja makes compilers disable colored output by default. +GENERATOR += -DVELOX_FORCE_COLORED_OUTPUT=ON +endif +endif + +ifndef USE_CCACHE +ifneq ($(shell which ccache), ) +USE_CCACHE=-DCMAKE_CXX_COMPILER_LAUNCHER=ccache +endif +endif + +NUM_THREADS ?= $(shell getconf _NPROCESSORS_CONF 2>/dev/null || echo 1) +CPU_TARGET ?= "avx" + +all: release #: Build the release version + +clean: #: Delete all build artifacts + rm -rf $(BUILD_BASE_DIR) + +submodules: + git submodule sync --recursive + git submodule update --init --recursive + +cmake: submodules #: Use CMake to create a Makefile build system + mkdir -p $(BUILD_BASE_DIR)/$(BUILD_DIR) && \ + cmake -B \ + "$(BUILD_BASE_DIR)/$(BUILD_DIR)" \ + ${CMAKE_FLAGS} \ + $(GENERATOR) \ + $(USE_CCACHE) \ + ${EXTRA_CMAKE_FLAGS} + +build: #: Build the software based in BUILD_DIR and BUILD_TYPE variables + cmake --build $(BUILD_BASE_DIR)/$(BUILD_DIR) -j $(NUM_THREADS) + +debug: #: Build with debugging symbols + $(MAKE) cmake BUILD_DIR=debug BUILD_TYPE=Debug + $(MAKE) build BUILD_DIR=debug -j ${NUM_THREADS} + +release: #: Build the release version + $(MAKE) cmake BUILD_DIR=release BUILD_TYPE=Release && \ + $(MAKE) build BUILD_DIR=release + +unittest: debug #: Build with debugging and run unit tests + cd $(BUILD_BASE_DIR)/debug && ctest -j ${NUM_THREADS} -VV --output-on-failure + +format-fix: #: Fix formatting issues in the main branch + scripts/format-check.py format main --fix + +format-check: #: Check for formatting issues on the main branch + clang-format --version + scripts/format-check.py format main + +header-fix: #: Fix license header issues in the current branch + scripts/format-check.py header main --fix + +header-check: #: Check for license header issues on the main branch + scripts/format-check.py header main + +help: #: Show the help messages + @cat $(firstword $(MAKEFILE_LIST)) | \ + awk '/^[-a-z]+:/' | \ + awk -F: '{ printf("%-20s %s\n", $$1, $$NF) }' diff --git a/dwio/nimble/common/BitEncoder.h b/dwio/nimble/common/BitEncoder.h new file mode 100644 index 0000000..aad2561 --- /dev/null +++ b/dwio/nimble/common/BitEncoder.h @@ -0,0 +1,75 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include + +namespace facebook::nimble { + +// Class to put bits into, and read bits from, a buffer. +class BitEncoder { + public: + // It's up to the user to not write/read beyond the provided buffer. + // Note that you probably want to ensure the memory is zero'd out. + explicit BitEncoder(char* buffer) + : buffer_(buffer), + writeWord_(reinterpret_cast(buffer)), + readWord_(reinterpret_cast(buffer)) {} + // We don't prevent you from using the putBits call after using this + // constructor, but you really shouldn't. + explicit BitEncoder(const char* buffer) + : BitEncoder(const_cast(buffer)) {} + + // Places |value| occupying |numBits| into the stream. Behavior + // is undefined if |value| is >= 2^numBits. + void putBits(uint64_t value, int numBits) { + *writeWord_ |= value << writeOffset_; + const int nextOffset = writeOffset_ + numBits; + if (nextOffset >= 64) { + ++writeWord_; + const int spilloverBits = nextOffset - 64; + if (numBits == 64 && spilloverBits == 0) { + return; + } + *writeWord_ |= value >> (numBits - spilloverBits); + writeOffset_ = spilloverBits; + } else { + writeOffset_ = nextOffset; + } + } + + uint64_t bitsWritten() { + return ((reinterpret_cast(writeWord_) - buffer_) << 3) + + writeOffset_; + } + + uint64_t getBits(int numBits) { + const int nextOffset = readOffset_ + numBits; + if (nextOffset >= 64) { + const uint64_t lowBits = *readWord_ >> readOffset_; + ++readWord_; + const int spilloverBits = nextOffset - 64; + readOffset_ = spilloverBits; + if (spilloverBits == 0) { + return lowBits; + } + return lowBits | + (*readWord_ & ((1ULL << spilloverBits) - 1ULL)) + << (numBits - spilloverBits); + } else { + const uint64_t result = + (*readWord_ >> readOffset_) & ((1ULL << numBits) - 1ULL); + readOffset_ = nextOffset; + return result; + } + } + + private: + char* buffer_; + uint64_t* writeWord_; + const uint64_t* readWord_; + int writeOffset_ = 0; + int readOffset_ = 0; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/Bits.cpp b/dwio/nimble/common/Bits.cpp new file mode 100644 index 0000000..3458436 --- /dev/null +++ b/dwio/nimble/common/Bits.cpp @@ -0,0 +1,183 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include "dwio/nimble/common/Bits.h" +#include "velox/common/base/BitUtil.h" + +namespace facebook::nimble::bits { + +void packBitmap(std::span bools, char* bitmap) { + uint64_t* word = reinterpret_cast(bitmap); + const uint64_t loopCount = bools.size() >> 6; + const uint64_t remainder = bools.size() - (loopCount << 6); + const bool* rawBools = bools.data(); + // TODO: use SIMD? + for (uint64_t i = 0; i < loopCount; ++i) { + for (int j = 0; j < 64; ++j) { + *word |= static_cast(*rawBools++) << j; + } + ++word; + } + for (int j = 0; j < remainder; ++j) { + *word |= static_cast(*rawBools++) << j; + } +} + +uint32_t countSetBits(uint32_t row, uint32_t rowCount, const char* bitmap) { + // Count bits remaining in current partial word, then count bits in + // each word, then finish with partial final word. + uint32_t nonNullCount = 0; + const uint32_t currentWord = row >> 6; + const uint32_t currentUsedBits = row - (currentWord << 6); + const uint64_t* word = + reinterpret_cast(bitmap) + currentWord; + if (currentUsedBits) { + // Do we only need to read bits from this one word? + if (row + rowCount <= (currentWord << 6) + 64) { + const uint64_t mask = (1ULL << rowCount) - 1; + return __builtin_popcountll((*word++ >> currentUsedBits) & mask); + } + nonNullCount += __builtin_popcountll(*word++ >> currentUsedBits); + rowCount -= (64 - currentUsedBits); + } + // We're now trued up to a word boundary. + const uint32_t wordCount = rowCount >> 6; + const uint32_t remainder = rowCount - (wordCount << 6); + for (uint32_t i = 0; i < wordCount; ++i) { + nonNullCount += __builtin_popcountll(*word++); + } + // Add in the bits in the remaining partial word. + if (remainder > 0) { + const uint64_t mask = (1ULL << remainder) - 1; + nonNullCount += __builtin_popcountll(*word & mask); + } + return nonNullCount; +} + +uint32_t findNonNullIndices( + uint32_t row, + uint32_t rowCount, + const char* bitmap, + uint32_t* indices) { + uint32_t* nextIndex = indices; + const uint32_t currentWord = row >> 6; + const uint32_t currentIndex = currentWord << 6; + const uint32_t currentUsedBits = row - currentIndex; + const uint64_t* nextWord = + reinterpret_cast(bitmap) + currentWord; + // Consume the remaining bits in the current word. + int index = -1; + if (currentUsedBits) { + uint64_t word = *nextWord++ >> currentUsedBits; + // Do we only need to read bits from this one word? + if (row + rowCount <= currentIndex + 64) { + word &= ((1ULL << rowCount) - 1); + while (int setBit = __builtin_ffsll(word)) { + index += setBit; + word >>= setBit; + *nextIndex++ = index; + } + return nextIndex - indices; + } + while (int setBit = __builtin_ffsll(word)) { + index += setBit; + word >>= setBit; + *nextIndex++ = index; + } + index = 63 - currentUsedBits; + rowCount -= (64 - currentUsedBits); + } + // We're now trued up to a word boundary. + const uint32_t wordCount = rowCount >> 6; + const uint32_t remainder = rowCount - (wordCount << 6); + for (uint32_t i = 0; i < wordCount; ++i) { + uint64_t word = *nextWord++; + uint32_t tempIndex = index; + constexpr uint64_t shiftEdgeCase = 1ULL << 63; + if (word == shiftEdgeCase) { + *nextIndex++ = tempIndex + 64; + } else { + while (int setBit = __builtin_ffsll(word)) { + tempIndex += setBit; + word >>= setBit; + *nextIndex++ = tempIndex; + } + } + index += 64; + } + // Parse last partial word. + if (remainder > 0) { + const uint64_t mask = (1ULL << remainder) - 1; + uint64_t word = *nextWord & mask; + while (int setBit = __builtin_ffsll(word)) { + index += setBit; + word >>= setBit; + *nextIndex++ = index; + } + } + return nextIndex - indices; +} + +uint32_t findNullIndices( + uint32_t row, + uint32_t rowCount, + const char* bitmap, + uint32_t* indices) { + uint32_t* nextIndex = indices; + for (uint32_t i = 0; i < rowCount; ++i) { + // TODO: optimize this like FindNonNullIndices. + if (getBit(row + i, bitmap)) { + *nextIndex++ = i; + } + } + return nextIndex - indices; +} + +void setBits(uint64_t begin, uint64_t end, char* bitmap) { + auto wordPtr = reinterpret_cast(bitmap); + velox::bits::forEachWord( + begin, + end, + [wordPtr](int32_t idx, uint64_t mask) { + wordPtr[idx] |= static_cast(-1) & mask; + }, + [wordPtr](int32_t idx) { wordPtr[idx] = static_cast(-1); }); +} + +void clearBits(uint64_t begin, uint64_t end, char* bitmap) { + auto wordPtr = reinterpret_cast(bitmap); + velox::bits::forEachWord( + begin, + end, + [wordPtr](int32_t idx, uint64_t mask) { wordPtr[idx] &= ~mask; }, + [wordPtr](int32_t idx) { wordPtr[idx] = 0; }); +} + +uint32_t +findSetBit(const char* bitmap, uint32_t begin, uint32_t end, uint32_t n) { + // TODO: maybe there is something fancy in bmi2 + while (begin < end) { + if (getBit(begin, bitmap) && --n == 0) { + break; + } + ++begin; + } + return begin; +} + +void BitmapBuilder::copy(const Bitmap& other, uint32_t begin, uint32_t end) { + auto source = static_cast(other.bits()); + auto dest = static_cast(bitmap_); + auto firstByte = begin / 8; + if (begin % 8) { + uint8_t mask = (1 << (begin % 8)) - 1; + dest[firstByte] = (dest[firstByte] & mask) | (source[firstByte] & ~mask); + ++firstByte; + } + // @lint-ignore CLANGSECURITY facebook-security-vulnerable-memcpy + std::memcpy( + dest + firstByte, + source + firstByte, + bits::bytesRequired(end) - firstByte); +} + +} // namespace facebook::nimble::bits diff --git a/dwio/nimble/common/Bits.h b/dwio/nimble/common/Bits.h new file mode 100644 index 0000000..24fa267 --- /dev/null +++ b/dwio/nimble/common/Bits.h @@ -0,0 +1,193 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include +#include + +// Common functions related to bits/bit-packing. +// +// The inlined methods are too simple to test or dont need tests +// as they are for debugging. The non-inlined methods are tested +// indirectly through the nullableColumnTests.cpp. + +namespace facebook::nimble::bits { + +// Returns the number of bits required to store the value. +// For a value of 0 we return 1. +inline int bitsRequired(uint64_t value) noexcept { + return 64 - __builtin_clzll(value | 1); +} + +// How many buckets are required to hold some elements? +inline uint64_t bucketsRequired( + uint64_t elementCount, + uint64_t bucketSize) noexcept { + return elementCount / bucketSize + (elementCount % bucketSize != 0); +} + +// Num bytes required to hold X bits. +inline uint64_t bytesRequired(uint64_t bitCount) noexcept { + return bucketsRequired(bitCount, 8); +} + +// Sets the |n|th bit to true. +inline void setBit(int64_t n, char* c) { + const int64_t byte = n >> 3; + const int64_t remainder = n & 7; + c[byte] |= (1 << remainder); +} + +// Sets the |n|th bit to |bit|. Note that if bit is false, this +// is a no-op. The intention is to let you avoid a branch. +inline void maybeSetBit(int64_t n, char* c, bool bit) { + const int64_t byte = n >> 3; + const int64_t remainder = n & 7; + c[byte] |= (bit << remainder); +} + +// Sets the |n|th bit to false. +inline void clearBit(int64_t n, char* c) { + const int64_t byte = n >> 3; + const int64_t remainder = n & 7; + c[byte] &= ~(1 << remainder); +} + +// Retrieves the |n|th bit. +inline bool getBit(int64_t n, const char* c) { + const int64_t byte = n >> 3; + const int64_t remainder = n & 7; + return c[byte] & (1 << remainder); +} + +// Sets bits [|begin|, |end|) in |bitmap| to true. Expects |bitmap| to be word +// aligned. +void setBits(uint64_t begin, uint64_t end, char* bitmap); + +// Sets bits [|begin|, |end|) in |bitmap| to false. Expects |bitmap| to be word +// aligned. +void clearBits(uint64_t begin, uint64_t end, char* bitmap); + +// Packs |bools| in bitmap. bitmap must point to a region of at least +// FixedBitArray::BufferSize(nulls.size(), 1) bytes. Note that we do not +// first clear the bytes being written to in |bitmap|, so if |bitmap| does +// not point to a zeroed out region after this call bit i will be set if +// bools[i] is true OR bit i was originally true. +void packBitmap(std::span bools, char* bitmap); + +// Returns how many of the bits in the |bitmap| starting at index |row| and +// going |rowCount| forward are are non-null. +uint32_t countSetBits(uint32_t row, uint32_t rowCount, const char* bitmap); + +// Finds the indices (relative to the row) of the non-null indices (i.e. set +// bits) among the next |rowCount| rows. |indices| should point to a buffer +// Y enough to hold rowCount values. We return the number of non-nulls. +uint32_t findNonNullIndices( + uint32_t row, + uint32_t rowCount, + const char* bitmap, + uint32_t* indices); + +// Same as above, but returns the indices of the nulls (i.e. unset bits). +uint32_t findNullIndices( + uint32_t row, + uint32_t rowCount, + const char* bitmap, + uint32_t* indices); + +// Finds the index of |n|th set bit in [|begin|, |end|) in |bitmap|. If not +// found, return |end|. +uint32_t +findSetBit(const char* bitmap, uint32_t begin, uint32_t end, uint32_t n); + +// Prints out the bits in numeric type in nibbles. Not efficient, +// only should be used for debugging/prototyping. +template +std::string printBits(T c) { + std::string result; + for (int i = 0; i < (sizeof(T) << 3); ++i) { + if (i > 0 && i % 4 == 0) { + result += ' '; + } + if (c & 1) { + result += '1'; + } else { + result += '0'; + } + c >>= 1; + } + // We actually want little endian order. + std::reverse(result.begin(), result.end()); + return result; +} + +template +std::string printHex(T c) { + std::string result(2 + 2 * sizeof(T), 0); + result[0] = '0'; + result[1] = 'x'; + char* pos = result.data() + result.size() - 1; + for (int i = 0; i < 2 * sizeof(T); ++i) { + const uint32_t nibble = c & 15; + if (nibble < 10) { + *pos = nibble + '0'; + } else { + *pos = (nibble - 10) + 'a'; + } + c >>= 4; + --pos; + } + return result; +} + +class Bitmap { + public: + Bitmap(const void* bitmap, uint32_t size) + : bitmap_{static_cast(const_cast(bitmap))}, size_{size} {} + + bool test(uint32_t pos) const { + return getBit(pos, bitmap_); + } + + uint32_t size() const { + return size_; + } + + const void* bits() const { + return bitmap_; + } + + protected: + char* bitmap_; + uint32_t size_; +}; + +class BitmapBuilder : public Bitmap { + public: + BitmapBuilder(void* bitmap, uint32_t size) : Bitmap{bitmap, size} {} + + void set(uint32_t pos) { + setBit(pos, bitmap_); + } + + void maybeSet(uint32_t pos, bool bit) { + maybeSetBit(pos, bitmap_, bit); + } + + void set(uint32_t begin, uint32_t end) { + setBits(begin, end, bitmap_); + } + + void clear(uint32_t begin, uint32_t end) { + clearBits(begin, end, bitmap_); + } + + // Copy the specified range from the source bitmap into this one. It + // guarantees |begin| is the beginning bit offset, but may copy more beyond + // |end|. + void copy(const Bitmap& other, uint32_t begin, uint32_t end); +}; + +} // namespace facebook::nimble::bits diff --git a/dwio/nimble/common/Buffer.h b/dwio/nimble/common/Buffer.h new file mode 100644 index 0000000..af7d2d1 --- /dev/null +++ b/dwio/nimble/common/Buffer.h @@ -0,0 +1,95 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include "velox/buffer/Buffer.h" +#include "velox/common/memory/Memory.h" + +#include +#include +#include +#include + +// Basic memory buffer interface, aka arena. +// +// Standard usage: +// char* pos = buffer->reserve(100); +// write data to [pos, pos + 100) +// pos = buffer->Reserve(27); +// write data to [pos, pos + 27) +// and so on + +namespace facebook::nimble { + +// Internally manages memory in chunks. Releases memory only upon destruction. +// Buffer is NOT threadsafe: external locking is required. +class Buffer { + using MemoryPool = facebook::velox::memory::MemoryPool; + + public: + explicit Buffer( + MemoryPool& memoryPool, + uint64_t initialChunkSize = kMinChunkSize) + : memoryPool_(memoryPool) { + addChunk(initialChunkSize); + reserveEnd_ = pos_; + } + + // Returns a pointer to a block of memory of size bytes that can be written + // to, and guarantees for the lifetime of *this that that region will remain + // valid. Does NOT guarantee that the region is initially 0'd. + char* reserve(uint64_t bytes) { + std::scoped_lock l(mutex_); + if (reserveEnd_ + bytes <= chunkEnd_) { + pos_ = reserveEnd_; + reserveEnd_ += bytes; + } else { + addChunk(bytes); + } + return pos_; + } + + // Copies |data| into the chunk, returning a view to the copied data. + std::string_view writeString(std::string_view data) { + char* pos = reserve(data.size()); + // @lint-ignore CLANGSECURITY facebook-security-vulnerable-memcpy + std::memcpy(pos, data.data(), data.size()); + return {pos, data.size()}; + } + + MemoryPool& getMemoryPool() { + return memoryPool_; + } + + std::string_view takeOwnership(velox::BufferPtr&& bufferPtr) { + std::string_view chunk{bufferPtr->as(), bufferPtr->size()}; + chunks_.push_back(std::move(bufferPtr)); + return chunk; + } + + private: + static constexpr uint64_t kMinChunkSize = 1LL << 20; + + void addChunk(uint64_t bytes) { + const uint64_t chunkSize = std::max(bytes, kMinChunkSize); + auto bufferPtr = + velox::AlignedBuffer::allocate(chunkSize, &memoryPool_); + pos_ = bufferPtr->asMutable(); + chunkEnd_ = pos_ + chunkSize; + reserveEnd_ = pos_ + bytes; + chunks_.push_back(std::move(bufferPtr)); + } + + char* chunkEnd_; + char* pos_; + char* reserveEnd_; + std::vector chunks_; + MemoryPool& memoryPool_; + // NOTE: this is temporary fix, to quickly enable parallel access to the + // buffer class. In the near future, we are going to templetize this class to + // produce a concurrent and a non-concurrent variants, and change the call + // sites to use each variant when needed. + std::mutex mutex_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/CMakeLists.txt b/dwio/nimble/common/CMakeLists.txt new file mode 100644 index 0000000..def7944 --- /dev/null +++ b/dwio/nimble/common/CMakeLists.txt @@ -0,0 +1,13 @@ +add_subdirectory(tests) + +add_library( + nimble_common + Bits.cpp + Checksum.cpp + FixedBitArray.cpp + Huffman.cpp + MetricsLogger.cpp + Types.cpp + Varint.cpp) + +target_link_libraries(nimble_common velox_memory Folly::folly) diff --git a/dwio/nimble/common/Checksum.cpp b/dwio/nimble/common/Checksum.cpp new file mode 100644 index 0000000..1b07dec --- /dev/null +++ b/dwio/nimble/common/Checksum.cpp @@ -0,0 +1,60 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#include "dwio/nimble/common/Checksum.h" +#include "dwio/nimble/common/Exceptions.h" + +#define XXH_INLINE_ALL +#include + +namespace facebook::nimble { + +namespace { +class Xxh3_64Checksum : public Checksum { + public: + Xxh3_64Checksum() : state_{XXH3_createState()} { + NIMBLE_DASSERT(state_ != nullptr, "Failed to initialize Xxh3_64Checksum."); + reset(); + } + + ~Xxh3_64Checksum() override { + XXH3_freeState(state_); + } + + void update(std::string_view data) override { + auto result = XXH3_64bits_update(state_, data.data(), data.size()); + NIMBLE_ASSERT(result != XXH_ERROR, "XXH3_64bits_update error."); + } + + uint64_t getChecksum(bool reset) override { + auto ret = static_cast(XXH3_64bits_digest(state_)); + if (UNLIKELY(reset)) { + this->reset(); + } + return ret; + } + + ChecksumType getType() const override { + return ChecksumType::XXH3_64; + } + + private: + XXH3_state_t* state_; + + void reset() { + auto result = XXH3_64bits_reset(state_); + NIMBLE_ASSERT(result != XXH_ERROR, "XXH3_64bits_reset error."); + } +}; +} // namespace + +std::unique_ptr ChecksumFactory::create(ChecksumType type) { + switch (type) { + case ChecksumType::XXH3_64: + return std::make_unique(); + default: + NIMBLE_NOT_SUPPORTED( + fmt::format("Unsupported checksum type: {}", toString(type))); + } +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/Checksum.h b/dwio/nimble/common/Checksum.h new file mode 100644 index 0000000..d7af1bc --- /dev/null +++ b/dwio/nimble/common/Checksum.h @@ -0,0 +1,25 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/common/Types.h" + +#include +#include + +namespace facebook::nimble { + +class Checksum { + public: + virtual ~Checksum() = default; + virtual void update(std::string_view data) = 0; + virtual uint64_t getChecksum(bool reset = false) = 0; + virtual ChecksumType getType() const = 0; +}; + +class ChecksumFactory { + public: + static std::unique_ptr create(ChecksumType type); +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/DefaultMetricsLogger.cpp b/dwio/nimble/common/DefaultMetricsLogger.cpp new file mode 100644 index 0000000..e699e57 --- /dev/null +++ b/dwio/nimble/common/DefaultMetricsLogger.cpp @@ -0,0 +1,88 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/common/DefaultMetricsLogger.h" +#include +#include "dwio/nimble/common/Exceptions.h" + +namespace facebook::nimble { + +DefaultMetricsLogger::DefaultMetricsLogger( + const std::string& ns, + const std::string& table, + const std::string& hostName, + const std::string& clientId, + std::string queryId) + : ns_{ns}, + table_{table}, + hostName_{hostName}, + clientId_{clientId}, + queryId_{std::move(queryId)} { + NIMBLE_DASSERT(!queryId_.empty(), "Empty query id passed in!"); +} + +void DefaultMetricsLogger::populateAccessorInfo( + logger::XldbAlphaLogger& log) const { + // Do not set the unset fields for better queries. + if (LIKELY(!ns_.empty())) { + log.setNS(ns_); + } + if (LIKELY(!table_.empty())) { + log.setTable(table_); + } + if (LIKELY(!clientId_.empty())) { + log.setClient(clientId_); + } + log.setQueryID(queryId_); +} + +void DefaultMetricsLogger::logException( + std::string_view operation, + const std::string& errorMessage) const { + logger::XldbAlphaLogger log; + populateAccessorInfo(log); + log.setOperationSV(operation); + log.setError(errorMessage); + LOG_VIA_LOGGER_ASYNC(log); +} + +void DefaultMetricsLogger::logStripeLoad( + const StripeLoadMetrics& metrics) const { + logger::XldbAlphaLogger log; + populateAccessorInfo(log); + log.setOperationSV(kStripeLoadOperation); + log.setCPUTime(metrics.cpuUsec); + log.setWallTime(metrics.wallTimeUsec); + log.setSerializedRunStats(folly::toJson(metrics.serialize())); + LOG_VIA_LOGGER_ASYNC(log); +} + +void DefaultMetricsLogger::logStripeFlush( + const StripeFlushMetrics& metrics) const { + logger::XldbAlphaLogger log; + populateAccessorInfo(log); + log.setOperationSV(kStripeFlushOperation); + log.setCPUTime(metrics.flushCpuUsec); + log.setWallTime(metrics.flushWallTimeUsec); + log.setSerializedRunStats(folly::toJson(metrics.serialize())); + LOG_VIA_LOGGER_ASYNC(log); +} + +void DefaultMetricsLogger::logFileClose(const FileCloseMetrics& metrics) const { + logger::XldbAlphaLogger log; + populateAccessorInfo(log); + log.setOperationSV(kFileCloseOperation); + log.setCPUTime(metrics.totalFlushCpuUsec); + log.setWallTime(metrics.totalFlushWallTimeUsec); + log.setSerializedRunStats(folly::toJson(metrics.serialize())); + LOG_VIA_LOGGER_ASYNC(log); +} + +void DefaultMetricsLogger::logZstrongContext(const std::string& context) const { + logger::XldbAlphaLogger log; + populateAccessorInfo(log); + log.setOperationSV(kZstrong); + log.setSerializedDebugStats(context); + LOG_VIA_LOGGER_ASYNC(log); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/DefaultMetricsLogger.h b/dwio/nimble/common/DefaultMetricsLogger.h new file mode 100644 index 0000000..9063e8c --- /dev/null +++ b/dwio/nimble/common/DefaultMetricsLogger.h @@ -0,0 +1,38 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "common/strings/UUID.h" +#include "dsi/logger/configs/XldbAlphaLoggerConfig/Logger.h" +#include "dwio/nimble/common/MetricsLogger.h" + +namespace facebook::nimble { + +class DefaultMetricsLogger : public MetricsLogger { + public: + DefaultMetricsLogger( + const std::string& ns, + const std::string& table, + const std::string& hostName, + const std::string& clientId, + std::string queryId = strings::generateUUID()); + + void logException(std::string_view operation, const std::string& errorMessage) + const override; + + void logStripeLoad(const StripeLoadMetrics& metrics) const override; + void logStripeFlush(const StripeFlushMetrics& metrics) const override; + void logFileClose(const FileCloseMetrics& metrics) const override; + void logZstrongContext(const std::string& context) const override; + + private: + void populateAccessorInfo(logger::XldbAlphaLogger& log) const; + + std::string ns_; + std::string table_; + std::string hostName_; + std::string clientId_; + std::string queryId_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/EncodingPrimitives.h b/dwio/nimble/common/EncodingPrimitives.h new file mode 100644 index 0000000..5ae3161 --- /dev/null +++ b/dwio/nimble/common/EncodingPrimitives.h @@ -0,0 +1,121 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include + +#include "folly/io/IOBuf.h" + +// Primitives for reading and writing to a stream of bytes. + +namespace facebook::nimble::encoding { + +// The write functions serialize a datum into a buffer, advancing the +// buffer pointer appropriately. + +// Write a type in native format for numeric types, via WriteString +// for string types. +template +inline void write(T value, char*& pos) { + *reinterpret_cast(pos) = value; + pos += sizeof(T); +} + +inline void writeChar(char value, char*& pos) { + write(value, pos); +} + +inline void writeUint32(uint32_t value, char*& pos) { + write(value, pos); +} + +inline void writeUint64(uint64_t value, char*& pos) { + write(value, pos); +} + +// Just the chars, no leading length. +inline void writeBytes(std::string_view value, char*& pos) { + std::copy(value.cbegin(), value.cend(), pos); + pos += value.size(); +} + +// Just the buffers, no leading length. +inline void writeBuffers(const folly::IOBuf& buffers, char*& pos) { + for (const auto buffer : buffers) { + std::copy(buffer.cbegin(), buffer.cend(), pos); + pos += buffer.size(); + } +} + +// 4 byte length followed by chars. +inline void writeString(std::string_view value, char*& pos) { + writeUint32(value.size(), pos); + writeBytes(value, pos); +} + +template <> +inline void write(std::string_view value, char*& pos) { + writeString(value, pos); +} + +template <> +inline void write(const std::string& value, char*& pos) { + writeString({value.data(), value.size()}, pos); +} + +// The read functions extract a datum from the buffer (stored in the +// format output by the Write functions), advancing the buffer +// pointer appropriately. + +// Reads a type written via Write. +template +inline TReturn read(const char*& pos) { + const T value = *reinterpret_cast(pos); + pos += sizeof(T); + return static_cast(value); +} + +inline char readChar(const char*& pos) { + return read(pos); +} + +inline uint32_t readUint32(const char*& pos) { + return read(pos); +} + +inline uint64_t readUint64(const char*& pos) { + return read(pos); +} + +inline std::string_view readString(const char*& pos) { + const uint32_t size = readUint32(pos); + std::string_view result(pos, size); + pos += size; + return result; +} + +inline std::string readOwnedString(const char*& pos) { + const uint32_t size = readUint32(pos); + std::string result(pos, size); + pos += size; + return result; +} + +template <> +inline std::string_view read( + const char*& pos) { + return readString(pos); +} + +template <> +inline std::string read(const char*& pos) { + return readOwnedString(pos); +} + +template +inline TReturn peek(const char* pos) { + return static_cast(*reinterpret_cast(pos)); +} + +} // namespace facebook::nimble::encoding diff --git a/dwio/nimble/common/EncodingType.h b/dwio/nimble/common/EncodingType.h new file mode 100644 index 0000000..5979812 --- /dev/null +++ b/dwio/nimble/common/EncodingType.h @@ -0,0 +1,42 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include + +#include "dwio/nimble/common/Types.h" + +namespace facebook::nimble { + +// This class is used to manage semantic type and its corresponding encoding +// type. For most data types, they are the same. They will be different for +// floating point types. The main reason it is different for floating point +// numbers are because of +/-0.0 ( +0.0 == -0.0 but their sign bits are +// different), and NaNs can have different payload bits but even if NaNs +// with the same bits, == will return false, that causes trouble for encoding. +// So they will be treated as integers during encoding (their type will +// still be encoded as double/float) +template +struct EncodingPhysicalType { + using type = typename TypeTraits::physicalType; + + static std::span asEncodingPhysicalTypeSpan( + std::span values) { + return std::span( + reinterpret_cast(values.data()), values.size()); + } + + static type asEncodingPhysicalType(T v) { + return *reinterpret_cast(&v); + } + + static T asEncodingLogicalType(type v) { + return *reinterpret_cast(&v); + } +}; + +#define AS(D, v) *(reinterpret_cast(&(v))) + +#define AS_CONST(D, v) *(reinterpret_cast(&(v))) +} // namespace facebook::nimble diff --git a/dwio/nimble/common/Entropy.h b/dwio/nimble/common/Entropy.h new file mode 100644 index 0000000..3a8b2c9 --- /dev/null +++ b/dwio/nimble/common/Entropy.h @@ -0,0 +1,37 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include "folly/container/F14Map.h" + +// Entropy encoding. + +namespace facebook::nimble::entropy { + +// Returns the Shannon entropy from symbol counts. +inline double entropy(std::span counts, int sumOfCount) { + double entropy = 0; + for (int count : counts) { + const double p = double(count) / sumOfCount; + entropy -= p * log2(p); + } + return entropy; +} + +// Calculates the Shannon entropy of a data set. +template +double computeEntropy(std::span data) { + folly::F14FastMap counts; + for (auto datum : data) { + ++counts[datum]; + } + std::vector finalCounts; + finalCounts.reserve(counts.size()); + for (const auto& p : counts) { + finalCounts.push_back(p.second); + } + return entropy(finalCounts, data.size()); +} + +} // namespace facebook::nimble::entropy diff --git a/dwio/nimble/common/Exceptions.h b/dwio/nimble/common/Exceptions.h new file mode 100644 index 0000000..1033137 --- /dev/null +++ b/dwio/nimble/common/Exceptions.h @@ -0,0 +1,470 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include "folly/FixedString.h" +#include "folly/experimental/symbolizer/Symbolizer.h" +#include "folly/synchronization/CallOnce.h" + +// Standard errors used throughout the codebase. + +namespace facebook::nimble { + +namespace error_source { +using namespace folly::string_literals; + +// Errors where the root cause of the problem is either because of bad input +// or an unsupported pattern of use are classified with source USER. +inline constexpr auto User = "USER"_fs; + +// Errors where the root cause of the problem is an unexpected internal state in +// the system. +inline constexpr auto Internal = "INTERNAL"_fs; + +// Errors where the root cause of the problem is the result of a dependency or +// an environment failures. +inline constexpr auto External = "EXTERNAL"_fs; +} // namespace error_source + +namespace error_code { +using namespace folly::string_literals; + +// An error raised when an argument verification fails +inline constexpr auto InvalidArgument = "INVALID_ARGUMENT"_fs; + +// An error raised when the current state of a component is invalid. +inline constexpr auto InvalidState = "INVALID_STATE"_fs; + +// An error raised when unreachable code point was executed. +inline constexpr auto UnreachableCode = "UNREACHABLE_CODE"_fs; + +// An error raised when a requested operation is not yet implemented. +inline constexpr auto NotImplemented = "NOT_IMPLEMENTED"_fs; + +// An error raised when a requested operation is not supported. +inline constexpr auto NotSupported = "NOT_SUPPORTED"_fs; + +// As error raised during encoding optimization, when incompatible encoding is +// attempted. +inline constexpr auto IncompatibleEncoding = "INCOMPATIBLE_ENCODING"_fs; + +// Am error raised if a corrupted file is detected +inline constexpr auto CorruptedFile = "CORRUPTED_FILE"_fs; + +// We do not know how to classify it yet. +inline constexpr auto Unknown = "UNKNOWN"_fs; +} // namespace error_code + +namespace external_source { +using namespace folly::string_literals; + +// Warm Storage +inline constexpr auto WarmStorage = "WARM_STORAGE"_fs; + +// Local File System +inline constexpr auto LocalFileSystem = "FILE_SYSTEM"_fs; + +} // namespace external_source + +// Based exception, used by all other Nimble exceptions. Provides common +// functionality for all other exception types. +class NimbleException : public std::exception { + public: + explicit NimbleException( + std::string_view exceptionName, + const char* fileName, + size_t fileLine, + const char* functionName, + std::string_view failingExpression, + std::string_view errorMessage, + std::string_view errorCode, + bool retryable) + : exceptionName_{std::move(exceptionName)}, + fileName_{fileName}, + fileLine_{fileLine}, + functionName_{functionName}, + failingExpression_{std::move(failingExpression)}, + errorMessage_{std::move(errorMessage)}, + errorCode_{std::move(errorCode)}, + retryable_{retryable} { + captureStackTraceFrames(); + } + + const char* what() const noexcept { + try { + folly::call_once(once_, [&] { finalizeMessage(); }); + return finalizedMessage_.c_str(); + } catch (...) { + return ""; + } + } + + const std::string& exceptionName() const { + return exceptionName_; + } + + const char* fileName() const { + return fileName_; + } + + size_t fileLine() const { + return fileLine_; + } + + const char* functionName() const { + return functionName_; + } + + const std::string& failingExpression() const { + return failingExpression_; + } + + const std::string& errorMessage() const { + return errorMessage_; + } + + virtual const std::string_view errorSource() const = 0; + + const std::string& errorCode() const { + return errorCode_; + } + + bool retryable() const { + return retryable_; + } + + protected: + virtual void appendMessage(std::string& /* message */) const {} + + private: + void captureStackTraceFrames() { + try { + constexpr size_t skipFrames = 2; + constexpr size_t maxFrames = 200; + uintptr_t addresses[maxFrames]; + ssize_t n = folly::symbolizer::getStackTrace(addresses, maxFrames); + + if (n < skipFrames) { + return; + } + + exceptionFrames_.assign(addresses + skipFrames, addresses + n); + } catch (const std::exception& ex) { + LOG(WARNING) << "Unable to capture stack trace: " << ex.what(); + } catch (...) { // Should never happen, catchall + LOG(WARNING) << "Unable to capture stack trace."; + } + } + + void finalizeMessage() const { + finalizedMessage_ += exceptionName_; + finalizedMessage_ += "\nError Source: "; + finalizedMessage_ += errorSource(); + finalizedMessage_ += "\nError Code: "; + finalizedMessage_ += errorCode_; + if (!errorMessage_.empty()) { + finalizedMessage_ += "\nError Message: "; + finalizedMessage_ += errorMessage_; + } + finalizedMessage_ += "\n"; + appendMessage(finalizedMessage_); + finalizedMessage_ += "Retryable: "; + finalizedMessage_ += retryable_ ? "True" : "False"; + finalizedMessage_ += "\nLocation: "; + finalizedMessage_ += functionName_; + finalizedMessage_ += "@"; + finalizedMessage_ += fileName_; + finalizedMessage_ += ":"; + finalizedMessage_ += folly::to(fileLine_); + + if (!failingExpression_.empty()) { + finalizedMessage_ += "\nExpression: "; + finalizedMessage_ += failingExpression_; + } + + if (LIKELY(!exceptionFrames_.empty())) { + std::vector symbolizedFrames; + symbolizedFrames.resize(exceptionFrames_.size()); + + folly::symbolizer::Symbolizer symbolizer{ + folly::symbolizer::LocationInfoMode::FULL}; + symbolizer.symbolize( + exceptionFrames_.data(), + symbolizedFrames.data(), + symbolizedFrames.size()); + + folly::symbolizer::StringSymbolizePrinter printer{ + folly::symbolizer::StringSymbolizePrinter::COLOR}; + printer.println(symbolizedFrames.data(), symbolizedFrames.size()); + + finalizedMessage_ += "\nStack Trace:\n"; + finalizedMessage_ += printer.str(); + } + } + + std::vector exceptionFrames_; + const std::string stackTrace_; + const std::string exceptionName_; + const char* fileName_; + const size_t fileLine_; + const char* functionName_; + const std::string failingExpression_; + const std::string errorMessage_; + const std::string errorCode_; + const bool retryable_; + + mutable folly::once_flag once_; + mutable std::string finalizedMessage_; +}; + +// Exception representing an error originating by a user misusing Nimble. +class NimbleUserError : public NimbleException { + public: + NimbleUserError( + const char* fileName, + size_t fileLine, + const char* functionName, + std::string_view failingExpression, + std::string_view errorMessage, + std::string_view errorCode, + bool retryable) + : NimbleException( + "NimbleUserError", + fileName, + fileLine, + functionName, + failingExpression, + errorMessage, + errorCode, + retryable) {} + + const std::string_view errorSource() const override { + return error_source::User; + } +}; + +// Excpetion representing unexpected behavior in Nimble. This usually means a +// bug in Nimble. +class NimbleInternalError : public NimbleException { + public: + NimbleInternalError( + const char* fileName, + size_t fileLine, + const char* functionName, + std::string_view failingExpression, + std::string_view errorMessage, + std::string_view errorCode, + bool retryable) + : NimbleException( + "NimbleInternalError", + fileName, + fileLine, + functionName, + failingExpression, + errorMessage, + errorCode, + retryable) {} + + const std::string_view errorSource() const override { + return error_source::Internal; + } +}; + +// Exception representing an issue originating from an Nimble external +// dependency (for example, Warm Storage or file system). These exceptions +// should not affect Nimble's SLA. +class NimbleExternalError : public NimbleException { + public: + NimbleExternalError( + const char* fileName, + const size_t fileLine, + const char* functionName, + std::string_view failingExpression, + std::string_view errorMessage, + std::string_view errorCode, + bool retryable, + std::string_view externalSource) + : NimbleException( + "NimbleExternalError", + fileName, + fileLine, + functionName, + failingExpression, + errorMessage, + errorCode, + retryable), + externalSource_{externalSource} {} + + const std::string_view errorSource() const override { + return error_source::External; + } + + void appendMessage(std::string& message) const override { + message += "External Source: "; + message += externalSource_; + message += "\n"; + } + + private: + const std::string externalSource_; +}; + +#define _NIMBLE_RAISE_EXCEPTION( \ + exception, expression, message, code, retryable) \ + throw exception( \ + __FILE__, __LINE__, __FUNCTION__, expression, message, code, retryable) + +#define _NIMBLE_RAISE_EXCEPTION_EXTENDED( \ + exception, expression, message, code, retryable, ...) \ + throw exception( \ + __FILE__, \ + __LINE__, \ + __FUNCTION__, \ + expression, \ + message, \ + code, \ + retryable, \ + __VA_ARGS__) + +#define NIMBLE_RAISE_USER_ERROR(expression, message, code, retryable) \ + _NIMBLE_RAISE_EXCEPTION( \ + ::facebook::nimble::NimbleUserError, \ + expression, \ + message, \ + code, \ + retryable) + +#define NIMBLE_RAISE_INTERNAL_ERROR(expression, message, code, retryable) \ + _NIMBLE_RAISE_EXCEPTION( \ + ::facebook::nimble::NimbleInternalError, \ + expression, \ + message, \ + code, \ + retryable) + +#define NIMBLE_RAISE_EXTERNAL_ERROR( \ + expression, source, message, code, retryable) \ + _NIMBLE_RAISE_EXCEPTION_EXTENDED( \ + ::facebook::nimble::NimbleExternalError, \ + expression, \ + message, \ + code, \ + retryable, \ + source) + +// Check user related conditions. Failure of this condition means the user +// misused Nimble and will trigger a user error. +#define NIMBLE_CHECK(condition, message) \ + if (UNLIKELY(!(condition))) { \ + NIMBLE_RAISE_USER_ERROR( \ + #condition, \ + message, \ + ::facebook::nimble::error_code::InvalidArgument, \ + /* retryable */ false); \ + } + +// Assert an internal Nimble expected behavior. Failure of this condition means +// Nimble encountered an unexpected behavior and will trigger an internal error. +#define NIMBLE_ASSERT(condition, message) \ + if (UNLIKELY(!(condition))) { \ + NIMBLE_RAISE_INTERNAL_ERROR( \ + #condition, \ + message, \ + ::facebook::nimble::error_code::InvalidState, \ + /* retryable */ false); \ + } + +// Verify a result from an external Nimble dependency. Failure of this condition +// means an external dependency returned an error and all retries (if +// applicable) were exhausted. This will trigger an external error. +#define NIMBLE_VERIFY_EXTERNAL(condition, source, code, retryable, message) \ + if (UNLIKELY(!(condition))) { \ + NIMBLE_RAISE_EXTERNAL_ERROR( \ + #condition, \ + ::facebook::nimble::external_source::source, \ + message, \ + code, \ + retryable); \ + } + +// Verify an expected file format conditions. Failure of this condition means +// the file is corrupted (e.g. passed magic number and version verification, but +// got unexpected format). This will trigger a user error. +#define NIMBLE_CHECK_FILE(condition, message) \ + if (UNLIKELY(!(condition))) { \ + NIMBLE_RAISE_USER_ERROR( \ + #condition, \ + message, \ + ::facebook::nimble::error_code::CorruptedFile, \ + /* retryable */ false); \ + } + +// Should be raised when we don't expect to hit a code path, but we did. This +// means a bug in Nimble. +#define NIMBLE_UNREACHABLE(message) \ + NIMBLE_RAISE_INTERNAL_ERROR( \ + "", \ + message, \ + ::facebook::nimble::error_code::UnreachableCode, \ + /* retryable */ false); + +// Should be raised in places where we still didn't implement the required +// functionality, but intend to do so in the future. This raises an internal +// error to indicate users needs this functionality, but we don't provide it. +#define NIMBLE_NOT_IMPLEMENTED(message) \ + NIMBLE_RAISE_INTERNAL_ERROR( \ + "", \ + message, \ + ::facebook::nimble::error_code::NotImplemented, \ + /* retryable */ false); + +// Should be raised in places where we don't support a functionality, and have +// no intention to support it in the future. This raises a user error, as the +// user should not expect this functionality to exist in the first place. +#define NIMBLE_NOT_SUPPORTED(message) \ + NIMBLE_RAISE_USER_ERROR( \ + "", \ + message, \ + ::facebook::nimble::error_code::NotSupported, \ + /* retryable */ false); + +// Incompatible Encoding errors are used in Nimble's encoding optimization, to +// indicate that an attempted encoding is incompatible with the data and should +// be avoided. +#define NIMBLE_INCOMPATIBLE_ENCODING(message) \ + NIMBLE_RAISE_USER_ERROR( \ + "", \ + message, \ + ::facebook::nimble::error_code::IncompatibleEncoding, \ + /* retryable */ false); + +// Should be used in "catch all" exception handlers, where we can't classify the +// error correctly. These errors mean that we are missing error classification. +#define NIMBLE_UNKNOWN(message) \ + NIMBLE_RAISE_INTERNAL_ERROR( \ + "", \ + message, \ + ::facebook::nimble::error_code::Unknown, \ + /* retryable */ true); + +// Should be used in "catch all" exception handlers wrapping external Nimble +// dependencies, where we can't classify the error correctly. These errors mean +// that we are missing error classification. +#define NIMBLE_UNKNOWN_EXTERNAL(source, message) \ + NIMBLE_RAISE_EXTERNAL_ERROR( \ + "", \ + ::facebook::nimble::external_source::source, \ + message, \ + ::facebook::nimble::error_code::Unknown, \ + /* retryable */ true); + +#ifndef NDEBUG +#define NIMBLE_DCHECK(condition, message) NIMBLE_CHECK(condition, message) +#define NIMBLE_DASSERT(condition, message) NIMBLE_ASSERT(condition, message) +#else +#define NIMBLE_DCHECK(condition, message) NIMBLE_CHECK(true, message) +#define NIMBLE_DASSERT(condition, message) NIMBLE_ASSERT(true, message) +#endif + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/FixedBitArray.cpp b/dwio/nimble/common/FixedBitArray.cpp new file mode 100644 index 0000000..f348058 --- /dev/null +++ b/dwio/nimble/common/FixedBitArray.cpp @@ -0,0 +1,701 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include "dwio/nimble/common/FixedBitArray.h" +#include + +namespace facebook::nimble { + +// Warning: do not change this function or lots of horrible data corruption +// will probably happen. +uint64_t FixedBitArray::bufferSize(uint64_t elementCount, int bitWidth) { + // We may read up to 7 bytes beyond the last theoretically needed byte, + // as we access whole machine words. + constexpr int kSlopSize = 7; + return bits::bytesRequired(elementCount * bitWidth) + kSlopSize; +} + +FixedBitArray::FixedBitArray(char* buffer, int bitWidth) + : buffer_(buffer), bitWidth_(bitWidth) { + DCHECK(bitWidth_ >= 0); + DCHECK(bitWidth_ <= 64); + // mask_ has the first bitWidth bits set to 1, rest 0. + mask_ = bitWidth == 64 ? (~0ULL) : ((1ULL << bitWidth_) - 1); +} + +uint64_t FixedBitArray::get(uint64_t index) const { + const uint64_t bits = index * bitWidth_; + const uint64_t offset = bits >> 3; + const uint64_t remainder = bits & 7; + const uint64_t word = *reinterpret_cast(buffer_ + offset); + // For widths > 58 bits, the value may overflow into the next word. + if (bitWidth_ > 58) { + const int overflow = bitWidth_ + remainder - 64; + if (overflow > 0) { + const uint64_t nextWord = + *reinterpret_cast(buffer_ + offset + 8); + return ((word >> remainder) | (nextWord << (bitWidth_ - overflow))) & + mask_; + } + } + return (word >> remainder) & mask_; +} + +uint32_t FixedBitArray::get32(uint64_t index) const { + const uint64_t bits = index * bitWidth_; + const uint64_t offset = bits >> 3; + const uint64_t remainder = bits & 7; + const uint64_t word = *reinterpret_cast(buffer_ + offset); + // Don't have to worry about overflow here since bitWidth_ <= 32. + return (word >> remainder) & mask_; +} + +void FixedBitArray::set(uint64_t index, uint64_t value) { + const uint64_t bits = index * bitWidth_; + const uint64_t offset = bits >> 3; + const uint64_t remainder = bits & 7; + uint64_t& word = *reinterpret_cast(buffer_ + offset); + // For widths > 58 bits, the value may overflow into the next word. + if (bitWidth_ > 58) { + const int overflow = bitWidth_ + remainder - 64; + if (overflow > 0) { + uint64_t& nextWord = *reinterpret_cast(buffer_ + offset + 8); + nextWord |= value >> (bitWidth_ - overflow); + } + } + word |= value << remainder; +} + +void FixedBitArray::set32(uint64_t index, uint32_t value) { + const uint64_t bits = index * bitWidth_; + const uint64_t offset = bits >> 3; + const uint64_t remainder = bits & 7; + uint64_t& word = *reinterpret_cast(buffer_ + offset); + // Don't have to worry about overflow here since bitWidth_ <= 32. + word |= static_cast(value) << remainder; +} + +void FixedBitArray::zeroAndSet(uint64_t index, uint64_t value) { + const uint64_t bits = index * bitWidth_; + const uint64_t offset = bits >> 3; + const uint64_t remainder = bits & 7; + uint64_t& word = *reinterpret_cast(buffer_ + offset); + word &= ~(mask_ << remainder); + // For widths > 58 bits, the value may overflow into the next word. + if (bitWidth_ > 58) { + const int overflow = bitWidth_ + remainder - 64; + if (overflow > 0) { + uint64_t& nextWord = *reinterpret_cast(buffer_ + offset + 8); + nextWord &= 0xFFFFFFFFFFFFFFFF << overflow; + nextWord |= value >> (bitWidth_ - overflow); + } + } + word |= value << remainder; +} + +namespace internal { + +// T is the output data types, namely uint32_t or uint64_t. +template +void bulkGet32Loop( + uint64_t& word, + const uint64_t** nextWord, + T** values, + T baseline) { + constexpr uint64_t kMask = (1ULL << bitWidth) - 1ULL; + // Some bits for the next value may still be present in word. + constexpr uint64_t spillover = (loopPosition * bitWidth) % 64 == 0 + ? 0 + : ((loopPosition + 1) * bitWidth) % 64; + // Note that this isn't a real branch as its on a constexpr. + if constexpr (spillover > 0) { + uint64_t remainder = word; + word = **nextWord; + ++(*nextWord); + if constexpr (withBaseline) { + **values = + ((remainder | word << (bitWidth - spillover)) & kMask) + baseline; + } else { + **values = (remainder | word << (bitWidth - spillover)) & kMask; + } + word >>= spillover; + ++(*values); + } else { + word = **nextWord; + ++(*nextWord); + } + // How many remaining values are in this word? + constexpr int valueCount = (64 - spillover) / bitWidth; + for (int i = 0; i < valueCount; ++i) { + if constexpr (withBaseline) { + **values = (word & kMask) + baseline; + } else { + **values = word & kMask; + } + ++(*values); + word >>= bitWidth; + } + constexpr int nextLoopPosition = loopPosition + valueCount + (spillover > 0); + bulkGet32Loop( + word, nextWord, values, baseline); +} + +// Unfortunately we cannot partially specialize the template for the +// terminal case of loopPosition = 64 so we must explicitly specify them. +#define BULK_GET32_LOOP_TERMINAL_CASE(bitWidth) \ + template <> \ + void bulkGet32Loop( \ + uint64_t & word, \ + const uint64_t** nextWord, \ + uint32_t** values, \ + uint32_t baseline) {} \ + template <> \ + void bulkGet32Loop( \ + uint64_t & word, \ + const uint64_t** nextWord, \ + uint64_t** values, \ + uint64_t baseline) {} \ + template <> \ + void bulkGet32Loop( \ + uint64_t & word, \ + const uint64_t** nextWord, \ + uint32_t** values, \ + uint32_t baseline) {} \ + template <> \ + void bulkGet32Loop( \ + uint64_t & word, \ + const uint64_t** nextWord, \ + uint64_t** values, \ + uint64_t baseline) {} + +BULK_GET32_LOOP_TERMINAL_CASE(1) +BULK_GET32_LOOP_TERMINAL_CASE(2) +BULK_GET32_LOOP_TERMINAL_CASE(3) +BULK_GET32_LOOP_TERMINAL_CASE(4) +BULK_GET32_LOOP_TERMINAL_CASE(5) +BULK_GET32_LOOP_TERMINAL_CASE(6) +BULK_GET32_LOOP_TERMINAL_CASE(7) +BULK_GET32_LOOP_TERMINAL_CASE(8) +BULK_GET32_LOOP_TERMINAL_CASE(9) +BULK_GET32_LOOP_TERMINAL_CASE(10) +BULK_GET32_LOOP_TERMINAL_CASE(11) +BULK_GET32_LOOP_TERMINAL_CASE(12) +BULK_GET32_LOOP_TERMINAL_CASE(13) +BULK_GET32_LOOP_TERMINAL_CASE(14) +BULK_GET32_LOOP_TERMINAL_CASE(15) +BULK_GET32_LOOP_TERMINAL_CASE(16) +BULK_GET32_LOOP_TERMINAL_CASE(17) +BULK_GET32_LOOP_TERMINAL_CASE(18) +BULK_GET32_LOOP_TERMINAL_CASE(19) +BULK_GET32_LOOP_TERMINAL_CASE(20) +BULK_GET32_LOOP_TERMINAL_CASE(21) +BULK_GET32_LOOP_TERMINAL_CASE(22) +BULK_GET32_LOOP_TERMINAL_CASE(23) +BULK_GET32_LOOP_TERMINAL_CASE(24) +BULK_GET32_LOOP_TERMINAL_CASE(25) +BULK_GET32_LOOP_TERMINAL_CASE(26) +BULK_GET32_LOOP_TERMINAL_CASE(27) +BULK_GET32_LOOP_TERMINAL_CASE(28) +BULK_GET32_LOOP_TERMINAL_CASE(29) +BULK_GET32_LOOP_TERMINAL_CASE(30) +BULK_GET32_LOOP_TERMINAL_CASE(31) +BULK_GET32_LOOP_TERMINAL_CASE(32) + +#undef BULK_GET32_LOOP_TERMINAL_CASE + +template +void bulkGet32Internal( + const FixedBitArray& fixedBitArray, + const char* buffer, + uint64_t start, + uint64_t length, + T* values, + T baseline) { + // Every 64 elements we know the slot will end on a word boundary. + // (It might end on a boundary before that, which should technically + // let us be more efficient if we used that instead because we can use + // the non-bulk API less, but the benefit is pretty small so we do the + // simple thing for now). + // + // We first use the non-bulk method to align ourselves to a 64-element + // boundary, then dispatch to the appropriate bit-width-specific code + // for the 64-element loops, then finish with the non-bulk method. + const uint64_t alignedStart = bits::bucketsRequired(start, 64) << 6; + if (start + length < alignedStart) { + for (uint64_t i = start; i < start + length; ++i) { + if constexpr (withBaseline) { + // TODO: An alternative would be to have a separate (constexpr) loop + // adding baselines at the end and hopefully the compiler will vectorize + // it (with -mavx2?). But it might be slower due to the extra loop + // condition cvhecks. Needto benchmark it. + *values = fixedBitArray.get32(i) + baseline; + } else { + *values = fixedBitArray.get32(i); + } + ++values; + } + return; + } + for (uint64_t i = start; i < alignedStart; ++i) { + if constexpr (withBaseline) { + *values = fixedBitArray.get32(i) + baseline; + } else { + *values = fixedBitArray.get32(i); + } + ++values; + } + const uint64_t loopCount = (length - (alignedStart - start)) >> 6; + switch (fixedBitArray.bitWidth()) { +#define BULK_GET32_SWITCH_CASE(bitWidth) \ + case bitWidth: { \ + const uint64_t* nextWord = reinterpret_cast( \ + buffer + ((alignedStart * bitWidth) >> 3)); \ + uint64_t word; \ + for (uint64_t i = 0; i < loopCount; ++i) { \ + bulkGet32Loop( \ + word, &nextWord, &values, baseline); \ + } \ + break; \ + } + + BULK_GET32_SWITCH_CASE(1) + BULK_GET32_SWITCH_CASE(2) + BULK_GET32_SWITCH_CASE(3) + BULK_GET32_SWITCH_CASE(4) + BULK_GET32_SWITCH_CASE(5) + BULK_GET32_SWITCH_CASE(6) + BULK_GET32_SWITCH_CASE(7) + BULK_GET32_SWITCH_CASE(8) + BULK_GET32_SWITCH_CASE(9) + BULK_GET32_SWITCH_CASE(10) + BULK_GET32_SWITCH_CASE(11) + BULK_GET32_SWITCH_CASE(12) + BULK_GET32_SWITCH_CASE(13) + BULK_GET32_SWITCH_CASE(14) + BULK_GET32_SWITCH_CASE(15) + BULK_GET32_SWITCH_CASE(16) + BULK_GET32_SWITCH_CASE(17) + BULK_GET32_SWITCH_CASE(18) + BULK_GET32_SWITCH_CASE(19) + BULK_GET32_SWITCH_CASE(20) + BULK_GET32_SWITCH_CASE(21) + BULK_GET32_SWITCH_CASE(22) + BULK_GET32_SWITCH_CASE(23) + BULK_GET32_SWITCH_CASE(24) + BULK_GET32_SWITCH_CASE(25) + BULK_GET32_SWITCH_CASE(26) + BULK_GET32_SWITCH_CASE(27) + BULK_GET32_SWITCH_CASE(28) + BULK_GET32_SWITCH_CASE(29) + BULK_GET32_SWITCH_CASE(30) + BULK_GET32_SWITCH_CASE(31) + BULK_GET32_SWITCH_CASE(32) + +#undef BULK_GET32_SWITCH_CASE + + default: + LOG(FATAL) << "bit width must lie in [1, 32], got: " + << fixedBitArray.bitWidth(); + } + const uint64_t remainderStart = alignedStart + (loopCount << 6); + const uint64_t remainderEnd = start + length; + for (uint64_t i = remainderStart; i < remainderEnd; ++i) { + if constexpr (withBaseline) { + *values = fixedBitArray.get32(i) + baseline; + } else { + *values = fixedBitArray.get32(i); + } + ++values; + } + return; +} + +} // namespace internal + +void FixedBitArray::bulkGet32(uint64_t start, uint64_t length, uint32_t* values) + const { + internal::bulkGet32Internal( + *this, buffer_, start, length, values, 0); +} + +void FixedBitArray::bulkGet32Into64( + uint64_t start, + uint64_t length, + uint64_t* values) const { + internal::bulkGet32Internal( + *this, buffer_, start, length, values, 0); +} + +void FixedBitArray::bulkGetWithBaseline32( + uint64_t start, + uint64_t length, + uint32_t* values, + uint32_t baseline) const { + internal::bulkGet32Internal( + *this, buffer_, start, length, values, baseline); +} + +void FixedBitArray::bulkGetWithBaseline32Into64( + uint64_t start, + uint64_t length, + uint64_t* values, + uint64_t baseline) const { + internal::bulkGet32Internal( + *this, buffer_, start, length, values, baseline); +} + +template +void bulkSet32Loop( + uint64_t** nextWord, + const uint32_t** values, + uint32_t baseline) { + // Some bits for the next value may need to put be in the current word. + constexpr int spillover = (loopPosition * bitWidth) % 64 == 0 + ? 0 + : ((loopPosition + 1) * bitWidth) % 64; + // Note that this isn't a real branch as its on a constexpr. + if constexpr (spillover > 0) { + if constexpr (withBaseline) { + **nextWord |= static_cast(**values - baseline) + << (64 - bitWidth + spillover); + ++(*nextWord); + **nextWord |= + static_cast(**values - baseline) >> (bitWidth - spillover); + ++(*values); + } else { + **nextWord |= static_cast(**values) + << (64 - bitWidth + spillover); + ++(*nextWord); + **nextWord |= static_cast(**values) >> (bitWidth - spillover); + ++(*values); + } + } else { + ++(*nextWord); + } + // How many remaining values are in this word? + constexpr int valueCount = (64 - spillover) / bitWidth; + int offset = spillover; + for (int i = 0; i < valueCount; ++i) { + if constexpr (withBaseline) { + **nextWord |= static_cast(**values - baseline) << offset; + } else { + **nextWord |= static_cast(**values) << offset; + } + offset += bitWidth; + ++(*values); + } + constexpr int nextLoopPosition = loopPosition + valueCount + (spillover > 0); + bulkSet32Loop( + nextWord, values, baseline); +} + +// Unfortunately we cannot partially specialize the template for the +// terminal case of loopPosition = 64 so we must explicitly specify them. +#define BULK_SET32_LOOP_TERMINAL_CASE(bitWidth) \ + template <> \ + void bulkSet32Loop( \ + uint64_t * *nextWord, const uint32_t** values, uint32_t baseline) {} \ + template <> \ + void bulkSet32Loop( \ + uint64_t * *nextWord, const uint32_t** values, uint32_t baseline) {} + +BULK_SET32_LOOP_TERMINAL_CASE(1) +BULK_SET32_LOOP_TERMINAL_CASE(2) +BULK_SET32_LOOP_TERMINAL_CASE(3) +BULK_SET32_LOOP_TERMINAL_CASE(4) +BULK_SET32_LOOP_TERMINAL_CASE(5) +BULK_SET32_LOOP_TERMINAL_CASE(6) +BULK_SET32_LOOP_TERMINAL_CASE(7) +BULK_SET32_LOOP_TERMINAL_CASE(8) +BULK_SET32_LOOP_TERMINAL_CASE(9) +BULK_SET32_LOOP_TERMINAL_CASE(10) +BULK_SET32_LOOP_TERMINAL_CASE(11) +BULK_SET32_LOOP_TERMINAL_CASE(12) +BULK_SET32_LOOP_TERMINAL_CASE(13) +BULK_SET32_LOOP_TERMINAL_CASE(14) +BULK_SET32_LOOP_TERMINAL_CASE(15) +BULK_SET32_LOOP_TERMINAL_CASE(16) +BULK_SET32_LOOP_TERMINAL_CASE(17) +BULK_SET32_LOOP_TERMINAL_CASE(18) +BULK_SET32_LOOP_TERMINAL_CASE(19) +BULK_SET32_LOOP_TERMINAL_CASE(20) +BULK_SET32_LOOP_TERMINAL_CASE(21) +BULK_SET32_LOOP_TERMINAL_CASE(22) +BULK_SET32_LOOP_TERMINAL_CASE(23) +BULK_SET32_LOOP_TERMINAL_CASE(24) +BULK_SET32_LOOP_TERMINAL_CASE(25) +BULK_SET32_LOOP_TERMINAL_CASE(26) +BULK_SET32_LOOP_TERMINAL_CASE(27) +BULK_SET32_LOOP_TERMINAL_CASE(28) +BULK_SET32_LOOP_TERMINAL_CASE(29) +BULK_SET32_LOOP_TERMINAL_CASE(30) +BULK_SET32_LOOP_TERMINAL_CASE(31) +BULK_SET32_LOOP_TERMINAL_CASE(32) + +#undef BULK_SET32_LOOP_TERMINAL_CASE + +template +void bulkSetInternal32( + FixedBitArray& fixedBitArray, + char* buffer, + uint64_t start, + uint64_t length, + const uint32_t* values, + uint32_t baseline) { + // Same general logic as BulkGet32. See the comments there. + switch (fixedBitArray.bitWidth()) { +#define BULK_SET32_SWITCH_CASE(bitWidth) \ + case bitWidth: { \ + const uint64_t alignedStart = bits::bucketsRequired(start, 64) << 6; \ + if (start + length < alignedStart) { \ + for (uint64_t i = start; i < start + length; ++i) { \ + if constexpr (withBaseline) { \ + fixedBitArray.set32(i, (*values) - baseline); \ + } else { \ + fixedBitArray.set32(i, *values); \ + } \ + ++values; \ + } \ + return; \ + } \ + for (uint64_t i = start; i < alignedStart; ++i) { \ + if constexpr (withBaseline) { \ + fixedBitArray.set32(i, (*values) - baseline); \ + } else { \ + fixedBitArray.set32(i, *values); \ + } \ + ++values; \ + } \ + const uint64_t loopCount = (length - (alignedStart - start)) >> 6; \ + uint64_t* nextWord = reinterpret_cast( \ + buffer + ((alignedStart * bitWidth) >> 3) - 8); \ + for (uint64_t i = 0; i < loopCount; ++i) { \ + bulkSet32Loop(&nextWord, &values, baseline); \ + } \ + const uint64_t remainderStart = alignedStart + (loopCount << 6); \ + const uint64_t remainderEnd = start + length; \ + for (uint64_t i = remainderStart; i < remainderEnd; ++i) { \ + if constexpr (withBaseline) { \ + fixedBitArray.set32(i, (*values) - baseline); \ + } else { \ + fixedBitArray.set32(i, *values); \ + } \ + ++values; \ + } \ + return; \ + } + + BULK_SET32_SWITCH_CASE(1) + BULK_SET32_SWITCH_CASE(2) + BULK_SET32_SWITCH_CASE(3) + BULK_SET32_SWITCH_CASE(4) + BULK_SET32_SWITCH_CASE(5) + BULK_SET32_SWITCH_CASE(6) + BULK_SET32_SWITCH_CASE(7) + BULK_SET32_SWITCH_CASE(8) + BULK_SET32_SWITCH_CASE(9) + BULK_SET32_SWITCH_CASE(10) + BULK_SET32_SWITCH_CASE(11) + BULK_SET32_SWITCH_CASE(12) + BULK_SET32_SWITCH_CASE(13) + BULK_SET32_SWITCH_CASE(14) + BULK_SET32_SWITCH_CASE(15) + BULK_SET32_SWITCH_CASE(16) + BULK_SET32_SWITCH_CASE(17) + BULK_SET32_SWITCH_CASE(18) + BULK_SET32_SWITCH_CASE(19) + BULK_SET32_SWITCH_CASE(20) + BULK_SET32_SWITCH_CASE(21) + BULK_SET32_SWITCH_CASE(22) + BULK_SET32_SWITCH_CASE(23) + BULK_SET32_SWITCH_CASE(24) + BULK_SET32_SWITCH_CASE(25) + BULK_SET32_SWITCH_CASE(26) + BULK_SET32_SWITCH_CASE(27) + BULK_SET32_SWITCH_CASE(28) + BULK_SET32_SWITCH_CASE(29) + BULK_SET32_SWITCH_CASE(30) + BULK_SET32_SWITCH_CASE(31) + BULK_SET32_SWITCH_CASE(32) + +#undef BULK_SET32_SWITCH_CASE + + default: + LOG(FATAL) << "bit width must lie in [1, 32], got: " + << fixedBitArray.bitWidth(); + } +} + +void FixedBitArray::bulkSet32( + uint64_t start, + uint64_t length, + const uint32_t* values) { + bulkSetInternal32(*this, buffer_, start, length, values, 0); +} + +void FixedBitArray::bulkSet32WithBaseline( + uint64_t start, + uint64_t length, + const uint32_t* values, + uint32_t baseline) { + bulkSetInternal32(*this, buffer_, start, length, values, baseline); +} + +template +void equals32Loop( + const uint32_t value, + const uint64_t equalsMask, + uint64_t word, + uint64_t** nextWord, + uint64_t* outputWord) { + constexpr uint64_t kMask = (1ULL << bitWidth) - 1ULL; + constexpr uint64_t spillover = (loopPosition * bitWidth) % 64 == 0 + ? 0 + : ((loopPosition + 1) * bitWidth) % 64; + if (spillover > 0) { + uint64_t remainder = word; + word = **nextWord; + ++(*nextWord); + if (((remainder | word << (bitWidth - spillover)) & kMask) == value) { + *outputWord |= (1ULL << loopPosition); + } + word >>= spillover; + } else { + word = **nextWord; + ++(*nextWord); + } + constexpr int valueCount = (64 - spillover) / bitWidth; + uint64_t maskedWord = word ^ equalsMask; + constexpr int offset = loopPosition + (spillover > 0); + for (int i = 0; i < valueCount; ++i) { + *outputWord |= static_cast((maskedWord & kMask) == 0) + << (offset + i); + maskedWord >>= bitWidth; + } + constexpr int nextWordShift = valueCount * bitWidth; + const uint64_t nextSpillWord = + nextWordShift == 64 ? 0 : (word >> nextWordShift); + constexpr int nextLoopPosition = offset + valueCount; + equals32Loop( + value, equalsMask, nextSpillWord, nextWord, outputWord); +} + +#define EQUALS32_TERMINAL_CASE(bitWidth) \ + template <> \ + void equals32Loop( \ + uint32_t value, \ + uint64_t equalsMask, \ + uint64_t word, \ + uint64_t * *nextWord, \ + uint64_t * outputWord) {} + +EQUALS32_TERMINAL_CASE(1) +EQUALS32_TERMINAL_CASE(2) +EQUALS32_TERMINAL_CASE(3) +EQUALS32_TERMINAL_CASE(4) +EQUALS32_TERMINAL_CASE(5) +EQUALS32_TERMINAL_CASE(6) +EQUALS32_TERMINAL_CASE(7) +EQUALS32_TERMINAL_CASE(8) +EQUALS32_TERMINAL_CASE(9) +EQUALS32_TERMINAL_CASE(10) +EQUALS32_TERMINAL_CASE(11) +EQUALS32_TERMINAL_CASE(12) +EQUALS32_TERMINAL_CASE(13) +EQUALS32_TERMINAL_CASE(14) +EQUALS32_TERMINAL_CASE(15) +EQUALS32_TERMINAL_CASE(16) +EQUALS32_TERMINAL_CASE(17) +EQUALS32_TERMINAL_CASE(18) +EQUALS32_TERMINAL_CASE(19) +EQUALS32_TERMINAL_CASE(20) +EQUALS32_TERMINAL_CASE(21) +EQUALS32_TERMINAL_CASE(22) +EQUALS32_TERMINAL_CASE(23) +EQUALS32_TERMINAL_CASE(24) +EQUALS32_TERMINAL_CASE(25) +EQUALS32_TERMINAL_CASE(26) +EQUALS32_TERMINAL_CASE(27) +EQUALS32_TERMINAL_CASE(28) +EQUALS32_TERMINAL_CASE(29) +EQUALS32_TERMINAL_CASE(30) +EQUALS32_TERMINAL_CASE(31) +EQUALS32_TERMINAL_CASE(32) + +#undef EQUALS32_TERMINAL_CASE + +void FixedBitArray::equals32( + uint64_t start, + uint64_t length, + uint32_t value, + char* bitVector) const { + // Per the header comment we require that start be a multiple of 64. + CHECK_EQ(start & 63, 0); + // First build the equality mask we'll use during the loop. + uint64_t equalsMask = 0; + FixedBitArray maskFixedBitArray((char*)&equalsMask, bitWidth_); + const int maskSlots = 64 / bitWidth_; + for (int i = 0; i < maskSlots; ++i) { + maskFixedBitArray.set(i, value); + } + const uint64_t loopCount = length >> 6; + uint64_t* nextWord = + reinterpret_cast(buffer_ + ((start * bitWidth_) >> 3)); + uint64_t* outputWord = reinterpret_cast(bitVector); + switch (bitWidth_) { +#define EQUALS32_SWITCH_CASE(bitWidth) \ + case bitWidth: { \ + for (uint64_t i = 0; i < loopCount; ++i) { \ + equals32Loop(value, equalsMask, 0, &nextWord, outputWord); \ + ++outputWord; \ + } \ + break; \ + } + + EQUALS32_SWITCH_CASE(1) + EQUALS32_SWITCH_CASE(2) + EQUALS32_SWITCH_CASE(3) + EQUALS32_SWITCH_CASE(4) + EQUALS32_SWITCH_CASE(5) + EQUALS32_SWITCH_CASE(6) + EQUALS32_SWITCH_CASE(7) + EQUALS32_SWITCH_CASE(8) + EQUALS32_SWITCH_CASE(9) + EQUALS32_SWITCH_CASE(10) + EQUALS32_SWITCH_CASE(11) + EQUALS32_SWITCH_CASE(12) + EQUALS32_SWITCH_CASE(13) + EQUALS32_SWITCH_CASE(14) + EQUALS32_SWITCH_CASE(15) + EQUALS32_SWITCH_CASE(16) + EQUALS32_SWITCH_CASE(17) + EQUALS32_SWITCH_CASE(18) + EQUALS32_SWITCH_CASE(19) + EQUALS32_SWITCH_CASE(20) + EQUALS32_SWITCH_CASE(21) + EQUALS32_SWITCH_CASE(22) + EQUALS32_SWITCH_CASE(23) + EQUALS32_SWITCH_CASE(24) + EQUALS32_SWITCH_CASE(25) + EQUALS32_SWITCH_CASE(26) + EQUALS32_SWITCH_CASE(27) + EQUALS32_SWITCH_CASE(28) + EQUALS32_SWITCH_CASE(29) + EQUALS32_SWITCH_CASE(30) + EQUALS32_SWITCH_CASE(31) + EQUALS32_SWITCH_CASE(32) + +#undef EQUALS32_SWITCH_CASE + } + const uint64_t remainderStart = start + (loopCount >> 6); + const uint64_t remainderEnd = start + length; + // Hrm actually in the case we are talking about here the final piece + // could be a written as a single word itself. + for (uint64_t i = remainderStart; i < remainderEnd; ++i) { + if (get32(i) == value) { + bits::setBit(i - start, bitVector); + } + } + return; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/FixedBitArray.h b/dwio/nimble/common/FixedBitArray.h new file mode 100644 index 0000000..97e661a --- /dev/null +++ b/dwio/nimble/common/FixedBitArray.h @@ -0,0 +1,125 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/common/Bits.h" + +// Packs integers into a fixed number (1-64) of bits and retrieves them. +// The 'bulk' API provided is particularly efficient at setting/getting +// ranges of values with small bit widths. +// +// Typical usage: +// std::vector data = ... +// int bitsRequired = BitsRequired(*maxElement(data.begin(), data.end()); +// auto buffer = std::make_unique( +// FixedBitArray::BufferSize(data.size(), bitsRequired)); +// FixedBitArray fixedBitArray(buffer.get(), bitsRequired); +// fixedBitArray.bulkSet32(0, data.size(), data.data()); +// +// And then you'd store the buffer somewhere, and later read the values back, +// recovering the info stored in data. Note that the FBA does not own +// any data. + +namespace facebook::nimble { + +class FixedBitArray { + public: + // Computes the buffer size needed to hold |elementCount| values with the + // specified |bitWidth|. Note that we allocate a few bytes of 'slop' space + // beyond what is mathematically required to speed things up. + static uint64_t bufferSize(uint64_t elementCount, int bitWidth); + + // Not legal to use; included so we can default construct on the stack. + FixedBitArray() = default; + + // Creates a fixed bit array stored at buffer whose elements can lie within + // the range [0, 2^|bitWidth|). The |buffer| must already be preallocated to + // the appropriate size, as given by BufferSize. + FixedBitArray(char* buffer, int bitWidth); + + // Convenience constructor for reading read-only data. Note that you are NOT + // prevented by the code from using the non-const calls on a FBA constructed + // via this method, but to do so is undefined behavior. + FixedBitArray(std::string_view buffer, int bitWidth) + : FixedBitArray(const_cast(buffer.data()), bitWidth) {} + + // Sets the |index|'th slot to the given |value|. If value + // is >= 2^|bitWidth| behavior is undefined. + // + // IMPORTANT NOTE: this does NOT function as normal assignment. + // We do NOT first 0 out the slot, i.e. we use or semantics. + // If you need to first 0 it out, use ZeroAndSet. + void set(uint64_t index, uint64_t value); + + // Zeroes a slot and then sets it (like normal assignment). + void zeroAndSet(uint64_t index, uint64_t value); + + // Gets the |index|'th value. + uint64_t get(uint64_t index) const; + + // Versions of set/get that only work with bitWidth <= 32, but are slightly + // faster. Same assignment caveat applies to set32 as to set. + void set32(uint64_t index, uint32_t value); + uint32_t get32(uint64_t index) const; + + // Retrieves a contiguous subarray from slots [start, start + length). + // Considerably faster than looping a get call. Only callable when + // bit width <= 32. + void bulkGet32(uint64_t start, uint64_t length, uint32_t* values) const; + + // Same as above, but outputs into 8-bytes values. Still requires that bit + // width <= 32. + void bulkGet32Into64(uint64_t start, uint64_t length, uint64_t* values) const; + + // Same as above. Add baseline to every value. + void bulkGetWithBaseline32( + uint64_t start, + uint64_t length, + uint32_t* values, + uint32_t baseline) const; + + // Same as above. Add baseline to every value. + void bulkGetWithBaseline32Into64( + uint64_t start, + uint64_t length, + uint64_t* values, + uint64_t baseline) const; + + // Sets a contiguous subarray of slots from [start, start + length). + // Considerably faster than looping/ a get call. Only callable when bitWidth + // <= 32. Same semantics as set -- see the warning there. + void bulkSet32(uint64_t start, uint64_t length, const uint32_t* values); + + // Same as above. Subtracts baseline from every value. + void bulkSet32WithBaseline( + uint64_t start, + uint64_t length, + const uint32_t* values, + uint32_t baseline); + + // Fills in the bit-packed |bitVector| with whether each value in + // [start, start + length) equals |value|. Should be faster than + // doing bulkGet32 and then doing your own equality check + bitpack. + // + // bitVector must point to a buffer of at least size BufferSize(length, 1); + // + // REQUIRES: start is a multiple of 64. If you need to run equals on + // a range where that isn't true, you'll have to manually do the part + // up to the first multiple of 64 yourself + adjust the results accordingly. + void equals32( + uint64_t start, + uint64_t length, + uint32_t value, + char* bitVector) const; + + int bitWidth() const { + return bitWidth_; + } + + private: + char* buffer_; + int bitWidth_; + uint64_t mask_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/Huffman.cpp b/dwio/nimble/common/Huffman.cpp new file mode 100644 index 0000000..967f089 --- /dev/null +++ b/dwio/nimble/common/Huffman.cpp @@ -0,0 +1,159 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include + +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/Huffman.h" + +namespace facebook::nimble::huffman { + +std::vector generateHuffmanEncodingTable( + std::span counts, + int* treeDepth) { + struct Element { + int symbol; + uint32_t count; + Element* leftChild; + Element* rightChild; + uint32_t encoding; + uint32_t encodingLength; + }; + // TODO: Ugh this is actually terrible with unique ptrs. Just have an owning + // vector of 2x length counts Elements and use raw pointers. Duh. + std::vector> pendingElements(counts.size()); + for (int i = 0; i < counts.size(); ++i) { + pendingElements[i] = std::make_unique(); + Element* e = pendingElements[i].get(); + e->symbol = i; + e->count = counts[i]; + e->leftChild = nullptr; + e->rightChild = nullptr; + } + auto compare = [](const std::unique_ptr& a, + const std::unique_ptr& b) { + return a->count > b->count; + }; + std::make_heap(pendingElements.begin(), pendingElements.end(), compare); + + // Get 2 smallest elements, merge them into a new one, repeat. + std::vector> doneElements; + while (pendingElements.size() >= 2) { + std::pop_heap(pendingElements.begin(), pendingElements.end(), compare); + std::unique_ptr leftChild = std::move(pendingElements.back()); + pendingElements.pop_back(); + std::pop_heap(pendingElements.begin(), pendingElements.end(), compare); + std::unique_ptr rightChild = std::move(pendingElements.back()); + pendingElements.pop_back(); + std::unique_ptr merged = std::make_unique(); + merged->symbol = -1; + merged->count = leftChild->count + rightChild->count; + merged->leftChild = leftChild.get(); + merged->rightChild = rightChild.get(); + doneElements.push_back(std::move(leftChild)); + doneElements.push_back(std::move(rightChild)); + pendingElements.push_back(std::move(merged)); + std::push_heap(pendingElements.begin(), pendingElements.end(), compare); + } + + // Now traverse the tree and fill out the encodings. + std::vector entries(counts.size()); + std::queue traverse; + traverse.push(pendingElements.front().get()); + traverse.front()->encoding = 0; + traverse.front()->encodingLength = 0; + *treeDepth = 0; + while (!traverse.empty()) { + Element* next = traverse.front(); + *treeDepth = std::max(*treeDepth, (int)next->encodingLength); + traverse.pop(); + if (next->leftChild) { + next->leftChild->encoding = + next->encoding | (1UL << next->encodingLength); + next->leftChild->encodingLength = next->encodingLength + 1; + traverse.push(next->leftChild); + } + if (next->rightChild) { + next->rightChild->encoding = next->encoding; + next->rightChild->encodingLength = next->encodingLength + 1; + traverse.push(next->rightChild); + } + // For leaf elements go ahead and output the entries. Remember that + // depth (encoding length) is lowest 5 bits. + if (next->symbol != -1) { + entries[next->symbol] = (next->encoding << 5) | next->encodingLength; + } + } + return entries; +} + +namespace { + +int minTreeDepth(uint32_t size) { + int minDepth = 1; + uint64_t capacity = 2; + while (capacity < size) { + ++minDepth; + capacity <<= 1; + } + return minDepth; +} + +} // namespace + +std::vector generateHuffmanEncodingTableWithMaxDepth( + std::span counts, + int maxDepth, + int* treeDepth) { + CHECK_GE(maxDepth, minTreeDepth(counts.size())); + auto table = generateHuffmanEncodingTable(counts, treeDepth); + if (*treeDepth <= maxDepth) { + return table; + } + std::vector ownedCounts; + ownedCounts.reserve(counts.size()); + for (int i = 0; i < counts.size(); ++i) { + if (counts[i] > 1) { + ownedCounts.push_back(counts[i] >> 1); + } else { + ownedCounts.push_back(1); + } + } + while (true) { + table = generateHuffmanEncodingTable(ownedCounts, treeDepth); + if (*treeDepth <= maxDepth) { + return table; + } + // TODO: Might want to shift dynamically based on how far we are from the + // max depth. + for (auto& count : ownedCounts) { + if (count > 1) { + count >>= 1; + } + } + } +} + +DecodingTable::DecodingTable( + facebook::velox::memory::MemoryPool& memoryPool, + std::span encodingTable, + int treeDepth) + : treeDepth_(treeDepth), lookupTable_(&memoryPool, 1 << treeDepth_) { + CHECK_LE(treeDepth_, 15); + CHECK_LE(encodingTable.size(), 4096); + for (uint32_t i = 0; i < encodingTable.size(); ++i) { + const uint32_t entry = encodingTable[i]; + const uint16_t entryDepth = entry & 15; + const uint16_t slotBits = treeDepth - entryDepth; + const uint16_t entrySlots = 1 << slotBits; + // Remember entry still uses 5 bits for depth even though we only use 4 + // in the in memory table. + const uint16_t encodingSuffix = entry >> 5; + const uint16_t lookupTableEntry = (i << 4) | entryDepth; + for (uint16_t j = 0; j < entrySlots; ++j) { + lookupTable_.data()[(j << entryDepth) | encodingSuffix] = + lookupTableEntry; + } + } +} + +} // namespace facebook::nimble::huffman diff --git a/dwio/nimble/common/Huffman.h b/dwio/nimble/common/Huffman.h new file mode 100644 index 0000000..610f184 --- /dev/null +++ b/dwio/nimble/common/Huffman.h @@ -0,0 +1,289 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include + +#include "dwio/nimble/common/BitEncoder.h" +#include "dwio/nimble/common/Bits.h" +#include "dwio/nimble/common/Vector.h" +#include "velox/common/memory/Memory.h" + +// Huffman encoder/decoder specialized for speed on small numbers (up to a few +// thousand) of symbols. +// +// A huffman encoding is an entropy encoding of k symbols in the range [0, n). +// We require that the range be dense, i.e. each symbol in the range appears +// at least once. (So clearly k >= n.) +// +// The basic idea is we use the frequencies of those symbols to generate the +// optimal encoding table, which can be easily serialized. We use this table +// to encode the data into a bit stream. On the decompression side we transform +// the serialized encoding table into an lookup table structure and then use +// that structure to quickly pull bits from the stream and reconstruct the +// original entries. +// +// Example usage: +// std::vector data = /* your data on dense range [0, n) */ +// std::vector counts = /* calculate the counts of your data */ +// int treeDepth; +// auto encodingTable = GenerateHuffmanEncodingTable(counts, &treeDepth); +// std::string encodedStream = HuffmanEncode(encodingTable, data); +// /* store the stream somewhere, store the encoding table with it +// /* then load it when you want to read the data... +// DecodingTable decodingTable(encodingTable, treeDepth); +// std::vector recoveredData(data.size()); +// decodingTable.Decode( +// data.size(), encodedStream.data(), recoveredData.data()); +// /* data and recoveredData are now the same */ +// +// This code is pretty fast. For the streamed decode on a small (tree depth 10) +// encoding table, we see decode speeds as fast as 5-6 cycles/symbol. + +namespace facebook::nimble::huffman { + +// The length of the string returned from HuffmanEncode is the number of +// bytes required to hold the bits, plus 7 for slop space. +constexpr int kHuffmanSlopBytes = 7; + +// Each entry in our encoding table is a uint32_t, where the bottom 5 bits +// give the depth of the encoding and the top bits gives the path to the leaf +// representing the entry. +std::vector generateHuffmanEncodingTable( + std::span counts, + int* treeDepth); + +// Similar to above, but enforces that tree depth be at most |maxDepth| by +// normalizing the counts (if necessary). Note that the minimal max depth is, +// of course, ceil(log2(counts.size()) -- don't ask for a lesser one! +std::vector generateHuffmanEncodingTableWithMaxDepth( + std::span counts, + int maxDepth, + int* treeDepth); + +// Uses the provided |encodingTable| to huffman encode the |symbols|. +template +std::string huffmanEncode( + std::span encodingTable, + std::span symbols); + +// This decoding table only allows a max depth of up to 15 and a max encoding +// table size of 4096. Note that a depth of 15 will be considerably slower +// than <= 14 due to not fitting in the 32kb L1 cache. +class DecodingTable { + public: + DecodingTable( + facebook::velox::memory::MemoryPool& memoryPool, + std::span encodingTable, + int treeDepth); + + // Decodes |n| symbols from |data| into |output|. + template + void decode(uint32_t n, const char* data, T* output) const; + + // Similar to Decode, but simultaneously decodes two independent data streams + // at once, extracting n symbols from each and placing them in the respective + // |output|. This was measured to decode 50-80% more symbols/sec than Decode. + template + void decodeStreamed( + uint32_t n, + const char* data1, + const char* data2, + T* output1, + T* output2) const; + + private: + const int treeDepth_; + Vector lookupTable_; +}; + +// +// End of pubic API. Implementations follow. +// + +template +std::string huffmanEncode( + std::span encodingTable, + std::span symbols) { + // We could iterate through the symbols and get the total length, then + // allocate the string. Or we could just over allocate and shrink. We'll do + // the latter for now. 2 because our max tree depth is 15 and hence can + // obviously fit within 2 bytes per symbol. + std::string buffer(2 * symbols.size() + kHuffmanSlopBytes, 0); + BitEncoder bitEncoder(buffer.data()); + for (uint32_t symbol : symbols) { + const uint32_t entry = encodingTable[symbol]; + bitEncoder.putBits(entry >> 5, entry & 31UL); + } + // Don't forget the slop bytes! + const int bytes = + bits::bytesRequired(bitEncoder.bitsWritten()) + kHuffmanSlopBytes; + buffer.resize(bytes); + return buffer; +} + +namespace internal { + +template +void huffmanDecode( + uint32_t n, + const char* data, + const uint16_t* lookupTable, + T* output) { + constexpr uint64_t mask = (1ULL << treeDepth) - 1ULL; + constexpr int symbolsPerLoop = 64 / treeDepth; + const uint32_t loopCount = n / symbolsPerLoop; + const uint32_t remainder = n - symbolsPerLoop * loopCount; + BitEncoder bitEncoder(data); + uint64_t currentWord = bitEncoder.getBits(64); + for (int i = 0; i < loopCount; ++i) { + int bitsUsed = 0; + for (int j = 0; j < symbolsPerLoop; ++j) { + const uint64_t tableSlot = currentWord & mask; + const uint16_t tableEntry = lookupTable[tableSlot]; + *output++ = tableEntry >> 4; + bitsUsed += tableEntry & 15; + currentWord >>= tableEntry & 15; + } + currentWord |= (bitEncoder.getBits(bitsUsed) << (64 - bitsUsed)); + } + for (int i = 0; i < remainder; ++i) { + const uint64_t tableSlot = currentWord & mask; + const uint16_t tableEntry = lookupTable[tableSlot]; + *output++ = tableEntry >> 4; + currentWord >>= tableEntry & 15; + } +} + +} // namespace internal + +template +void DecodingTable::decode(uint32_t n, const char* data, T* output) const { + switch (treeDepth_) { +#define CASE(treeDepth) \ + case treeDepth: { \ + return internal::huffmanDecode( \ + n, data, lookupTable_.data(), output); \ + } + CASE(1) + CASE(2) + CASE(3) + CASE(4) + CASE(5) + CASE(6) + CASE(7) + CASE(8) + CASE(9) + CASE(10) + CASE(11) + CASE(12) + CASE(13) + CASE(14) + CASE(15) +#undef CASE + + default: { + LOG(FATAL) << "programming error: encountered DecodingTable w/tree depth " + << "output the allowed range of [1, 15]: " << treeDepth_; + } + } +} + +namespace internal { + +// Note that the zstd checks for the "bmi2" instruction set, and if available +// dispatches to a function with __attribute__((__target__("bmi2"))) set. Might +// be worth doing later. +template +void huffmanDecodeStreamed( + uint32_t n, + const char* data1, + const char* data2, + const uint16_t* lookupTable, + T* output1, + T* output2) { + constexpr uint64_t mask = (1ULL << treeDepth) - 1ULL; + constexpr int symbolsPerLoop = 64 / treeDepth; + const uint32_t loopCount = n / symbolsPerLoop; + const uint32_t remainder = n - symbolsPerLoop * loopCount; + BitEncoder bitEncoder1(data1); + uint64_t currentWord1 = bitEncoder1.getBits(64); + BitEncoder bitEncoder2(data2); + uint64_t currentWord2 = bitEncoder2.getBits(64); + for (int i = 0; i < loopCount; ++i) { + int bitsUsed1 = 0; + int bitsUsed2 = 0; + for (int j = 0; j < symbolsPerLoop; ++j) { + const uint64_t tableSlot1 = currentWord1 & mask; + const uint16_t tableEntry1 = lookupTable[tableSlot1]; + *output1++ = tableEntry1 >> 4; + bitsUsed1 += tableEntry1 & 15; + currentWord1 >>= tableEntry1 & 15; + const uint64_t tableSlot2 = currentWord2 & mask; + const uint16_t tableEntry2 = lookupTable[tableSlot2]; + *output2++ = tableEntry2 >> 4; + bitsUsed2 += tableEntry2 & 15; + currentWord2 >>= tableEntry2 & 15; + } + currentWord1 |= (bitEncoder1.getBits(bitsUsed1) << (64 - bitsUsed1)); + currentWord2 |= (bitEncoder2.getBits(bitsUsed2) << (64 - bitsUsed2)); + } + for (int i = 0; i < remainder; ++i) { + const uint64_t tableSlot1 = currentWord1 & mask; + const uint16_t tableEntry1 = lookupTable[tableSlot1]; + *output1++ = tableEntry1 >> 4; + currentWord1 >>= tableEntry1 & 15; + const uint64_t tableSlot2 = currentWord2 & mask; + const uint16_t tableEntry2 = lookupTable[tableSlot2]; + *output2++ = tableEntry2 >> 4; + currentWord2 >>= tableEntry2 & 15; + } +} + +} // namespace internal + +// It's actually at first a bit mysterious why this is faster than the single +// stream version. The answer seems to be that having the two streams better +// saturates the cpu pipeline (thanks to the OoO optimization window). It's +// interesting to note that 2 streams seems to be optimal: 3 at once is roughly +// the same speed (perhaps a few % faster), but 4+ is slower. +template +void DecodingTable::decodeStreamed( + uint32_t n, + const char* data1, + const char* data2, + T* output1, + T* output2) const { + switch (treeDepth_) { +#define CASE(treeDepth) \ + case treeDepth: { \ + return internal::huffmanDecodeStreamed( \ + n, data1, data2, lookupTable_.data(), output1, output2); \ + } + CASE(1) + CASE(2) + CASE(3) + CASE(4) + CASE(5) + CASE(6) + CASE(7) + CASE(8) + CASE(9) + CASE(10) + CASE(11) + CASE(12) + CASE(13) + CASE(14) + CASE(15) +#undef CASE + + default: { + LOG(FATAL) << "programming error: encountered DecodingTable w/tree depth " + << "output the allowed range of [1, 15]: " << treeDepth_; + } + } +} + +} // namespace facebook::nimble::huffman diff --git a/dwio/nimble/common/IndexMap.h b/dwio/nimble/common/IndexMap.h new file mode 100644 index 0000000..625bf7b --- /dev/null +++ b/dwio/nimble/common/IndexMap.h @@ -0,0 +1,68 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/common/NimbleCompare.h" +#include "folly/container/F14Map.h" + +// An index map maintain a bijection between a set and each element's unique +// index in [0, n). E.g. the first key introduced to the index map gets index 0, +// the second 1, and so on. Looking up key by index and index by key are both +// possible. + +namespace facebook::nimble { + +// Don't insert more than INT32_MAX symbols into the map. +template +class IndexMap { + public: + // Returns the unique index associated with the key, adding the new key to + // the internal T<->int32_t mapping if it hasn't been seen before. + int32_t index(const T& key) noexcept { + auto it = indices_.find(key); + if (it == indices_.end()) { + indices_.emplace(key, indices_.size()); + keys_.push_back(key); + return indices_.size() - 1; + } + return it->second; + } + + // Returns the index of an existing key, or -1 if the key has not been + // seen previously. + int32_t readOnlyIndex(const T& key) const noexcept { + auto it = indices_.find(key); + if (it == indices_.end()) { + return -1; + } + return it->second; + } + + // Retrieves a previously inserted key via its index. index must lie in + // [0, size()). + const T& key(int32_t index) noexcept { + DCHECK_LT(index, keys_.size()); + return keys_[index]; + } + + // The current size size will be the index of the next previously unseen key + // passed to index. + int32_t size() noexcept { + return indices_.size(); + } + + // Transfers ownership of the keys_ to the caller. *this should not + // be used after this is called. + std::vector&& releaseKeys() { + indices_.clear(); + return std::move(keys_); + } + + private: + folly:: + F14FastMap, NimbleComparator> + indices_; + std::vector keys_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/MetricsLogger.cpp b/dwio/nimble/common/MetricsLogger.cpp new file mode 100644 index 0000000..9ee44dd --- /dev/null +++ b/dwio/nimble/common/MetricsLogger.cpp @@ -0,0 +1,41 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/common/MetricsLogger.h" + +namespace facebook::nimble { + +folly::dynamic StripeLoadMetrics::serialize() const { + folly::dynamic obj = folly::dynamic::object; + obj["stripeIndex"] = stripeIndex; + obj["rowsInStripe"] = rowsInStripe; + obj["streamCount"] = streamCount; + obj["totalStreamSize"] = totalStreamSize; + return obj; +} + +folly::dynamic StripeFlushMetrics::serialize() const { + folly::dynamic obj = folly::dynamic::object; + obj["inputSize"] = inputSize; + obj["rowCount"] = rowCount; + obj["stripeSize"] = stripeSize; + obj["trackedMemory"] = trackedMemory; + return obj; +} + +folly::dynamic FileCloseMetrics::serialize() const { + folly::dynamic obj = folly::dynamic::object; + obj["rowCount"] = rowCount; + obj["inputSize"] = inputSize; + obj["stripeCount"] = stripeCount; + obj["fileSize"] = fileSize; + obj["totalFlushCpuUsec"] = totalFlushCpuUsec; + obj["totalFlushWallTimeUsec"] = totalFlushWallTimeUsec; + return obj; +} + +LoggingScope::Context& LoggingScope::Context::get() { + thread_local static Context ctx; + return ctx; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/MetricsLogger.h b/dwio/nimble/common/MetricsLogger.h new file mode 100644 index 0000000..55c7c53 --- /dev/null +++ b/dwio/nimble/common/MetricsLogger.h @@ -0,0 +1,104 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include + +namespace facebook::nimble { + +struct StripeLoadMetrics { + uint32_t stripeIndex; + uint32_t rowsInStripe; + uint32_t streamCount{0}; + uint32_t totalStreamSize{0}; + // TODO: add IO sizes. + + // TODO: add encoding summary + + size_t cpuUsec; + size_t wallTimeUsec; + + folly::dynamic serialize() const; +}; + +// Might be good to capture via some kind of run stats struct. +// We can then adapt the run stats to file writer run stats. +struct StripeFlushMetrics { + // Stripe shape summary. + uint64_t inputSize; + uint64_t rowCount; + uint64_t stripeSize; + + // We would add flush policy states here when wired up in the future. + // uint64_t inputChunkSize_; + + // TODO: add some type of encoding summary in a follow-up diff. + + // Memory footprint + uint64_t trackedMemory; + // uint64_t residentMemory; + + // Perf stats. + uint64_t flushCpuUsec; + uint64_t flushWallTimeUsec; + // Add IOStatistics when we have finished WS api consolidations. + + folly::dynamic serialize() const; +}; + +struct FileCloseMetrics { + uint64_t rowCount; + uint64_t inputSize; + uint64_t stripeCount; + uint64_t fileSize; + + // Perf stats. + uint64_t totalFlushCpuUsec; + uint64_t totalFlushWallTimeUsec; + // Add IOStatistics when we have finished WS api consolidations. + + folly::dynamic serialize() const; +}; + +class MetricsLogger { + public: + constexpr static std::string_view kStripeLoadOperation{"STRIPE_LOAD"}; + constexpr static std::string_view kStripeFlushOperation{"STRIPE_FLUSH"}; + constexpr static std::string_view kFileCloseOperation{"FILE_CLOSE"}; + constexpr static std::string_view kZstrong{"ZSTRONG"}; + + virtual ~MetricsLogger() = default; + + virtual void logException( + std::string_view /* operation */, + const std::string& /* errorMessage */) const {} + + virtual void logStripeLoad(const StripeLoadMetrics& /* metrics */) const {} + virtual void logStripeFlush(const StripeFlushMetrics& /* metrics */) const {} + virtual void logFileClose(const FileCloseMetrics& /* metrics */) const {} + virtual void logZstrongContext(const std::string&) const {} +}; + +class LoggingScope { + public: + explicit LoggingScope(const MetricsLogger& logger) { + Context::get().logger = &logger; + } + + ~LoggingScope() { + Context::get().logger = nullptr; + } + + static const MetricsLogger* getLogger() { + return Context::get().logger; + } + + private: + struct Context { + const MetricsLogger* logger; + + static Context& get(); + }; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/NimbleCompare.h b/dwio/nimble/common/NimbleCompare.h new file mode 100644 index 0000000..08d251c --- /dev/null +++ b/dwio/nimble/common/NimbleCompare.h @@ -0,0 +1,72 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include + +namespace facebook::nimble { + +template +class NimbleCompare { + public: + static bool equals(const T& a, const T& b); +}; + +template +inline bool NimbleCompare::equals(const T& a, const T& b) { + return a == b; +} + +template +class NimbleCompare< + T, + std::enable_if_t || std::is_same_v>> { + public: + // double or float + using FloatingType = + typename std::conditional, double, float>::type; + // 64bit integer or 32 bit integer + using IntegralType = typename std:: + conditional, int64_t, int32_t>::type; + + static_assert(sizeof(FloatingType) == sizeof(IntegralType)); + + static bool equals(const T& a, const T& b); + + // This will be convenient when for debug, logging. + static IntegralType asInteger(const T& a); +}; + +template +inline bool NimbleCompare< + T, + std::enable_if_t || std::is_same_v>>:: + equals(const T& a, const T& b) { + // For floating point types, we do bit-wise comparison, for other types, + // just use the original ==. + // TODO: handle NaN. + return *(reinterpret_cast(&(a))) == + *(reinterpret_cast(&(b))); +} + +template +inline typename NimbleCompare< + T, + std::enable_if_t< + std::is_same_v || std::is_same_v>>::IntegralType +NimbleCompare< + T, + std::enable_if_t || std::is_same_v>>:: + asInteger(const T& a) { + return *(reinterpret_cast(&(a))); +} + +template +struct NimbleComparator { + constexpr bool operator()(const T& a, const T& b) const { + return NimbleCompare::equals(a, b); + } +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/Rle.h b/dwio/nimble/common/Rle.h new file mode 100644 index 0000000..fd9dddb --- /dev/null +++ b/dwio/nimble/common/Rle.h @@ -0,0 +1,42 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include "dwio/nimble/common/NimbleCompare.h" +#include "dwio/nimble/common/Vector.h" + +#include + +// Functions related to run length encoding. + +namespace facebook::nimble::rle { + +// TODO: This utility is only used by RLEEncoding. Consider moving it there. + +template +void computeRuns( + std::span data, + Vector* runLengths, + Vector* runValues) { + static_assert(!std::is_floating_point_v); + if (data.empty()) { + return; + } + uint32_t runLength = 1; + T last = data[0]; + for (int i = 1; i < data.size(); ++i) { + if (data[i] == last) { + ++runLength; + } else { + runLengths->push_back(runLength); + runValues->push_back(last); + last = data[i]; + runLength = 1; + } + } + runLengths->push_back(runLength); + runValues->push_back(last); +} + +} // namespace facebook::nimble::rle diff --git a/dwio/nimble/common/StopWatch.cpp b/dwio/nimble/common/StopWatch.cpp new file mode 100644 index 0000000..3a57706 --- /dev/null +++ b/dwio/nimble/common/StopWatch.cpp @@ -0,0 +1,47 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include "dwio/nimble/common/StopWatch.h" +#include "common/time/ClockGettimeNS.h" + +namespace facebook::nimble { + +void StopWatch::start() { + if (!running_) { + startNs_ = fb_clock_gettime_ns(CLOCK_MONOTONIC); + running_ = true; + } +} + +void StopWatch::stop() { + if (running_) { + elapsedNs_ += fb_clock_gettime_ns(CLOCK_MONOTONIC) - startNs_; + running_ = false; + } +} + +void StopWatch::reset() { + running_ = false; + elapsedNs_ = 0; +} + +double StopWatch::elapsed() { + return static_cast(elapsedNsec()) / (1000 * 1000 * 1000); +} + +int64_t StopWatch::elapsedNsec() { + if (running_) { + stop(); + start(); + } + return elapsedNs_; +} + +int64_t StopWatch::elapsedUsec() { + return elapsedNsec() / 1000; +} + +int64_t StopWatch::elapsedMsec() { + return elapsedNsec() / (1000 * 1000); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/StopWatch.h b/dwio/nimble/common/StopWatch.h new file mode 100644 index 0000000..ef21a2a --- /dev/null +++ b/dwio/nimble/common/StopWatch.h @@ -0,0 +1,50 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include + +// A simple stopwatch measuring to the nanosecond level. +// Useful for timing code while prototyping. The accuracy should be 'good', +// but its not clear what that actually means, and in any case it will depend +// on the underlying details of the system. + +namespace facebook::nimble { + +class StopWatch { + public: + StopWatch() = default; + + // Start timing. A no-op if we're already timing. + void start(); + + // Stop timing. A no-op if we aren't currently timing. + void stop(); + + // Restore *this to its newly constructed state. + void reset(); + + // Reset and then Start in one. + void restart() { + reset(); + start(); + } + + // Returns the elapsed time in seconds without rounding. Does not stop + // the stopwatch if it is running, but does reduce the accuracy a bit + // so don't call the Elapsed* functions unnecessarily. + double elapsed(); + + // Returns the elapsed time on the stopwatch in the appropriate units, + // rounding down. Does not stop the stopwatch if it is running. + int64_t elapsedNsec(); + int64_t elapsedUsec(); + int64_t elapsedMsec(); + + private: + int64_t startNs_; + int64_t elapsedNs_ = 0; + bool running_ = false; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/Types.cpp b/dwio/nimble/common/Types.cpp new file mode 100644 index 0000000..693f20a --- /dev/null +++ b/dwio/nimble/common/Types.cpp @@ -0,0 +1,120 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include "dwio/nimble/common/Types.h" + +#include +#include "dwio/nimble/common/Exceptions.h" + +namespace facebook::nimble { + +std::ostream& operator<<(std::ostream& out, EncodingType encodingType) { + return out << toString(encodingType); +} + +std::string toString(EncodingType encodingType) { + switch (encodingType) { + case EncodingType::Trivial: + return "Trivial"; + case EncodingType::RLE: + return "RLE"; + case EncodingType::Dictionary: + return "Dictionary"; + case EncodingType::FixedBitWidth: + return "FixedBitWidth"; + case EncodingType::Nullable: + return "Nullable"; + case EncodingType::SparseBool: + return "SparseBool"; + case EncodingType::Varint: + return "Varint"; + case EncodingType::Delta: + return "Delta"; + case EncodingType::Constant: + return "Constant"; + case EncodingType::MainlyConstant: + return "MainlyConstant"; + case EncodingType::Sentinel: + return "Sentinel"; + } + return fmt::format( + "Unknown encoding type: {}", static_cast(encodingType)); +} + +std::ostream& operator<<(std::ostream& out, DataType dataType) { + return out << toString(dataType); +} + +std::string toString(DataType dataType) { + switch (dataType) { + case DataType::Int8: + return "Int8"; + case DataType::Int16: + return "Int16"; + case DataType::Uint8: + return "Uint8"; + case DataType::Uint16: + return "Uint16"; + case DataType::Int32: + return "Int32"; + case DataType::Int64: + return "Int64"; + case DataType::Uint32: + return "Uint32"; + case DataType::Uint64: + return "Uint64"; + case DataType::Float: + return "Float"; + case DataType::Double: + return "Double"; + case DataType::Bool: + return "Bool"; + case DataType::String: + return "String"; + default: + return fmt::format( + "Unknown data type: {}", static_cast(dataType)); + } +} + +std::string toString(CompressionType compressionType) { + switch (compressionType) { + case CompressionType::Uncompressed: + return "Uncompressed"; + case CompressionType::Zstd: + return "Zstd"; + case CompressionType::Zstrong: + return "Zstrong"; + default: + return fmt::format( + "Unknown compression type: {}", + static_cast(compressionType)); + } +} + +std::ostream& operator<<(std::ostream& out, CompressionType compressionType) { + return out << toString(compressionType); +} + +template <> +void Variant::set( + VariantType& target, + std::string_view source) { + target = std::string(source); +} + +template <> +std::string_view Variant::get(VariantType& source) { + return std::get(source); +} + +std::string toString(ChecksumType type) { + switch (type) { + case ChecksumType::XXH3_64: + return "XXH3_64"; + default: + return fmt::format( + "Unknown checksum type: {}", static_cast(type)); + } +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/Types.h b/dwio/nimble/common/Types.h new file mode 100644 index 0000000..a6dd05e --- /dev/null +++ b/dwio/nimble/common/Types.h @@ -0,0 +1,298 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include +#include + +// Single file containing all the types and enums in nimble, as well as some +// templates for mapping between those types and C++ types. +// +// A note on types: For all of our enum classes, we assume the number of +// types will stay small, so that each one is representable by a single byte. +// Don't violate this! + +namespace facebook::nimble { + +using VariantType = std::variant< + int8_t, + uint8_t, + int16_t, + uint16_t, + int32_t, + uint32_t, + int64_t, + uint64_t, + float, + double, + bool, + std::string>; + +template +class Variant { + public: + static void set(VariantType& target, T source) { + target = source; + } + + static T get(VariantType& source) { + return std::get(source); + } +}; + +template <> +void Variant::set( + VariantType& target, + std::string_view source); + +template <> +std::string_view Variant::get(VariantType& source); + +enum class EncodingType { + // Native encoding for numerics, simple packed chars with offsets for strings, + // bitpacked for bools. All data types supported. + Trivial = 0, + // Run length encoded data. The runs lengths are bit packed, and the run + // values are encoded like the trivial encoding. All data types supported. + RLE = 1, + // Data with the uniques encoded in a dictionary and the indices into that + // dictionary. All data types except bools supported. + Dictionary = 2, + // Stores integer types packed into a fixed number of bits (namely, the + // smallest required to represent the largest element). Currently only + // works with non-negative values, we may add ZigZag encoding later. + FixedBitWidth = 3, + // Stores nullable data using a 'sentinel' value to represent nulls in a + // single non-nullable encoding. + Sentinel = 4, + // Stores nullable data by wrapping one subencoding representing the non-nulls + // with another subencoding marking which rows are null. + Nullable = 5, + // Stores indices to set (or unset) bits. Useful for storing sparse data, such + // as when only a few rows in a encoding are non-null. + SparseBool = 6, + // Stores integer types via varint encoding. Currently only + // works with non-negative values, we may add ZigZag encoding later. + Varint = 7, + // Stores integer types with a delta encoding. Currently only supports + // positive deltas. + Delta = 8, + // Stores constant (i.e. only 1 unique value) data. + Constant = 9, + // Stores 'mainly constant' data, i.e. treats one particular value as special, + // using a bool child vector to store whether each row is that special value, + // and stores the non-special values as a separate encoding. + MainlyConstant = 10, +}; +std::string toString(EncodingType encodingType); +std::ostream& operator<<(std::ostream& out, EncodingType encodingType); + +enum class DataType : uint8_t { + Undefined = 0, + Int8 = 1, + Uint8 = 2, + Int16 = 3, + Uint16 = 4, + Int32 = 5, + Uint32 = 6, + Int64 = 7, + Uint64 = 8, + Float = 9, + Double = 10, + Bool = 11, + String = 12, +}; + +std::string toString(DataType dataType); +std::ostream& operator<<(std::ostream& out, DataType dataType); + +// General string compression. Make sure values here match those in the footer +// specification +enum class CompressionType : uint8_t { + Uncompressed = 0, + // Zstd doesn't require us to externally store level or any other info. + Zstd = 1, + Zstrong = 2, +}; + +std::string toString(CompressionType compressionType); +std::ostream& operator<<(std::ostream& out, CompressionType compressionType); + +enum class ChecksumType : uint8_t { XXH3_64 = 0 }; + +std::string toString(ChecksumType type); + +// A CompresionType and any type-specific configuration params. +struct CompressionParams { + CompressionType type; + + // For zstd. + int zstdLevel = 1; +}; + +// Parameters controlling the search for the optimal encoding on a data set. +struct OptimalSearchParams { + // Whether recursive structures are allowed. E.g. a encoding may use another + // encoding as a subencoding, and may use the encoding factory to find the + // best subencoding. We must terminate the recursion at some depth. With the + // default of 1 allowed recursion the top level encoding may use recursive + // encodings, but its subencodings may not. + int allowedRecursions = 1; + + // Whether to log debug info during the search (such as estimated sizes of + // each encoding considered, etc.); + bool logSearch = false; + + // Helps align log messages when log_search=true to help distinguish + // subencoding log messages from higher-level ones. + int logDepth = 0; + + // Entropy encodings, such as HuffmanEncoding, can be much more compact than + // others, but are quite a bit slower. However, compared to applying a general + // string compression on top of another encoding they are relatively fast. + // In general the entropy encodings will also be smaller than GSC on top of + // a non-entropy encoding. + bool enableEntropyEncodings = true; + + // For some dimension encodings that will frequently be grouped by, we may + // want to force dictionary enabled encodings so grouping by that can will be + // fast. + bool requireDictionaryEnabled = false; +}; + +template +struct TypeTraits {}; + +template <> +struct TypeTraits { + using physicalType = uint8_t; + using sumType = int64_t; + static constexpr DataType dataType = DataType::Int8; +}; + +template <> +struct TypeTraits { + using physicalType = uint8_t; + using sumType = uint64_t; + static constexpr DataType dataType = DataType::Uint8; +}; + +template <> +struct TypeTraits { + using physicalType = uint16_t; + using sumType = int64_t; + static constexpr DataType dataType = DataType::Int16; +}; + +template <> +struct TypeTraits { + using physicalType = uint16_t; + using sumType = uint64_t; + static constexpr DataType dataType = DataType::Uint16; +}; + +template <> +struct TypeTraits { + using physicalType = uint32_t; + using sumType = int64_t; + static constexpr DataType dataType = DataType::Int32; +}; + +template <> +struct TypeTraits { + using physicalType = uint32_t; + using sumType = uint64_t; + static constexpr DataType dataType = DataType::Uint32; +}; + +template <> +struct TypeTraits { + using physicalType = uint64_t; + using sumType = int64_t; + static constexpr DataType dataType = DataType::Int64; +}; + +template <> +struct TypeTraits { + using physicalType = uint64_t; + using sumType = uint64_t; + static constexpr DataType dataType = DataType::Uint64; +}; + +template <> +struct TypeTraits { + using physicalType = uint32_t; + using sumType = double; + static constexpr DataType dataType = DataType::Float; +}; + +template <> +struct TypeTraits { + using physicalType = uint64_t; + using sumType = double; + static constexpr DataType dataType = DataType::Double; +}; + +template <> +struct TypeTraits { + using physicalType = bool; + static constexpr DataType dataType = DataType::Bool; +}; + +template <> +struct TypeTraits { + using physicalType = std::string; + static constexpr DataType dataType = DataType::String; +}; + +template <> +struct TypeTraits { + using physicalType = std::string_view; + static constexpr DataType dataType = DataType::String; +}; + +template +constexpr bool isFourByteIntegralType() { + return std::is_same_v || std::is_same_v; +} + +template +constexpr bool isSignedIntegralType() { + return std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v; +} + +template +constexpr bool isUnsignedIntegralType() { + return std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v; +} + +template +constexpr bool isIntegralType() { + return isSignedIntegralType() || isUnsignedIntegralType(); +} + +template +constexpr bool isFloatingPointType() { + return std::is_same_v || std::is_same_v; +} + +template +constexpr bool isNumericType() { + return isIntegralType() || isFloatingPointType(); +} + +template +constexpr bool isStringType() { + return std::is_same_v || std::is_same_v; +} + +template +constexpr bool isBoolType() { + return std::is_same_v; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/Varint.cpp b/dwio/nimble/common/Varint.cpp new file mode 100644 index 0000000..503422e --- /dev/null +++ b/dwio/nimble/common/Varint.cpp @@ -0,0 +1,828 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#ifdef __x86_64__ +#include +#endif //__x86_64__ + +#ifdef __aarch64__ +#include "common/aarch64/compat.h" +#endif //__aarch64__ + +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/Varint.h" +#include "folly/CpuId.h" + +namespace facebook::nimble::varint { + +const char* bulkVarintSkip(uint64_t n, const char* pos) { + const uint64_t* word = reinterpret_cast(pos); + while (n >= 8) { + // Zeros in the 8 * ith bits indicate termination of a varint. + n -= __builtin_popcountll(~(*word++) & 0x8080808080808080ULL); + } + pos = reinterpret_cast(word); + while (n--) { + skipVarint(&pos); + } + return pos; +} + +uint64_t bulkVarintSize32(std::span values) { + constexpr uint8_t kLookupSizeTable32[32] = {5, 5, 5, 5, 4, 4, 4, 4, 4, 4, 4, + 3, 3, 3, 3, 3, 3, 3, 2, 2, 2, 2, + 2, 2, 2, 1, 1, 1, 1, 1, 1, 1}; + uint64_t size = 0; + for (uint32_t value : values) { + size += kLookupSizeTable32[__builtin_clz(value | 1U)]; + } + return size; +} + +uint64_t bulkVarintSize64(std::span values) { + constexpr uint8_t kLookupSizeTable64[64] = { + 10, 9, 9, 9, 9, 9, 9, 9, 8, 8, 8, 8, 8, 8, 8, 7, 7, 7, 7, 7, 7, 7, + 6, 6, 6, 6, 6, 6, 6, 5, 5, 5, 5, 5, 5, 5, 4, 4, 4, 4, 4, 4, 4, 3, + 3, 3, 3, 3, 3, 3, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1}; + uint64_t size = 0; + for (uint64_t value : values) { + size += kLookupSizeTable64[__builtin_clzll(value | 1ULL)]; + } + return size; +} + +// Declaration of the function we build via generated code below. +template +__attribute__((__target__("bmi2"))) +// __attribute__ ((optimize("Os"))) +const char* +bulkVarintDecodeBmi2(uint64_t n, const char* pos, T* output); + +const char* bulkVarintDecode32(uint64_t n, const char* pos, uint32_t* output) { + static bool hasBmi2 = folly::CpuId().bmi2(); + if (hasBmi2) { + return bulkVarintDecodeBmi2(n, pos, output); + } + for (uint64_t i = 0; i < n; ++i) { + *output++ = readVarint32(&pos); + } + return pos; +} + +const char* bulkVarintDecode64(uint64_t n, const char* pos, uint64_t* output) { + static bool hasBmi2 = folly::CpuId().bmi2(); + if (hasBmi2) { + return bulkVarintDecodeBmi2(n, pos, output); + } + for (uint64_t i = 0; i < n; ++i) { + *output++ = readVarint64(&pos); + } + return pos; +} + +// Codegen for the cases below. Useful if we want to try to tweak something +// (different mask length, etc) in the future. + +// std::string codegenVarintMask(int endByte, int len) { +// CHECK_GE(endByte + 1, len); +// std::string s = "0x0000000000000000ULL"; +// int offset = 5 + endByte * 2; +// for (int i = 0; i < len; ++i) { +// *(s.end() - offset) = '7'; +// *(s.end() - offset + 1) = 'f'; +// offset -= 2; +// } +// return s; +// } + +// std::string codegen32(uint64_t controlBits, int maskLength) { +// CHECK(controlBits < (1 << maskLength)); +// std::string s = absl::Substitute( +// " case $0ULL: {", controlBits); +// int lastZero = -1; +// int numVariants = 0; +// bool carryoverUsed = false; +// for (int nextBit = 0; nextBit < maskLength; ++nextBit) { +// // A zero control bit means we detected the end of a varint, so +// // we can construct a mask of the bottom 7 bits starting at the end +// // of the nextBit byte and going back (nextBit - lastZero) bytes. +// if ((controlBits & (1ULL << nextBit)) == 0) { +// if (carryoverUsed) { +// s += absl::Substitute("\n *output++ = _pext_u64(word, $0);", +// CodegenVarintMask(nextBit, nextBit - +// lastZero)); +// } else { +// s += absl::Substitute("\n const uint64_t firstValue = " +// "_pext_u64(word, $0);", +// CodegenVarintMask(nextBit, nextBit - +// lastZero)); +// s += "\n *output++ = (firstValue << carryoverBits) | +// carryover;"; carryoverUsed = true; +// } +// lastZero = nextBit; +// ++numVariants; +// } +// } +// // Ending on a complete varint, not completing any varint, and completing +// // at least 1 varint but no ending on one are all distinct cases. +// if (lastZero == -1) { +// s += absl::Substitute("\n carryover |= " +// "_pext_u64(word, $0) << carryoverBits;", +// CodegenVarintMask(maskLength - 1, maskLength)); +// s += absl::Substitute("\n carryoverBits += $0;", 7 * maskLength); +// } else if (lastZero == maskLength - 1) { +// s += "\n carryover = 0ULL;"; +// s += "\n carryoverBits = 0;"; +// s += absl::Substitute("\n n -= $0;", numVariants); +// } else { +// s += absl::Substitute("\n carryover = _pext_u64(word, $0);", +// CodegenVarintMask(maskLength - 1, +// maskLength - 1 - lastZero)); +// s += absl::Substitute("\n carryoverBits = $0;", +// 7 * (maskLength - 1 - lastZero)); +// s += absl::Substitute("\n n -= $0;", numVariants); +// } +// // s += absl::Substitute("\n pos += $0;", maskLength); +// s += "\n continue;"; +// s += "\n }"; +// return s; +// } + +template +const char* bulkVarintDecodeBmi2(uint64_t n, const char* pos, T* output) { + constexpr uint64_t mask = 0x0000808080808080; + // Note that we could of course use a maskLength of up to 8. But I found + // that with maskLength > 6 we start to spill out of the l1i cache in + // opt mode and that counterbalances the gain. Plus the first run and/or + // small n are more expensive as we have to load more instructions. + constexpr int maskLength = 6; + uint64_t carryover = 0; + int carryoverBits = 0; + pos -= maskLength; + // Also note that a handful of these cases are impossible for 32-bit varints. + // We coould save a tiny bit of program size by pruning them out. + while (n >= 8) { + pos += maskLength; + uint64_t word = *reinterpret_cast(pos); + const uint64_t controlBits = _pext_u64(word, mask); + switch (controlBits) { + case 0ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + *output++ = _pext_u64(word, 0x000000007f000000ULL); + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 6; + continue; + } + case 1ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + *output++ = _pext_u64(word, 0x000000007f000000ULL); + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 5; + continue; + } + case 2ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f7f00ULL); + *output++ = _pext_u64(word, 0x000000007f000000ULL); + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 5; + continue; + } + case 3ULL: { + const uint64_t firstValue = _pext_u64(word, 0x00000000007f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x000000007f000000ULL); + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 4; + continue; + } + case 4ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x000000007f7f0000ULL); + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 5; + continue; + } + case 5ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x000000007f7f0000ULL); + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 4; + continue; + } + case 6ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x000000007f7f7f00ULL); + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 4; + continue; + } + case 7ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000007f7f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 3; + continue; + } + case 8ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + *output++ = _pext_u64(word, 0x0000007f7f000000ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 5; + continue; + } + case 9ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + *output++ = _pext_u64(word, 0x0000007f7f000000ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 4; + continue; + } + case 10ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f7f00ULL); + *output++ = _pext_u64(word, 0x0000007f7f000000ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 4; + continue; + } + case 11ULL: { + const uint64_t firstValue = _pext_u64(word, 0x00000000007f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000007f7f000000ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 3; + continue; + } + case 12ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x0000007f7f7f0000ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 4; + continue; + } + case 13ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000007f7f7f0000ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 3; + continue; + } + case 14ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000007f7f7f7f00ULL); + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 3; + continue; + } + case 15ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000007f7f7f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00007f0000000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 2; + continue; + } + case 16ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + *output++ = _pext_u64(word, 0x000000007f000000ULL); + *output++ = _pext_u64(word, 0x00007f7f00000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 5; + continue; + } + case 17ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + *output++ = _pext_u64(word, 0x000000007f000000ULL); + *output++ = _pext_u64(word, 0x00007f7f00000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 4; + continue; + } + case 18ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f7f00ULL); + *output++ = _pext_u64(word, 0x000000007f000000ULL); + *output++ = _pext_u64(word, 0x00007f7f00000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 4; + continue; + } + case 19ULL: { + const uint64_t firstValue = _pext_u64(word, 0x00000000007f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x000000007f000000ULL); + *output++ = _pext_u64(word, 0x00007f7f00000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 3; + continue; + } + case 20ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x000000007f7f0000ULL); + *output++ = _pext_u64(word, 0x00007f7f00000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 4; + continue; + } + case 21ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x000000007f7f0000ULL); + *output++ = _pext_u64(word, 0x00007f7f00000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 3; + continue; + } + case 22ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x000000007f7f7f00ULL); + *output++ = _pext_u64(word, 0x00007f7f00000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 3; + continue; + } + case 23ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000007f7f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00007f7f00000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 2; + continue; + } + case 24ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + *output++ = _pext_u64(word, 0x00007f7f7f000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 4; + continue; + } + case 25ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + *output++ = _pext_u64(word, 0x00007f7f7f000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 3; + continue; + } + case 26ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f7f00ULL); + *output++ = _pext_u64(word, 0x00007f7f7f000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 3; + continue; + } + case 27ULL: { + const uint64_t firstValue = _pext_u64(word, 0x00000000007f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00007f7f7f000000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 2; + continue; + } + case 28ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x00007f7f7f7f0000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 3; + continue; + } + case 29ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00007f7f7f7f0000ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 2; + continue; + } + case 30ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00007f7f7f7f7f00ULL); + carryover = 0ULL; + carryoverBits = 0; + n -= 2; + continue; + } + case 31ULL: { + const uint64_t firstValue = _pext_u64(word, 0x00007f7f7f7f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + carryover = 0ULL; + carryoverBits = 0; + n -= 1; + continue; + } + case 32ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + *output++ = _pext_u64(word, 0x000000007f000000ULL); + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 5; + continue; + } + case 33ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + *output++ = _pext_u64(word, 0x000000007f000000ULL); + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 4; + continue; + } + case 34ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f7f00ULL); + *output++ = _pext_u64(word, 0x000000007f000000ULL); + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 4; + continue; + } + case 35ULL: { + const uint64_t firstValue = _pext_u64(word, 0x00000000007f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x000000007f000000ULL); + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 3; + continue; + } + case 36ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x000000007f7f0000ULL); + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 4; + continue; + } + case 37ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x000000007f7f0000ULL); + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 3; + continue; + } + case 38ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x000000007f7f7f00ULL); + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 3; + continue; + } + case 39ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000007f7f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000007f00000000ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 2; + continue; + } + case 40ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + *output++ = _pext_u64(word, 0x0000007f7f000000ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 4; + continue; + } + case 41ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + *output++ = _pext_u64(word, 0x0000007f7f000000ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 3; + continue; + } + case 42ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f7f00ULL); + *output++ = _pext_u64(word, 0x0000007f7f000000ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 3; + continue; + } + case 43ULL: { + const uint64_t firstValue = _pext_u64(word, 0x00000000007f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000007f7f000000ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 2; + continue; + } + case 44ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x0000007f7f7f0000ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 3; + continue; + } + case 45ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000007f7f7f0000ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 2; + continue; + } + case 46ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000007f7f7f7f00ULL); + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 2; + continue; + } + case 47ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000007f7f7f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + carryover = _pext_u64(word, 0x00007f0000000000ULL); + carryoverBits = 7; + n -= 1; + continue; + } + case 48ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + *output++ = _pext_u64(word, 0x000000007f000000ULL); + carryover = _pext_u64(word, 0x00007f7f00000000ULL); + carryoverBits = 14; + n -= 4; + continue; + } + case 49ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + *output++ = _pext_u64(word, 0x000000007f000000ULL); + carryover = _pext_u64(word, 0x00007f7f00000000ULL); + carryoverBits = 14; + n -= 3; + continue; + } + case 50ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f7f00ULL); + *output++ = _pext_u64(word, 0x000000007f000000ULL); + carryover = _pext_u64(word, 0x00007f7f00000000ULL); + carryoverBits = 14; + n -= 3; + continue; + } + case 51ULL: { + const uint64_t firstValue = _pext_u64(word, 0x00000000007f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x000000007f000000ULL); + carryover = _pext_u64(word, 0x00007f7f00000000ULL); + carryoverBits = 14; + n -= 2; + continue; + } + case 52ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x000000007f7f0000ULL); + carryover = _pext_u64(word, 0x00007f7f00000000ULL); + carryoverBits = 14; + n -= 3; + continue; + } + case 53ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x000000007f7f0000ULL); + carryover = _pext_u64(word, 0x00007f7f00000000ULL); + carryoverBits = 14; + n -= 2; + continue; + } + case 54ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x000000007f7f7f00ULL); + carryover = _pext_u64(word, 0x00007f7f00000000ULL); + carryoverBits = 14; + n -= 2; + continue; + } + case 55ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000007f7f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + carryover = _pext_u64(word, 0x00007f7f00000000ULL); + carryoverBits = 14; + n -= 1; + continue; + } + case 56ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + carryover = _pext_u64(word, 0x00007f7f7f000000ULL); + carryoverBits = 21; + n -= 3; + continue; + } + case 57ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f0000ULL); + carryover = _pext_u64(word, 0x00007f7f7f000000ULL); + carryoverBits = 21; + n -= 2; + continue; + } + case 58ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x00000000007f7f00ULL); + carryover = _pext_u64(word, 0x00007f7f7f000000ULL); + carryoverBits = 21; + n -= 2; + continue; + } + case 59ULL: { + const uint64_t firstValue = _pext_u64(word, 0x00000000007f7f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + carryover = _pext_u64(word, 0x00007f7f7f000000ULL); + carryoverBits = 21; + n -= 1; + continue; + } + case 60ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + *output++ = _pext_u64(word, 0x0000000000007f00ULL); + carryover = _pext_u64(word, 0x00007f7f7f7f0000ULL); + carryoverBits = 28; + n -= 2; + continue; + } + case 61ULL: { + const uint64_t firstValue = _pext_u64(word, 0x0000000000007f7fULL); + *output++ = (firstValue << carryoverBits) | carryover; + carryover = _pext_u64(word, 0x00007f7f7f7f0000ULL); + carryoverBits = 28; + n -= 1; + continue; + } + case 62ULL: { + const uint64_t firstValue = _pext_u64(word, 0x000000000000007fULL); + *output++ = (firstValue << carryoverBits) | carryover; + carryover = _pext_u64(word, 0x00007f7f7f7f7f00ULL); + carryoverBits = 35; + n -= 1; + continue; + } + case 63ULL: { + carryover |= _pext_u64(word, 0x00007f7f7f7f7f7fULL) << carryoverBits; + carryoverBits += 42; + continue; + } + default: { + NIMBLE_UNREACHABLE("Control bits must be < 64"); + } + } + } + pos += maskLength; + if (n > 0) { + if constexpr (std::is_same::value) { + *output++ = readVarint32(&pos) << carryoverBits | carryover; + for (uint64_t i = 1; i < n; ++i) { + *output++ = readVarint32(&pos); + } + } else { + *output++ = readVarint64(&pos) << carryoverBits | carryover; + for (uint64_t i = 1; i < n; ++i) { + *output++ = readVarint64(&pos); + } + } + } + return pos; +} + +} // namespace facebook::nimble::varint diff --git a/dwio/nimble/common/Varint.h b/dwio/nimble/common/Varint.h new file mode 100644 index 0000000..07a66e3 --- /dev/null +++ b/dwio/nimble/common/Varint.h @@ -0,0 +1,110 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include + +// Varint-related encoding methods. Same binary encoding as folly/Varint.h +// but with a different API, a faster decode method thanks to +// to not checking bounds, and bulk decoding methods. No functions in this +// library check bounds or deal with buffer overflow. +// +// The bulk decoding methods are particularly interesting. They range from +// 4x faster than the non-bulk in the best case (all 1 byte varints) to +// about 50% faster in the worst case (random byte lengths), assuming that +// bmi2 is available. See varintBenchmark.h for some details. + +namespace facebook::nimble::varint { + +// Decode n varints at once. Makes use of bmi2 instruction set if its +// available. Returns pos updated past the n varints. +const char* bulkVarintDecode32(uint64_t n, const char* pos, uint32_t* output); +const char* bulkVarintDecode64(uint64_t n, const char* pos, uint64_t* output); + +// Skips n varints, returning pos updated past the n varints. +const char* bulkVarintSkip(uint64_t n, const char* pos); + +// Returns the number of bytes the |values| will occupy after varint encoding. +uint64_t bulkVarintSize32(std::span values); +uint64_t bulkVarintSize64(std::span values); + +// Inline non-bulk methods follow below. + +template +inline void writeVarint(T val, char** pos) noexcept { + while (val >= 128) { + *((*pos)++) = 0x80 | (val & 0x7f); + val >>= 7; + } + *((*pos)++) = val; +} + +inline void skipVarint(const char** pos) noexcept { + while (*((*pos)++) & 128) { + } +} + +inline uint32_t readVarint32(const char** pos) noexcept { + uint32_t value = (**pos) & 127; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= (**pos & 127) << 7; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= (**pos & 127) << 14; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= (**pos & 127) << 21; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= (*((*pos)++) & 127) << 28; + return value; +} + +inline uint64_t readVarint64(const char** pos) noexcept { + uint64_t value = (**pos) & 127; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= (**pos & 127) << 7; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= (**pos & 127) << 14; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= (**pos & 127) << 21; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= static_cast(**pos & 127) << 28; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= static_cast(**pos & 127) << 35; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= static_cast(**pos & 127) << 42; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= static_cast(**pos & 127) << 49; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= static_cast(**pos & 127) << 56; + if (!(*((*pos)++) & 128)) { + return value; + } + value |= static_cast(*((*pos)++) & 127) << 63; + return value; +} + +} // namespace facebook::nimble::varint diff --git a/dwio/nimble/common/Vector.h b/dwio/nimble/common/Vector.h new file mode 100644 index 0000000..8adbe79 --- /dev/null +++ b/dwio/nimble/common/Vector.h @@ -0,0 +1,286 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/common/Exceptions.h" +#include "velox/buffer/Buffer.h" +#include "velox/common/memory/Memory.h" + +#include +#include + +// Basically equivalent to std::vector, but without the edge case for booleans, +// i.e. data() returns T* for all T. This allows for implicit conversion to +// std::span for all T. + +namespace facebook::nimble { + +template +class Vector { + using InnerType = + typename std::conditional, uint8_t, T>::type; + + public: + Vector(velox::memory::MemoryPool* memoryPool, size_t size, T value) + : memoryPool_{memoryPool} { + init(size); + std::fill(dataRawPtr_, dataRawPtr_ + size_, value); + } + + Vector(velox::memory::MemoryPool* memoryPool, size_t size) + : memoryPool_{memoryPool} { + init(size); + } + + explicit Vector(velox::memory::MemoryPool* memoryPool) + : memoryPool_{memoryPool} { + capacity_ = 0; + size_ = 0; + data_ = nullptr; + dataRawPtr_ = nullptr; +#ifndef NDEBUG + dataRawPtr_ = placeholder_.data(); +#endif + } + + template + Vector(velox::memory::MemoryPool* memoryPool, It first, It last) + : memoryPool_{memoryPool} { + auto size = last - first; + init(size); + std::copy(first, last, dataRawPtr_); + } + + Vector(const Vector& other) { + *this = other; + } + + Vector& operator=(const Vector& other) { + if (this != &other) { + size_ = other.size(); + capacity_ = other.capacity_; + memoryPool_ = other.memoryPool_; + allocateBuffer(); + std::copy(other.dataRawPtr_, other.dataRawPtr_ + size_, dataRawPtr_); + } + return *this; + } + + Vector(Vector&& other) noexcept { + *this = std::move(other); + } + + Vector& operator=(Vector&& other) noexcept { + if (this != &other) { + size_ = other.size(); + capacity_ = other.capacity_; + data_ = std::move(other.data_); +#ifndef NDEBUG + dataRawPtr_ = placeholder_.data(); +#endif + if (data_ != nullptr) { + dataRawPtr_ = reinterpret_cast(data_->asMutable()); + } + memoryPool_ = other.memoryPool_; + other.size_ = 0; + other.capacity_ = 0; + } + return *this; + } + + Vector(velox::memory::MemoryPool* memoryPool, std::initializer_list l) + : memoryPool_{memoryPool} { + init(l.size()); + std::copy(l.begin(), l.end(), dataRawPtr_); + } + + inline velox::memory::MemoryPool* memoryPool() { + return memoryPool_; + } + + uint64_t size() const { + return size_; + } + bool empty() const { + return size_ == 0; + } + uint64_t capacity() const { + return capacity_; + } + T& operator[](uint64_t i) { + return dataRawPtr_[i]; + } + const T& operator[](uint64_t i) const { + return dataRawPtr_[i]; + } + T* begin() { + return dataRawPtr_; + } + T* end() { + return dataRawPtr_ + size_; + } + const T* begin() const { + return dataRawPtr_; + } + const T* end() const { + return dataRawPtr_ + size_; + } + T& back() { + return dataRawPtr_[size_ - 1]; + } + const T& back() const { + return dataRawPtr_[size_ - 1]; + } + + // Directly updates the size_ to |size|. Useful if you've filled in some data + // directly using the underlying raw pointers. + void update_size(uint64_t size) { + size_ = size; + } + + // Fills all of data_ with T(). + void zero_out() { + std::fill(dataRawPtr_, dataRawPtr_ + capacity_, T()); + } + + // Fills all of data_ with the given value. + void fill(T value) { + std::fill(dataRawPtr_, dataRawPtr_ + capacity_, value); + } + + // Resets *this to a newly constructed empty state. + void clear() { + capacity_ = 0; + size_ = 0; + data_.reset(); + dataRawPtr_ = nullptr; + } + + void insert(T* output, const T* inputStart, const T* inputEnd) { + const uint64_t inputSize = inputEnd - inputStart; + const uint64_t distanceToEnd = end() - output; + if (inputSize > distanceToEnd) { + const uint64_t sizeChange = inputSize - distanceToEnd; + const uint64_t distanceFromBegin = output - begin(); + resize(size_ + sizeChange); + std::move(inputStart, inputEnd, begin() + distanceFromBegin); + } else { + std::move(inputStart, inputEnd, output); + } + } + + // Add |copies| copies of |value| to the end of the vector. + void extend(uint64_t copies, T value) { + reserve(size_ + copies); + std::fill(end(), end() + copies, value); + size_ += copies; + } + + T* data() noexcept { + return dataRawPtr_; + } + + const T* data() const noexcept { + return dataRawPtr_; + } + + void push_back(T value) { + if (size_ == capacity_) { + reserve(calculateNewSize(capacity_)); + } + dataRawPtr_[size_] = value; + ++size_; + } + + template + void emplace_back(Args&&... args) { + if (size_ == capacity_) { + reserve(calculateNewSize(capacity_)); + } + // use placement new to construct the object. + new (dataRawPtr_ + size_) InnerType(std::forward(args)...); + ++size_; + } + + // Ensures that *this can hold |size| elements. Does NOT shrink to + // fit if |size| is less than size(), and does NOT initialize any new + // values. + void reserve(uint64_t size) { + if (size > capacity_) { + capacity_ = size; + auto newData = + velox::AlignedBuffer::allocate(capacity_, memoryPool_); + if (data_ != nullptr && size_ > 0) { + std::move( + dataRawPtr_, + dataRawPtr_ + size_, + reinterpret_cast(newData->template asMutable())); + } + data_ = std::move(newData); + dataRawPtr_ = reinterpret_cast(data_->asMutable()); + } + } + + // Changes size_ to |size|. Does NOT shrink to fit the new size, and + // does NOT initialize any new elements if |size| is greater than size_. + void resize(uint64_t size) { + reserve(size); + size_ = size; + } + + // Changes size_ to |newSize|. Does NOT shrink to fit the new size. Initialize + // any new elements to |value| if |newSize| is greater than size_. + void resize(uint64_t newSize, const T& value) { + auto initialSize = size_; + resize(newSize); + + if (size_ > initialSize) { + std::fill(dataRawPtr_ + initialSize, end(), value); + } + } + + velox::BufferPtr releaseOwnership() { + velox::BufferPtr tmp = std::move(data_); + tmp->setSize(size_); + capacity_ = 0; + size_ = 0; + data_ = nullptr; + dataRawPtr_ = nullptr; +#ifndef NDEBUG + dataRawPtr_ = placeholder_.data(); +#endif + return tmp; + } + + private: + inline void init(size_t size) { + capacity_ = size; + size_ = size; + allocateBuffer(); + } + + inline size_t calculateNewSize(size_t size) { + auto newSize = size <<= 1; + if (newSize == 0) { + return velox::AlignedBuffer::kSizeofAlignedBuffer; + } + + return newSize; + } + + inline void allocateBuffer() { + data_ = velox::AlignedBuffer::allocate(capacity_, memoryPool_); + dataRawPtr_ = reinterpret_cast(data_->asMutable()); + } + + velox::memory::MemoryPool* memoryPool_; + velox::BufferPtr data_; + uint64_t capacity_; + uint64_t size_; + T* dataRawPtr_; +#ifndef NDEBUG + inline static std::array placeholder_; +#endif +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/common/benchmarks/VarintBenchmark.cpp b/dwio/nimble/common/benchmarks/VarintBenchmark.cpp new file mode 100644 index 0000000..964cd60 --- /dev/null +++ b/dwio/nimble/common/benchmarks/VarintBenchmark.cpp @@ -0,0 +1,325 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include +#include + +#include "dwio/nimble/common/Varint.h" +#include "folly/Benchmark.h" +#include "folly/Random.h" +#include "folly/Varint.h" + +using namespace ::facebook; + +const int kNumElements = 1000 * 1000; + +// Basically same code as dwrf::IntDecoder::readVuLong. +uint64_t DwrfRead(const char** bufferStart, const char* bufferEnd) { + if (LIKELY(bufferEnd - *bufferStart >= folly::kMaxVarintLength64)) { + const char* p = *bufferStart; + uint64_t val; + do { + int64_t b; + b = *p++; + val = (b & 0x7f); + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 7; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 14; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 21; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 28; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 35; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 42; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 49; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x7f) << 56; + if (UNLIKELY(b >= 0)) { + break; + } + b = *p++; + val |= (b & 0x01) << 63; + if (LIKELY(b >= 0)) { + break; + } else { + throw std::runtime_error{"invalid encoding: likely corrupt data"}; + } + } while (false); + *bufferStart = p; + return val; + } else { + // this part isn't the same, but doesn't measurably effect time. + return nimble::varint::readVarint64(bufferStart); + } +} + +// Makes random data uniform over in bit width over 32 bits. +std::vector MakeUniformData(int num_elements = kNumElements) { + std::vector data(num_elements); + for (int i = 0; i < num_elements; ++i) { + const int bit_shift = 1 + folly::Random::secureRand32() % 32; + data[i] = folly::Random::secureRand32() % (1 << bit_shift); + } + return data; +} + +// Makes 95% 1 byte, 5% 2 byte data. +std::vector MakeSkewedData(int num_elements = kNumElements) { + std::vector data(num_elements); + for (int i = 0; i < num_elements; ++i) { + if (folly::Random::secureRand32() % 20) { + data[i] = folly::Random::secureRand32() % (1 << 7); + } else { + data[i] = folly::Random::secureRand32() % (1 << 14); + } + } + return data; +} + +BENCHMARK(Encode, iters) { + std::vector data; + std::unique_ptr buf; + BENCHMARK_SUSPEND { + data = MakeUniformData(); + buf = std::make_unique(kNumElements * folly::kMaxVarintLength32); + } + while (iters--) { + char* pos = buf.get(); + for (int i = 0; i < kNumElements; ++i) { + nimble::varint::writeVarint(data[i], &pos); + } + CHECK_GE(pos - buf.get(), kNumElements); + } +} + +BENCHMARK(FollyEncode, iters) { + std::vector data; + std::unique_ptr buf; + BENCHMARK_SUSPEND { + data = MakeUniformData(); + buf = std::make_unique(kNumElements * folly::kMaxVarintLength32); + } + while (iters--) { + uint8_t* pos = buf.get(); + for (int i = 0; i < kNumElements; ++i) { + pos += folly::encodeVarint(data[i], pos); + } + CHECK_GE(pos - buf.get(), kNumElements); + } +} + +BENCHMARK(NimbleDecodeUniform, iters) { + std::vector data; + std::unique_ptr buf; + std::vector recovered; + BENCHMARK_SUSPEND { + recovered.resize(kNumElements); + data = MakeUniformData(); + buf = std::make_unique(kNumElements * folly::kMaxVarintLength32); + char* pos = buf.get(); + for (int i = 0; i < kNumElements; ++i) { + nimble::varint::writeVarint(data[i], &pos); + } + } + while (iters--) { + const char* cpos = buf.get(); + for (int i = 0; i < kNumElements; ++i) { + recovered[i] = nimble::varint::readVarint32(&cpos); + } + CHECK_EQ(recovered.back(), data.back()); + } +} + +BENCHMARK(NimbleBulkDecodeUniform, iters) { + std::vector data; + std::unique_ptr buf; + std::vector recovered; + BENCHMARK_SUSPEND { + recovered.resize(kNumElements); + data = MakeUniformData(); + buf = std::make_unique(kNumElements * folly::kMaxVarintLength32); + char* pos = buf.get(); + for (int i = 0; i < kNumElements; ++i) { + nimble::varint::writeVarint(data[i], &pos); + } + } + while (iters--) { + const char* cpos = buf.get(); + nimble::varint::bulkVarintDecode32(kNumElements, cpos, recovered.data()); + CHECK_EQ(recovered.back(), data.back()); + } +} + +BENCHMARK(FollyDecodeUniform, iters) { + std::vector data; + std::unique_ptr buf; + uint8_t* pos; + std::vector recovered; + BENCHMARK_SUSPEND { + recovered.resize(kNumElements); + data = MakeUniformData(); + buf = std::make_unique(kNumElements * folly::kMaxVarintLength32); + pos = buf.get(); + for (int i = 0; i < kNumElements; ++i) { + pos += folly::encodeVarint(data[i], pos); + } + } + while (iters--) { + const uint8_t* fstart = buf.get(); + const uint8_t* fend = buf.get() + (pos - buf.get()); + folly::Range frange(fstart, fend); + for (int i = 0; i < kNumElements; ++i) { + recovered[i] = folly::decodeVarint(frange); + } + CHECK_EQ(recovered.back(), data.back()); + } +} + +BENCHMARK(DwrfDecodeUniform, iters) { + std::vector data; + std::unique_ptr buf; + std::vector recovered; + uint64_t varint_bytes; + BENCHMARK_SUSPEND { + recovered.resize(kNumElements); + data = MakeUniformData(); + buf = std::make_unique(kNumElements * folly::kMaxVarintLength32); + char* pos = buf.get(); + for (int i = 0; i < kNumElements; ++i) { + nimble::varint::writeVarint(data[i], &pos); + } + varint_bytes = pos - buf.get(); + } + while (iters--) { + const char* cpos = buf.get(); + const char* end = cpos + varint_bytes; + for (int i = 0; i < kNumElements; ++i) { + recovered[i] = DwrfRead(&cpos, end); + } + CHECK_EQ(recovered.back(), data.back()); + } +} + +BENCHMARK(NimbleDecodeSkewed, iters) { + std::vector data; + std::unique_ptr buf; + std::vector recovered; + BENCHMARK_SUSPEND { + recovered.resize(kNumElements); + data = MakeSkewedData(); + buf = std::make_unique(kNumElements * folly::kMaxVarintLength32); + char* pos = buf.get(); + for (int i = 0; i < kNumElements; ++i) { + nimble::varint::writeVarint(data[i], &pos); + } + } + while (iters--) { + const char* cpos = buf.get(); + for (int i = 0; i < kNumElements; ++i) { + recovered[i] = nimble::varint::readVarint32(&cpos); + } + CHECK_EQ(recovered.back(), data.back()); + } +} + +BENCHMARK(NimbleBulkDecodeSkewed, iters) { + std::vector data; + std::unique_ptr buf; + std::vector recovered; + BENCHMARK_SUSPEND { + recovered.resize(kNumElements); + data = MakeSkewedData(); + buf = std::make_unique(kNumElements * folly::kMaxVarintLength32); + char* pos = buf.get(); + for (int i = 0; i < kNumElements; ++i) { + nimble::varint::writeVarint(data[i], &pos); + } + } + while (iters--) { + const char* cpos = buf.get(); + nimble::varint::bulkVarintDecode32(kNumElements, cpos, recovered.data()); + CHECK_EQ(recovered.back(), data.back()); + } +} + +BENCHMARK(FollyDecodeSkewed, iters) { + std::vector data; + std::unique_ptr buf; + uint8_t* pos; + std::vector recovered; + BENCHMARK_SUSPEND { + recovered.resize(kNumElements); + data = MakeSkewedData(); + buf = std::make_unique(kNumElements * folly::kMaxVarintLength32); + pos = buf.get(); + for (int i = 0; i < kNumElements; ++i) { + pos += folly::encodeVarint(data[i], pos); + } + } + while (iters--) { + const uint8_t* fstart = buf.get(); + const uint8_t* fend = buf.get() + (pos - buf.get()); + folly::Range frange(fstart, fend); + for (int i = 0; i < kNumElements; ++i) { + recovered[i] = folly::decodeVarint(frange); + } + CHECK_EQ(recovered.back(), data.back()); + } +} + +BENCHMARK(DwrfDecodeSkewed, iters) { + std::vector data; + std::unique_ptr buf; + std::vector recovered; + uint64_t varint_bytes; + BENCHMARK_SUSPEND { + recovered.resize(kNumElements); + data = MakeSkewedData(); + buf = std::make_unique(kNumElements * folly::kMaxVarintLength32); + char* pos = buf.get(); + for (int i = 0; i < kNumElements; ++i) { + nimble::varint::writeVarint(data[i], &pos); + } + varint_bytes = pos - buf.get(); + } + while (iters--) { + const char* cpos = buf.get(); + const char* end = cpos + varint_bytes; + for (int i = 0; i < kNumElements; ++i) { + recovered[i] = DwrfRead(&cpos, end); + } + CHECK_EQ(recovered.back(), data.back()); + } +} + +int main() { + folly::runBenchmarks(); +} diff --git a/dwio/nimble/common/tests/BitEncoderTests.cpp b/dwio/nimble/common/tests/BitEncoderTests.cpp new file mode 100644 index 0000000..49328ee --- /dev/null +++ b/dwio/nimble/common/tests/BitEncoderTests.cpp @@ -0,0 +1,65 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include +#include "dwio/nimble/common/BitEncoder.h" +#include "folly/Random.h" + +using namespace ::facebook; + +TEST(BitEncoderTests, WriteThenReadDifferentBitLengths) { + constexpr int elementCount = 1000; + std::vector buffer(4 * elementCount); + nimble::BitEncoder bitEncoder(buffer.data()); + for (int i = 0; i < elementCount; ++i) { + bitEncoder.putBits(i % 31, (i % 31) + 1); + } + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(i % 31, bitEncoder.getBits((i % 31) + 1)); + } +} + +TEST(BitEncoderTests, WriteAndReadIntermixed) { + constexpr int elementCount = 1000; + std::vector data; + for (int i = 0; i < elementCount; ++i) { + data.push_back(2 * i); + } + std::vector bitLengths; + for (int i = 0; i < elementCount; ++i) { + bitLengths.push_back(11 + (i % 10)); + } + std::vector buffer(4 * elementCount); + nimble::BitEncoder bitEncoder(buffer.data()); + for (int i = 0; i < elementCount; i += 2) { + bitEncoder.putBits(data[i], bitLengths[i]); + bitEncoder.putBits(data[i + 1], bitLengths[i + 1]); + ASSERT_EQ(data[i >> 1], bitEncoder.getBits(bitLengths[i >> 1])); + } + for (int i = elementCount / 2; i < elementCount; ++i) { + ASSERT_EQ(data[i], bitEncoder.getBits(bitLengths[i])); + } +} + +TEST(BitEncoderTests, WriteThenReadFullBitRange) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + constexpr int elementCount = 1000; + std::vector data; + for (int i = 0; i < elementCount; ++i) { + data.push_back(folly::Random::rand64(rng)); + } + for (int bits = 1; bits < 65; ++bits) { + const uint64_t mask = bits == 64 ? (~0ULL) : ((1ULL << bits) - 1); + std::vector buffer(8 * elementCount); + nimble::BitEncoder bitEncoder(buffer.data()); + for (int i = 0; i < elementCount; ++i) { + bitEncoder.putBits(data[i] & mask, bits); + } + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(bitEncoder.getBits(bits), data[i] & mask); + } + } +} diff --git a/dwio/nimble/common/tests/BitsTests.cpp b/dwio/nimble/common/tests/BitsTests.cpp new file mode 100644 index 0000000..4be7d58 --- /dev/null +++ b/dwio/nimble/common/tests/BitsTests.cpp @@ -0,0 +1,90 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include +#include +#include "dwio/nimble/common/Bits.h" +#include "folly/Random.h" + +using namespace facebook::nimble::bits; + +template +void repeat(int32_t times, const T& t) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + while (times-- > 0) { + t(rng); + } +} + +TEST(BitsTests, setBits) { + repeat(10, [](auto& rng) { + auto size = folly::Random::rand32(64 * 1024, rng) + 1; + auto begin = folly::Random::rand32(size, rng); + auto end = folly::Random::rand32(begin, size, rng); + std::vector bitmap(bucketsRequired(size, 64) * 8, 0); + setBits(begin, end, bitmap.data()); + for (auto i = 0; i < size; ++i) { + bool expected = (i >= begin && i < end); + EXPECT_EQ(expected, getBit(i, bitmap.data())) << i; + } + }); +} + +TEST(BitsTests, clearBits) { + repeat(10, [](auto& rng) { + auto size = folly::Random::rand32(64 * 1024, rng) + 1; + auto begin = folly::Random::rand32(size, rng); + auto end = folly::Random::rand32(begin, size, rng); + std::vector bitmap(bucketsRequired(size, 64) * 8, 0xff); + clearBits(begin, end, bitmap.data()); + for (auto i = 0; i < size; ++i) { + bool expected = (i < begin || i >= end); + EXPECT_EQ(expected, getBit(i, bitmap.data())) << i; + } + }); +} + +TEST(BitsTests, findSetBit) { + repeat(10, [](auto& rng) { + auto size = folly::Random::rand32(64 * 1024, rng) + 1; + std::vector bitmap(bucketsRequired(size, 64) * 8, 0); + auto begin = folly::Random::rand32(size, rng); + auto n = (size - begin) / 3; + auto setBits = 0; + auto pos = size; + for (auto i = begin; i < size; ++i) { + if (folly::Random::oneIn(3, rng)) { + setBit(i, bitmap.data()); + if (++setBits == n) { + pos = i; + } + } + } + EXPECT_EQ(pos, findSetBit(bitmap.data(), begin, size, n)); + }); +} + +TEST(BitsTests, copy) { + repeat(1, [](auto& rng) { + auto size = folly::Random::rand32(64 * 1024, rng) + 1; + auto begin = folly::Random::rand32(size, rng); + auto end = folly::Random::oneIn(2) ? folly::Random::rand32(begin, size, rng) + : size; + std::vector src(bucketsRequired(size, 64) * 8, 0); + std::vector dst(src.size(), 0); + for (auto i = begin; i < end; ++i) { + maybeSetBit(i, src.data(), folly::Random::oneIn(2, rng)); + } + Bitmap srcBitmap{src.data(), size}; + BitmapBuilder dstBitmap{dst.data(), size}; + dstBitmap.copy(srcBitmap, begin, end); + for (auto i = 0; i < end; ++i) { + if (i < begin) { + EXPECT_FALSE(getBit(i, dst.data())); + } else { + EXPECT_EQ(getBit(i, src.data()), getBit(i, dst.data())) << i; + } + } + }); +} diff --git a/dwio/nimble/common/tests/CMakeLists.txt b/dwio/nimble/common/tests/CMakeLists.txt new file mode 100644 index 0000000..cca4034 --- /dev/null +++ b/dwio/nimble/common/tests/CMakeLists.txt @@ -0,0 +1,20 @@ +add_library(nimble_common_file_writer NimbleFileWriter.cpp) + +target_link_libraries(nimble_common_file_writer nimble_common + nimble_velox_writer velox_vector) + +add_executable( + nimble_common_tests + BitEncoderTests.cpp + BitsTests.cpp + ExceptionTests.cpp + FixedBitArrayTests.cpp + HuffmanTests.cpp + IndexMapTests.cpp + VarintTests.cpp + VectorTests.cpp) + +add_test(nimble_common_tests nimble_common_tests) + +target_link_libraries(nimble_common_tests nimble_common gtest gtest_main + glog::glog Folly::folly) diff --git a/dwio/nimble/common/tests/ExceptionTests.cpp b/dwio/nimble/common/tests/ExceptionTests.cpp new file mode 100644 index 0000000..4eed18e --- /dev/null +++ b/dwio/nimble/common/tests/ExceptionTests.cpp @@ -0,0 +1,242 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include "dwio/nimble/common/Exceptions.h" + +using namespace ::facebook; + +template +void verifyException( + const T& e, + const std::string& exceptionName, + const std::string& fileName, + const std::string& fileLine, + const std::string& functionName, + const std::string& failingExpression, + const std::string& errorMessage, + const std::string& errorSource, + const std::string& errorCode, + const std::string& retryable, + const std::string& additionalMessage = "") { + EXPECT_EQ(fileName, e.fileName()); + if (!fileLine.empty()) { + EXPECT_EQ(fileLine, folly::to(e.fileLine())); + } + EXPECT_EQ(functionName, e.functionName()); + EXPECT_EQ(failingExpression, e.failingExpression()); + EXPECT_EQ(errorMessage, e.errorMessage()); + EXPECT_EQ(errorSource, e.errorSource()); + EXPECT_EQ(errorCode, e.errorCode()); + EXPECT_EQ(retryable, e.retryable() ? "True" : "False"); + + EXPECT_NE( + std::string(e.what()).find(exceptionName + "\n"), std::string::npos); + EXPECT_NE( + std::string(e.what()).find("Error Source: " + errorSource + "\n"), + std::string::npos); + EXPECT_NE( + std::string(e.what()).find("Error Code: " + errorCode + "\n"), + std::string::npos); + if (!errorMessage.empty()) { + EXPECT_NE( + std::string(e.what()).find("Error Message: " + errorMessage + "\n"), + std::string::npos); + } + EXPECT_NE( + std::string(e.what()).find("Retryable: " + retryable + "\n"), + std::string::npos); + EXPECT_NE( + std::string(e.what()).find( + "Location: " + + folly::to(functionName, '@', fileName, ':', fileLine)), + std::string::npos); + if (!failingExpression.empty()) { + EXPECT_NE( + std::string(e.what()).find("Expression: " + failingExpression + "\n"), + std::string::npos); + } + EXPECT_NE(std::string(e.what()).find("Stack Trace:\n"), std::string::npos); + + if (!additionalMessage.empty()) { + EXPECT_NE(std::string(e.what()).find(additionalMessage), std::string::npos); + } +} + +TEST(ExceptionTests, Format) { + verifyException( + nimble::NimbleUserError( + "file1", 23, "func1", "expr1", "err1", "code1", true), + "NimbleUserError", + "file1", + "23", + "func1", + "expr1", + "err1", + "USER", + "code1", + "True"); + + verifyException( + nimble::NimbleInternalError( + "file2", 24, "func2", "expr2", "err2", "code2", false), + "NimbleInternalError", + "file2", + "24", + "func2", + "expr2", + "err2", + "INTERNAL", + "code2", + "False"); + + verifyException( + nimble::NimbleExternalError( + "file3", 25, "func3", "expr3", "err3", "code3", true, "source"), + "NimbleExternalError", + "file3", + "25", + "func3", + "expr3", + "err3", + "EXTERNAL", + "code3", + "True"); +} + +TEST(ExceptionTests, Check) { + int a = 5; + try { + NIMBLE_CHECK(a < 3, "error message1"); + } catch (const nimble::NimbleUserError& e) { + verifyException( + e, + "NimbleUserError", + __FILE__, + "", + "TestBody", + "a < 3", + "error message1", + "USER", + "INVALID_ARGUMENT", + "False"); + } +} + +TEST(ExceptionTests, Assert) { + int a = 5; + try { + NIMBLE_ASSERT(a > 8, "error message2"); + } catch (const nimble::NimbleInternalError& e) { + verifyException( + e, + "NimbleInternalError", + __FILE__, + "", + "TestBody", + "a > 8", + "error message2", + "INTERNAL", + "INVALID_STATE", + "False"); + } +} + +TEST(ExceptionTests, Verify) { + try { + NIMBLE_VERIFY_EXTERNAL( + 1 == 2, + LocalFileSystem, + nimble::error_code::NotSupported, + true, + "error message3"); + } catch (const nimble::NimbleExternalError& e) { + verifyException( + e, + "NimbleExternalError", + __FILE__, + "", + "TestBody", + "1 == 2", + "error message3", + "EXTERNAL", + "NOT_SUPPORTED", + "True", + "External Source: FILE_SYSTEM"); + } +} + +TEST(ExceptionTests, Unreachable) { + try { + NIMBLE_UNREACHABLE("error message"); + } catch (const nimble::NimbleInternalError& e) { + verifyException( + e, + "NimbleInternalError", + __FILE__, + "", + "TestBody", + "", + "error message", + "INTERNAL", + "UNREACHABLE_CODE", + "False"); + } +} + +TEST(ExceptionTests, NotImplemented) { + try { + NIMBLE_NOT_IMPLEMENTED("error message7"); + } catch (const nimble::NimbleInternalError& e) { + verifyException( + e, + "NimbleInternalError", + __FILE__, + "", + "TestBody", + "", + "error message7", + "INTERNAL", + "NOT_IMPLEMENTED", + "False"); + } +} + +TEST(ExceptionTests, NotSupported) { + try { + NIMBLE_NOT_SUPPORTED("error message6"); + } catch (const nimble::NimbleUserError& e) { + verifyException( + e, + "NimbleUserError", + __FILE__, + "", + "TestBody", + "", + "error message6", + "USER", + "NOT_SUPPORTED", + "False"); + } +} + +TEST(ExceptionTests, StackTraceThreads) { + // Make sure captured stack trace doesn't need anything from thread local + // storage + std::exception_ptr e; + auto throwFunc = []() { NIMBLE_CHECK(false, "Test."); }; + std::thread t([&]() { + try { + throwFunc(); + } catch (...) { + e = std::current_exception(); + } + }); + + t.join(); + + ASSERT_NE(nullptr, e); + EXPECT_NE( + std::string::npos, + folly::exceptionStr(e).find( + "facebook::nimble::NimbleException::NimbleException")); +} diff --git a/dwio/nimble/common/tests/FixedBitArrayTests.cpp b/dwio/nimble/common/tests/FixedBitArrayTests.cpp new file mode 100644 index 0000000..fa6bfe1 --- /dev/null +++ b/dwio/nimble/common/tests/FixedBitArrayTests.cpp @@ -0,0 +1,340 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include + +#include "dwio/nimble/common/FixedBitArray.h" +#include "folly/Benchmark.h" +#include "folly/Random.h" + +using namespace ::facebook; + +TEST(FixedBitArrayTests, SetThenGet) { + constexpr int bitWidth = 10; + constexpr int elementCount = 10000; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.set(i, (i + i * i) % (1 << bitWidth)); + } + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(fixedBitArray.get(i), (i + i * i) % (1 << bitWidth)); + } +} + +TEST(FixedBitArrayTests, zeroAndSet) { + constexpr int bitWidth = 4; + constexpr int elementCount = 1234; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.set(i, (i + i * i) % (1 << bitWidth)); + } + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(fixedBitArray.get(i), (i + i * i) % (1 << bitWidth)); + } + for (int i = 0; i < elementCount; ++i) { + if (i % 3 == 0) { + fixedBitArray.zeroAndSet(i, i % (1 << bitWidth)); + } + } + for (int i = 0; i < elementCount; ++i) { + if (i % 3 == 0) { + ASSERT_EQ(fixedBitArray.get(i), i % (1 << bitWidth)); + } else { + ASSERT_EQ(fixedBitArray.get(i), (i + i * i) % (1 << bitWidth)); + } + } +} + +constexpr int kNumTestsPerBitWidth = 10; +constexpr int kMaxElements = 1000; + +TEST(FixedBitArrayTests, SetThenGetRandom) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int bitWidth = 1; bitWidth <= 64; ++bitWidth) { + const uint64_t elementMask = + bitWidth == 64 ? (~0ULL) : ((1ULL << bitWidth) - 1); + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = folly::Random::rand32(rng) % kMaxElements; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = folly::Random::rand64(rng) & elementMask; + } + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.set(i, randomValues[i]); + } + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(fixedBitArray.get(i), randomValues[i]); + } + } + } +} + +TEST(FixedBitArrayTests, zeroAndSetThenGetRandom) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int bitWidth = 1; bitWidth <= 64; ++bitWidth) { + const uint64_t elementMask = + bitWidth == 64 ? (~0ULL) : ((1ULL << bitWidth) - 1); + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = folly::Random::rand32(rng) % kMaxElements; + std::vector buffer( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth), 0xFF); + nimble::FixedBitArray fixedBitArray(buffer.data(), bitWidth); + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = folly::Random::rand64(rng) & elementMask; + } + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.zeroAndSet(i, randomValues[i]); + } + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(fixedBitArray.get(i), randomValues[i]) << bitWidth; + } + } + } +} + +TEST(FixedBitArrayTests, Set32ThenGet32Random) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int bitWidth = 1; bitWidth <= 32; ++bitWidth) { + const uint64_t elementMask = (1ULL << bitWidth) - 1; + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = folly::Random::rand32(rng) % kMaxElements; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = folly::Random::rand32(rng) & elementMask; + } + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.set32(i, randomValues[i]); + } + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(fixedBitArray.get32(i), randomValues[i]) + << i << " " << bitWidth; + } + } + } +} + +TEST(FixedBitArrayTests, bulkGet32) { + constexpr int kBitWidth = 7; + constexpr int kNumElements = 27; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(kNumElements, kBitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), kBitWidth); + for (int i = 0; i < kNumElements; ++i) { + fixedBitArray.set32(i, 100 - i); + } + std::vector values(kNumElements); + fixedBitArray.bulkGet32(0, kNumElements, values.data()); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(values[i], 100 - i); + } + std::vector values64(kNumElements); + fixedBitArray.bulkGet32Into64(0, kNumElements, values64.data()); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(values64[i], 100 - i); + } +} + +TEST(FixedBitArrayTests, BulkGet32Random) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int bitWidth = 1; bitWidth <= 32; ++bitWidth) { + const uint64_t elementMask = (1ULL << bitWidth) - 1; + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = folly::Random::rand32(rng) % kMaxElements; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = folly::Random::rand32(rng) & elementMask; + } + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.set32(i, randomValues[i]); + } + std::vector values(elementCount); + fixedBitArray.bulkGet32(0, elementCount, values.data()); + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(values[i], randomValues[i]); + } + for (int i = 0; i < elementCount; ++i) { + uint32_t element; + fixedBitArray.bulkGet32(i, 1, &element); + ASSERT_EQ(element, values[i]); + } + std::vector values64(elementCount); + fixedBitArray.bulkGet32Into64(0, elementCount, values64.data()); + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(values64[i], randomValues[i]); + } + for (int i = 0; i < elementCount; ++i) { + uint64_t element; + fixedBitArray.bulkGet32Into64(i, 1, &element); + ASSERT_EQ(element, values64[i]); + } + } + } +} + +TEST(FixedBitArrayTests, BulkGetWithBaseline32Random) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int bitWidth = 4; bitWidth <= 32; ++bitWidth) { + const uint64_t maxElement = (1ULL << bitWidth); + const uint64_t baseline = folly::Random::rand32(rng) % maxElement; + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = folly::Random::rand32(rng) % kMaxElements; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = folly::Random::rand32(rng) % (maxElement - baseline); + } + for (int i = 0; i < elementCount; ++i) { + fixedBitArray.set32(i, randomValues[i]); + } + std::vector values(elementCount); + fixedBitArray.bulkGetWithBaseline32( + 0, elementCount, values.data(), baseline); + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(values[i], randomValues[i] + baseline) + << "Bit Width: " << bitWidth << ", i: " << i; + } + for (int i = 0; i < elementCount; ++i) { + uint32_t element; + fixedBitArray.bulkGetWithBaseline32(i, 1, &element, baseline); + ASSERT_EQ(element, randomValues[i] + baseline); + } + std::vector values64(elementCount); + fixedBitArray.bulkGetWithBaseline32Into64( + 0, elementCount, values64.data(), baseline); + for (int i = 0; i < elementCount; ++i) { + ASSERT_EQ(values64[i], randomValues[i] + baseline); + } + for (int i = 0; i < elementCount; ++i) { + uint64_t element; + fixedBitArray.bulkGetWithBaseline32Into64(i, 1, &element, baseline); + ASSERT_EQ(element, randomValues[i] + baseline); + } + } + } +} + +TEST(FixedBitArrayTests, BulkSet32Random) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + for (int bitWidth = 1; bitWidth <= 32; ++bitWidth) { + const uint64_t maxElement = (1ULL << bitWidth); + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = 1 + folly::Random::rand32(rng) % kMaxElements; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = folly::Random::rand32(rng) % maxElement; + } + const int offset = folly::Random::rand32(rng) % elementCount; + const int size = elementCount - offset; + fixedBitArray.bulkSet32(offset, size, randomValues.data() + offset); + for (int i = 0; i < size; ++i) { + ASSERT_EQ(fixedBitArray.get32(offset + i), randomValues[offset + i]); + } + } + } +} + +TEST(FixedBitArrayTests, BulkSet32WithBaselineRandom) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + for (int bitWidth = 4; bitWidth <= 32; ++bitWidth) { + const uint64_t maxElement = (1ULL << bitWidth); + const uint64_t baseline = folly::Random::rand32(rng) % maxElement; + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = 1 + folly::Random::rand32(rng) % kMaxElements; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + std::vector randomValues(elementCount); + std::vector randomValuesWithBaseline(elementCount); + for (int i = 0; i < elementCount; ++i) { + randomValues[i] = folly::Random::rand32(rng) % (maxElement - baseline); + randomValuesWithBaseline[i] = randomValues[i] + baseline; + } + const int offset = folly::Random::rand32(rng) % elementCount; + const int size = elementCount - offset; + fixedBitArray.bulkSet32WithBaseline( + offset, size, randomValuesWithBaseline.data() + offset, baseline); + for (int i = 0; i < size; ++i) { + ASSERT_EQ(fixedBitArray.get32(offset + i), randomValues[offset + i]) + << "Bit Width: " << bitWidth << ", i: " << i; + } + } + } +} + +TEST(FixedBitArrayTests, Equals32Random) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int bitWidth = 1; bitWidth <= 32; ++bitWidth) { + const uint64_t elementMask = (1ULL << bitWidth) - 1; + for (int test = 0; test < kNumTestsPerBitWidth; ++test) { + const int elementCount = 1 + folly::Random::rand32(rng) % kMaxElements; + auto buffer = std::make_unique( + nimble::FixedBitArray::bufferSize(elementCount, bitWidth)); + nimble::FixedBitArray fixedBitArray(buffer.get(), bitWidth); + std::vector randomValues(elementCount); + for (int i = 0; i < elementCount; ++i) { + // Test both the full range and also a case where equals are actually + // likely. + randomValues[i] = folly::Random::rand32(rng) & elementMask; + if (test % 2 == 1) { + randomValues[i] = randomValues[i] % 10; + } + fixedBitArray.set32(i, randomValues[i]); + } + // Remember start needs to be a multiple of 64. + const int start = + (folly::Random::rand32(rng) % elementCount) & (0xFFFFFF00); + const int length = elementCount - start; + auto equalsBuffer = std::make_unique( + nimble::FixedBitArray::bufferSize(length, 1)); + const uint32_t equalsValue = + randomValues[folly::Random::rand32(rng) % elementCount]; + fixedBitArray.equals32(start, length, equalsValue, equalsBuffer.get()); + for (int i = 0; i < length; ++i) { + ASSERT_EQ( + nimble::bits::getBit(i, equalsBuffer.get()), + randomValues[start + i] == equalsValue); + } + } + } +} diff --git a/dwio/nimble/common/tests/HuffmanTests.cpp b/dwio/nimble/common/tests/HuffmanTests.cpp new file mode 100644 index 0000000..1fb06bc --- /dev/null +++ b/dwio/nimble/common/tests/HuffmanTests.cpp @@ -0,0 +1,85 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include "dwio/nimble/common/Huffman.h" + +using namespace ::facebook; + +class HuffmanTests : public ::testing::Test { + protected: + void SetUp() override { + pool_ = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + } + + std::shared_ptr pool_; +}; + +TEST_F(HuffmanTests, EndToEnd) { + const std::vector counts = {64, 32, 16, 8, 4, 1, 1, 1, 1}; + // Note that the counts aren't reflective of the data. That's okay. It will + // still work just fine as long as the data lies in the range + // [0, counts.size()). + const std::vector data = {0, 0, 1, 3, 0, 6, 7, 8, 0, 0, 5}; + int treeDepth = 0; + auto encodingTable = + nimble::huffman::generateHuffmanEncodingTable(counts, &treeDepth); + ASSERT_EQ(treeDepth, 7); + std::string huffmanEncoded = + nimble::huffman::huffmanEncode(encodingTable, data); + std::vector recovered(data.size()); + nimble::huffman::DecodingTable decodingTable( + *pool_, encodingTable, treeDepth); + decodingTable.decode(data.size(), huffmanEncoded.data(), recovered.data()); + for (int i = 0; i < data.size(); ++i) { + ASSERT_EQ(recovered[i], data[i]); + } +} + +TEST_F(HuffmanTests, MaxDepthEndToEnd) { + std::vector counts = {64, 32, 16, 8, 4, 1, 1, 1, 1}; + const std::vector data = {4, 2, 1, 0, 0, 6, 7, 8, 8}; + int treeDepth; + auto encodingTable = + nimble::huffman::generateHuffmanEncodingTableWithMaxDepth( + counts, 5, &treeDepth); + ASSERT_LE(treeDepth, 5); + std::string huffmanEncoded = + nimble::huffman::huffmanEncode(encodingTable, data); + std::vector recovered(data.size()); + nimble::huffman::DecodingTable decodingTable( + *pool_, encodingTable, treeDepth); + decodingTable.decode(data.size(), huffmanEncoded.data(), recovered.data()); + for (int i = 0; i < data.size(); ++i) { + ASSERT_EQ(recovered[i], data[i]); + } +} + +TEST_F(HuffmanTests, StreamedEndToEnd) { + const std::vector counts = {7, 8, 9, 10, 11, 12, 13}; + const std::vector data1 = {0, 1, 2, 3, 4, 5, 6}; + const std::vector data2 = {6, 6, 5, 5, 4, 0, 1}; + int treeDepth; + auto encodingTable = + nimble::huffman::generateHuffmanEncodingTable(counts, &treeDepth); + ASSERT_LE(treeDepth, counts.size() - 1); + std::string huffmanEncoded1 = + nimble::huffman::huffmanEncode(encodingTable, data1); + std::string huffmanEncoded2 = + nimble::huffman::huffmanEncode(encodingTable, data2); + std::vector recovered1(data1.size()); + std::vector recovered2(data2.size()); + nimble::huffman::DecodingTable decodingTable( + *pool_, encodingTable, treeDepth); + decodingTable.decodeStreamed( + data1.size(), + huffmanEncoded1.data(), + huffmanEncoded2.data(), + recovered1.data(), + recovered2.data()); + for (int i = 0; i < data1.size(); ++i) { + ASSERT_EQ(recovered1[i], data1[i]); + } + for (int i = 0; i < data2.size(); ++i) { + ASSERT_EQ(recovered2[i], data2[i]); + } +} diff --git a/dwio/nimble/common/tests/IndexMapTests.cpp b/dwio/nimble/common/tests/IndexMapTests.cpp new file mode 100644 index 0000000..ed05a43 --- /dev/null +++ b/dwio/nimble/common/tests/IndexMapTests.cpp @@ -0,0 +1,52 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include "dwio/nimble/common/IndexMap.h" + +using namespace ::facebook; + +TEST(IndexMapTests, DoubleIndexMap) { + { + // First, verify that flat_hash_map doesn't work with double/float + folly::F14FastMap indices; + indices.emplace(+0.0, 1); + int count = indices.size(); + EXPECT_EQ(1, count); + indices.emplace(-0.0, 2); + count = indices.size(); + EXPECT_EQ(1, count); + } + + { + nimble::IndexMap dIndices; + dIndices.index(+0.0); + int count = dIndices.size(); + EXPECT_EQ(1, count); + dIndices.index(-0.0); + count = dIndices.size(); + EXPECT_EQ(2, count); + } +} + +TEST(IndexMapTests, FloatIndexMap) { + { + // First, verify that flat_hash_map doesn't work with double/float + folly::F14FastMap indices; + indices.emplace(+0.0f, 1); + int count = indices.size(); + EXPECT_EQ(1, count); + indices.emplace(-0.0f, 2); + count = indices.size(); + EXPECT_EQ(1, count); + } + + { + nimble::IndexMap dIndices; + dIndices.index(+0.0f); + int count = dIndices.size(); + EXPECT_EQ(1, count); + dIndices.index(-0.0f); + count = dIndices.size(); + EXPECT_EQ(2, count); + } +} diff --git a/dwio/nimble/common/tests/NimbleFileWriter.cpp b/dwio/nimble/common/tests/NimbleFileWriter.cpp new file mode 100644 index 0000000..de8f9e1 --- /dev/null +++ b/dwio/nimble/common/tests/NimbleFileWriter.cpp @@ -0,0 +1,52 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/common/tests/NimbleFileWriter.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/velox/VeloxWriter.h" +#include "dwio/nimble/velox/VeloxWriterOptions.h" + +namespace facebook::nimble::test { + +std::string createNimbleFile( + velox::memory::MemoryPool& memoryPool, + const velox::VectorPtr& vector, + nimble::VeloxWriterOptions writerOptions, + bool flushAfterWrite) { + return createNimbleFile( + memoryPool, + std::vector{vector}, + std::move(writerOptions), + flushAfterWrite); +} + +std::string createNimbleFile( + velox::memory::MemoryPool& memoryPool, + const std::vector& vectors, + nimble::VeloxWriterOptions writerOptions, + bool flushAfterWrite) { + std::string file; + auto writeFile = std::make_unique(&file); + + NIMBLE_ASSERT(vectors.size() > 0, "Expecting at least one input vector."); + auto& type = vectors[0]->type(); + + for (int i = 1; i < vectors.size(); ++i) { + NIMBLE_ASSERT( + vectors[i]->type()->kindEquals(type), + "All vectors should have the same schema."); + } + + nimble::VeloxWriter writer( + memoryPool, type, std::move(writeFile), std::move(writerOptions)); + for (const auto& vector : vectors) { + writer.write(vector); + if (flushAfterWrite) { + writer.flush(); + } + } + writer.close(); + + return file; +} + +} // namespace facebook::nimble::test diff --git a/dwio/nimble/common/tests/NimbleFileWriter.h b/dwio/nimble/common/tests/NimbleFileWriter.h new file mode 100644 index 0000000..d211ba8 --- /dev/null +++ b/dwio/nimble/common/tests/NimbleFileWriter.h @@ -0,0 +1,22 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/velox/VeloxWriterOptions.h" +#include "velox/vector/BaseVector.h" + +namespace facebook::nimble::test { + +std::string createNimbleFile( + velox::memory::MemoryPool& memoryPool, + const std::vector& vectors, + nimble::VeloxWriterOptions writerOptions = {}, + bool flushAfterWrite = true); + +std::string createNimbleFile( + velox::memory::MemoryPool& memoryPool, + const velox::VectorPtr& vector, + nimble::VeloxWriterOptions writerOptions = {}, + bool flushAfterWrite = true); + +} // namespace facebook::nimble::test diff --git a/dwio/nimble/common/tests/StopWatchTests.cpp b/dwio/nimble/common/tests/StopWatchTests.cpp new file mode 100644 index 0000000..0bbfc75 --- /dev/null +++ b/dwio/nimble/common/tests/StopWatchTests.cpp @@ -0,0 +1,57 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include "dwio/nimble/common/StopWatch.h" +#include "folly/Benchmark.h" + +using namespace ::facebook; + +namespace { +void WasteSomeTime() { + int tot = 0; + for (int i = 0; i < 100; ++i) { + tot += i; + folly::doNotOptimizeAway(tot); + } +} +} // namespace + +TEST(StopWatchTests, BasicCorrectness) { + nimble::StopWatch watch; + watch.start(); + WasteSomeTime(); + watch.stop(); + + ASSERT_GE(watch.elapsedNsec(), 0LL); + ASSERT_GE(watch.elapsed(), 0.0); + + watch.reset(); + ASSERT_EQ(watch.elapsedNsec(), 0); + ASSERT_EQ(watch.elapsed(), 0.0); + + watch.start(); + WasteSomeTime(); + watch.stop(); + int64_t elapsed1 = watch.elapsedNsec(); + ASSERT_GT(elapsed1, 0LL); + watch.start(); + WasteSomeTime(); + int64_t elapsed2 = watch.elapsedNsec(); + ASSERT_GT(elapsed2, elapsed1); + watch.stop(); + int64_t elapsed3 = watch.elapsedNsec(); + WasteSomeTime(); + int64_t elapsed4 = watch.elapsedNsec(); + ASSERT_EQ(elapsed3, elapsed4); + + ASSERT_EQ(elapsed4 / 1000, watch.elapsedUsec()); + ASSERT_EQ(elapsed4 / (1000 * 1000), watch.elapsedMsec()); + + watch.reset(); + watch.start(); + watch.start(); + WasteSomeTime(); + watch.stop(); + watch.stop(); + ASSERT_GT(watch.elapsedNsec(), 0LL); +} diff --git a/dwio/nimble/common/tests/TestUtils.h b/dwio/nimble/common/tests/TestUtils.h new file mode 100644 index 0000000..4e9fe79 --- /dev/null +++ b/dwio/nimble/common/tests/TestUtils.h @@ -0,0 +1,362 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include + +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" +#include "folly/Random.h" +#include "velox/common/file/File.h" +#include "velox/common/memory/Memory.h" + +// Utilities to support testing in nimble. + +namespace facebook::nimble::testing { + +// Adds random data covering the whole range of the data type. +template +void addRandomData(RNG&& rng, int rowCount, Vector* data, Buffer* buffer); + +// Draw an int uniformly from [min, max). +class Util { + public: + Util(velox::memory::MemoryPool& memoryPool) : memoryPool_(memoryPool) {} + + // Makes between 1 and maxRows random values over T's data range. + // For strings we improvise a bit. + template + Vector makeRandomData(RNG&& rng, uint32_t maxRows, Buffer* buffer); + + // A random length vector of all the same data. + template + Vector makeConstantData(RNG&& rng, uint32_t maxRows, Buffer* buffer); + + // A random length vector of runs of various lengths. + template + Vector makeRLEData(RNG&& rng, uint32_t maxRows, Buffer* buffer); + + // Makes data compatible with our HuffmanColumn (namely in the + // range [0, 4096)) for integer data only. + template + Vector makeHuffmanData(RNG&& rng, uint32_t maxRows, Buffer* buffer); + + // A Vector of all the Make* Vectors. + template + std::vector> + makeDataPatterns(RNG&& rng, uint32_t maxRows, Buffer* buffer); + + // Trims data down so that it will be friendly to sums, i.e. so it won't + // overflow the sum type. This mainly applies to 64-bit integers. + template + Vector sumFriendlyData(const Vector& data); + + inline uint64_t uniformRandom(uint64_t min, uint64_t max) { + CHECK_LT(min, max); + return min + (folly::Random::rand32() % (max - min)); + } + + // Draw an int uniformly from [0, max); + inline uint64_t uniformRandom(uint64_t max) { + return uniformRandom(0, max); + } + + private: + velox::memory::MemoryPool& memoryPool_; +}; + +// +// End of public API. Implementation follows. +// + +template +inline void +addRandomData(RNG&& rng, int rowCount, Vector* data, Buffer* /* buffer */) { + // Half the time only add positive data. + if (!std::is_signed_v || folly::Random::rand32() % 2) { + for (int i = 0; i < rowCount; ++i) { + if constexpr (sizeof(T) > 4) { + const uint64_t rand = folly::Random::rand64(std::forward(rng)); + data->push_back(*reinterpret_cast(&rand)); + + } else { + const uint32_t rand = folly::Random::rand32(std::forward(rng)); + data->push_back(*reinterpret_cast(&rand)); + } + } + } else { + for (int i = 0; i < rowCount; ++i) { + if constexpr (sizeof(T) > 4) { + const uint64_t rand = + folly::Random::rand64(std::forward(rng)) & ((1ULL << 63) - 1); + data->push_back(*reinterpret_cast(&rand)); + } else { + const uint32_t rand = + folly::Random::rand32(std::forward(rng)) & ((1U << 31) - 1); + data->push_back(*reinterpret_cast(&rand)); + } + } + } +} + +template +inline void addRandomData( + RNG&& rng, + int rowCount, + Vector* data, + Buffer* /* buffer */) { + for (int i = 0; i < rowCount; ++i) { + const uint32_t rand = folly::Random::rand32(std::forward(rng)); + data->push_back(static_cast(rand)); + } +} + +template +inline void addRandomData( + RNG&& rng, + int rowCount, + Vector* data, + Buffer* /* buffer */) { + for (int i = 0; i < rowCount; ++i) { + const uint64_t rand = folly::Random::rand64(std::forward(rng)); + data->push_back(static_cast(rand)); + } +} + +template +inline void addRandomData( + RNG&& rng, + int rowCount, + Vector* data, + Buffer* buffer) { + for (int i = 0; i < rowCount; ++i) { + // This is a bit arbitrary, but lets stick to 50 char max string length. + // We can have some separate tests for large strings if we want. + const int len = folly::Random::rand32(std::forward(rng)) % 50; + char* pos = buffer->reserve(len); + for (int j = 0; j < len; ++j) { + pos[j] = folly::Random::rand32(std::forward(rng)) % 256; + } + data->emplace_back(pos, len); + } +} + +template +inline void addRandomData( + RNG&& rng, + int rowCount, + Vector* data, + Buffer* /* buffer */) { + for (int i = 0; i < rowCount; ++i) { + data->push_back(folly::Random::rand32(std::forward(rng)) & 1); + } +} + +template +Vector Util::makeRandomData(RNG&& rng, uint32_t maxRows, Buffer* buffer) { + Vector randomData(&memoryPool_); + const int rowCount = + 1 + folly::Random::rand32(std::forward(rng)) % maxRows; + randomData.reserve(rowCount); + addRandomData(std::forward(rng), rowCount, &randomData, buffer); + return randomData; +} + +template +Vector Util::makeConstantData(RNG&& rng, uint32_t maxRows, Buffer* buffer) { + Vector data = makeRandomData(std::forward(rng), maxRows, buffer); + for (int i = 1; i < data.size(); ++i) { + data[i] = data[0]; + } + return data; +} + +template +Vector Util::makeRLEData(RNG&& rng, uint32_t maxRows, Buffer* buffer) { + Vector data = makeRandomData(std::forward(rng), maxRows, buffer); + int index = 0; + while (index < data.size()) { + const uint32_t runLength = data.size() - index == 1 ? 1 + : 1 + + folly::Random::rand32(std::forward(rng)) % + (data.size() - index - 1); + for (int i = 0; i < runLength; ++i) { + data[index + i] = data[index]; + } + index += runLength; + } + return data; +} + +template +Vector +Util::makeHuffmanData(RNG&& rng, uint32_t maxRows, Buffer* /* buffer */) { + const int rowCount = + 1 + folly::Random::rand32(std::forward(rng)) % maxRows; + Vector huffmanData(&memoryPool_); + huffmanData.reserve(rowCount); + // Half the time draw all the symbols from within [0, symbolCount), + // emulating encoding a dictionary index, and the other half of the time + // emulate encoding a normal stream by drawing from the full [0, 4096) range. + const int symbolCount = + 1 + folly::Random::rand32(std::forward(rng)) % rowCount; + const int maxValue = + (folly::Random::rand32(std::forward(rng)) & 1) ? symbolCount : 4096; + for (uint32_t i = 0; i < rowCount; ++i) { + huffmanData.push_back( + folly::Random::rand32(std::forward(rng)) % maxValue); + } + return huffmanData; +} + +template +std::vector> +Util::makeDataPatterns(RNG&& rng, uint32_t maxRows, Buffer* buffer) { + std::vector> patterns; + patterns.push_back( + makeRandomData(std::forward(rng), maxRows, buffer)); + patterns.push_back( + makeConstantData(std::forward(rng), maxRows, buffer)); + patterns.push_back(makeRLEData(std::forward(rng), maxRows, buffer)); + if constexpr (isIntegralType()) { + patterns.push_back( + makeHuffmanData(std::forward(rng), maxRows, buffer)); + } + return patterns; +} + +template +Vector sumFriendlyDataHelper( + velox::memory::MemoryPool&, + const Vector& data) { + return data; +} + +template +Vector Util::sumFriendlyData(const Vector& data) { + return data; +} + +template <> +inline Vector Util::sumFriendlyData(const Vector& data) { + Vector sumData(&memoryPool_); + // 15 is somewhat arbitrary shift, but this is testing code, so meh. + for (int64_t datum : data) { + sumData.push_back(datum >> 15); + } + return sumData; +} + +template <> +inline Vector Util::sumFriendlyData(const Vector& data) { + Vector sumData(&memoryPool_); + // 15 is somewhat arbitrary shift, but this is testing code, so meh. + for (uint64_t datum : data) { + sumData.push_back(datum >> 15); + } + return sumData; +} + +struct Chunk { + uint64_t offset; + uint64_t size; +}; + +// Wrapper around InMemoryReadFile (can't inherit, as InMemoryReadFile is final) +// which tracks all offsets and sizes being read. This is used to verify Nimble +// reader coalese behavior. +class InMemoryTrackableReadFile final : public velox::ReadFile { + public: + explicit InMemoryTrackableReadFile( + std::string_view file, + bool shouldProduceChainedBuffers) + : file_{file}, + shouldProduceChainedBuffers_{shouldProduceChainedBuffers} {} + + std::string_view pread(uint64_t offset, uint64_t length, void* buf) + const final { + chunks_.push_back({offset, length}); + return file_.pread(offset, length, buf); + } + + std::string pread(uint64_t offset, uint64_t length) const final { + chunks_.push_back({offset, length}); + return file_.pread(offset, length); + } + + uint64_t preadv( + uint64_t /* offset */, + const std::vector>& /* buffers */) const final { + NIMBLE_NOT_SUPPORTED("Not used by Nimble"); + } + + void preadv( + folly::Range regions, + folly::Range iobufs) const override { + VELOX_CHECK_EQ(regions.size(), iobufs.size()); + for (size_t i = 0; i < regions.size(); ++i) { + const auto& region = regions[i]; + auto& output = iobufs[i]; + if (shouldProduceChainedBuffers_) { + chunks_.push_back({region.offset, region.length}); + uint64_t splitPoint = region.length / 2; + output = folly::IOBuf(folly::IOBuf::CREATE, splitPoint); + file_.pread(region.offset, splitPoint, output.writableData()); + output.append(splitPoint); + const uint64_t nextLength = region.length - splitPoint; + auto next = folly::IOBuf::create(nextLength); + file_.pread( + region.offset + splitPoint, nextLength, next->writableData()); + next->append(nextLength); + output.appendChain(std::move(next)); + } else { + output = folly::IOBuf(folly::IOBuf::CREATE, region.length); + pread(region.offset, region.length, output.writableData()); + output.append(region.length); + } + } + } + + uint64_t size() const final { + return file_.size(); + } + + uint64_t memoryUsage() const final { + return file_.memoryUsage(); + } + + // Mainly for testing. Coalescing isn't helpful for in memory data. + void setShouldCoalesce(bool shouldCoalesce) { + file_.setShouldCoalesce(shouldCoalesce); + } + + bool shouldCoalesce() const final { + return file_.shouldCoalesce(); + } + + const std::vector& chunks() { + return chunks_; + } + + void resetChunks() { + chunks_.clear(); + } + + std::string getName() const override { + return ""; + } + + uint64_t getNaturalReadSize() const override { + return 1024; + } + + private: + velox::InMemoryReadFile file_; + bool shouldProduceChainedBuffers_; + mutable std::vector chunks_; +}; + +} // namespace facebook::nimble::testing diff --git a/dwio/nimble/common/tests/VarintTests.cpp b/dwio/nimble/common/tests/VarintTests.cpp new file mode 100644 index 0000000..26612f6 --- /dev/null +++ b/dwio/nimble/common/tests/VarintTests.cpp @@ -0,0 +1,115 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include + +#include "dwio/nimble/common/Varint.h" +#include "folly/Random.h" +#include "folly/Range.h" +#include "folly/Varint.h" + +using namespace ::facebook; + +namespace { +const int kNumElements = 10000; +} + +TEST(VarintTests, WriteRead32) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + std::vector data; + // Generate data with uniform bit length. + for (int i = 0; i < kNumElements; ++i) { + const int bitShift = folly::Random::rand32(rng) % 32; + data.push_back(folly::Random::rand32(rng) >> bitShift); + } + + auto buffer = + std::make_unique(kNumElements * folly::kMaxVarintLength32); + char* pos = buffer.get(); + for (int i = 0; i < kNumElements; ++i) { + nimble::varint::writeVarint(data[i], &pos); + } + + auto follyBuffer = + std::make_unique(kNumElements * folly::kMaxVarintLength32); + uint8_t* fpos = follyBuffer.get(); + for (int i = 0; i < kNumElements; ++i) { + fpos += folly::encodeVarint(data[i], fpos); + } + + ASSERT_EQ(pos - buffer.get(), fpos - follyBuffer.get()); + + ASSERT_EQ(nimble::varint::bulkVarintSize32(data), pos - buffer.get()); + + const char* cpos = buffer.get(); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(data[i], nimble::varint::readVarint32(&cpos)); + } + + const uint8_t* fstart = follyBuffer.get(); + const uint8_t* fend = follyBuffer.get() + (fpos - follyBuffer.get()); + folly::Range frange(fstart, fend); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(data[i], folly::decodeVarint(frange)); + } + + std::vector bulk(kNumElements); + cpos = buffer.get(); + nimble::varint::bulkVarintDecode32(kNumElements, cpos, bulk.data()); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(data[i], bulk[i]); + } +} + +TEST(VarintTests, WriteRead64) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + std::vector data; + // Generate data with uniform bit length. + for (int i = 0; i < kNumElements; ++i) { + const int bitShift = folly::Random::rand64(rng) % 64; + data.push_back(folly::Random::rand64(rng) >> bitShift); + } + + auto buffer = + std::make_unique(kNumElements * folly::kMaxVarintLength64); + char* pos = buffer.get(); + for (int i = 0; i < kNumElements; ++i) { + nimble::varint::writeVarint(data[i], &pos); + } + + auto follyBuffer = + std::make_unique(kNumElements * folly::kMaxVarintLength64); + uint8_t* fpos = follyBuffer.get(); + for (int i = 0; i < kNumElements; ++i) { + fpos += folly::encodeVarint(data[i], fpos); + } + + ASSERT_EQ(pos - buffer.get(), fpos - follyBuffer.get()); + + ASSERT_EQ(nimble::varint::bulkVarintSize64(data), pos - buffer.get()); + + const char* cpos = buffer.get(); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(data[i], nimble::varint::readVarint64(&cpos)); + } + + const uint8_t* fstart = follyBuffer.get(); + const uint8_t* fend = follyBuffer.get() + (fpos - follyBuffer.get()); + folly::Range frange(fstart, fend); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(data[i], folly::decodeVarint(frange)); + } + + std::vector bulk(kNumElements); + cpos = buffer.get(); + nimble::varint::bulkVarintDecode64(kNumElements, cpos, bulk.data()); + for (int i = 0; i < kNumElements; ++i) { + ASSERT_EQ(data[i], bulk[i]); + } +} diff --git a/dwio/nimble/common/tests/VectorTests.cpp b/dwio/nimble/common/tests/VectorTests.cpp new file mode 100644 index 0000000..286321b --- /dev/null +++ b/dwio/nimble/common/tests/VectorTests.cpp @@ -0,0 +1,230 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include "dwio/nimble/common/Vector.h" +#include "velox/common/memory/Memory.h" + +DECLARE_bool(velox_enable_memory_usage_track_in_default_memory_pool); + +using namespace ::facebook; + +class VectorTests : public ::testing::Test { + protected: + static void SetUpTestCase() { + FLAGS_velox_enable_memory_usage_track_in_default_memory_pool = true; + } + + void SetUp() override { + pool_ = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + } + + std::shared_ptr pool_; +}; + +TEST_F(VectorTests, InitializerList) { + nimble::Vector v1(pool_.get(), {3, 4}); + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(3, v1[0]); + EXPECT_EQ(4, v1[1]); +} + +TEST_F(VectorTests, FromRange) { + std::vector source{4, 5, 6}; + nimble::Vector v1(pool_.get(), source.begin(), source.end()); + EXPECT_EQ(3, v1.size()); + EXPECT_EQ(4, v1[0]); + EXPECT_EQ(5, v1[1]); + EXPECT_EQ(6, v1[2]); +} + +TEST_F(VectorTests, EqualOp1) { + nimble::Vector v1(pool_.get()); + v1.push_back(1); + v1.emplace_back(2); + v1.push_back(3); + + nimble::Vector v2(pool_.get()); + v2.push_back(4); + v2.emplace_back(5); + + EXPECT_EQ(3, v1.size()); + EXPECT_EQ(1, v1[0]); + EXPECT_EQ(2, v1[1]); + EXPECT_EQ(3, v1[2]); + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(4, v2[0]); + EXPECT_EQ(5, v2[1]); + + v1 = v2; + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(4, v1[0]); + EXPECT_EQ(5, v1[1]); + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(4, v2[0]); + EXPECT_EQ(5, v2[1]); +} + +TEST_F(VectorTests, ExplicitMoveEqualOp) { + nimble::Vector v1(pool_.get()); + v1.push_back(1); + v1.emplace_back(2); + v1.push_back(3); + + nimble::Vector v2(pool_.get()); + v2.push_back(4); + v2.emplace_back(5); + + EXPECT_EQ(3, v1.size()); + ASSERT_FALSE(v1.empty()); + EXPECT_EQ(1, v1[0]); + EXPECT_EQ(2, v1[1]); + EXPECT_EQ(3, v1[2]); + EXPECT_EQ(2, v2.size()); + ASSERT_FALSE(v2.empty()); + EXPECT_EQ(4, v2[0]); + EXPECT_EQ(5, v2[1]); + + v1 = std::move(v2); + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(4, v1[0]); + EXPECT_EQ(5, v1[1]); + // @lint-ignore CLANGTIDY bugprone-use-after-move + EXPECT_EQ(0, v2.size()); + ASSERT_TRUE(v2.empty()); +} + +TEST_F(VectorTests, MoveEqualOp1) { + nimble::Vector v1(pool_.get()); + v1.push_back(1); + v1.emplace_back(2); + v1.push_back(3); + EXPECT_EQ(3, v1.size()); + EXPECT_EQ(1, v1[0]); + EXPECT_EQ(2, v1[1]); + EXPECT_EQ(3, v1[2]); + v1 = nimble::Vector(pool_.get(), {4, 5}); + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(4, v1[0]); + EXPECT_EQ(5, v1[1]); +} + +TEST_F(VectorTests, CopyCtr) { + nimble::Vector v2(pool_.get()); + v2.push_back(3); + v2.emplace_back(4); + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(3, v2[0]); + EXPECT_EQ(4, v2[1]); + nimble::Vector v1(v2); + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(3, v1[0]); + EXPECT_EQ(4, v1[1]); + + // make sure they do not share buffer + v1[0] = 1; + v1[1] = 2; + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(1, v1[0]); + EXPECT_EQ(2, v1[1]); + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(3, v2[0]); + EXPECT_EQ(4, v2[1]); +} + +TEST_F(VectorTests, BoolInitializerList) { + nimble::Vector v1(pool_.get(), {true, false, true}); + EXPECT_EQ(3, v1.size()); + EXPECT_EQ(true, v1[0]); + EXPECT_EQ(false, v1[1]); + EXPECT_EQ(true, v1[2]); +} + +TEST_F(VectorTests, BoolEqualOp1) { + nimble::Vector v1(pool_.get()); + v1.push_back(false); + v1.emplace_back(true); + v1.push_back(true); + + EXPECT_EQ(3, v1.size()); + EXPECT_EQ(false, v1[0]); + EXPECT_EQ(true, v1[1]); + EXPECT_EQ(true, v1[2]); + + nimble::Vector v2(pool_.get()); + v2.push_back(true); + v2.emplace_back(false); + + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(true, v2[0]); + EXPECT_EQ(false, v2[1]); + + v1 = v2; + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(true, v1[0]); + EXPECT_EQ(false, v1[1]); + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(true, v2[0]); + EXPECT_EQ(false, v2[1]); +} + +TEST_F(VectorTests, BoolMoveEqualOp1) { + nimble::Vector v1(pool_.get()); + v1.push_back(true); + v1.emplace_back(false); + v1.push_back(false); + + EXPECT_EQ(3, v1.size()); + EXPECT_EQ(true, v1[0]); + EXPECT_EQ(false, v1[1]); + EXPECT_EQ(false, v1[2]); + + v1 = nimble::Vector(pool_.get(), {false, true}); + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(false, v1[0]); + EXPECT_EQ(true, v1[1]); +} + +TEST_F(VectorTests, BoolCopyCtr) { + nimble::Vector v2(pool_.get()); + v2.push_back(true); + v2.emplace_back(false); + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(true, v2[0]); + EXPECT_EQ(false, v2[1]); + nimble::Vector v1(v2); + EXPECT_EQ(2, v1.size()); + EXPECT_EQ(true, v1[0]); + EXPECT_EQ(false, v1[1]); + EXPECT_EQ(2, v2.size()); + EXPECT_EQ(true, v2[0]); + EXPECT_EQ(false, v2[1]); +} + +TEST_F(VectorTests, MemoryCleanup) { + EXPECT_EQ(0, pool_->currentBytes()); + { + nimble::Vector v(pool_.get()); + EXPECT_EQ(0, pool_->currentBytes()); + v.resize(1000, 10); + EXPECT_NE(0, pool_->currentBytes()); + } + EXPECT_EQ(0, pool_->currentBytes()); + { + nimble::Vector v(pool_.get()); + EXPECT_EQ(0, pool_->currentBytes()); + v.resize(1000, 10); + EXPECT_NE(0, pool_->currentBytes()); + + auto vCopy(v); + } + EXPECT_EQ(0, pool_->currentBytes()); + { + nimble::Vector v(pool_.get()); + EXPECT_EQ(0, pool_->currentBytes()); + v.resize(1000, 10); + EXPECT_NE(0, pool_->currentBytes()); + + auto vCopy(std::move(v)); + } + EXPECT_EQ(0, pool_->currentBytes()); +} diff --git a/dwio/nimble/encodings/CMakeLists.txt b/dwio/nimble/encodings/CMakeLists.txt new file mode 100644 index 0000000..15bbc92 --- /dev/null +++ b/dwio/nimble/encodings/CMakeLists.txt @@ -0,0 +1,15 @@ +add_subdirectory(tests) + +add_library( + nimble_encodings + Compression.cpp + Encoding.cpp + EncodingFactoryNew.cpp + EncodingLayout.cpp + EncodingLayoutCapture.cpp + RleEncoding.cpp + SparseBoolEncoding.cpp + Statistics.cpp + TrivialEncoding.cpp) + +target_link_libraries(nimble_encodings Folly::folly absl::flat_hash_map) diff --git a/dwio/nimble/encodings/Compression.cpp b/dwio/nimble/encodings/Compression.cpp new file mode 100644 index 0000000..d6b56e5 --- /dev/null +++ b/dwio/nimble/encodings/Compression.cpp @@ -0,0 +1,75 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/encodings/Compression.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/encodings/CompressionInternal.h" + +namespace facebook::nimble { + +/* static */ CompressionResult Compression::compress( + velox::memory::MemoryPool& memoryPool, + std::string_view data, + DataType dataType, + int bitWidth, + const CompressionPolicy& compressionPolicy) { + auto compression = compressionPolicy.compression(); + switch (compression.compressionType) { +#ifdef NIMBLE_HAS_ZSTD + case CompressionType::Zstd: + return compressZstd( + memoryPool, + data, + dataType, + bitWidth, + compressionPolicy, + compression.parameters.zstd); +#endif + +#ifdef NIMBLE_HAS_ZSTRONG + case CompressionType::Zstrong: + return compressZstrong( + memoryPool, + data, + dataType, + bitWidth, + compressionPolicy, + compression.parameters.zstrong); +#endif + + default: + break; + } + NIMBLE_NOT_SUPPORTED(fmt::format( + "Unsupported compression type: {}.", + toString(compression.compressionType))); +} + +/* static */ Vector Compression::uncompress( + velox::memory::MemoryPool& memoryPool, + CompressionType compressionType, + std::string_view data) { + switch (compressionType) { + case CompressionType::Uncompressed: { + NIMBLE_UNREACHABLE( + "uncompress() shouldn't be called on uncompressed buffer."); + } + +#ifdef NIMBLE_HAS_ZSTD + case CompressionType::Zstd: + return uncompressZstd(memoryPool, compressionType, data); +#endif + +#ifdef NIMBLE_HAS_ZSTRONG + case CompressionType::Zstrong: + return uncompressZstrong(memoryPool, compressionType, data); +#endif + + default: + break; + } + NIMBLE_NOT_SUPPORTED(fmt::format( + "Unsupported decompression type: {}.", toString(compressionType))); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/Compression.h b/dwio/nimble/encodings/Compression.h new file mode 100644 index 0000000..9c5f6c1 --- /dev/null +++ b/dwio/nimble/encodings/Compression.h @@ -0,0 +1,158 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/encodings/EncodingSelection.h" +#include "folly/io/IOBuf.h" + +namespace facebook::nimble { + +struct CompressionResult { + CompressionType compressionType; + std::optional> buffer; +}; + +class Compression { + public: + static CompressionResult compress( + velox::memory::MemoryPool& memoryPool, + std::string_view data, + DataType dataType, + int bitWidth, + const CompressionPolicy& compressionPolicy); + + static Vector uncompress( + velox::memory::MemoryPool& memoryPool, + CompressionType compressionType, + std::string_view data); +}; + +// Encodings using compression repeat the same pattern, which involves trying to +// compress the data, and throwing it away if compression policy decides to +// throw it away. This class tries to extract this common logic into one place. +// Note: There are actually two sub-patterns, therefore, there are two CTors in +// this class. More on this below. +template +class CompressionEncoder { + public: + // This CTor handles the sub-pattern where the source data is already encoded + // correctly, so no extra encoding is needed. + CompressionEncoder( + velox::memory::MemoryPool& memoryPool, + const CompressionPolicy& compressionPolicy, + DataType dataType, + std::string_view uncompressedBuffer, + int bitWidth = 0) + : dataSize_{uncompressedBuffer.size()}, + compressionType_{CompressionType::Uncompressed} { + if (uncompressedBuffer.size() == 0 || + compressionPolicy.compression().compressionType == + CompressionType::Uncompressed) { + // No compression, just use the original buffer. + data_ = uncompressedBuffer; + return; + } + + auto compressionResult = Compression::compress( + memoryPool, uncompressedBuffer, dataType, bitWidth, compressionPolicy); + + if (compressionResult.compressionType == CompressionType::Uncompressed) { + // Compression declined. Use the original buffer. + data_ = uncompressedBuffer; + return; + } + + // Compression accepted. Use the compressed buffer. + compressed_ = std::move(compressionResult.buffer); + data_ = {compressed_->data(), compressed_->size()}; + dataSize_ = compressed_->size(); + compressionType_ = compressionResult.compressionType; + } + + // This CTor handles the sub-pattern where the source data requires special + // encoding before it is compressed/written. + // Note that in this case, the target buffer for the newly encoded data is + // different if compression is applied (it is written to a temp buffer), or if + // compression is skipped (written directly to the stream buffer). + CompressionEncoder( + velox::memory::MemoryPool& memoryPool, + const CompressionPolicy& compressionPolicy, + DataType dataType, + int bitWidth, + size_t uncompressedSize, + std::function()> allocateUncompressedBuffer, + std::function encoder) + : dataSize_{uncompressedSize}, + compressionType_{CompressionType::Uncompressed}, + encoder_{encoder} { + if (uncompressedSize == 0 || + compressionPolicy.compression().compressionType == + CompressionType::Uncompressed) { + // No compression. Do not encode the data yet. It will be encoded later + // (in write()) directly into the output buffer. + return; + } + + // Compression is attempted. Encode the data before compressing it, into a + // temp buffer. + auto uncompressed = allocateUncompressedBuffer(); + char* pos = uncompressed.data(); + encoder_(pos); + auto compressionResult = Compression::compress( + memoryPool, + {uncompressed.data(), uncompressed.size()}, + dataType, + bitWidth, + compressionPolicy); + + if (compressionResult.compressionType == CompressionType::Uncompressed) { + // Compression declined. Since we already encoded the data, remember the + // temp buffer for later. + // Note: data size is still the same uncompressed size. + data_ = uncompressed; + return; + } + + // Compression accepted. Use the compressed buffer. + compressed_ = std::move(compressionResult.buffer); + data_ = {compressed_->data(), compressed_->size()}; + dataSize_ = compressed_->size(); + compressionType_ = compressionResult.compressionType; + } + + size_t getSize() { + return dataSize_; + } + + void write(char*& pos) { + if (!data_.has_value()) { + // If we are here, it means we handle uncompressed data that needs to be + // encoded directly into the target buffer. + encoder_(pos); + return; + } + + if (data_->data() == nullptr) { + return; + } + + std::copy(data_->begin(), data_->end(), pos); + pos += data_->size(); + } + + CompressionType compressionType() { + return compressionType_; + } + + private: + size_t dataSize_; + std::optional> data_; + std::optional> compressed_; + CompressionType compressionType_; + std::function encoder_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/CompressionInternal.h b/dwio/nimble/encodings/CompressionInternal.h new file mode 100644 index 0000000..0529538 --- /dev/null +++ b/dwio/nimble/encodings/CompressionInternal.h @@ -0,0 +1,36 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/encodings/Compression.h" +#include "dwio/nimble/encodings/EncodingSelection.h" + +namespace facebook::nimble { + +CompressionResult compressZstd( + velox::memory::MemoryPool& memoryPool, + std::string_view data, + DataType dataType, + int bitWidth, + const CompressionPolicy& compressionPolicy, + const ZstdCompressionParameters& zstdParams); + +CompressionResult compressZstrong( + velox::memory::MemoryPool& memoryPool, + std::string_view data, + DataType dataType, + int bitWidth, + const CompressionPolicy& compressionPolicy, + const ZstrongCompressionParameters& zstrongParams); + +Vector uncompressZstd( + velox::memory::MemoryPool& memoryPool, + CompressionType compressionType, + std::string_view data); + +Vector uncompressZstrong( + velox::memory::MemoryPool& memoryPool, + CompressionType compressionType, + std::string_view data); + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/CompressionZstd.cpp b/dwio/nimble/encodings/CompressionZstd.cpp new file mode 100644 index 0000000..3c01163 --- /dev/null +++ b/dwio/nimble/encodings/CompressionZstd.cpp @@ -0,0 +1,55 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include +#include +#include +#include "dwio/nimble/encodings/CompressionInternal.h" + +namespace facebook::nimble { + +CompressionResult compressZstd( + velox::memory::MemoryPool& memoryPool, + std::string_view data, + DataType dataType, + int /* bitWidth */, + const CompressionPolicy& /* compressionPolicy */, + const ZstdCompressionParameters& zstdParams) { + Vector buffer{&memoryPool, data.size() + sizeof(uint32_t)}; + auto pos = buffer.data(); + encoding::writeUint32(data.size(), pos); + auto ret = ZSTD_compress( + pos, data.size(), data.data(), data.size(), zstdParams.compressionLevel); + if (ZSTD_isError(ret)) { + NIMBLE_ASSERT( + ZSTD_getErrorCode(ret) == ZSTD_ErrorCode::ZSTD_error_dstSize_tooSmall, + fmt::format( + "Error while compressing data: {}", ZSTD_getErrorName(ret))); + return { + .compressionType = CompressionType::Uncompressed, + .buffer = std::nullopt, + }; + } + + buffer.resize(ret + sizeof(uint32_t)); + return { + .compressionType = CompressionType::Zstd, + .buffer = std::move(buffer), + }; +} + +Vector uncompressZstd( + velox::memory::MemoryPool& memoryPool, + CompressionType compressionType, + std::string_view data) { + auto pos = data.data(); + const uint32_t uncompressedSize = encoding::readUint32(pos); + Vector buffer{&memoryPool, uncompressedSize}; + auto ret = ZSTD_decompress( + buffer.data(), buffer.size(), pos, data.size() - sizeof(uint32_t)); + NIMBLE_CHECK( + !ZSTD_isError(ret), + fmt::format("Error uncompressing data: {}", ZSTD_getErrorName(ret))); + return buffer; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/CompressionZstrong.cpp b/dwio/nimble/encodings/CompressionZstrong.cpp new file mode 100644 index 0000000..e8f692e --- /dev/null +++ b/dwio/nimble/encodings/CompressionZstrong.cpp @@ -0,0 +1,214 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/encodings/CompressionInternal.h" + +#include +#include +#include +#include +#include +#include "data_compression/experimental/zstrong_compressors/xldb/lib/xldb_compressor.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/MetricsLogger.h" + +// Temporary flags to enable testing. +DEFINE_int32( + nimble_zstrong_override_compression_level, + -1, + "Override the compression level passed to Zstrong."); +DEFINE_int32( + nimble_zstrong_override_decompression_level, + -1, + "Override the decompression level passed to Zstrong."); + +namespace facebook::nimble { +namespace { + +// Rate limit the error logs because they can explode in volume. +// We don't rate limit the version logs in this layer because it has its own +// logic for rate limiting. +class ErrorLogRateLimiter { + using Clock = folly::chrono::coarse_steady_clock; + static Clock::duration constexpr kDefaultResetInterval = + std::chrono::hours(1); + static constexpr long kDefaultLogQuota = 100; + + public: + bool shouldLog() { + const long lastResetTime = lastResetTime_.load(std::memory_order_acquire); + const long currentTime = Clock::now().time_since_epoch().count(); + + if (currentTime - lastResetTime >= kDefaultResetInterval.count()) { + LOG(WARNING) << "Resetting log quota"; + logCount_ = 0; + lastResetTime_.store(currentTime, std::memory_order_release); + } + return ++logCount_ <= kDefaultLogQuota; + } + + private: + std::atomic_long logCount_{0}; + std::atomic lastResetTime_{0}; +}; + +static bool shouldLogError() { + static ErrorLogRateLimiter limiter; + return limiter.shouldLog(); +} + +class ZstrongLogger : public zstrong::compressors::xldb::XLDBLogger { + void log(std::string_view msg) { + try { + auto logger = LoggingScope::getLogger(); + // It's possible for logger to be nullptr in unit test context. + if (LIKELY(logger != nullptr)) { + logger->logZstrongContext(std::string(msg)); + } + } catch (...) { + LOG(WARNING) << "Failed to log Zstrong context."; + } + } + + public: + void logCompressionVersion(int formatVersion) override { + log(generateJSONLog( + LogType::VERSION_LOG, Operation::COMPRESS, formatVersion)); + } + + void logDecompressionVersion(int formatVersion) override { + log(generateJSONLog( + LogType::VERSION_LOG, Operation::DECOMPRESS, formatVersion)); + } + + void logCompressionError( + std::string_view msg, + std::optional formatVersion = std::nullopt) override { + if (shouldLogError()) { + LOG(WARNING) << fmt::format("Logging zstrong compression error: {}", msg); + log(generateJSONLog( + LogType::ERROR, Operation::COMPRESS, formatVersion, msg)); + } + } + + void logDecompressionError( + std::string_view msg, + std::optional formatVersion = std::nullopt) override { + if (shouldLogError()) { + LOG(WARNING) << fmt::format( + "Logging zstrong decompression error: {}", msg); + log(generateJSONLog( + LogType::ERROR, Operation::DECOMPRESS, formatVersion, msg)); + } + } + + ~ZstrongLogger() override = default; +}; + +} // namespace + +CompressionResult compressZstrong( + velox::memory::MemoryPool& memoryPool, + std::string_view data, + DataType dataType, + int bitWidth, + const CompressionPolicy& compressionPolicy, + const ZstrongCompressionParameters& zstrongParams) { + zstrong::compressors::xldb::DataType dt = + zstrong::compressors::xldb::DataType::Unknown; + switch (dataType) { + case DataType::Uint64: + case DataType::Int64: + bitWidth = 64; + dt = zstrong::compressors::xldb::DataType::U64; + break; + case DataType::Uint32: + case DataType::Int32: + bitWidth = 32; + dt = zstrong::compressors::xldb::DataType::Int; + break; + case DataType::Uint16: + case DataType::Int16: + bitWidth = 16; + dt = zstrong::compressors::xldb::DataType::Int; + break; + case DataType::Uint8: + case DataType::Int8: + bitWidth = 8; + dt = zstrong::compressors::xldb::DataType::Int; + break; + // Used for bit packed data. + case DataType::Undefined: { + if (zstrongParams.useVariableBitWidthCompressor) { + dt = zstrong::compressors::xldb::DataType::Int; + } + break; + } + case DataType::Float: + dt = zstrong::compressors::xldb::DataType::F32; + break; + default: + // TODO: support other datatypes + break; + } + zstrong::compressors::xldb::Compressor compressor( + std::make_unique()); + const int32_t compressionLevel = + FLAGS_nimble_zstrong_override_compression_level > 0 + ? FLAGS_nimble_zstrong_override_compression_level + : zstrongParams.compressionLevel; + const int32_t decompressionLevel = + FLAGS_nimble_zstrong_override_decompression_level > 0 + ? FLAGS_nimble_zstrong_override_decompression_level + : zstrongParams.decompressionLevel; + zstrong::compressors::xldb::CompressParams params = { + .level = zstrong::compressors::xldb::Level( + compressionLevel, decompressionLevel), + .datatype = dt, + .integerBitWidth = folly::to(bitWidth), + .bruteforce = false, + }; + + const auto uncompressedSize = data.size(); + const auto bound = compressor.compressBound(uncompressedSize); + Vector buffer{&memoryPool, bound}; + size_t compressedSize = 0; + compressedSize = compressor.compress( + buffer.data(), buffer.size(), data.data(), data.size(), params); + + if (!compressionPolicy.shouldAccept( + CompressionType::Zstrong, uncompressedSize, compressedSize) || + compressedSize >= bound) { + return { + .compressionType = CompressionType::Uncompressed, + .buffer = std::nullopt, + }; + } + + buffer.resize(compressedSize); + return { + .compressionType = CompressionType::Zstrong, + .buffer = std::move(buffer), + }; +} + +Vector uncompressZstrong( + velox::memory::MemoryPool& memoryPool, + CompressionType compressionType, + std::string_view data) { + auto compressor = + zstrong::compressors::xldb::Compressor(std::make_unique()); + const auto uncompressedSize = + compressor.decompressedSize(data.data(), data.length()); + Vector buffer{&memoryPool, uncompressedSize}; + const auto actualSize = compressor.decompress( + buffer.data(), buffer.size(), data.data(), data.size()); + NIMBLE_CHECK( + actualSize == uncompressedSize, + fmt::format( + "Corrupted stream. Decompressed object does not match expected size. Expected: {}, Actual: {}.", + uncompressedSize, + actualSize)); + return buffer; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/ConstantEncoding.h b/dwio/nimble/encodings/ConstantEncoding.h new file mode 100644 index 0000000..861c507 --- /dev/null +++ b/dwio/nimble/encodings/ConstantEncoding.h @@ -0,0 +1,115 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include "dwio/nimble/common/Bits.h" +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/EncodingType.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/encodings/EncodingSelection.h" + +// Encodes data that is constant, i.e. there is a single unique value. + +namespace facebook::nimble { + +// Data layout is: +// Encoding::kPrefixSize bytes: standard Encoding prefix +// X bytes: the constant value via encoding primitive. +template +class ConstantEncoding final + : public TypedEncoding::physicalType> { + public: + using cppDataType = T; + using physicalType = typename TypeTraits::physicalType; + + ConstantEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data); + + void reset() final; + void skip(uint32_t rowCount) final; + void materialize(uint32_t rowCount, void* buffer) final; + + static std::string_view encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer); + + std::string debugString(int offset) const final; + + private: + physicalType value_; +}; + +// +// End of public API. Implementation follows. +// + +template +ConstantEncoding::ConstantEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data) + : TypedEncoding(memoryPool, data) { + const char* pos = data.data() + Encoding::kPrefixSize; + value_ = encoding::read(pos); + NIMBLE_CHECK(pos == data.end(), "Unexpected constant encoding end"); +} + +template +void ConstantEncoding::reset() {} + +template +void ConstantEncoding::skip(uint32_t /* rowCount */) {} + +template +void ConstantEncoding::materialize(uint32_t rowCount, void* buffer) { + physicalType* castBuffer = static_cast(buffer); + for (uint32_t i = 0; i < rowCount; ++i) { + castBuffer[i] = value_; + } +} + +template +std::string_view ConstantEncoding::encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer) { + if (values.empty()) { + NIMBLE_INCOMPATIBLE_ENCODING("ConstantEncoding cannot be empty."); + } + + if (selection.statistics().uniqueCounts().size() != 1) { + NIMBLE_INCOMPATIBLE_ENCODING("ConstantEncoding requires constant data."); + } + + const uint32_t rowCount = values.size(); + uint32_t encodingSize = Encoding::kPrefixSize; + if constexpr (isStringType()) { + encodingSize += 4 + values[0].size(); + } else { + encodingSize += sizeof(physicalType); + } + char* reserved = buffer.reserve(encodingSize); + char* pos = reserved; + Encoding::serializePrefix( + EncodingType::Constant, TypeTraits::dataType, rowCount, pos); + encoding::write(values[0], pos); + NIMBLE_DASSERT(pos - reserved == encodingSize, "Encoding size mismatch."); + return {reserved, encodingSize}; +} + +template +std::string ConstantEncoding::debugString(int offset) const { + return fmt::format( + "{}{}<{}> rowCount={} value={}", + std::string(offset, ' '), + toString(Encoding::encodingType()), + toString(Encoding::dataType()), + Encoding::rowCount(), + AS_CONST(T, value_)); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/DeltaEncoding.h b/dwio/nimble/encodings/DeltaEncoding.h new file mode 100644 index 0000000..7122345 --- /dev/null +++ b/dwio/nimble/encodings/DeltaEncoding.h @@ -0,0 +1,331 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include + +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/EncodingType.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" + +// Stores integer data in a delta encoding. We use three child encodings: +// one for whether each row is a delta from the last or a restatement, +// one for the deltas, and one one for the restatements. For now we +// only support positive deltas. +// +// As an example, consider the data +// +// 1 2 4 1 2 3 4 1 2 4 8 8 +// +// The is-restatement bool vector is +// T F F T F F F T F F F F +// +// The delta vector is +// 1 2 1 1 1 1 2 4 0 +// +// The restatement vector is +// 1 1 1 + +namespace facebook::nimble { + +// Data layout is: +// Encoding::kPrefixSize bytes: standard Encoding prefix +// 4 bytes: restatement relative offset (X) +// 4 bytes: is-restatement relative offset (Y) +// X bytes: delta encoding bytes +// Y bytes: restatement encoding bytes +// Z bytes: is-restatement encoding bytes +template +class DeltaEncoding final + : public TypedEncoding::physicalType> { + public: + using cppDataType = T; + using physicalType = typename TypeTraits::physicalType; + + DeltaEncoding(velox::memory::MemoryPool& memoryPool, std::string_view data); + + void reset() final; + void skip(uint32_t rowCount) final; + void materialize(uint32_t rowCount, void* buffer) final; + + std::string debugString(int offset) const final; + + private: + physicalType currentValue_; + std::unique_ptr deltas_; + std::unique_ptr restatements_; + std::unique_ptr isRestatements_; + // Temporary bufs. + Vector deltasBuffer_; + Vector restatementsBuffer_; + Vector isRestatementsBuffer_; +}; + +// +// End of public API. Implementation follows. +// + +template +DeltaEncoding::DeltaEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data) + : TypedEncoding(memoryPool, data), + deltasBuffer_(&memoryPool), + restatementsBuffer_(&memoryPool), + isRestatementsBuffer_(&memoryPool) { + auto pos = data.data() + Encoding::kPrefixSize; + const uint32_t restatementsOffset = encoding::readUint32(pos); + const uint32_t isRestatementsOffset = encoding::readUint32(pos); + deltas_ = EncodingFactory::decode(memoryPool, {pos, restatementsOffset}); + pos += restatementsOffset; + restatements_ = + EncodingFactory::decode(memoryPool, {pos, isRestatementsOffset}); + pos += isRestatementsOffset; + isRestatements_ = EncodingFactory::decode( + memoryPool, {pos, static_cast(data.end() - pos)}); +} + +template +void DeltaEncoding::reset() { + deltas_->reset(); + restatements_->reset(); + isRestatements_->reset(); +} + +template +void DeltaEncoding::skip(uint32_t rowCount) { + if (rowCount == 0) { + return; + } + isRestatementsBuffer_.resize(rowCount); + isRestatements_->materialize(rowCount, isRestatementsBuffer_.data()); + const uint32_t numRestatements = std::accumulate( + isRestatementsBuffer_.begin(), isRestatementsBuffer_.end(), 0UL); + // Find the last restatement, then accumulate deltas forward from it. + int64_t lastRestatement = rowCount - 1; + while (lastRestatement >= 0) { + if (isRestatementsBuffer_[lastRestatement]) { + break; + } + --lastRestatement; + } + if (lastRestatement >= 0) { + restatements_->skip(numRestatements - 1); + restatements_->materialize(1, ¤tValue_); + const uint32_t deltasToSkip = lastRestatement - + std::accumulate(isRestatementsBuffer_.begin(), + isRestatementsBuffer_.begin() + lastRestatement, + 0UL); + deltas_->skip(deltasToSkip); + } + const uint32_t deltasToAccumulate = rowCount - 1 - lastRestatement; + deltasBuffer_.resize(deltasToAccumulate); + deltas_->materialize(deltasToAccumulate, deltasBuffer_.data()); + currentValue_ += std::accumulate( + deltasBuffer_.begin(), deltasBuffer_.end(), physicalType()); +} + +template +void DeltaEncoding::materialize(uint32_t rowCount, void* buffer) { + isRestatementsBuffer_.resize(rowCount); + isRestatements_->materialize(rowCount, isRestatementsBuffer_.data()); + const uint32_t numRestatements = std::accumulate( + isRestatementsBuffer_.begin(), isRestatementsBuffer_.end(), 0UL); + restatementsBuffer_.reserve(numRestatements); + restatements_->materialize(numRestatements, restatementsBuffer_.data()); + deltasBuffer_.reserve(rowCount - numRestatements); + deltas_->materialize(rowCount - numRestatements, deltasBuffer_.data()); + physicalType* castValue = static_cast(buffer); + physicalType* nextRestatement = restatementsBuffer_.begin(); + physicalType* nextDelta = deltasBuffer_.begin(); + for (uint32_t i = 0; i < rowCount; ++i) { + if (isRestatementsBuffer_[i]) { + currentValue_ = *nextRestatement++; + } else { + currentValue_ += *nextDelta++; + } + *castValue++ = currentValue_; + } +} + +// namespace internal { + +// template +// void computeDeltas( +// std::span values, +// Vector* deltas, +// Vector* restatements, +// Vector* isRestatements) { +// isRestatements->push_back(true); +// restatements->push_back(values[0]); +// // For signed integer types we avoid the potential overflow in the +// // delta by restating whenever the last value was negative and the +// // next is positive. We could be more elegant by storing the +// // deltas as the appropriate unsigned type. +// if constexpr (isSignedIntegralType()) { +// for (uint32_t i = 1; i < values.size(); ++i) { +// const bool crossesZero = values[i] > 0 && values[i - 1] < 0; +// if (values[i] >= values[i - 1] && !crossesZero) { +// isRestatements->push_back(false); +// deltas->push_back(values[i] - values[i - 1]); +// } else { +// isRestatements->push_back(true); +// restatements->push_back(values[i]); +// } +// } +// } else { +// for (uint32_t i = 1; i < values.size(); ++i) { +// if (values[i] >= values[i - 1]) { +// isRestatements->push_back(false); +// deltas->push_back(values[i] - values[i - 1]); +// } else { +// isRestatements->push_back(true); +// restatements->push_back(values[i]); +// } +// } +// } +// } + +// } // namespace internal + +// template +// bool DeltaEncoding::estimateSize( +// velox::memory::MemoryPool& memoryPool, +// std::span dataValues, +// OptimalSearchParams searchParams, +// encodings::EncodingParameters& encodingParameters, +// uint32_t* size) { +// auto values = +// EncodingPhysicalType::asEncodingPhysicalTypeSpan(dataValues); if +// (values.empty()) { +// return false; +// } +// Vector deltas(&memoryPool); +// Vector restatements(&memoryPool); +// Vector isRestatements(&memoryPool); +// internal::computeDeltas(values, &deltas, &restatements, &isRestatements); +// uint32_t deltasSize; +// auto& deltaEncodingParameters = encodingParameters.set_delta(); +// estimateOptimalEncodingSize( +// memoryPool, +// deltas, +// searchParams, +// &deltasSize, +// deltaEncodingParameters.deltasParameters().ensure()); +// uint32_t restatementsSize; +// estimateOptimalEncodingSize( +// memoryPool, +// restatements, +// searchParams, +// &restatementsSize, +// deltaEncodingParameters.restatementsParameters().ensure()); +// uint32_t isRestatementsSize; +// estimateOptimalEncodingSize( +// memoryPool, +// isRestatements, +// searchParams, +// &isRestatementsSize, +// deltaEncodingParameters.isRestatementsParameters().ensure()); +// *size = Encoding::kPrefixSize + 8 + deltasSize + restatementsSize + +// isRestatementsSize; +// return true; +// } + +// template +// std::string_view DeltaEncoding::serialize( +// std::span values, +// Buffer* buffer) { +// uint32_t unusedSize; +// encodings::EncodingParameters encodingParameters; +// // Hrm should we pass these in? This call won't normally be used outside of +// // testing. +// OptimalSearchParams searchParams; +// estimateSize( +// buffer->getMemoryPool(), +// values, +// searchParams, +// encodingParameters, +// &unusedSize); +// return serialize(values, encodingParameters, buffer); +// } + +// template +// std::string_view DeltaEncoding::serialize( +// std::span dataValues, +// const encodings::EncodingParameters& encodingParameters, +// Buffer* buffer) { +// auto values = +// EncodingPhysicalType::asEncodingPhysicalTypeSpan(dataValues); if +// (values.empty()) { +// NIMBLE_INCOMPATIBLE_ENCODING("DeltaEncoding can't be used with 0 rows."); +// } +// NIMBLE_CHECK( +// encodingParameters.getType() == +// encodings::EncodingParameters::Type::delta && +// encodingParameters.delta_ref().has_value() && +// encodingParameters.delta_ref()->deltasParameters().has_value() && +// encodingParameters.delta_ref() +// ->restatementsParameters() +// .has_value() && +// encodingParameters.delta_ref() +// ->isRestatementsParameters() +// .has_value(), +// "Incomplete or incompatible Delta encoding parameters."); + +// auto& memoryPool = buffer->getMemoryPool(); +// const uint32_t rowCount = values.size(); +// Vector deltas(&memoryPool); +// Vector restatements(&memoryPool); +// Vector isRestatements(&memoryPool); +// auto& deltaEncodingParameters = encodingParameters.delta_ref().value(); +// internal::computeDeltas(values, &deltas, &restatements, &isRestatements); +// std::string_view serializedDeltas = serializeEncoding( +// deltas, deltaEncodingParameters.deltasParameters().value(), buffer); +// std::string_view serializedRestatements = serializeEncoding( +// restatements, +// deltaEncodingParameters.restatementsParameters().value(), +// buffer); +// std::string_view serializedIsRestatements = serializeEncoding( +// isRestatements, +// deltaEncodingParameters.isRestatementsParameters().value(), +// buffer); +// const uint32_t encodingSize = Encoding::kPrefixSize + 8 + +// serializedDeltas.size() + serializedRestatements.size() + +// serializedIsRestatements.size(); +// char* reserved = buffer->reserve(encodingSize); +// char* pos = reserved; +// Encoding::serializePrefix( +// EncodingType::Delta, TypeTraits::dataType, rowCount, pos); +// encoding::writeUint32(serializedDeltas.size(), pos); +// encoding::writeUint32(serializedRestatements.size(), pos); +// encoding::writeBytes(serializedDeltas, pos); +// encoding::writeBytes(serializedRestatements, pos); +// encoding::writeBytes(serializedIsRestatements, pos); +// NIMBLE_DASSERT(pos - reserved == encodingSize, "Encoding size mismatch."); +// return {reserved, encodingSize}; +// } + +template +std::string DeltaEncoding::debugString(int offset) const { + std::string log = Encoding::debugString(offset); + log += fmt::format( + "\n{}deltas child:\n{}", + std::string(offset + 2, ' '), + deltas_->debugString(offset + 4)); + log += fmt::format( + "\n{}restatements child:\n{}", + std::string(offset + 2, ' '), + restatements_->debugString(offset + 4)); + log += fmt::format( + "\n{}isRestatements child:\n{}", + std::string(offset + 2, ' '), + isRestatements_->debugString(offset + 4)); + return log; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/DictionaryEncoding.h b/dwio/nimble/encodings/DictionaryEncoding.h new file mode 100644 index 0000000..df2cffa --- /dev/null +++ b/dwio/nimble/encodings/DictionaryEncoding.h @@ -0,0 +1,175 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/EncodingType.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" +#include "dwio/nimble/encodings/EncodingIdentifier.h" +#include "dwio/nimble/encodings/EncodingSelection.h" +#include "dwio/nimble/encodings/TrivialEncoding.h" +#include "folly/container/F14Map.h" +#include "velox/common/memory/Memory.h" + +// A dictionary encoded stream is comprised of two pieces: a mapping from the +// n unique values in a stream to the integers [0, n) and the vector of indices +// that scatter those uniques back to the original ordering. + +namespace facebook::nimble { + +// The layout for a dictionary encoding is: +// Encoding::kPrefixSize bytes: standard Encoding prefix +// 4 bytes: alphabet size +// XX bytes: alphabet encoding bytes +// YY bytes: indices encoding bytes +template +class DictionaryEncoding + : public TypedEncoding::physicalType> { + public: + using cppDataType = T; + using physicalType = typename TypeTraits::physicalType; + + static const int kAlphabetSizeOffset = Encoding::kPrefixSize; + + DictionaryEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data); + + void reset() final; + void skip(uint32_t rowCount) final; + void materialize(uint32_t rowCount, void* buffer) final; + + static std::string_view encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer); + + std::string debugString(int offset) const final; + + private: + // Stores pre-loaded alphabet + Vector alphabet_; + + // Indices are uint32_t. + std::unique_ptr indicesEncoding_; + std::unique_ptr alphabetEncoding_; + Vector buffer_; // Temporary buffer. +}; + +// +// End of public API. Implementation follows. +// + +template +DictionaryEncoding::DictionaryEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data) + : TypedEncoding(memoryPool, data), + alphabet_{&memoryPool}, + buffer_(&memoryPool) { + auto pos = data.data() + kAlphabetSizeOffset; + const uint32_t alphabetSize = encoding::readUint32(pos); + + alphabetEncoding_ = + EncodingFactory::decode(this->memoryPool_, {pos, alphabetSize}); + const uint32_t alphabetCount = alphabetEncoding_->rowCount(); + alphabet_.resize(alphabetCount); + alphabetEncoding_->materialize(alphabetCount, alphabet_.data()); + + pos += alphabetSize; + indicesEncoding_ = EncodingFactory::decode( + this->memoryPool_, {pos, static_cast(data.end() - pos)}); +} + +template +void DictionaryEncoding::reset() { + indicesEncoding_->reset(); +} + +template +void DictionaryEncoding::skip(uint32_t rowCount) { + indicesEncoding_->skip(rowCount); +} + +template +void DictionaryEncoding::materialize(uint32_t rowCount, void* buffer) { + buffer_.resize(rowCount); + indicesEncoding_->materialize(rowCount, buffer_.data()); + + T* output = static_cast(buffer); + for (uint32_t index : buffer_) { + *output = alphabet_[index]; + ++output; + } +} + +template +std::string_view DictionaryEncoding::encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer) { + const uint32_t valueCount = values.size(); + const uint32_t alphabetCount = selection.statistics().uniqueCounts().size(); + + folly::F14FastMap alphabetMapping; + alphabetMapping.reserve(alphabetCount); + Vector alphabet{&buffer.getMemoryPool()}; + alphabet.reserve(alphabetCount); + uint32_t index = 0; + for (const auto& pair : selection.statistics().uniqueCounts()) { + alphabet.push_back(pair.first); + alphabetMapping.emplace(pair.first, index++); + } + + Vector indices{&buffer.getMemoryPool()}; + indices.reserve(valueCount); + for (const auto& value : values) { + auto it = alphabetMapping.find(value); + NIMBLE_DASSERT( + it != alphabetMapping.end(), + "Statistics corruption. Missing alphabet entry."); + indices.push_back(it->second); + } + + Buffer tempBuffer{buffer.getMemoryPool()}; + std::string_view serializedAlphabet = + selection.template encodeNested( + EncodingIdentifiers::Dictionary::Alphabet, {alphabet}, tempBuffer); + std::string_view serializedIndices = + selection.template encodeNested( + EncodingIdentifiers::Dictionary::Indices, {indices}, tempBuffer); + + const uint32_t encodingSize = Encoding::kPrefixSize + 4 + + serializedAlphabet.size() + serializedIndices.size(); + char* reserved = buffer.reserve(encodingSize); + char* pos = reserved; + Encoding::serializePrefix( + EncodingType::Dictionary, TypeTraits::dataType, valueCount, pos); + encoding::writeUint32(serializedAlphabet.size(), pos); + encoding::writeBytes(serializedAlphabet, pos); + encoding::writeBytes(serializedIndices, pos); + NIMBLE_DASSERT(pos - reserved == encodingSize, "Encoding size mismatch."); + return {reserved, encodingSize}; +} + +template +std::string DictionaryEncoding::debugString(int offset) const { + std::string log = Encoding::debugString(offset); + log += fmt::format( + "\n{}indices child:\n{}", + std::string(offset, ' '), + indicesEncoding_->debugString(offset + 2)); + log += fmt::format( + "\n{}alphabet:\n{}entries={}", + std::string(offset, ' '), + std::string(offset + 2, ' '), + alphabet_.size()); + return log; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/Encoding.cpp b/dwio/nimble/encodings/Encoding.cpp new file mode 100644 index 0000000..0d26ce1 --- /dev/null +++ b/dwio/nimble/encodings/Encoding.cpp @@ -0,0 +1,51 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" + +namespace facebook::nimble { + +EncodingType Encoding::encodingType() const { + return static_cast(data_[kEncodingTypeOffset]); +} + +DataType Encoding::dataType() const { + return static_cast(data_[kDataTypeOffset]); +} + +uint32_t Encoding::rowCount() const { + return *reinterpret_cast(data_.data() + kRowCountOffset); +} + +/* static */ void Encoding::copyIOBuf(char* pos, const folly::IOBuf& buf) { + [[maybe_unused]] size_t length = buf.computeChainDataLength(); + for (auto data : buf) { + std::copy(data.cbegin(), data.cend(), pos); + pos += data.size(); + length -= data.size(); + } + NIMBLE_DASSERT(length == 0, "IOBuf chain length corruption"); +} + +void Encoding::serializePrefix( + EncodingType encodingType, + DataType dataType, + uint32_t rowCount, + char*& pos) { + encoding::writeChar(static_cast(encodingType), pos); + encoding::writeChar(static_cast(dataType), pos); + encoding::writeUint32(rowCount, pos); +} + +std::string Encoding::debugString(int offset) const { + return fmt::format( + "{}{}<{}> rowCount={}", + std::string(offset, ' '), + toString(encodingType()), + toString(dataType()), + rowCount()); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/Encoding.h b/dwio/nimble/encodings/Encoding.h new file mode 100644 index 0000000..85f852e --- /dev/null +++ b/dwio/nimble/encodings/Encoding.h @@ -0,0 +1,286 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/common/Bits.h" +#include "dwio/nimble/common/EncodingType.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" +#include "velox/common/memory/Memory.h" + +#include + +// The Encoding class defines an interface for interacting with encodings +// (aka vectors, aka arrays) of encoded data. The API is tailored for +// typical usage patterns within query engines, and is designed to be +// easily extensible. +// +// Some general notes: +// 1. Output bitmaps are compatible with the FixedBitArray class. In +// particular this means you must use FixedBitArray::bufferSize to +// determine the space needed for the bitmaps you pass to, e.g., +// the output buffer of Encoding::Equals. +// 2. When passing types into functions via the void* interfaces, +// and when materializing data, numeric types are input/output via +// their native types, while string are passed via std::string_view +// -- NOT std::string! +// 3. We refer to the 'row pointer' in various places. Think of this as +// the iterator over the encoding's data. Encoding::Reset() is like +// rp = vec.begin(), advancing N rows is like rp = rp + N, etc. +// 4. An Encoding doesn't own the underlying data. If you want to create two +// copies for some reason on the same data go ahead, they are totally +// independent. +// 5. Each subclass of Encoding will provide its own serialization methods. +// The resulting string should be passed to the subclass's constructor +// at read time. +// 6. Although serialized Encodings can be stored independently if desired +// (as they are just strings), when written to disk they are normally +// stored together in a Tablet (see tablet.h). +// 7. We consistently use 4 byte unsigned integers for things like offsets +// and row counts. Encodings, and indeed entire Tablets, are required to +// fit within a uint32_t in size. This is enforced in the normal +// ingestion pathways. If you directly use the serialization methods, +// be careful that this remains true! +// 8. See the notes in the Encoding class about array encodings. These +// encodings +// have slightly different semantics than other encodings. + +namespace facebook::nimble { + +class Encoding { + public: + Encoding(velox::memory::MemoryPool& memoryPool, std::string_view data) + : memoryPool_{memoryPool}, data_{data} {} + virtual ~Encoding() = default; + + EncodingType encodingType() const; + DataType dataType() const; + uint32_t rowCount() const; + + static void copyIOBuf(char* pos, const folly::IOBuf& buf); + + virtual uint32_t nullCount() const { + return 0; + } + + // Resets the internal state (e.g. row pointer) to newly constructed form. + virtual void reset() = 0; + + // Advances the row pointer N rows. Note that we don't provide a 'backwards' + // iterator; if you need to move your row pointer back, reset() and skip(). + virtual void skip(uint32_t rowCount) = 0; + + // Materializes the next |rowCount| rows into buffer. Advances + // the row pointer |rowCount|. + // + // Remember that buffer is interpreted as a physicalType*, with + // std::string_view* for strings and bool* (not a bitmap) for bool. For + // non-POD types, namely string data, note that the materialized values are + // only guaranteed valid until the next non-const call to *this. + // + // If this->isNullable(), null rows are materialized as physicalType(). To be + // able to distinguish between nulls and non-null physicalType() values, use + // MaterializeNullable. + virtual void materialize(uint32_t rowCount, void* buffer) = 0; + + // Nullable method. + // Like Materialize, but also sets the ith bit of bitmap to reflect whether + // ith row was null or not. 0 means null, 1 means not null. Null rows are left + // untouched instead of being filled with default values. + // + // If |scatterBitmap| is provided, |rowCount| still indicates how many items + // to be read from the encoding. However, instead of placing them sequentially + // in the output |buffer|, the items are scattered. This means that item will + // be place into the slot where the corresponding positional bit is set to 1 + // in |scatterBitmap|, (note that the value being read may be null). For every + // positional scatter bit set to 0, it will fill a null in the poisition in + // the output |buffer|. |rowCount| should match the number of bits set to 1 in + // |scatterBitmap|. For scattered reads, |buffer| and |nulls| should + // have enough space to accommodate |scatterBitmap.size()| items. When + // |offset| is specified, use the |scatterBitmap| starting from |offset| and + // scatter to |buffer| and |nulls| starting from |offset|. + // + // Returns number of items that are not null. In the case when all values are + // non null, |nulls| will not be filled with all 1s. It's expected that + // caller explicitly checks for that condition. + virtual uint32_t materializeNullable( + uint32_t rowCount, + void* buffer, + std::function nulls, + const bits::Bitmap* scatterBitmap = nullptr, + uint32_t offset = 0) = 0; + + // Whether this encoding is nullable, i.e. contains any nulls. This property + // modifies how engines need to interpret many of the function results, and a + // number of functions are only callable if isNullable() returns true. + virtual bool isNullable() const { + return false; + } + + // A number of functions are legal to call only if the encoding is dictionary + // enabled, such as retrieving the dictionary indices or looking up a value + // by dictionary index. + virtual bool isDictionaryEnabled() const { + return false; + } + + // Dictionary method. + // The size of the dictionary, which is equal to the number of unique values. + virtual uint32_t dictionarySize() const { + NIMBLE_NOT_SUPPORTED("Data is not dictionary encoded."); + } + + // Dictionary method. + // Returns the value at the given index, which must be in [0, num_entries). + // This pointer is only guaranteed valid until the next Entry call. + virtual const void* dictionaryEntry(uint32_t /* index */) const { + NIMBLE_NOT_SUPPORTED("Data is not dictionary encoded."); + } + + // Dictionary method. + // Materializes the next |rowCount| dictionary indices into buffer. Advances + // The row pointer |rowCount|. + virtual void materializeIndices( + uint32_t /* rowCount */, + uint32_t* /* buffer */) { + NIMBLE_NOT_SUPPORTED("Data is not dictionary encoded."); + } + + // A string for debugging/iteration that gives details about *this. + // Offset adds that many spaces before the msg (useful for children + // encodings). + virtual std::string debugString(int offset = 0) const; + + protected: + // The binary layout for each Encoding begins with the same prefix: + // 1 byte: EncodingType + // 1 byte: DataType + // 4 bytes: uint32_t num rows + static constexpr int kEncodingTypeOffset = 0; + static constexpr int kDataTypeOffset = 1; + static constexpr int kRowCountOffset = 2; + static constexpr int kPrefixSize = 6; + + static void serializePrefix( + EncodingType encodingType, + DataType dataType, + uint32_t rowCount, + char*& pos); + + velox::memory::MemoryPool& memoryPool_; + const std::string_view data_; +}; + +// The TypedEncoding class exposes the same interface as the base +// Encoding class but provides common typed implementation for some apis exposed +// by the Encoding class. T means semantic data type, physicalType means data +// type used for encoding +template +class TypedEncoding : public Encoding { + static_assert( + std::is_same_v::physicalType>); + + public: + TypedEncoding(velox::memory::MemoryPool& memoryPool, std::string_view data) + : Encoding{memoryPool, data} {} + + // Similar to materialize(), but scatters values to output buffer according to + // scatterBitmap. When scatterBitmap is nullptr or all 1's, the output + // nullBitmap will not be set. It's expected that caller explicitly checks + // against the return value and handle such all 1's cases properly. + uint32_t materializeNullable( + uint32_t rowCount, + void* buffer, + std::function nulls, + const bits::Bitmap* scatterBitmap, + uint32_t offset) override { + // 1. Read X items from the encoding. + // 2. Spread the items read in #1 into positions in |buffer| where + // |scatterBitmap| has a matching positional bit set to 1. + if (offset > 0) { + buffer = static_cast(buffer) + offset; + } + + materialize(rowCount, buffer); + + // TODO: check rowCount matches the number of set bits in scatterBitmap. + auto scatterCount = + scatterBitmap ? scatterBitmap->size() - offset : rowCount; + if (scatterCount == rowCount) { + // No need to scatter. Avoid setting nullBitmap since caller is expected + // to explicitly handle such non-null cases. + return rowCount; + } + + NIMBLE_CHECK( + rowCount < scatterCount, + fmt::format("Unexpected count {} vs {}", rowCount, scatterCount)); + + void* nullBitmap = nulls(); + + bits::BitmapBuilder nullBits{nullBitmap, offset + scatterCount}; + nullBits.copy(*scatterBitmap, offset, offset + scatterCount); + + // Scatter backward + uint32_t pos = offset + scatterCount - 1; + physicalType* output = + static_cast(buffer) + scatterCount - 1; + const physicalType* source = + static_cast(buffer) + rowCount - 1; + while (output != source) { + if (scatterBitmap->test(pos)) { + *output = *source; + --source; + } + --output; + --pos; + } + + return rowCount; + } +}; + +namespace detail { +template +class BufferedEncoding { + public: + explicit BufferedEncoding(std::unique_ptr encoding) + : bufferPosition_{BufferSize}, + encodingPosition_{0}, + encoding_{std::move(encoding)} {} + + T nextValue() { + if (UNLIKELY(bufferPosition_ == BufferSize)) { + auto rows = std::min( + BufferSize, encoding_->rowCount() - encodingPosition_); + encoding_->materialize(rows, buffer_.data()); + bufferPosition_ = 0; + } + ++encodingPosition_; + return buffer_[bufferPosition_++]; + } + + void reset() { + bufferPosition_ = BufferSize; + encodingPosition_ = 0; + encoding_->reset(); + } + + uint32_t position() const noexcept { + return encodingPosition_ - 1; + } + + uint32_t rowCount() const noexcept { + return encoding_->rowCount(); + } + + private: + uint16_t bufferPosition_; + uint32_t encodingPosition_; + std::unique_ptr encoding_; + std::array buffer_; +}; + +} // namespace detail +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/EncodingFactoryNew.cpp b/dwio/nimble/encodings/EncodingFactoryNew.cpp new file mode 100644 index 0000000..0a061fd --- /dev/null +++ b/dwio/nimble/encodings/EncodingFactoryNew.cpp @@ -0,0 +1,368 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/encodings/EncodingFactoryNew.h" +#include "dwio/nimble/encodings/ConstantEncoding.h" +#include "dwio/nimble/encodings/DictionaryEncoding.h" +#include "dwio/nimble/encodings/EncodingSelection.h" +#include "dwio/nimble/encodings/FixedBitWidthEncoding.h" +#include "dwio/nimble/encodings/MainlyConstantEncoding.h" +#include "dwio/nimble/encodings/NullableEncoding.h" +#include "dwio/nimble/encodings/RleEncoding.h" +#include "dwio/nimble/encodings/SparseBoolEncoding.h" +#include "dwio/nimble/encodings/TrivialEncoding.h" +#include "dwio/nimble/encodings/VarintEncoding.h" + +namespace facebook::nimble { + +namespace { + +template +static std::span::physicalType> toPhysicalSpan( + std::span values) { + return std::span::physicalType>( + reinterpret_cast::physicalType*>( + values.data()), + values.size()); +} +} // namespace + +std::unique_ptr EncodingFactory::decode( + velox::memory::MemoryPool& memoryPool, + std::string_view data) { + // Maybe we should have a magic number of encodings too? Hrm. + const EncodingType encodingType = static_cast(data[0]); + const DataType dataType = static_cast(data[1]); +#define RETURN_ENCODING_BY_LEAF_TYPE(Encoding, dataType) \ + switch (dataType) { \ + case DataType::Int8: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint8: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Int16: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint16: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Int32: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint32: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Int64: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint64: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Float: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Double: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Bool: \ + return std::make_unique>(memoryPool, data); \ + case DataType::String: \ + return std::make_unique>(memoryPool, data); \ + default: \ + NIMBLE_UNREACHABLE(fmt::format("Unknown encoding type {}.", dataType)) \ + } + +#define RETURN_ENCODING_BY_INTEGER_TYPE(Encoding, dataType) \ + switch (dataType) { \ + case DataType::Int8: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint8: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Int16: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint16: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Int32: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint32: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Int64: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint64: \ + return std::make_unique>(memoryPool, data); \ + default: \ + NIMBLE_UNREACHABLE(fmt::format( \ + "Trying to deserialize an integer-only stream for " \ + "a nonintegral data type {}.", \ + dataType)); \ + } + +#define RETURN_ENCODING_BY_VARINT_TYPE(Encoding, dataType) \ + switch (dataType) { \ + case DataType::Int32: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint32: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Int64: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint64: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Float: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Double: \ + return std::make_unique>(memoryPool, data); \ + default: \ + NIMBLE_UNREACHABLE(fmt::format( \ + "Trying to deserialize a varint stream for " \ + "an incompatible data type {}.", \ + dataType)); \ + } + +#define RETURN_ENCODING_BY_NON_BOOL_TYPE(Encoding, dataType) \ + switch (dataType) { \ + case DataType::Int8: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint8: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Int16: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint16: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Int32: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint32: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Int64: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint64: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Float: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Double: \ + return std::make_unique>(memoryPool, data); \ + case DataType::String: \ + return std::make_unique>(memoryPool, data); \ + default: \ + NIMBLE_UNREACHABLE(fmt::format( \ + "Trying to deserialize a non-bool stream for " \ + "the bool data type {}.", \ + dataType)); \ + } + +#define RETURN_ENCODING_BY_NUMERIC_TYPE(Encoding, dataType) \ + switch (dataType) { \ + case DataType::Int8: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint8: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Int16: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint16: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Int32: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint32: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Int64: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Uint64: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Float: \ + return std::make_unique>(memoryPool, data); \ + case DataType::Double: \ + return std::make_unique>(memoryPool, data); \ + default: \ + NIMBLE_UNREACHABLE(fmt::format( \ + "Trying to deserialize a non-numeric stream for " \ + "a numeric data type {}.", \ + dataType)); \ + } + + switch (encodingType) { + case EncodingType::Trivial: { + RETURN_ENCODING_BY_LEAF_TYPE(TrivialEncoding, dataType); + } + case EncodingType::RLE: { + RETURN_ENCODING_BY_LEAF_TYPE(RLEEncoding, dataType); + } + case EncodingType::Dictionary: { + RETURN_ENCODING_BY_NON_BOOL_TYPE(DictionaryEncoding, dataType); + } + case EncodingType::FixedBitWidth: { + RETURN_ENCODING_BY_NUMERIC_TYPE(FixedBitWidthEncoding, dataType); + } + case EncodingType::Nullable: { + RETURN_ENCODING_BY_LEAF_TYPE(NullableEncoding, dataType); + } + case EncodingType::SparseBool: { + NIMBLE_ASSERT( + dataType == DataType::Bool, + "Trying to deserialize a SparseBoolEncoding with a non-bool data type."); + return std::make_unique(memoryPool, data); + } + case EncodingType::Varint: { + RETURN_ENCODING_BY_VARINT_TYPE(VarintEncoding, dataType); + } + case EncodingType::Constant: { + RETURN_ENCODING_BY_LEAF_TYPE(ConstantEncoding, dataType); + } + case EncodingType::MainlyConstant: { + RETURN_ENCODING_BY_NON_BOOL_TYPE(MainlyConstantEncoding, dataType); + } + default: { + NIMBLE_UNREACHABLE( + "Trying to deserialize invalid EncodingType -- garbage input?"); + } + } +} + +template +std::string_view EncodingFactory::encode( + std::unique_ptr>&& selectorPolicy, + std::span values, + Buffer& buffer) { + using physicalType = typename TypeTraits::physicalType; + auto physicalValues = toPhysicalSpan(values); + auto statistics = Statistics::create(physicalValues); + auto selectionResult = selectorPolicy->select(physicalValues, statistics); + EncodingSelection selection{ + std::move(selectionResult), + std::move(statistics), + std::move(selectorPolicy)}; + return EncodingFactory::encode( + std::move(selection), physicalValues, buffer); +} + +template +std::string_view EncodingFactory::encodeNullable( + std::unique_ptr>&& selectorPolicy, + std::span values, + std::span nulls, + Buffer& buffer) { + using physicalType = typename TypeTraits::physicalType; + auto physicalValues = toPhysicalSpan(values); + auto statistics = Statistics::create(physicalValues); + auto selectionResult = + selectorPolicy->selectNullable(physicalValues, nulls, statistics); + EncodingSelection selection{ + std::move(selectionResult), + std::move(statistics), + std::move(selectorPolicy)}; + return EncodingFactory::encodeNullable( + std::move(selection), physicalValues, nulls, buffer); +} + +template +std::string_view EncodingFactory::encode( + EncodingSelection::physicalType>&& selection, + std::span::physicalType> values, + Buffer& buffer) { + using physicalType = typename TypeTraits::physicalType; + auto castedValues = toPhysicalSpan(values); + switch (selection.encodingType()) { + case EncodingType::Constant: { + return ConstantEncoding::encode(selection, castedValues, buffer); + } + case EncodingType::Trivial: { + return TrivialEncoding::encode(selection, castedValues, buffer); + } + case EncodingType::RLE: { + return RLEEncoding::encode(selection, castedValues, buffer); + } + case EncodingType::Dictionary: { + NIMBLE_DASSERT( + (!std::is_same::value && !castedValues.empty()), + "Invalid DictionaryEncoding selection."); + return DictionaryEncoding::encode(selection, castedValues, buffer); + } + case EncodingType::FixedBitWidth: { + if constexpr (isNumericType()) { + return FixedBitWidthEncoding::encode( + selection, castedValues, buffer); + } else { + NIMBLE_INCOMPATIBLE_ENCODING( + "FixedBitWidth encoding should not be selected for non-numeric data types."); + } + } + case EncodingType::Varint: { + // TODO: we can support floating point types, but currently Statistics + // doesn't calculate buckets for floating point types. We should convert + // floating point types to their physical type, and then Statistics and + // Varint encoding will just work. + if constexpr ( + isNumericType() && + (sizeof(physicalType) == 4 || sizeof(T) == 8)) { + return VarintEncoding::encode(selection, castedValues, buffer); + } else { + NIMBLE_INCOMPATIBLE_ENCODING( + "Varint encoding can only be selected for large numeric data types."); + } + } + case EncodingType::MainlyConstant: { + if constexpr (std::is_same::value) { + NIMBLE_INCOMPATIBLE_ENCODING( + "MainlyConstant encoding should not be selected for bool data types."); + } else { + return MainlyConstantEncoding::encode( + selection, castedValues, buffer); + } + } + case EncodingType::SparseBool: { + if constexpr (!std::is_same::value) { + NIMBLE_INCOMPATIBLE_ENCODING( + "SparseBool encoding should not be selected for non-bool data types."); + } else { + return SparseBoolEncoding::encode(selection, castedValues, buffer); + } + } + default: { + NIMBLE_NOT_SUPPORTED(fmt::format( + "Encoding {} is not supported.", toString(selection.encodingType()))); + } + } +} + +template +std::string_view EncodingFactory::encodeNullable( + EncodingSelection::physicalType>&& selection, + std::span::physicalType> values, + std::span nulls, + Buffer& buffer) { + auto physicalValues = toPhysicalSpan(values); + switch (selection.encodingType()) { + case EncodingType::Nullable: { + return NullableEncoding::encodeNullable( + selection, physicalValues, nulls, buffer); + } + default: { + NIMBLE_NOT_SUPPORTED(fmt::format( + "Encoding {} is not supported for nullable data.", + toString(selection.encodingType()))); + } + } +} + +#define DEFINE_TEMPLATES(type) \ + template std::string_view EncodingFactory::encode( \ + std::unique_ptr> && selectorPolicy, \ + std::span values, \ + Buffer & buffer); \ + template std::string_view EncodingFactory::encodeNullable( \ + std::unique_ptr> && selectorPolicy, \ + std::span values, \ + std::span nulls, \ + Buffer & buffer); \ + template std::string_view EncodingFactory::encode( \ + EncodingSelection::physicalType> && selection, \ + std::span::physicalType> values, \ + Buffer & buffer); \ + template std::string_view EncodingFactory::encodeNullable( \ + EncodingSelection::physicalType> && selection, \ + std::span::physicalType> values, \ + std::span nulls, \ + Buffer & buffer); + +DEFINE_TEMPLATES(int8_t); +DEFINE_TEMPLATES(uint8_t); +DEFINE_TEMPLATES(int16_t); +DEFINE_TEMPLATES(uint16_t); +DEFINE_TEMPLATES(int32_t); +DEFINE_TEMPLATES(uint32_t); +DEFINE_TEMPLATES(int64_t); +DEFINE_TEMPLATES(uint64_t); +DEFINE_TEMPLATES(float); +DEFINE_TEMPLATES(double); +DEFINE_TEMPLATES(bool); +DEFINE_TEMPLATES(std::string_view); + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/EncodingFactoryNew.h b/dwio/nimble/encodings/EncodingFactoryNew.h new file mode 100644 index 0000000..de44463 --- /dev/null +++ b/dwio/nimble/encodings/EncodingFactoryNew.h @@ -0,0 +1,66 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include + +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/encodings/Encoding.h" + +namespace facebook::nimble { + +template +class EncodingSelection; + +template +class EncodingSelectionPolicy; + +class EncodingFactory { + public: + static std::unique_ptr decode( + velox::memory::MemoryPool& memoryPool, + std::string_view data); + + template + static std::string_view encode( + std::unique_ptr>&& selectorPolicy, + std::span values, + Buffer& buffer); + + template + static std::string_view encodeNullable( + std::unique_ptr>&& selectorPolicy, + std::span values, + std::span nulls, + Buffer& buffer); + + private: + template + static std::string_view encode( + EncodingSelection::physicalType>&& selection, + std::span::physicalType> values, + Buffer& buffer); + + template + static std::string_view encodeNullable( + EncodingSelection::physicalType>&& selection, + std::span::physicalType> values, + std::span nulls, + Buffer& buffer); + + friend class EncodingSelection; + friend class EncodingSelection; + friend class EncodingSelection; + friend class EncodingSelection; + friend class EncodingSelection; + friend class EncodingSelection; + friend class EncodingSelection; + friend class EncodingSelection; + friend class EncodingSelection; + friend class EncodingSelection; + friend class EncodingSelection; + friend class EncodingSelection; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/EncodingIdentifier.h b/dwio/nimble/encodings/EncodingIdentifier.h new file mode 100644 index 0000000..34f0e27 --- /dev/null +++ b/dwio/nimble/encodings/EncodingIdentifier.h @@ -0,0 +1,43 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include + +namespace facebook::nimble { + +// When encoding contains nested encodings, each nested encoding has its +// own identifier. +using NestedEncodingIdentifier = uint8_t; + +struct EncodingIdentifiers { + struct Dictionary { + static constexpr NestedEncodingIdentifier Alphabet = 0; + static constexpr NestedEncodingIdentifier Indices = 1; + }; + + struct MainlyConstant { + static constexpr NestedEncodingIdentifier IsCommon = 0; + static constexpr NestedEncodingIdentifier OtherValues = 1; + }; + + struct Nullable { + static constexpr NestedEncodingIdentifier Data = 0; + static constexpr NestedEncodingIdentifier Nulls = 1; + }; + + struct RunLength { + static constexpr NestedEncodingIdentifier RunLengths = 0; + static constexpr NestedEncodingIdentifier RunValues = 1; + }; + + struct SparseBool { + static constexpr NestedEncodingIdentifier Indices = 0; + }; + + struct Trivial { + static constexpr NestedEncodingIdentifier Lengths = 0; + }; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/EncodingLayout.cpp b/dwio/nimble/encodings/EncodingLayout.cpp new file mode 100644 index 0000000..ac80cd7 --- /dev/null +++ b/dwio/nimble/encodings/EncodingLayout.cpp @@ -0,0 +1,119 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/encodings/EncodingLayout.h" +#include +#include +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/Exceptions.h" + +namespace facebook::nimble { + +namespace { +constexpr uint32_t kMinEncodingLayoutBufferSize = 5; +} + +EncodingLayout::EncodingLayout( + EncodingType encodingType, + CompressionType compressionType, + std::vector> children) + : encodingType_{encodingType}, + compressionType_{compressionType}, + children_{std::move(children)} {} + +int32_t EncodingLayout::serialize(std::span output) const { + // Serialized encoding layout is as follows: + // 1 byte - Encoding Type + // 1 byte - CompressionType + // 1 byte - Children (nested encoding) Count + // 2 bytes - Extra data size (currently always set to zero. Reserved for futre + // exra args for compression or encoding) + // Extra data size bytes - Extra data (currently not used) + // 1 bytes - 1st (nested encoding) child exists + // X bytes - 1st (nested encoding) child + // 1 bytes - 2nd (nested encoding) child exists + // Y bytes - 2nd (nested encoding) child + // ... + + // We store at least kMinEncodingLayoutBufferSize bytes: encoding type, + // compression type, children count and extra data size (2 bytes), plus one + // byte per child. + NIMBLE_CHECK( + output.size() >= kMinEncodingLayoutBufferSize + children_.size(), + "Captured encoding layout buffer too small."); + + output[0] = static_cast(encodingType_); + output[1] = static_cast(compressionType_); + output[2] = static_cast(children_.size()); + // Currently, extra data is not used and always set to zero. + output[3] = output[4] = 0; + + int32_t size = kMinEncodingLayoutBufferSize; + + for (auto i = 0; i < children_.size(); ++i) { + const auto& child = children_[i]; + if (child.has_value()) { + output[size++] = 1; + size += child->serialize(output.subspan(size)); + } else { + // Set child size to 0 + output[size++] = 0; + } + } + + return size; +} + +std::pair EncodingLayout::create( + std::string_view encoding) { + NIMBLE_CHECK( + encoding.size() >= kMinEncodingLayoutBufferSize, + "Invalid captured encoding layout. Buffer too small."); + + auto pos = encoding.data(); + const auto encodingType = encoding::read(pos); + const auto compressionType = encoding::read(pos); + const auto childrenCount = encoding::read(pos); + [[maybe_unused]] const auto extraDataSize = encoding::read(pos); + + NIMBLE_DASSERT(extraDataSize == 0, "Extra data currently not supported."); + + uint32_t offset = kMinEncodingLayoutBufferSize; + std::vector> children; + children.reserve(childrenCount); + for (auto i = 0; i < childrenCount; ++i) { + auto childExists = encoding::peek(encoding.data() + offset); + ++offset; + if (childExists > 0) { + auto encodingLayout = EncodingLayout::create(encoding.substr(offset)); + offset += encodingLayout.second; + children.emplace_back(std::move(encodingLayout.first)); + } else { + children.emplace_back(std::nullopt); + } + } + + return {{encodingType, compressionType, std::move(children)}, offset}; +} + +EncodingType EncodingLayout::encodingType() const { + return encodingType_; +} + +CompressionType EncodingLayout::compressionType() const { + return compressionType_; +} + +uint8_t EncodingLayout::childrenCount() const { + return children_.size(); +} + +const std::optional& EncodingLayout::child( + NestedEncodingIdentifier identifier) const { + NIMBLE_DCHECK( + identifier < children_.size(), + "Encoding layout identifier is out of range."); + + return children_[identifier]; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/EncodingLayout.h b/dwio/nimble/encodings/EncodingLayout.h new file mode 100644 index 0000000..f9ad074 --- /dev/null +++ b/dwio/nimble/encodings/EncodingLayout.h @@ -0,0 +1,37 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include +#include + +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/encodings/EncodingIdentifier.h" + +namespace facebook::nimble { + +class EncodingLayout { + public: + EncodingLayout( + EncodingType encodingType, + CompressionType compressionType, + std::vector> children = {}); + + EncodingType encodingType() const; + CompressionType compressionType() const; + uint8_t childrenCount() const; + const std::optional& child( + NestedEncodingIdentifier identifier) const; + + int32_t serialize(std::span output) const; + static std::pair create(std::string_view encoding); + + private: + EncodingType encodingType_; + CompressionType compressionType_; + std::vector> children_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/EncodingLayoutCapture.cpp b/dwio/nimble/encodings/EncodingLayoutCapture.cpp new file mode 100644 index 0000000..474d147 --- /dev/null +++ b/dwio/nimble/encodings/EncodingLayoutCapture.cpp @@ -0,0 +1,152 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/encodings/EncodingLayoutCapture.h" + +#include + +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/Types.h" + +namespace facebook::nimble { + +namespace { + +constexpr uint32_t kEncodingPrefixSize = 6; + +} // namespace + +EncodingLayout EncodingLayoutCapture::capture(std::string_view encoding) { + NIMBLE_CHECK( + encoding.size() >= kEncodingPrefixSize, "Encoding size too small."); + + const auto encodingType = + encoding::peek(encoding.data()); + CompressionType compressionType = CompressionType::Uncompressed; + + if (encodingType == EncodingType::FixedBitWidth || + encodingType == EncodingType::Trivial) { + compressionType = encoding::peek( + encoding.data() + kEncodingPrefixSize); + } + + std::vector> children; + switch (encodingType) { + case EncodingType::FixedBitWidth: + case EncodingType::Varint: + case EncodingType::Constant: { + // Non nested encodings have zero children + break; + } + case EncodingType::Trivial: { + const auto dataType = + encoding::peek(encoding.data() + 1); + if (dataType == DataType::String) { + const char* pos = encoding.data() + kEncodingPrefixSize + 1; + const uint32_t lengthsBytes = encoding::readUint32(pos); + + children.reserve(1); + children.emplace_back( + EncodingLayoutCapture::capture({pos, lengthsBytes})); + } + break; + } + case EncodingType::SparseBool: { + children.reserve(1); + children.emplace_back(EncodingLayoutCapture::capture( + encoding.substr(kEncodingPrefixSize + 1))); + break; + } + case EncodingType::MainlyConstant: { + children.reserve(2); + + const char* pos = encoding.data() + kEncodingPrefixSize; + const uint32_t isCommonBytes = encoding::readUint32(pos); + + children.emplace_back( + EncodingLayoutCapture::capture({pos, isCommonBytes})); + + pos += isCommonBytes; + const uint32_t otherValuesBytes = encoding::readUint32(pos); + + children.emplace_back( + EncodingLayoutCapture::capture({pos, otherValuesBytes})); + break; + } + case EncodingType::Dictionary: { + children.reserve(2); + const char* pos = encoding.data() + kEncodingPrefixSize; + const uint32_t alphabetBytes = encoding::readUint32(pos); + + children.emplace_back( + EncodingLayoutCapture::capture({pos, alphabetBytes})); + + pos += alphabetBytes; + + children.emplace_back(EncodingLayoutCapture::capture( + {pos, encoding.size() - (pos - encoding.data())})); + break; + } + case EncodingType::RLE: { + const auto dataType = + encoding::peek(encoding.data() + 1); + + children.reserve(dataType == DataType::Bool ? 1 : 2); + + const char* pos = encoding.data() + kEncodingPrefixSize; + const uint32_t runLengthBytes = encoding::readUint32(pos); + + children.emplace_back( + EncodingLayoutCapture::capture({pos, runLengthBytes})); + + if (dataType != DataType::Bool) { + pos += runLengthBytes; + + children.emplace_back(EncodingLayoutCapture::capture( + {pos, encoding.size() - (pos - encoding.data())})); + } + break; + } + case EncodingType::Delta: { + children.reserve(3); + + const char* pos = encoding.data() + kEncodingPrefixSize; + const uint32_t deltaBytes = encoding::readUint32(pos); + const uint32_t restatementBytes = encoding::readUint32(pos); + + children.emplace_back(EncodingLayoutCapture::capture({pos, deltaBytes})); + + pos += deltaBytes; + + children.emplace_back( + EncodingLayoutCapture::capture({pos, restatementBytes})); + + pos += restatementBytes; + + children.emplace_back(EncodingLayoutCapture::capture( + {pos, encoding.size() - (pos - encoding.data())})); + break; + } + case EncodingType::Nullable: { + const char* pos = encoding.data() + kEncodingPrefixSize; + const uint32_t dataBytes = encoding::readUint32(pos); + + // For nullable encodings we only capture the data encoding part, so we + // are "overwriting" the current captured node with the nested data node. + return EncodingLayoutCapture::capture({pos, dataBytes}); + + break; + } + case EncodingType::Sentinel: { + // For sentinel encodings we only capture the data encoding part, so we + // are "overwriting" the current captured node with the nested data node. + return EncodingLayoutCapture::capture( + encoding.substr(kEncodingPrefixSize + 8)); + break; + } + } + + return {encodingType, compressionType, std::move(children)}; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/EncodingLayoutCapture.h b/dwio/nimble/encodings/EncodingLayoutCapture.h new file mode 100644 index 0000000..6a44372 --- /dev/null +++ b/dwio/nimble/encodings/EncodingLayoutCapture.h @@ -0,0 +1,19 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include "dwio/nimble/encodings/EncodingLayout.h" + +namespace facebook::nimble { + +class EncodingLayoutCapture { + public: + // Captures an encoding tree from an encoded stream. + // It traverses the encoding headers in the stream and produces a serialized + // encoding tree layout. + // |encoding| - The serialized encoding + static EncodingLayout capture(std::string_view encoding); +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/EncodingSelection.h b/dwio/nimble/encodings/EncodingSelection.h new file mode 100644 index 0000000..721cf12 --- /dev/null +++ b/dwio/nimble/encodings/EncodingSelection.h @@ -0,0 +1,245 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include + +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" +#include "dwio/nimble/encodings/EncodingIdentifier.h" +#include "dwio/nimble/encodings/Statistics.h" + +namespace facebook::nimble { + +// Overview: +// Nimble has many encoding to chose from. Each encoding yields different +// results on different shapes of data. In addition, Nimble's encodings are +// nested (e.g. encodings may contain other encodings, to encode sub streams of +// data). To be able to successfully encode an Nimble stream, we must build an +// encoding tree, optimized for the given data. Nimble is using a top-down +// (greedy) selection approach, where, by inspecting the top level data, we pick +// an encoding, and use it to encode the data. If this is a nested encoding, we +// then break down the data, and apply the same algorithm on each nested +// encoding. Once we pick an encoding for a specific level, we commit to this +// selection and never go back and change it. +// +// Encoding selection allows plugging in different policy implementations for +// selecting an encoding tree. +// Encoding selection provides access to a rich set of statistics on the data, +// which allows selection policies to apply smart decisions on which encoding +// fits this data best. +// +// To implement an encoding selection policy, one must inherit from +// EncodingSelectionPolicy and implement two methods: +// 1. select(): Given the data and its statistics, return a selected encoding +// (and compression policy) +// 2. createImpl(): Create a nested encoding selection policy to be used for a +// nested encoding +// +// During the encoding process, the encode() method of the selected encoding is +// given access to an EncodingSelection class. This class allows the encode() +// method to access/reuse the pre-calculated statistics, and also allows +// continuing the selection process on nested data streams, by invoking the +// encodeNested() method on the nested stream. +// Internally, encodeNested() is creating a child encoding selection policy, +// suitable for the nested stream type, and triggers the selection process and +// encoding for the nested stream. + +class EncodingSelectionPolicyBase; + +// +// Compression policy type definitions: +// A compression policy defines which compression algorithm to apply on the data +// (if any) and what parameters to use for this compression algorithm. +// In addition, once compression is applied to the data, the compression policy +// can decide if the compressed result is statisfactory or if it should be +// discarded. +struct ZstdCompressionParameters { + int16_t compressionLevel = 3; +}; + +struct ZstrongCompressionParameters { + int16_t compressionLevel = 0; + int16_t decompressionLevel = 0; + bool useVariableBitWidthCompressor = true; +}; + +union CompressionParameters { + ZstdCompressionParameters zstd; + ZstrongCompressionParameters zstrong; +}; + +struct CompressionInformation { + CompressionType compressionType; + CompressionParameters parameters; +}; + +class CompressionPolicy { + public: + virtual CompressionInformation compression() const = 0; + virtual bool shouldAccept( + CompressionType /* compressionType */, + uint64_t /* uncompressedSize */, + uint64_t /* compressedSize */) const = 0; + + virtual ~CompressionPolicy() = default; +}; + +// Default compression policy. Default behavior (if not compression policy is +// provided) is to not compress. +class NoCompressionPolicy : public CompressionPolicy { + public: + CompressionInformation compression() const override { + return {.compressionType = CompressionType::Uncompressed}; + } + + virtual bool shouldAccept( + CompressionType /* compressionType */, + uint64_t /* uncompressedSize */, + uint64_t /* compressedSize */) const override { + return false; + } +}; + +// Type representing a selected encoding. +// This is the result type returned from the select() method of an encoding +// selection policy. Also provides access to the compression policies for nested +// data streams. +struct EncodingSelectionResult { + EncodingType encodingType; + std::function()> compressionPolicyFactory = + []() { return std::make_unique(); }; +}; + +// The EncodingSelection class is passed in to the encode() method of each +// encoding. It provides access to the current encoding selection details, and +// allows triggering nested necodings of nested data streams. +template +class EncodingSelection { + public: + EncodingSelection( + EncodingSelectionResult selectionResult, + Statistics&& statistics, + std::unique_ptr&& selectionPolicy) + : selectionResult_{std::move(selectionResult)}, + statistics_{std::move(statistics)}, + selectionPolicy_{std::move(selectionPolicy)} {} + + EncodingType encodingType() const noexcept { + return selectionResult_.encodingType; + } + + const Statistics& statistics() const noexcept { + return statistics_; + } + + std::unique_ptr compressionPolicy() const noexcept { + auto policy = selectionResult_.compressionPolicyFactory(); + return policy; + } + + // This method encodes a nested data stream. This includes + // triggering a new encoding selection operation for the nested data and + // recursively encoding internal nested stream (if further nested encodings + // are selected). + template + std::string_view encodeNested( + NestedEncodingIdentifier identifier, + std::span values, + Buffer& buffer); + + private: + EncodingSelectionResult selectionResult_; + Statistics statistics_; + std::unique_ptr selectionPolicy_; +}; + +// HACK: When triggering encoding of a nested data stream, this data stream +// often has different data type than the parent (for example, when encoding +// string data, using dictionary encoding, the dictionary indices nested stream +// will most likely be integers, thus different from the parent string data +// type). This means, that we need to create a selection process, using a child +// encoding selection policy, but on a different template type. +// Since we don't have access to the nested type at compile time, we use this +// intermediate (non-templated) class, to allow constructing the correct +// templated encoding selection policy class, at runtime. +class EncodingSelectionPolicyBase { + public: + template + std::unique_ptr create( + EncodingType encodingType, + NestedEncodingIdentifier identifier) { + return createImpl(encodingType, identifier, TypeTraits::dataType); + } + + virtual ~EncodingSelectionPolicyBase() = default; + + protected: + // This method allows creating a child (nested) encoding selection policy + // instance, with a different data type than the parent instance. + // This method allows transferring state between parent and nested encoding + // selection policy instances. + // A child encoding selection policy instance will be created for every nested + // stream of the parent encoding. The identifier passed in allows to + // differentiate between the nested streams. + virtual std::unique_ptr createImpl( + EncodingType encodingType, + NestedEncodingIdentifier identifier, + DataType type) = 0; +}; + +// This is the main base class for defining an encoding selection policy. +template +class EncodingSelectionPolicy : public EncodingSelectionPolicyBase { + using physicalType = typename TypeTraits::physicalType; + + public: + // This is the main encoding selection method. + // Given the data, and its statistics, the encoding selection policy should + // return the selected encoding (and the matching compression policy). + virtual EncodingSelectionResult select( + std::span values, + const Statistics& statistics) = 0; + + // Same as the |select()| method above, but for nullable values. + virtual EncodingSelectionResult selectNullable( + std::span values, + std::span nulls, + const Statistics& statistics) = 0; + + virtual ~EncodingSelectionPolicy() = default; +}; + +namespace { +template +std::unique_ptr unique_ptr_cast(std::unique_ptr src) { + return std::unique_ptr(static_cast(src.release())); +} +} // namespace + +template +template +std::string_view EncodingSelection::encodeNested( + NestedEncodingIdentifier identifier, + std::span values, + Buffer& buffer) { + // Create the nested encoding selection policy instance, and cast it to the + // strongly templated type. + auto nestedPolicy = std::unique_ptr>( + static_cast*>( + selectionPolicy_->template create(encodingType(), identifier) + .release())); + auto statistics = Statistics::create(values); + auto selectionResult = nestedPolicy->select(values, statistics); + return EncodingFactory::encode( + EncodingSelection{ + std::move(selectionResult), + std::move(statistics), + std::move(nestedPolicy)}, + values, + buffer); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/EncodingSelectionPolicy.h b/dwio/nimble/encodings/EncodingSelectionPolicy.h new file mode 100644 index 0000000..6e7c569 --- /dev/null +++ b/dwio/nimble/encodings/EncodingSelectionPolicy.h @@ -0,0 +1,1060 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include +#include "dwio/nimble/common/Bits.h" +#include "dwio/nimble/common/EncodingType.h" +#include "dwio/nimble/encodings/EncodingIdentifier.h" +#include "dwio/nimble/encodings/EncodingLayout.h" +#include "dwio/nimble/encodings/EncodingSelection.h" + +namespace facebook::nimble { + +using EncodingSelectionPolicyFactory = + std::function(DataType)>; + +// The following enables encoding selection debug messages. By default, these +// logs are turned off (with zero overhead). In tests (or in debug sessions), we +// enable these logs. +#define RED "\033[31m" +#define GREEN "\033[32m" +#define YELLOW "\033[33m" +#define BLUE "\033[34m" +#define PURPLE "\033[35m" +#define CYAN "\033[36m" +#define RESET_COLOR "\033[0m" + +#ifndef NIMBLE_ENCODING_SELECTION_DEBUG_MAX_ITEMS +#define NIMBLE_ENCODING_SELECTION_DEBUG_MAX_ITEMS 50 +#endif + +#ifdef NIMBLE_ENCODING_SELECTION_DEBUG +#define NIMBLE_SELECTION_LOG(stream) LOG(INFO) << stream << RESET_COLOR +#else +#define NIMBLE_SELECTION_LOG(stream) +#endif + +#define COMMA , +#define UNIQUE_PTR_FACTORY_EXTRA(data_type, class, extra_types, ...) \ + switch (data_type) { \ + case ::facebook::nimble::DataType::Uint8: { \ + return std::make_unique::physicalType extra_types>>( \ + __VA_ARGS__); \ + } \ + case ::facebook::nimble::DataType::Int8: { \ + return std::make_unique::physicalType extra_types>>( \ + __VA_ARGS__); \ + } \ + case ::facebook::nimble::DataType::Uint16: { \ + return std::make_unique::physicalType extra_types>>(__VA_ARGS__); \ + } \ + case ::facebook::nimble::DataType::Int16: { \ + return std::make_unique::physicalType extra_types>>( \ + __VA_ARGS__); \ + } \ + case ::facebook::nimble::DataType::Uint32: { \ + return std::make_unique::physicalType extra_types>>(__VA_ARGS__); \ + } \ + case ::facebook::nimble::DataType::Int32: { \ + return std::make_unique::physicalType extra_types>>( \ + __VA_ARGS__); \ + } \ + case ::facebook::nimble::DataType::Uint64: { \ + return std::make_unique::physicalType extra_types>>(__VA_ARGS__); \ + } \ + case ::facebook::nimble::DataType::Int64: { \ + return std::make_unique::physicalType extra_types>>( \ + __VA_ARGS__); \ + } \ + case ::facebook::nimble::DataType::Float: { \ + return std::make_unique::physicalType extra_types>>( \ + __VA_ARGS__); \ + } \ + case ::facebook::nimble::DataType::Double: { \ + return std::make_unique::physicalType extra_types>>( \ + __VA_ARGS__); \ + } \ + case ::facebook::nimble::DataType::Bool: { \ + return std::make_unique::physicalType extra_types>>( \ + __VA_ARGS__); \ + } \ + case ::facebook::nimble::DataType::String: { \ + return std::make_unique::physicalType extra_types>>(__VA_ARGS__); \ + } \ + default: { \ + NIMBLE_UNREACHABLE( \ + fmt::format("Unsupported data type {}.", toString(data_type))); \ + } \ + } + +#define UNIQUE_PTR_FACTORY(data_type, class, ...) \ + UNIQUE_PTR_FACTORY_EXTRA(data_type, class, , __VA_ARGS__) + +struct CompressionOptions { + float compressionAcceptRatio = 0.98f; + uint32_t zstdCompressionLevel = 3; + uint32_t zstrongCompressionLevel = 4; + uint32_t zstrongDecompressionLevel = 2; + bool useVariableBitWidthCompressor = false; +}; + +// This is the manual encoding selection implementation. +// The manual selection is using a manually crafted model to choose the most +// appropriate encoding based on the provided statistics. +template +class ManualEncodingSelectionPolicy : public EncodingSelectionPolicy { + using physicalType = typename TypeTraits::physicalType; + + public: + ManualEncodingSelectionPolicy( + std::vector> readFactors, + CompressionOptions compressionOptions, + std::optional identifier) + : readFactors_{std::move(readFactors)}, + compressionOptions_{std::move(compressionOptions)}, + identifier_{identifier} {} + + std::unique_ptr createImpl( + EncodingType encodingType, + NestedEncodingIdentifier identifier, + DataType type) override { + // In each sub-level of the encoding selection, we exclude the encodings + // selected in parent levels. Although this is not required (as hopefully, + // the model will not pick a nested encoding of the same type as the + // parent), it provides an additional safty net, making sure the encoding + // selection will eventually converge, and also slightly speeds up nested + // encoding selection. + // TODO: validate the assumptions here compared to brute forcing, to see if + // the same encoding is selected multiple times in the tree (for example, + // should we allow trivial string lengths to be encoded using trivial + // encoding?) + std::vector> filteredReadFactors; + filteredReadFactors.reserve(readFactors_.size() - 1); + for (const auto& pair : readFactors_) { + if (pair.first != encodingType) { + filteredReadFactors.push_back(pair); + } + } + UNIQUE_PTR_FACTORY_EXTRA( + type, + ManualEncodingSelectionPolicy, + COMMA FixedByteWidth, + std::move(filteredReadFactors), + compressionOptions_, + identifier); + } + + EncodingSelectionResult select( + std::span values, + const Statistics& statistics) override { + if (values.empty()) { + return { + .encodingType = EncodingType::Trivial, + }; + } + + float minCost = std::numeric_limits::max(); + EncodingType selectedEncoding = EncodingType::Trivial; + constexpr size_t maxItems = NIMBLE_ENCODING_SELECTION_DEBUG_MAX_ITEMS; + + // Iterate on all candidate encodings, and pick the encoding with the + // minimal cost. + for (const auto& pair : readFactors_) { + auto encodingType = pair.first; + auto size = estimateSize(encodingType, values.size(), statistics); + if (!size.has_value()) { + NIMBLE_SELECTION_LOG( + PURPLE << encodingType << " encoding is incompatible."); + continue; + } + + // We use read factor weights to raise/lower the favorability of each + // encoding. + auto readFactor = pair.second; + auto cost = size.value() * readFactor; + NIMBLE_SELECTION_LOG( + YELLOW << "Encoding: " << encodingType << ", Size: " << size.value() + << ", Factor: " << readFactor << ", Cost: " << cost); + if (cost < minCost) { + minCost = cost; + selectedEncoding = encodingType; + } + } + + // Currently, we always attempt to compress leaf data streams (using + // Zstrong). The logic behind this is that encoding selection optimizes the + // in-memory layout of data, while compression provides extra saving for + // persistent storage (and bandwidth). + class AlwaysCompressZstrongPolicy : public CompressionPolicy { + public: + explicit AlwaysCompressZstrongPolicy( + CompressionOptions compressionOptions) + : compressionOptions_{std::move(compressionOptions)} {} + + CompressionInformation compression() const override { + CompressionInformation information{ + .compressionType = CompressionType::Zstrong}; + // See + // https://docs.google.com/spreadsheets/d/1tjkVol68s94_Z2R4ROp1kpOkMwlrbHI0Jlg9VpRnk9c/edit?usp=sharing + // for summary on why these compression levels are used. + information.parameters.zstrong.compressionLevel = + compressionOptions_.zstrongCompressionLevel; + information.parameters.zstrong.decompressionLevel = + compressionOptions_.zstrongDecompressionLevel; + information.parameters.zstrong.useVariableBitWidthCompressor = + compressionOptions_.useVariableBitWidthCompressor; + return information; + } + + virtual bool shouldAccept( + CompressionType compressionType, + uint64_t uncompressedSize, + uint64_t compressedSize) const override { + if (uncompressedSize * compressionOptions_.compressionAcceptRatio < + compressedSize) { + NIMBLE_SELECTION_LOG( + BLUE << compressionType + << " compression rejected. Original size: " + << uncompressedSize + << ", Compressed size: " << compressedSize); + return false; + } + + NIMBLE_SELECTION_LOG( + CYAN << compressionType + << " compression accepted. Original size: " << uncompressedSize + << ", Compressed size: " << compressedSize); + return true; + } + + private: + const CompressionOptions compressionOptions_; + }; + + NIMBLE_SELECTION_LOG( + "Selected Encoding" + << (identifier_.has_value() + ? folly::to( + " [NestedEncodingIdentifier: ", identifier_.value(), "]") + : "") + << ": " << GREEN << selectedEncoding << RESET_COLOR + << ", Sampled Data: " + << folly::join( + ",", + std::span{ + values.data(), std::min(maxItems, values.size())}) + << (values.size() > maxItems ? "..." : "")); + return { + .encodingType = selectedEncoding, + .compressionPolicyFactory = [compressionOptions = + compressionOptions_]() { + return std::make_unique( + compressionOptions); + }}; + } + + EncodingSelectionResult selectNullable( + std::span /* values */, + std::span /* nulls */, + const Statistics& /* statistics */) override { + return { + .encodingType = EncodingType::Nullable, + }; + } + + const std::vector>& readFactors() const { + return readFactors_; + } + + float compressionAcceptRatio() const { + return compressionOptions_.compressionAcceptRatio; + } + + private: + static inline uint32_t + bitPackedBytes(uint64_t minValue, uint64_t maxValue, uint32_t count) { + if constexpr (FixedByteWidth) { + // NOTE: We round up the bits required to the next byte boundary. + // This is to match a (temporary) hack we added to FixedBitWidthEncoding, + // to mitigare compression issue on non-rounded bit widths. See + // dwio/nimble/encodings/FixedBitWidthEncoding.h for full details. + return bits::bytesRequired( + ((bits::bitsRequired(maxValue - minValue) + 7) & ~7) * count); + } else { + return bits::bytesRequired( + bits::bitsRequired(maxValue - minValue) * count); + } + } + + static std::optional estimateNumericSize( + EncodingType encodingType, + uint64_t entryCount, + const Statistics& statistics) { + switch (encodingType) { + case EncodingType::Constant: { + return statistics.uniqueCounts().size() == 1 + ? std::optional{getEncodingOverhead< + EncodingType::Constant, + physicalType>()} + : std::nullopt; + } + case EncodingType::MainlyConstant: { + // Assumptions: + // We store one entry for the common value. + // Number of uncommon values is total item count minus the max unique + // count (the common value count). + // For each uncommon value we store the its value. We assume they will + // be stored bit-packed. + // We also store a bitmap for all rows (is-common bitmap). This bitmap + // will most likely be stored as SparseBool. Therefore, for each + // uncommon value there will be an index. These indices will most likely + // be stored bit-packed, with bit width of max(rowCount). + + // Find most common item count + const auto maxUniqueCount = std::max_element( + statistics.uniqueCounts().cbegin(), + statistics.uniqueCounts().cend(), + [](const auto& a, const auto& b) { return a.second < b.second; }); + // Deduce uncommon values count + const auto uncommonCount = entryCount - maxUniqueCount->second; + // Assuming the uncommon values will be stored bit packed. + const auto uncommonValueSize = + bitPackedBytes(statistics.min(), statistics.max(), uncommonCount); + // Uncommon (sparse bool) bitmap will have index per uncommon value, + // stored bit packed. + const auto uncommonIndicesSize = + bitPackedBytes(0, entryCount, uncommonCount); + uint32_t overhead = + getEncodingOverhead() + + // Overhead for storing uncommon values + getEncodingOverhead() + + // Overhead for storing uncommon bitmap + getEncodingOverhead() + + getEncodingOverhead(); + return overhead + sizeof(physicalType) + uncommonValueSize + + uncommonIndicesSize; + } + case EncodingType::Trivial: { + return getEncodingOverhead() + + (entryCount * sizeof(physicalType)); + } + case EncodingType::FixedBitWidth: { + return getEncodingOverhead< + EncodingType::FixedBitWidth, + physicalType>() + + bitPackedBytes(statistics.min(), statistics.max(), entryCount); + } + case EncodingType::Dictionary: { + // Assumptions: + // Alphabet stored trivially. + // Indices are stored bit-packed, with bit width needed to store max + // dictionary size (which is the unique value count). + const uint64_t indicesSize = + bitPackedBytes(0, statistics.uniqueCounts().size(), entryCount); + const uint64_t alphabetSize = + statistics.uniqueCounts().size() * sizeof(physicalType); + uint32_t overhead = + getEncodingOverhead() + + // Alphabet overhead + getEncodingOverhead() + + // Indices overhead + getEncodingOverhead(); + return overhead + alphabetSize + indicesSize; + } + case EncodingType::RLE: { + // Assumptions: + // Run values are stored bit-packed (with bit width needed to store max + // value). Run lengths are stored using bit-packing (with bit width + // needed to store max repetition count). + + const auto runValuesSize = bitPackedBytes( + statistics.min(), + statistics.max(), + statistics.consecutiveRepeatCount()); + const auto runLengthsSize = bitPackedBytes( + statistics.minRepeat(), + statistics.maxRepeat(), + statistics.consecutiveRepeatCount()); + uint32_t overhead = + getEncodingOverhead() + + // Overhead of run values + getEncodingOverhead() + + // Overhead of run lengths + getEncodingOverhead(); + return overhead + runValuesSize + runLengthsSize; + } + case EncodingType::Varint: { + // Note: the condition below actually support floating point numbers as + // well, as we use physicalType, which, for floating point numbers, is + // an integer. + if constexpr (isIntegralType() && sizeof(T) >= 4) { + // First (7 bit) bucket produces 1 byte number. Second bucket produce + // 2 byte number and so forth. + size_t i = 0; + const uint64_t dataSize = std::accumulate( + statistics.bucketCounts().cbegin(), + statistics.bucketCounts().cend(), + 0, + [&i](const uint64_t sum, const uint64_t bucketSize) { + return sum + (bucketSize * (++i)); + }); + return getEncodingOverhead() + + dataSize; + } else { + return std::nullopt; + } + } + default: { + return std::nullopt; + } + } + } + + static std::optional estimateBoolSize( + EncodingType encodingType, + size_t entryCount, + const Statistics& statistics) { + switch (encodingType) { + case EncodingType::Constant: { + return statistics.uniqueCounts().size() == 1 + ? std::optional{getEncodingOverhead< + EncodingType::Constant, + physicalType>()} + : std::nullopt; + } + case EncodingType::SparseBool: { + // Assumptions: + // Uncommon indices are stored bit-packed (with bit width capable of + // representing max entry count). + + const auto exceptionCount = std::min( + statistics.uniqueCounts().at(true), + statistics.uniqueCounts().at(false)); + uint32_t overhead = + getEncodingOverhead() + + // Overhead for storing exception indices + getEncodingOverhead(); + return overhead + sizeof(bool) + + bitPackedBytes(0, entryCount, exceptionCount); + } + case EncodingType::Trivial: { + return getEncodingOverhead() + + bitPackedBytes(0, 1, entryCount); + } + case EncodingType::RLE: { + // Assumptions: + // Run lengths are stored using bit-packing (with bit width + // needed to store max repetition count). + + const auto runLengthsSize = bitPackedBytes( + statistics.minRepeat(), + statistics.maxRepeat(), + statistics.consecutiveRepeatCount()); + uint32_t overhead = + getEncodingOverhead() + + // Overhead of run lengths + getEncodingOverhead(); + return overhead + sizeof(bool) + runLengthsSize; + } + default: { + return std::nullopt; + } + } + } + + static std::optional estimateStringSize( + EncodingType encodingType, + size_t entryCount, + const Statistics& statistics) { + const uint32_t maxStringSize = statistics.max().size(); + switch (encodingType) { + case EncodingType::Constant: { + return statistics.uniqueCounts().size() == 1 + ? std::optional{getEncodingOverhead< + EncodingType::Constant, + physicalType>(maxStringSize)} + : std::nullopt; + } + case EncodingType::MainlyConstant: { + // Assumptions: + // We store one entry for the common value. + // For each uncommon value we store the its value. + // Uncommon values will be stored trivially, or if there are + // repetitions, we assume they will be nested as a dictionary. Either + // way, it means each (unique) string value will be stored exactly once. + // (Note: for strings we store the blob size and the lengths, assuming + // bit-packed). + // We also store a bitmap for all rows (is-common bitmap). This bitmap + // will most likely be stored as SparseBool. Therefore, for each + // uncommon value there will be an index. These indices will most likely + // be stored bit-packed, with bit width of max(rowCount). + + // Find the most common item count + const auto maxUniqueCount = std::max_element( + statistics.uniqueCounts().cbegin(), + statistics.uniqueCounts().cend(), + [](const auto& a, const auto& b) { return a.second < b.second; }); + // Get the total blob size for all (unique) strings + const uint64_t alphabetByteSize = std::accumulate( + statistics.uniqueCounts().cbegin(), + statistics.uniqueCounts().cend(), + 0, + [](const uint32_t sum, const auto& unique) { + return sum + unique.first.size(); + }); + // Deduce uncommon values count + const auto uncommonCount = entryCount - maxUniqueCount->second; + // Remove common string from blob, and add the (bit-packed) string + // lengths to the size. + // Note: we remove the common string, as it is later added back when + // calculating the header overhead. + const uint64_t alphabetSize = alphabetByteSize - + maxUniqueCount->first.size() + + bitPackedBytes(statistics.min().size(), + statistics.max().size(), + statistics.uniqueCounts().size()); + // Uncommon values (sparse bool) bitmap will have index per value, + // stored bit packed. + const auto uncommonIndicesSize = + bitPackedBytes(0, entryCount, uncommonCount); + uint32_t overhead = + getEncodingOverhead( + maxUniqueCount->first.size()) + + // Overhead for storing uncommon values + getEncodingOverhead( + maxStringSize) + + // Overhead for storing uncommon bitmap + getEncodingOverhead(); + + return overhead + alphabetSize + uncommonIndicesSize; + } + case EncodingType::Trivial: { + // We assume string lengths will be stored bit packed. + return getEncodingOverhead( + maxStringSize) + + statistics.totalStringsLength() + + bitPackedBytes( + statistics.min().size(), + statistics.max().size(), + entryCount); + } + case EncodingType::Dictionary: { + // Assumptions: + // Alphabet stored trivially. + // Indices are stored bit-packed, with bit width needed to store max + // dictionary size (which is the unique value count). + const uint64_t indicesSize = + bitPackedBytes(0, statistics.uniqueCounts().size(), entryCount); + // Get the total blob size for all (unique) strings + const uint64_t alphabetByteSize = std::accumulate( + statistics.uniqueCounts().cbegin(), + statistics.uniqueCounts().cend(), + 0, + [](const uint32_t sum, const auto& unique) { + return sum + unique.first.size(); + }); + // Add (bit-packed) string lengths + const uint64_t alphabetSize = alphabetByteSize + + bitPackedBytes(statistics.min().size(), + statistics.max().size(), + statistics.uniqueCounts().size()); + uint32_t overhead = + getEncodingOverhead( + maxStringSize) + + // Alphabet overhead + getEncodingOverhead( + maxStringSize) + + // Indices overhead + getEncodingOverhead(); + return overhead + alphabetSize + indicesSize; + } + case EncodingType::RLE: { + // Assumptions: + // Run values are stored using dictionary (and inside, trivial + + // bit-packing). Run lengths are stored using bit-packing (with bit + // width needed to store max repetition count). + + uint64_t runValuesSize = + // (unique) strings blob size + std::accumulate( + statistics.uniqueCounts().cbegin(), + statistics.uniqueCounts().cend(), + 0, + [](const uint32_t sum, const auto& unique) { + return sum + unique.first.size(); + }) + + // (bit-packed) string lengths + bitPackedBytes( + statistics.min().size(), + statistics.max().size(), + statistics.consecutiveRepeatCount()) + + // dictionary indices + bitPackedBytes( + 0, + statistics.uniqueCounts().size(), + statistics.consecutiveRepeatCount()); + const auto runLengthsSize = bitPackedBytes( + statistics.minRepeat(), + statistics.maxRepeat(), + statistics.consecutiveRepeatCount()); + uint32_t overhead = + getEncodingOverhead() + + // Overhead of run values + getEncodingOverhead() + + getEncodingOverhead() + + getEncodingOverhead() + + // Overhead of run lengths + getEncodingOverhead(); + return overhead + runValuesSize + runLengthsSize; + } + default: { + return std::nullopt; + } + } + } + + static std::optional estimateSize( + EncodingType encodingType, + size_t entryCount, + const Statistics& statistics) { + if constexpr (isNumericType()) { + return estimateNumericSize(encodingType, entryCount, statistics); + } else if constexpr (isBoolType()) { + return estimateBoolSize(encodingType, entryCount, statistics); + } else if constexpr (isStringType()) { + return estimateStringSize(encodingType, entryCount, statistics); + } + + NIMBLE_UNREACHABLE(fmt::format( + "Unable to estimate size for type {}.", folly::demangle(typeid(T)))); + } + + template + static uint32_t getEncodingOverhead(uint32_t size = sizeof(dataType)) { + constexpr uint32_t commonPrefixSize = 6; + if constexpr (encodingType == EncodingType::Trivial) { + if constexpr (isStringType()) { + return commonPrefixSize + + // Trivial strings also have nested encoding for lengths. + getEncodingOverhead() + + 5; // CompressionType (1 byte), BlobOffset (4 bytes) + } else { + return commonPrefixSize + 1; // CompressionType (1 byte) + } + } else if constexpr (encodingType == EncodingType::Constant) { + if constexpr (isStringType()) { + return commonPrefixSize + size + + sizeof(uint32_t); // CommonValue (string bytes and string length) + } else { + return commonPrefixSize + sizeof(dataType); // CompressionType (1 byte) + } + } else if constexpr (encodingType == EncodingType::FixedBitWidth) { + return commonPrefixSize + 2 + + sizeof(dataType); // CompressionType (1 byte), Baseline (size of + // dataType), Bit Width (1 byte) + } else if constexpr (encodingType == EncodingType::MainlyConstant) { + if constexpr (isStringType()) { + return commonPrefixSize + (2 * sizeof(uint32_t)) + size + + sizeof(uint32_t); // IsCommon size (4 bytes), OtherValues size (4 + // bytes), Common value (string bytes and string + // length) + } else { + return commonPrefixSize + (2 * sizeof(uint32_t)) + + sizeof(dataType); // IsCommon size (4 bytes), OtherValues size (4 + // bytes), Common value (size of dataType) + } + } else if constexpr (encodingType == EncodingType::SparseBool) { + return commonPrefixSize + 1; // IsSet flag (1 byte) + } else if constexpr (encodingType == EncodingType::Dictionary) { + return commonPrefixSize + sizeof(uint32_t); // Alphabet size (4 byte) + } else if constexpr (encodingType == EncodingType::RLE) { + return commonPrefixSize + sizeof(uint32_t); // Run Lengths size (4 byte) + } else if constexpr (encodingType == EncodingType::Varint) { + return commonPrefixSize + sizeof(dataType); // Baseline (size of dataType) + } else if constexpr (encodingType == EncodingType::Nullable) { + return commonPrefixSize + sizeof(uint32_t); // Data size (4 bytes) + } + } + + // These read factors are used to raise or lower the chances of an encoding to + // be picked. Right now, these represent mostly the cpu cost to decode the + // values. Trivial is being boosed as we factor in the added benefit of + // applying compression. + // See ManualEncodingSelectionPolicyFactory::defaultReadFactors for the + // default. + const std::vector> readFactors_; + const CompressionOptions compressionOptions_; + const std::optional identifier_; +}; + +// This model is trained offline and parameters are updated here. +// The parameters are relatively robust and do not need to be updated +// unless we add / remove an encoding type. +template +struct EncodingPredictionModel { + using physicalType = typename TypeTraits::physicalType; + explicit EncodingPredictionModel() + : maxRepeatParam(1.52), minRepeatParam(1.13), uniqueParam(2.589) {} + float maxRepeatParam; + float minRepeatParam; + float uniqueParam; + float predict(const Statistics& statistics) { + // TODO: Utilize more features within statistics for prediction. + auto maxRepeat = statistics.maxRepeat(); + auto minRepeat = statistics.minRepeat(); + auto unique = statistics.uniqueCounts().size(); + return maxRepeatParam * maxRepeat + minRepeatParam * minRepeat + + uniqueParam * unique; + } +}; + +class ManualEncodingSelectionPolicyFactory { + public: + ManualEncodingSelectionPolicyFactory( + std::vector> readFactors = + defaultReadFactors(), + CompressionOptions compressionOptions = {}) + : readFactors_{std::move(readFactors)}, + compressionOptions_{std::move(compressionOptions)} {} + + std::unique_ptr createPolicy( + DataType dataType) const { + UNIQUE_PTR_FACTORY( + dataType, + ManualEncodingSelectionPolicy, + readFactors_, + compressionOptions_, + std::nullopt); + } + + static std::vector possibleEncodings() { + return { + EncodingType::Constant, + EncodingType::Trivial, + EncodingType::FixedBitWidth, + EncodingType::MainlyConstant, + EncodingType::SparseBool, + EncodingType::Dictionary, + EncodingType::RLE, + EncodingType::Varint, + }; + } + + static std::vector> defaultReadFactors() { + return { + {EncodingType::Constant, 1.0}, + {EncodingType::Trivial, 0.7}, + {EncodingType::FixedBitWidth, 0.9}, + {EncodingType::MainlyConstant, 1.0}, + {EncodingType::SparseBool, 1.0}, + {EncodingType::Dictionary, 1.0}, + {EncodingType::RLE, 1.0}, + {EncodingType::Varint, 1.0}, + }; + } + + static std::vector> parseReadFactors( + const std::string& val) { + std::vector> readFactors; + std::vector parts; + folly::split(';', folly::trimWhitespace(val), parts); + readFactors.reserve(parts.size()); + + std::vector possibleEncodings = + nimble::ManualEncodingSelectionPolicyFactory::possibleEncodings(); + std::vector possibleEncodingStrings; + possibleEncodingStrings.reserve(possibleEncodings.size()); + std::transform( + possibleEncodings.cbegin(), + possibleEncodings.cend(), + std::back_inserter(possibleEncodingStrings), + [](const auto& encoding) { return toString(encoding); }); + for (const auto& part : parts) { + std::vector kv; + folly::split('=', part, kv); + NIMBLE_CHECK( + kv.size() == 2, + fmt::format( + "Invalid read factor format. " + "Expected format is =;=. " + "Unable to parse '{}'.", + part)); + auto value = folly::tryTo(kv[1]); + NIMBLE_CHECK( + value.hasValue(), + fmt::format( + "Unable to parse read factor value '{}' in '{}'. Expected valid float value.", + kv[1], + part)); + bool found = false; + auto key = folly::trimWhitespace(kv[0]); + for (auto i = 0; i < possibleEncodings.size(); ++i) { + auto encoding = possibleEncodings[i]; + // @lint-ignore CLANGTIDY facebook-hte-LocalUncheckedArrayBounds + if (key == possibleEncodingStrings[i]) { + found = true; + readFactors.emplace_back(encoding, value.value()); + break; + } + } + NIMBLE_CHECK( + found, + fmt::format( + "Unknown or unexpected read factor encoding '{}'. Allowed values: {}", + key, + folly::join(",", possibleEncodingStrings))); + } + return readFactors; + } + + private: + const std::vector> readFactors_; + const CompressionOptions compressionOptions_; +}; + +// This is a learned encoding selection implementation. +template +class LearnedEncodingSelectionPolicy : public EncodingSelectionPolicy { + using physicalType = typename TypeTraits::physicalType; + + public: + LearnedEncodingSelectionPolicy( + std::vector encodingChoices, + std::optional identifier, + EncodingPredictionModel encodingModel = EncodingPredictionModel()) + : encodingChoices_{std::move(encodingChoices)}, + identifier_{identifier}, + mlModel_{encodingModel} {} + + LearnedEncodingSelectionPolicy() + : LearnedEncodingSelectionPolicy{ + possibleEncodingChoices(), + std::nullopt, + EncodingPredictionModel()} {} + + EncodingSelectionResult select( + std::span values, + const Statistics& statistics) override { + if (values.empty()) { + return { + .encodingType = EncodingType::Trivial, + }; + } + + float prediction = mlModel_.predict(statistics); + if (prediction > 0.1) { + return { + .encodingType = EncodingType::Trivial, + }; + } + // TODO: Implement a multi-class Encoding model so that we can predict not + // only a trivial encoding but also other encodings. + + return { + .encodingType = EncodingType::Trivial, + }; + } + + EncodingSelectionResult selectNullable( + std::span /* values */, + std::span /* nulls */, + const Statistics& /* statistics */) override { + return { + .encodingType = EncodingType::Nullable, + }; + } + + std::unique_ptr createImpl( + EncodingType encodingType, + NestedEncodingIdentifier identifier, + DataType type) override { + // In each sub-level of the encoding selection, we exclude the encodings + // selected in parent levels. Although this is not required (as hopefully, + // the model will not pick a nested encoding of the same type as the + // parent), it provides an additional safty net, making sure the encoding + // selection will eventually converge, and also slightly speeds up nested + // encoding selection. + // TODO: validate the assumptions here compared to brute forcing, to see if + // the same encoding is selected multiple times in the tree (for example, + // should we allow trivial string lengths to be encoded using trivial + // encoding?) + std::vector filteredReadFactors; + filteredReadFactors.reserve(encodingChoices_.size() - 1); + for (const auto& encodingType_ : encodingChoices_) { + if (encodingType_ != encodingType) { + filteredReadFactors.push_back(encodingType_); + } + } + UNIQUE_PTR_FACTORY( + type, + LearnedEncodingSelectionPolicy, + std::move(filteredReadFactors), + identifier); + } + + private: + static std::vector possibleEncodingChoices() { + return { + EncodingType::Constant, + EncodingType::Trivial, + EncodingType::FixedBitWidth, + EncodingType::MainlyConstant, + EncodingType::SparseBool, + EncodingType::Dictionary, + EncodingType::RLE, + EncodingType::Varint, + }; + } + + std::vector encodingChoices_; + std::optional identifier_; + EncodingPredictionModel mlModel_; +}; + +class ReplayedCompressionPolicy : public nimble::CompressionPolicy { + public: + explicit ReplayedCompressionPolicy( + nimble::CompressionType compressionType, + CompressionOptions compressionOptions) + : compressionType_{compressionType}, + compressionOptions_{std::move(compressionOptions)} {} + + nimble::CompressionInformation compression() const override { + if (compressionType_ == nimble::CompressionType::Uncompressed) { + return {.compressionType = nimble::CompressionType::Uncompressed}; + } + + if (compressionType_ == nimble::CompressionType::Zstd) { + nimble::CompressionInformation information{ + .compressionType = nimble::CompressionType::Zstd}; + information.parameters.zstd.compressionLevel = + compressionOptions_.zstdCompressionLevel; + return information; + } + + nimble::CompressionInformation information{ + .compressionType = nimble::CompressionType::Zstrong}; + information.parameters.zstrong.compressionLevel = + compressionOptions_.zstrongCompressionLevel; + information.parameters.zstrong.decompressionLevel = + compressionOptions_.zstrongDecompressionLevel; + information.parameters.zstrong.useVariableBitWidthCompressor = + compressionOptions_.useVariableBitWidthCompressor; + return information; + } + + virtual bool shouldAccept( + nimble::CompressionType /* compressionType */, + uint64_t uncompressedSize, + uint64_t compressedSize) const override { + if (uncompressedSize * compressionOptions_.compressionAcceptRatio < + compressedSize) { + NIMBLE_SELECTION_LOG( + BLUE << compressionType_ << " compression rejected. Original size: " + << uncompressedSize << ", Compressed size: " << compressedSize); + return false; + } + + NIMBLE_SELECTION_LOG( + CYAN << compressionType_ << " compression accepted. Original size: " + << uncompressedSize << ", Compressed size: " << compressedSize); + return true; + } + + private: + const nimble::CompressionType compressionType_; + const CompressionOptions compressionOptions_; +}; + +template +class ReplayedEncodingSelectionPolicy; + +template +class ReplayedEncodingSelectionPolicy + : public nimble::EncodingSelectionPolicy { + using physicalType = typename nimble::TypeTraits::physicalType; + + public: + ReplayedEncodingSelectionPolicy( + EncodingLayout encodingLayout, + CompressionOptions compressionOptions, + const EncodingSelectionPolicyFactory& encodingSelectionPolicyFactory) + : encodingLayout_{encodingLayout}, + compressionOptions_{std::move(compressionOptions)}, + encodingSelectionPolicyFactory_{encodingSelectionPolicyFactory} {} + + nimble::EncodingSelectionResult select( + std::span /* values */, + const nimble::Statistics& /* statistics */) override { + NIMBLE_SELECTION_LOG( + CYAN << "Replaying encoding " << encodingLayout_.encodingType()); + return { + .encodingType = encodingLayout_.encodingType(), + .compressionPolicyFactory = [this]() { + return std::make_unique( + encodingLayout_.compressionType(), compressionOptions_); + }}; + } + + EncodingSelectionResult selectNullable( + std::span /* values */, + std::span /* nulls */, + const Statistics& /* statistics */) override { + NIMBLE_SELECTION_LOG( + CYAN << "Replaying nullable encoding " + << encodingLayout_.encodingType()); + encodingLayout_ = EncodingLayout{ + EncodingType::Nullable, + CompressionType::Uncompressed, + { + /* Data */ std::move(encodingLayout_), + /* Nulls */ std::nullopt, + }}; + return { + .encodingType = EncodingType::Nullable, + }; + } + + std::unique_ptr createImpl( + nimble::EncodingType /* encodingType */, + nimble::NestedEncodingIdentifier identifier, + nimble::DataType type) override { + NIMBLE_ASSERT( + identifier < encodingLayout_.childrenCount(), + "Sub-encoding identifier out of range."); + auto child = encodingLayout_.child(identifier); + + if (child.has_value()) { + UNIQUE_PTR_FACTORY( + type, + ReplayedEncodingSelectionPolicy, + child.value(), + compressionOptions_, + encodingSelectionPolicyFactory_); + } else { + return encodingSelectionPolicyFactory_(type); + } + } + + private: + EncodingLayout encodingLayout_; + const CompressionOptions compressionOptions_; + const EncodingSelectionPolicyFactory& encodingSelectionPolicyFactory_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/FixedBitWidthEncoding.h b/dwio/nimble/encodings/FixedBitWidthEncoding.h new file mode 100644 index 0000000..1f96501 --- /dev/null +++ b/dwio/nimble/encodings/FixedBitWidthEncoding.h @@ -0,0 +1,216 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include + +#include "dwio/nimble/common/Bits.h" +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/EncodingType.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/FixedBitArray.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/encodings/Compression.h" +#include "dwio/nimble/encodings/Encoding.h" + +// The FixedBitWidthEncoding stores integer data in a fixed number of +// bits equal to the number of bits required to represent the largest value in +// the encoding. For now we only support encoding non-negative values, but +// we may later add an optional zigzag encoding that will let us handle +// negatives. + +namespace facebook::nimble { + +// Data layout is: +// Encoding::kPrefixSize bytes: standard Encoding prefix +// 1 byte: compression type +// sizeof(T) byte: baseline value +// 1 byte: bit width +// FixedBitArray::BufferSize(rowCount, bit_width) bytes: packed values. +template +class FixedBitWidthEncoding final + : public TypedEncoding::physicalType> { + public: + using cppDataType = T; + using physicalType = typename TypeTraits::physicalType; + + static const int kCompressionOffset = Encoding::kPrefixSize; + static const int kPrefixSize = 2 + sizeof(T); + + FixedBitWidthEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data); + + void reset() final; + void skip(uint32_t rowCount) final; + void materialize(uint32_t rowCount, void* buffer) final; + + static std::string_view encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer); + + std::string debugString(int offset) const final; + + private: + int bitWidth_; + physicalType baseline_; + FixedBitArray fixedBitArray_; + uint32_t row_ = 0; + Vector uncompressedData_; + Vector buffer_; +}; + +// +// End of public API. Implementations follow. +// + +template +FixedBitWidthEncoding::FixedBitWidthEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data) + : TypedEncoding{memoryPool, data}, + uncompressedData_{&memoryPool}, + buffer_{&memoryPool} { + auto pos = data.data() + kCompressionOffset; + auto compressionType = static_cast(encoding::readChar(pos)); + baseline_ = encoding::read(pos); + bitWidth_ = static_cast(encoding::readChar(pos)); + if (compressionType != CompressionType::Uncompressed) { + uncompressedData_ = Compression::uncompress( + memoryPool, + compressionType, + {pos, static_cast(data.end() - pos)}); + fixedBitArray_ = FixedBitArray{ + {uncompressedData_.data(), uncompressedData_.size()}, bitWidth_}; + } else { + fixedBitArray_ = + FixedBitArray{{pos, static_cast(data.end() - pos)}, bitWidth_}; + } +} + +template +void FixedBitWidthEncoding::reset() { + row_ = 0; +} + +template +void FixedBitWidthEncoding::skip(uint32_t rowCount) { + row_ += rowCount; +} + +template +void FixedBitWidthEncoding::materialize(uint32_t rowCount, void* buffer) { + if constexpr (isFourByteIntegralType()) { + fixedBitArray_.bulkGetWithBaseline32( + row_, rowCount, static_cast(buffer), baseline_); + } else { + if (sizeof(physicalType) == 8 && bitWidth_ <= 32) { + fixedBitArray_.bulkGetWithBaseline32Into64( + row_, rowCount, static_cast(buffer), baseline_); + } else { + const uint32_t start = row_; + const uint32_t end = start + rowCount; + physicalType* output = static_cast(buffer); + for (uint32_t i = start; i < end; ++i) { + *output++ = fixedBitArray_.get(i) + baseline_; + } + } + } + row_ += rowCount; +} + +template +std::string_view FixedBitWidthEncoding::encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer) { + static_assert( + std::is_same_v< + typename std::make_unsigned::type, + physicalType>, + "Physical type must be unsigned."); + const uint32_t rowCount = values.size(); + // NOTE: If we end up with bitsRequired not being a multiple of 8, it degrades + // compression efficiency a lot. To (temporarily) mitigate this, we revert to + // using "FixedByteWidth", meaning that we round the required bitness to the + // closest byte. This is a temporary solution. Going forward we should + // evaluate other options. For example: + // 1. Asymetric bit width, where we don't encode the data during write + // (instead we use Trivial + Compression), but we encode it when we read (for + // better in-memory representation). + // 2. Apply compression only if bit width is a multiple of 8 + // 3. Try both bit width and byte width and pick one. + // 4. etc... + const int bitsRequired = + (bits::bitsRequired( + selection.statistics().max() - selection.statistics().min()) + + 7) & + ~7; + + const uint32_t fixedBitArraySize = + FixedBitArray::bufferSize(values.size(), bitsRequired); + + Vector vector{&buffer.getMemoryPool()}; + + auto dataCompressionPolicy = selection.compressionPolicy(); + CompressionEncoder compressionEncoder{ + buffer.getMemoryPool(), + *dataCompressionPolicy, + DataType::Undefined, + bitsRequired, + fixedBitArraySize, + [&]() { + vector.resize(fixedBitArraySize); + return std::span{vector}; + }, + [&, baseline = selection.statistics().min()](char*& pos) { + memset(pos, 0, fixedBitArraySize); + FixedBitArray fba(pos, bitsRequired); + if constexpr (sizeof(physicalType) == 4) { + fba.bulkSet32WithBaseline( + 0, + rowCount, + reinterpret_cast(values.data()), + baseline); + } else { + // TODO: We may want to support 32-bit mode with (u)int64 here as + // well. + for (uint32_t i = 0; i < values.size(); ++i) { + fba.set(i, values[i] - baseline); + } + } + pos += fixedBitArraySize; + return pos; + }}; + + const uint32_t encodingSize = Encoding::kPrefixSize + + FixedBitWidthEncoding::kPrefixSize + compressionEncoder.getSize(); + char* reserved = buffer.reserve(encodingSize); + char* pos = reserved; + Encoding::serializePrefix( + EncodingType::FixedBitWidth, TypeTraits::dataType, rowCount, pos); + encoding::writeChar( + static_cast(compressionEncoder.compressionType()), pos); + encoding::write(selection.statistics().min(), pos); + encoding::writeChar(bitsRequired, pos); + compressionEncoder.write(pos); + + NIMBLE_DASSERT(encodingSize == pos - reserved, "Encoding size mismatch."); + return {reserved, encodingSize}; +} + +template +std::string FixedBitWidthEncoding::debugString(int offset) const { + return fmt::format( + "{}{}<{}> rowCount={} bit_width={}", + std::string(offset, ' '), + toString(Encoding::encodingType()), + toString(Encoding::dataType()), + Encoding::rowCount(), + bitWidth_); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/MainlyConstantEncoding.h b/dwio/nimble/encodings/MainlyConstantEncoding.h new file mode 100644 index 0000000..9a53336 --- /dev/null +++ b/dwio/nimble/encodings/MainlyConstantEncoding.h @@ -0,0 +1,238 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/EncodingType.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" +#include "dwio/nimble/encodings/EncodingIdentifier.h" +#include "dwio/nimble/encodings/EncodingSelection.h" +#include "velox/common/memory/Memory.h" + +// Encodes data that is 'mainly' a single value by using a bool child vectors +// to mark the rows that are that value, and another child encoding to encode +// the other values. We don't actually require that the single value be any +// given fraction of the data, but generally the encoding is effective when that +// constant fraction is large (say 50%+). All data types except bool are +// supported. +// +// E.g. if the data is +// 1 1 2 1 1 1 3 1 1 9 +// +// The single value would be 1 +// +// The bool vector would be +// T T F T T T F T T F +// +// The other-values child encoding would be +// 2 3 9 +// +// This construction is quite similar to that of NullableEncoding, but instead +// of marking nulls specially we mark the constant value. + +namespace facebook::nimble { + +// Data layout is: +// Encoding::kPrefixSize bytes: standard Encoding prefix +// 4 bytes: num isCommon encoding bytes (X) +// X bytes: isCommon encoding bytes +// 4 bytes: num otherValues encoding bytes (Y) +// Y bytes: otherValues encoding bytes +// Z bytes: the constant value via encoding primitive. +template +class MainlyConstantEncoding final + : public TypedEncoding::physicalType> { + public: + using cppDataType = T; + using physicalType = typename TypeTraits::physicalType; + + MainlyConstantEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data); + + void reset() final; + void skip(uint32_t rowCount) final; + void materialize(uint32_t rowCount, void* buffer) final; + + static std::string_view encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer); + + std::string debugString(int offset) const final; + + private: + std::unique_ptr isCommon_; + std::unique_ptr otherValues_; + physicalType commonValue_; + // Temporary bufs. + Vector isCommonBuffer_; + Vector otherValuesBuffer_; +}; + +// +// End of public API. Implementation follows. +// + +template +MainlyConstantEncoding::MainlyConstantEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data) + : TypedEncoding(memoryPool, data), + isCommonBuffer_(&memoryPool), + otherValuesBuffer_(&memoryPool) { + const char* pos = data.data() + Encoding::kPrefixSize; + const uint32_t isCommonBytes = encoding::readUint32(pos); + isCommon_ = EncodingFactory::decode(this->memoryPool_, {pos, isCommonBytes}); + pos += isCommonBytes; + const uint32_t otherValuesBytes = encoding::readUint32(pos); + otherValues_ = + EncodingFactory::decode(this->memoryPool_, {pos, otherValuesBytes}); + pos += otherValuesBytes; + commonValue_ = encoding::read(pos); + NIMBLE_CHECK(pos == data.end(), "Unexpected mainly constant encoding end"); +} + +template +void MainlyConstantEncoding::reset() { + isCommon_->reset(); + otherValues_->reset(); +} + +template +void MainlyConstantEncoding::skip(uint32_t rowCount) { + // Hrm this isn't ideal. We should return to this later -- a new + // encoding func? Encoding::Accumulate to add up next N rows? + isCommonBuffer_.resize(rowCount); + isCommon_->materialize(rowCount, isCommonBuffer_.data()); + const uint32_t commonCount = + std::accumulate(isCommonBuffer_.begin(), isCommonBuffer_.end(), 0U); + const uint32_t nonCommonCount = rowCount - commonCount; + if (nonCommonCount == 0) { + return; + } + + otherValues_->skip(nonCommonCount); +} + +template +void MainlyConstantEncoding::materialize(uint32_t rowCount, void* buffer) { + // This too isn't ideal. We will want an Encoding::Indices method or + // something our SparseBool can use, giving back just the set indices + // rather than a materialization. + isCommonBuffer_.resize(rowCount); + isCommon_->materialize(rowCount, isCommonBuffer_.data()); + const uint32_t commonCount = + std::accumulate(isCommonBuffer_.begin(), isCommonBuffer_.end(), 0U); + const uint32_t nonCommonCount = rowCount - commonCount; + + if (nonCommonCount == 0) { + physicalType* output = static_cast(buffer); + std::fill(output, output + rowCount, commonValue_); + return; + } + + otherValuesBuffer_.reserve(nonCommonCount); + otherValues_->materialize(nonCommonCount, otherValuesBuffer_.data()); + physicalType* output = static_cast(buffer); + const physicalType* nextOtherValue = otherValuesBuffer_.begin(); + // This is a generic scatter -- should we have a common scatter func? + for (uint32_t i = 0; i < rowCount; ++i) { + if (isCommonBuffer_[i]) { + *output++ = commonValue_; + } else { + *output++ = *nextOtherValue++; + } + } + NIMBLE_DASSERT( + nextOtherValue - otherValuesBuffer_.begin() == nonCommonCount, + "Encoding size mismatch."); +} + +namespace internal {} // namespace internal + +template +std::string_view MainlyConstantEncoding::encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer) { + if (values.empty()) { + NIMBLE_INCOMPATIBLE_ENCODING("MainlyConstantEncoding cannot be empty."); + } + + const auto commonElement = std::max_element( + selection.statistics().uniqueCounts().cbegin(), + selection.statistics().uniqueCounts().cend(), + [](const auto& a, const auto& b) { return a.second < b.second; }); + + const uint32_t entryCount = values.size(); + const uint32_t uncommonCount = entryCount - commonElement->second; + + Vector isCommon{&buffer.getMemoryPool(), values.size(), true}; + Vector otherValues(&buffer.getMemoryPool()); + otherValues.reserve(uncommonCount); + + physicalType commonValue = commonElement->first; + for (auto i = 0; i < values.size(); ++i) { + physicalType currentValue = values[i]; + if (currentValue != commonValue) { + isCommon[i] = false; + otherValues.push_back(std::move(currentValue)); + } + } + + Buffer tempBuffer{buffer.getMemoryPool()}; + std::string_view serializedIsCommon = selection.template encodeNested( + EncodingIdentifiers::MainlyConstant::IsCommon, isCommon, tempBuffer); + std::string_view serializedOtherValues = + selection.template encodeNested( + EncodingIdentifiers::MainlyConstant::OtherValues, + otherValues, + tempBuffer); + + uint32_t encodingSize = Encoding::kPrefixSize + 8 + + serializedIsCommon.size() + serializedOtherValues.size(); + if constexpr (isNumericType()) { + encodingSize += sizeof(physicalType); + } else { + encodingSize += 4 + commonValue.size(); + } + char* reserved = buffer.reserve(encodingSize); + char* pos = reserved; + Encoding::serializePrefix( + EncodingType::MainlyConstant, TypeTraits::dataType, entryCount, pos); + // TODO: Reorder these so that metadata is at the beginning. + encoding::writeString(serializedIsCommon, pos); + encoding::writeString(serializedOtherValues, pos); + encoding::write(commonValue, pos); + NIMBLE_DASSERT(pos - reserved == encodingSize, "Encoding size mismatch."); + return {reserved, encodingSize}; +} + +template +std::string MainlyConstantEncoding::debugString(int offset) const { + std::string log = fmt::format( + "{}{}<{}> rowCount={} commonValue={}", + std::string(offset, ' '), + toString(Encoding::encodingType()), + toString(Encoding::dataType()), + Encoding::rowCount(), + commonValue_); + log += fmt::format( + "\n{}isCommon child:\n{}", + std::string(offset + 2, ' '), + isCommon_->debugString(offset + 4)); + log += fmt::format( + "\n{}otherValues child:\n{}", + std::string(offset + 2, ' '), + otherValues_->debugString(offset + 4)); + return log; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/NullableEncoding.h b/dwio/nimble/encodings/NullableEncoding.h new file mode 100644 index 0000000..6ab4d63 --- /dev/null +++ b/dwio/nimble/encodings/NullableEncoding.h @@ -0,0 +1,259 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include "dwio/nimble/common/Bits.h" +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/EncodingType.h" +#include "dwio/nimble/common/FixedBitArray.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" +#include "dwio/nimble/encodings/EncodingIdentifier.h" +#include "dwio/nimble/encodings/EncodingSelection.h" + +// A nullable encoding holds a subencoding of non-null values and another +// subencoding of booleans representing whether each row was null. + +namespace facebook::nimble { + +// Data layout is: +// Encoding::kPrefixSize bytes: standard Encoding prefix +// 4 bytes: non-null child encoding size (X) +// X bytes: non-null child encoding bytes +// Y byes: null child encoding bytes +template +class NullableEncoding final + : public TypedEncoding::physicalType> { + public: + using cppDataType = T; + using physicalType = typename TypeTraits::physicalType; + + NullableEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data); + + uint32_t nullCount() const final; + bool isNullable() const final; + + void reset() final; + void skip(uint32_t rowCount) final; + void materialize(uint32_t rowCount, void* buffer) final; + uint32_t materializeNullable( + uint32_t rowCount, + void* buffer, + std::function nulls, + const bits::Bitmap* scatterBitmap = nullptr, + uint32_t offset = 0) final; + + static std::string_view encodeNullable( + EncodingSelection& selection, + std::span values, + std::span nulls, + Buffer& buffer); + + std::string debugString(int offset) const final; + + private: + // One bit for each row. A true bit represents a row with a non-null value. + const char* bitmap_; + std::unique_ptr nonNulls_; + std::unique_ptr nulls_; + uint32_t row_ = 0; + Vector indicesBuffer_; // Temporary buffer. + Vector charBuffer_; // Another temporary buffer. + Vector boolBuffer_; // Yet another temporary buffer. +}; + +// +// End of public API. Implementations follow. +// + +template +NullableEncoding::NullableEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data) + : TypedEncoding(memoryPool, data), + indicesBuffer_(&memoryPool), + charBuffer_(&memoryPool), + boolBuffer_(&memoryPool) { + const char* pos = data.data() + Encoding::kPrefixSize; + const uint32_t nonNullsBytes = encoding::readUint32(pos); + nonNulls_ = EncodingFactory::decode(this->memoryPool_, {pos, nonNullsBytes}); + pos += nonNullsBytes; + nulls_ = EncodingFactory::decode( + this->memoryPool_, {pos, static_cast(data.end() - pos)}); + NIMBLE_DASSERT( + Encoding::rowCount() == nulls_->rowCount(), "Nulls count mismatch."); +} + +template +uint32_t NullableEncoding::nullCount() const { + return nulls_->rowCount() - nonNulls_->rowCount(); +} + +template +bool NullableEncoding::isNullable() const { + return true; +} + +template +void NullableEncoding::reset() { + row_ = 0; + nonNulls_->reset(); + nulls_->reset(); +} + +template +void NullableEncoding::skip(uint32_t rowCount) { + // Hrm this isn't ideal. We should return to this later -- a new + // encoding func? Encoding::Accumulate to add up next N rows? + boolBuffer_.resize(rowCount); + nulls_->materialize(rowCount, boolBuffer_.data()); + const uint32_t nonNullCount = + std::accumulate(boolBuffer_.begin(), boolBuffer_.end(), 0U); + nonNulls_->skip(nonNullCount); +} + +template +void NullableEncoding::materialize(uint32_t rowCount, void* buffer) { + // This too isn't ideal. We will want an Encoding::Indices method or + // something our SparseBool can use, giving back just the set indices + // rather than a materialization. + boolBuffer_.resize(rowCount); + nulls_->materialize(rowCount, boolBuffer_.data()); + const uint32_t nonNullCount = + std::accumulate(boolBuffer_.begin(), boolBuffer_.end(), 0U); + nonNulls_->materialize(nonNullCount, buffer); + + if (nonNullCount != rowCount) { + physicalType* output = static_cast(buffer) + rowCount - 1; + const physicalType* lastNonNull = + static_cast(buffer) + nonNullCount - 1; + // This is a generic scatter -- should we have a common scatter func? + uint32_t pos = rowCount - 1; + while (output != lastNonNull) { + if (boolBuffer_[pos]) { + *output = *lastNonNull; + --lastNonNull; + } else { + *output = physicalType(); + } + --output; + --pos; + } + } + + row_ += rowCount; +} + +template +uint32_t NullableEncoding::materializeNullable( + uint32_t rowCount, + void* buffer, + std::function nulls, + const bits::Bitmap* scatterBitmap, + uint32_t offset) { + boolBuffer_.resize(rowCount); + nulls_->materialize(rowCount, boolBuffer_.data()); + const uint32_t nonNullCount = + std::accumulate(boolBuffer_.begin(), boolBuffer_.end(), 0U); + + if (offset > 0) { + buffer = static_cast(buffer) + offset; + } + nonNulls_->materialize(nonNullCount, buffer); + + auto scatterSize = scatterBitmap ? scatterBitmap->size() - offset : rowCount; + if (nonNullCount != scatterSize) { + void* nullBitmap = nulls(); + bits::BitmapBuilder nullBits{nullBitmap, offset + scatterSize}; + nullBits.clear(offset, offset + scatterSize); + + uint32_t pos = offset + scatterSize - 1; + physicalType* output = static_cast(buffer) + scatterSize - 1; + const physicalType* lastNonNull = + static_cast(buffer) + nonNullCount - 1; + auto nonNullIt = boolBuffer_.begin() + rowCount - 1; + + if (scatterSize != rowCount) { + // In scattered reads, spread the items into the right positions in + // |buffer| and |nullBitmap| based on the bits set to 1 in + // |scatterBitmap|. + while (output != lastNonNull) { + if (scatterBitmap->test(pos)) { + if (*nonNullIt--) { + nullBits.set(pos); + *output = *lastNonNull; + --lastNonNull; + } + } + --output; + --pos; + } + } else { + while (output != lastNonNull) { + if (*nonNullIt--) { + nullBits.set(pos); + *output = *lastNonNull; + --lastNonNull; + } + --output; + --pos; + } + } + + if (output >= buffer) { + nullBits.set(offset, pos + 1); + } + } + + row_ += rowCount; + return nonNullCount; +} + +template +std::string_view NullableEncoding::encodeNullable( + EncodingSelection& selection, + std::span values, + std::span nulls, + Buffer& buffer) { + const uint32_t rowCount = nulls.size(); + + Buffer tempBuffer{buffer.getMemoryPool()}; + std::string_view serializedValues = + selection.template encodeNested( + EncodingIdentifiers::Nullable::Data, values, tempBuffer); + std::string_view serializedNulls = selection.template encodeNested( + EncodingIdentifiers::Nullable::Nulls, nulls, tempBuffer); + + const uint32_t encodingSize = Encoding::kPrefixSize + 4 + + serializedValues.size() + serializedNulls.size(); + char* reserved = buffer.reserve(encodingSize); + char* pos = reserved; + Encoding::serializePrefix( + EncodingType::Nullable, TypeTraits::dataType, rowCount, pos); + encoding::writeString(serializedValues, pos); + encoding::writeBytes(serializedNulls, pos); + NIMBLE_DASSERT(pos - reserved == encodingSize, "Encoding size mismatch."); + return {reserved, encodingSize}; +} + +template +std::string NullableEncoding::debugString(int offset) const { + std::string log = Encoding::debugString(offset); + log += fmt::format( + "\n{}non-null child:\n{}", + std::string(offset + 2, ' '), + nonNulls_->debugString(offset + 4)); + log += fmt::format( + "\n{}null child:\n{}", + std::string(offset + 2, ' '), + nulls_->debugString(offset + 4)); + return log; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/RleEncoding.cpp b/dwio/nimble/encodings/RleEncoding.cpp new file mode 100644 index 0000000..082398e --- /dev/null +++ b/dwio/nimble/encodings/RleEncoding.cpp @@ -0,0 +1,29 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/encodings/RleEncoding.h" + +namespace facebook::nimble { + +RLEEncoding::RLEEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data) + : internal::RLEEncodingBase>(memoryPool, data) { + initialValue_ = *reinterpret_cast( + internal::RLEEncodingBase>::getValuesStart()); + NIMBLE_CHECK( + (internal::RLEEncodingBase>::getValuesStart() + + 1) == data.end(), + "Unexpected run length encoding end"); + internal::RLEEncodingBase>::reset(); +} + +bool RLEEncoding::nextValue() { + value_ = !value_; + return !value_; +} + +void RLEEncoding::resetValues() { + value_ = initialValue_; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/RleEncoding.h b/dwio/nimble/encodings/RleEncoding.h new file mode 100644 index 0000000..beba832 --- /dev/null +++ b/dwio/nimble/encodings/RleEncoding.h @@ -0,0 +1,231 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include +#include "dwio/nimble/common/Bits.h" +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/EncodingType.h" +#include "dwio/nimble/common/FixedBitArray.h" +#include "dwio/nimble/common/Rle.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" +#include "dwio/nimble/encodings/EncodingIdentifier.h" +#include "dwio/nimble/encodings/EncodingSelection.h" + +// Holds data in RLE format. Run lengths are bit packed, and the run values +// are stored trivially. +// +// Note: we might want to recursively use the encoding factory to encode the +// run values. This recursive use can lead to great compression, but also +// tends to slow things down, particularly write speed. + +namespace facebook::nimble { + +namespace internal { + +// Base case covers the datatype-independent functionality. We use the CRTP +// to avoid having to use virtual functions (namely on +// RLEEncodingBase::RunValue). +// Data layout is: +// Encoding::kPrefixSize bytes: standard Encoding data +// 4 bytes: runs size +// X bytes: runs encoding bytes +template +class RLEEncodingBase + : public TypedEncoding::physicalType> { + public: + using cppDataType = T; + using physicalType = typename TypeTraits::physicalType; + + RLEEncodingBase(velox::memory::MemoryPool& memoryPool, std::string_view data) + : TypedEncoding(memoryPool, data), + materializedRunLengths_{EncodingFactory::decode( + memoryPool, + {data.data() + Encoding::kPrefixSize + 4, + *reinterpret_cast( + data.data() + Encoding::kPrefixSize)})} {} + + void reset() { + materializedRunLengths_.reset(); + derived().resetValues(); + copiesRemaining_ = materializedRunLengths_.nextValue(); + currentValue_ = nextValue(); + } + + void skip(uint32_t rowCount) final { + uint32_t rowsLeft = rowCount; + // TODO: We should have skip blocks. + while (rowsLeft) { + if (rowsLeft < copiesRemaining_) { + copiesRemaining_ -= rowsLeft; + return; + } else { + rowsLeft -= copiesRemaining_; + copiesRemaining_ = materializedRunLengths_.nextValue(); + currentValue_ = nextValue(); + } + } + } + + void materialize(uint32_t rowCount, void* buffer) final { + uint32_t rowsLeft = rowCount; + physicalType* output = static_cast(buffer); + while (rowsLeft) { + if (rowsLeft < copiesRemaining_) { + std::fill(output, output + rowsLeft, currentValue_); + copiesRemaining_ -= rowsLeft; + return; + } else { + std::fill(output, output + copiesRemaining_, currentValue_); + output += copiesRemaining_; + rowsLeft -= copiesRemaining_; + copiesRemaining_ = materializedRunLengths_.nextValue(); + currentValue_ = nextValue(); + } + } + } + + static std::string_view encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer) { + const uint32_t valueCount = values.size(); + Vector runLengths(&buffer.getMemoryPool()); + Vector runValues(&buffer.getMemoryPool()); + rle::computeRuns(values, &runLengths, &runValues); + + Buffer tempBuffer{buffer.getMemoryPool()}; + std::string_view serializedRunLengths = + selection.template encodeNested( + EncodingIdentifiers::RunLength::RunLengths, runLengths, tempBuffer); + + std::string_view serializedRunValues = + getSerializedRunValues(selection, runValues, tempBuffer); + + const uint32_t encodingSize = Encoding::kPrefixSize + 4 + + serializedRunLengths.size() + serializedRunValues.size(); + char* reserved = buffer.reserve(encodingSize); + char* pos = reserved; + Encoding::serializePrefix( + EncodingType::RLE, TypeTraits::dataType, valueCount, pos); + encoding::writeString(serializedRunLengths, pos); + encoding::writeBytes(serializedRunValues, pos); + NIMBLE_DASSERT(pos - reserved == encodingSize, "Encoding size mismatch."); + return {reserved, encodingSize}; + } + + const char* getValuesStart() { + return this->data_.data() + Encoding::kPrefixSize + 4 + + *reinterpret_cast( + this->data_.data() + Encoding::kPrefixSize); + } + + RLEEncoding& derived() { + return *static_cast(this); + } + physicalType nextValue() { + return derived().nextValue(); + } + static std::string_view getSerializedRunValues( + EncodingSelection& selection, + const Vector& runValues, + Buffer& buffer) { + return RLEEncoding::getSerializedRunValues(selection, runValues, buffer); + } + + uint32_t copiesRemaining_ = 0; + physicalType currentValue_; + detail::BufferedEncoding materializedRunLengths_; +}; + +} // namespace internal + +// Handles the numeric cases. Bools and strings are templated below. +// Data layout is: +// RLEEncodingBase bytes +// 4 * sizeof(physicalType) bytes: run values +template +class RLEEncoding final : public internal::RLEEncodingBase> { + using physicalType = typename TypeTraits::physicalType; + + public: + explicit RLEEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data); + + physicalType nextValue(); + void resetValues(); + static std::string_view getSerializedRunValues( + EncodingSelection& selection, + const Vector& runValues, + Buffer& buffer) { + return selection.template encodeNested( + EncodingIdentifiers::RunLength::RunValues, runValues, buffer); + } + + private: + detail::BufferedEncoding values_; +}; + +// For the bool case we know the values will alternative between true +// and false, so in addition to the run lengths we need only store +// whether the first value is true or false. +// RLEEncodingBase bytes +// 1 byte: whether first row is true +template <> +class RLEEncoding final + : public internal::RLEEncodingBase> { + public: + RLEEncoding(velox::memory::MemoryPool& memoryPool, std::string_view data); + + bool nextValue(); + void resetValues(); + static std::string_view getSerializedRunValues( + EncodingSelection& /* selection */, + const Vector& runValues, + Buffer& buffer) { + char* reserved = buffer.reserve(sizeof(char)); + *reserved = runValues[0]; + return {reserved, 1}; + } + + private: + bool initialValue_; + bool value_; +}; + +// +// End of public API. Implementations follow. +// + +template +RLEEncoding::RLEEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data) + : internal::RLEEncodingBase>(memoryPool, data), + values_{EncodingFactory::decode( + memoryPool, + {internal::RLEEncodingBase>::getValuesStart(), + static_cast( + data.end() - + internal::RLEEncodingBase>:: + getValuesStart())})} { + internal::RLEEncodingBase>::reset(); +} + +template +typename RLEEncoding::physicalType RLEEncoding::nextValue() { + return values_.nextValue(); +} + +template +void RLEEncoding::resetValues() { + values_.reset(); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/SentinelEncoding.h b/dwio/nimble/encodings/SentinelEncoding.h new file mode 100644 index 0000000..1a7d63e --- /dev/null +++ b/dwio/nimble/encodings/SentinelEncoding.h @@ -0,0 +1,402 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include "dwio/nimble/common/Bits.h" +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/FixedBitArray.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "folly/container/F14Set.h" +#include "velox/common/memory/Memory.h" + +#include + +// A sentinel encoding chooses a value not present in the non-null values to +// represent a null value, and then encodes the data (with the sentinel value +// inserted into each slot that was null) as a normal encoding. The 'nullness' +// of the sentinel is handled in the read operations. + +namespace facebook::nimble { + +// Data layout is: +// Encoding::kPrefixSize bytes: standard Encoding prefix +// 4 bytes: num nulls +// 4 bytes: non-null child encoding size (X) +// X bytes: non-null child encoding bytes +// Y bytes: type-dependent sentinel value +template +class SentinelEncoding final + : public TypedEncoding::physicalType> { + public: + using cppDataType = T; + using physicalType = typename TypeTraits::physicalType; + + SentinelEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data); + + uint32_t nullCount() const final; + bool isNullable() const final; + + void reset() final; + void skip(uint32_t rowCount) final; + void materialize(uint32_t rowCount, void* buffer) final; + uint32_t materializeNullable( + uint32_t rowCount, + void* buffer, + std::function nulls, + const bits::Bitmap* scatterBitmap = nullptr, + uint32_t offset = 0) final; + + // // Our signature here for estimate size and serialize are a little + // different + // // than for non-nullable encodings, as we take in both the nulls and non + // // nulls. Remember that the size of the values must be equal to the number + // of + // // true values in |nulls|. + // // + // // Note that nulls[i] is set to true if the ith value is NOT null. + // static bool estimateSize( + // velox::memory::MemoryPool& memoryPool, + // std::span nonNullValues, + // std::span nulls, + // OptimalSearchParams optimalSearchParams, + // encodings::EncodingParameters& encodingParameters, + // uint32_t* size); + + // static std::string_view serialize( + // std::span nonNullValues, + // std::span nulls, + // const encodings::EncodingParameters& encodingParameters, + // Buffer* buffer); + + // // Estimates the best sentinel encoding using default search parameters and + // // then serializes. + // static std::string_view serialize( + // std::span nonNullValues, + // std::span nulls, + // Buffer* buffer); + + std::string debugString(int offset) const final; + + private: + std::unique_ptr sentineledData_; + physicalType sentinelValue_; + nimble::Vector buffer_; + uint32_t nullCount_; +}; + +template +SentinelEncoding::SentinelEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data) + : TypedEncoding(memoryPool, data), buffer_(&memoryPool) { + const char* pos = data.data() + Encoding::kPrefixSize; + nullCount_ = encoding::readUint32(pos); + const uint32_t sentineledBytes = encoding::readUint32(pos); + sentineledData_ = + deserializeEncoding(this->memoryPool_, {pos, sentineledBytes}); + pos += sentineledBytes; + sentinelValue_ = encoding::read(pos); + NIMBLE_CHECK(pos == data.end(), "Unexpected sentinel encoding end"); +} + +template +uint32_t SentinelEncoding::nullCount() const { + return nullCount_; +} + +template +bool SentinelEncoding::isNullable() const { + return true; +} + +template +void SentinelEncoding::reset() { + sentineledData_->reset(); +} + +template +void SentinelEncoding::skip(uint32_t rowCount) { + sentineledData_->skip(rowCount); +} + +template +void SentinelEncoding::materialize(uint32_t rowCount, void* buffer) { + sentineledData_->materialize(rowCount, buffer); + if (sentinelValue_ == physicalType()) { + return; + } + physicalType* castBuffer = static_cast(buffer); + for (uint32_t i = 0; i < rowCount; ++i) { + if (castBuffer[i] == sentinelValue_) { + castBuffer[i] = physicalType(); + } + } +} + +template +uint32_t SentinelEncoding::materializeNullable( + uint32_t rowCount, + void* buffer, + std::function nulls, + const bits::Bitmap* scatterBitmap, + uint32_t offset) { + if (offset > 0) { + buffer = static_cast(buffer) + offset; + } + sentineledData_->materialize(rowCount, buffer); + + // Member variables in tight loops are bad. + const physicalType localSentinel = sentinelValue_; + auto scatterCount = scatterBitmap ? scatterBitmap->size() - offset : rowCount; + void* nullBitmap = nulls(); + bits::BitmapBuilder nullBits{nullBitmap, offset + scatterCount}; + nullBits.clear(offset, offset + scatterCount); + + uint32_t nonNullCount = 0; + if (scatterCount != rowCount) { + physicalType* lastValue = static_cast(buffer) + rowCount; + physicalType* castBuffer = static_cast(buffer); + + for (int64_t i = offset + scatterCount - 1; i >= offset; --i) { + if (scatterBitmap->test(i)) { + --lastValue; + if (*lastValue != localSentinel) { + castBuffer[i] = *lastValue; + nullBits.set(i); + ++nonNullCount; + } + } + } + } else { + physicalType* castBuffer = static_cast(buffer); + for (uint32_t i = 0; i < rowCount; ++i) { + auto notNull = *castBuffer++ != localSentinel; + nonNullCount += static_cast(notNull); + nullBits.maybeSet(i, notNull); + } + } + + return nonNullCount; +} + +namespace { +template +inline uint64_t maxUniqueCount() { + return 1L << (8 * BytesInNumber); +} + +template <> +inline uint64_t maxUniqueCount<8L>() { + // can't do 1 << 64 because it will generate run time failure of: + // runtime error: shift exponent 64 is too large for 64-bit type 'long' + return std::numeric_limits::max(); +} +} // namespace + +template +std::optional findSentinelValue( + std::span nonNullValues, + std::string* /*unused*/) { + static_assert(isNumericType()); + folly::F14FastSet uniques( + nonNullValues.begin(), nonNullValues.end()); + if (UNLIKELY(uniques.size() >= maxUniqueCount())) { + return std::nullopt; + } + physicalType sentinel = physicalType(); + while (uniques.find(sentinel) != uniques.end()) { + ++sentinel; + } + return sentinel; +} + +template <> +inline std::optional findSentinelValue( + std::span nonNullValues, + std::string* sentinel) { + auto intToStringSentinel = [](int i) { + // Mildly inefficient, but really unlikely to ever matter. + std::string result; + while (i) { + result += char(i % 256); + i >>= 8; + } + return result; + }; + int next = 1; +loop_start: + for (const auto value : nonNullValues) { + if (value == *sentinel) { + *sentinel = intToStringSentinel(next++); + goto loop_start; + } + } + return *sentinel; +} + +template <> +inline std::optional findSentinelValue( + std::span nonNullValues, + std::string* sentinel) { + auto intToStringSentinel = [](int i) { + // Mildly inefficient, but really unlikely to ever matter. + std::string result; + while (i) { + result += char(i % 256); + i >>= 8; + } + return result; + }; + int next = 1; +loop_start: + for (const auto& value : nonNullValues) { + if (value == *sentinel) { + *sentinel = intToStringSentinel(next++); + goto loop_start; + } + } + return *sentinel; +} + +template +Vector createSentineledData( + velox::memory::MemoryPool& memoryPool, + std::span nonNullValues, + std::span nulls, + physicalType sentinelValue) { + Vector sentineledData(&memoryPool, nulls.size()); + auto it = nonNullValues.begin(); + for (int i = 0; i < nulls.size(); ++i) { + if (nulls[i]) { + sentineledData[i] = *it++; + } else { + sentineledData[i] = sentinelValue; + } + } + return sentineledData; +} + +// template +// std::string_view SentinelEncoding::serialize( +// std::span nonNullDataValues, +// std::span nulls, +// const encodings::EncodingParameters& encodingParameters, +// Buffer* buffer) { +// NIMBLE_CHECK( +// encodingParameters.getType() == +// encodings::EncodingParameters::Type::sentinel && +// encodingParameters.sentinel_ref().has_value() && +// encodingParameters.sentinel_ref()->valuesParameters().has_value(), +// "Incomplete or incompatible Sentinel encoding parameters."); + +// std::string sentinelHolder; // Used only for string-view type. +// auto& sentinelParameters = encodingParameters.sentinel_ref().value(); +// // TODO: Once we have a way to store cached values next to encodings, we +// // should store the previously calculated sentinel value and try to use it +// in +// // following serializations. +// auto nonNullValues = +// EncodingPhysicalType::asEncodingPhysicalTypeSpan(nonNullDataValues); +// auto sentinelOptional = findSentinelValue(nonNullValues, &sentinelHolder); +// if (!sentinelOptional.has_value()) { +// NIMBLE_INCOMPATIBLE_ENCODING( +// "Cannot use SentinelEncoding when no value is left for sentinel."); +// } +// auto sentinelValue = sentinelOptional.value(); +// auto& memoryPool = buffer->getMemoryPool(); +// const Vector sentineledData = +// createSentineledData(memoryPool, nonNullValues, nulls, sentinelValue); +// const uint32_t nullCount = +// nulls.size() - std::accumulate(nulls.begin(), nulls.end(), 0u); +// std::string_view sentineledEncoding = serializeEncoding( +// sentineledData, sentinelParameters.valuesParameters().value(), buffer); +// uint32_t encodingSize = Encoding::kPrefixSize + 8 + +// sentineledEncoding.size(); if constexpr (isNumericType()) { +// encodingSize += sizeof(physicalType); +// } else { +// encodingSize += 4 + sentinelValue.size(); +// } +// char* reserved = buffer->reserve(encodingSize); +// char* pos = reserved; +// Encoding::serializePrefix( +// EncodingType::Sentinel, +// TypeTraits::dataType, +// sentineledData.size(), +// pos); +// encoding::writeUint32(nullCount, pos); +// encoding::writeString(sentineledEncoding, pos); +// encoding::write(sentinelValue, pos); +// NIMBLE_DASSERT(pos - reserved == encodingSize, "Encoding size mismatch."); +// return {reserved, encodingSize}; +// } + +// template +// bool SentinelEncoding::estimateSize( +// velox::memory::MemoryPool& memoryPool, +// std::span nonNullDataValues, +// std::span nulls, +// OptimalSearchParams optimalSearchParams, +// encodings::EncodingParameters& encodingParameters, +// uint32_t* size) { +// std::string holder; +// auto nonNullValues = +// EncodingPhysicalType::asEncodingPhysicalTypeSpan(nonNullDataValues); +// auto sentinelOptional = findSentinelValue(nonNullValues, &holder); +// if (!sentinelOptional.has_value()) { +// NIMBLE_INCOMPATIBLE_ENCODING( +// "Cannot use SentinelEncoding when no value is left for sentinel."); +// } +// auto sentinelValue = sentinelOptional.value(); +// const Vector sentineledData = +// createSentineledData(memoryPool, nonNullValues, nulls, sentinelValue); +// uint32_t sentineledSize; +// auto& sentinelParameters = encodingParameters.set_sentinel(); +// estimateOptimalEncodingSize( +// memoryPool, +// sentineledData, +// optimalSearchParams, +// &sentineledSize, +// sentinelParameters.valuesParameters().ensure()); +// *size = Encoding::kPrefixSize + 8 + sentineledSize + sizeof(physicalType); +// return true; +// } + +// template +// std::string_view SentinelEncoding::serialize( +// std::span nonNullValues, +// std::span nulls, +// Buffer* buffer) { +// encodings::EncodingParameters encodingParameters; +// uint32_t unusedSize; +// SentinelEncoding::estimateSize( +// buffer->getMemoryPool(), +// nonNullValues, +// nulls, +// OptimalSearchParams(), +// encodingParameters, +// &unusedSize); +// return SentinelEncoding::serialize( +// nonNullValues, nulls, encodingParameters, buffer); +// } + +template +std::string SentinelEncoding::debugString(int offset) const { + std::string log = Encoding::debugString(offset); + log += fmt::format( + "\n{}sentineled child:\n{}", + std::string(offset + 2, ' '), + sentineledData_->debugString(offset + 4)); + log += fmt::format( + "\n{}num nulls: {}", std::string(offset + 2, ' '), nullCount_); + log += fmt::format( + "\n{}sentinel value: {}", std::string(offset + 2, ' '), sentinelValue_); + return log; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/SparseBoolEncoding.cpp b/dwio/nimble/encodings/SparseBoolEncoding.cpp new file mode 100644 index 0000000..f923e9d --- /dev/null +++ b/dwio/nimble/encodings/SparseBoolEncoding.cpp @@ -0,0 +1,114 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include + +#include "dwio/nimble/encodings/SparseBoolEncoding.h" + +#include "dwio/nimble/common/Bits.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/encodings/Compression.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" + +namespace facebook::nimble { + +SparseBoolEncoding::SparseBoolEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data) + : TypedEncoding{memoryPool, data}, + sparseValue_{static_cast(data[kSparseValueOffset])}, + indicesUncompressed_{&memoryPool}, + indices_{EncodingFactory::decode( + memoryPool, + {data.data() + kIndicesOffset, data.size() - kIndicesOffset})} { + reset(); +} + +void SparseBoolEncoding::reset() { + row_ = 0; + indices_.reset(); + nextIndex_ = indices_.nextValue(); +} + +void SparseBoolEncoding::skip(uint32_t rowCount) { + const uint32_t end = row_ + rowCount; + while (nextIndex_ < end) { + nextIndex_ = indices_.nextValue(); + } + row_ = end; +} + +void SparseBoolEncoding::materialize(uint32_t rowCount, void* buffer) { + const uint32_t end = row_ + rowCount; + if (sparseValue_) { + memset(buffer, 0, rowCount); + while (nextIndex_ < end) { + static_cast(buffer)[nextIndex_ - row_] = true; + nextIndex_ = indices_.nextValue(); + } + } else { + memset(buffer, 1, rowCount); + while (nextIndex_ < end) { + static_cast(buffer)[nextIndex_ - row_] = false; + nextIndex_ = indices_.nextValue(); + } + } + row_ = end; +} + +std::string_view SparseBoolEncoding::encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer) { + // Decide the polarity of the encoding. + const uint32_t valueCount = values.size(); + const uint32_t setCount = selection.statistics().uniqueCounts().at(true); + bool sparseValue; + uint32_t indexCount; + if (setCount > (valueCount >> 1)) { + sparseValue = false; + indexCount = valueCount - setCount; + } else { + sparseValue = true; + indexCount = setCount; + } + + Vector indices{&buffer.getMemoryPool()}; + indices.reserve(indexCount + 1); + if (sparseValue) { + for (auto i = 0; i < values.size(); ++i) { + if (values[i]) { + indices.push_back(i); + } + } + } else { + for (auto i = 0; i < values.size(); ++i) { + if (!values[i]) { + indices.push_back(i); + } + } + } + + // Pushing rowCount as the last item. Materialize relies on finding this value + // in order to stop looping as this value is greater than any possible index. + indices.push_back(valueCount); + + Buffer tempBuffer{buffer.getMemoryPool()}; + std::string_view serializedIndices = + selection.template encodeNested( + EncodingIdentifiers::SparseBool::Indices, indices, tempBuffer); + + const uint32_t encodingSize = Encoding::kPrefixSize + + SparseBoolEncoding::kPrefixSize + serializedIndices.size(); + char* reserved = buffer.reserve(encodingSize); + char* pos = reserved; + Encoding::serializePrefix( + EncodingType::SparseBool, DataType::Bool, valueCount, pos); + encoding::writeChar(sparseValue, pos); + encoding::writeBytes(serializedIndices, pos); + + NIMBLE_DASSERT(pos - reserved == encodingSize, "Encoding size mismatch."); + return {reserved, encodingSize}; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/SparseBoolEncoding.h b/dwio/nimble/encodings/SparseBoolEncoding.h new file mode 100644 index 0000000..533d91b --- /dev/null +++ b/dwio/nimble/encodings/SparseBoolEncoding.h @@ -0,0 +1,64 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/FixedBitArray.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/encodings/EncodingIdentifier.h" +#include "dwio/nimble/encodings/EncodingSelection.h" +#include "velox/common/memory/Memory.h" + +// Represents bools sparsely. Useful when only a small fraction are set (or not +// set). +// +// In terms of space, this will be better than the trivial bitmap encoding when +// the fraction of set (or unset) bits is smaller than 1 / lg(n), where n is the +// number of rows in the stream. On normal size files thats around 5%. +// +// This encoding will likely be of most use in representing the nulls +// subencoding inside a NullableEncoding. +// +// TODO: Should we generalize the SparseEncoding idea to other data types? +// Maybe. Or maybe in our mainly-constant encoding implementation we simple use +// a SparseBoolEncoding on top of the dense encoded data. + +namespace facebook::nimble { + +// Data layout is: +// Encoding::kPrefixSize bytes: standard Encoding prefix +// 1 byte: whether the sparse bits are set or unset +// XX bytes: indices encoding bytes +class SparseBoolEncoding final : public TypedEncoding { + public: + using cppDataType = bool; + + static constexpr int kPrefixSize = 1; + static constexpr int kSparseValueOffset = Encoding::kPrefixSize; + static constexpr int kIndicesOffset = Encoding::kPrefixSize + 1; + + SparseBoolEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data); + + void reset() final; + void skip(uint32_t rowCount) final; + void materialize(uint32_t rowCount, void* buffer) final; + + static std::string_view encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer); + + private: + // If true, then indices give the position of the set bits; if false they give + // the positions of the unset bits. + const bool sparseValue_; + Vector indicesUncompressed_; + detail::BufferedEncoding indices_; + uint32_t nextIndex_; // The current index (FBA value). + uint32_t row_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/Statistics.cpp b/dwio/nimble/encodings/Statistics.cpp new file mode 100644 index 0000000..6b891d9 --- /dev/null +++ b/dwio/nimble/encodings/Statistics.cpp @@ -0,0 +1,277 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/encodings/Statistics.h" +#include "dwio/nimble/common/Types.h" + +#include +#include +#include +#include + +namespace facebook::nimble { + +namespace { + +template +using MapType = typename UniqueValueCounts::MapType; + +} // namespace + +template +void Statistics::populateRepeats() const { + uint64_t consecutiveRepeatCount = 0; + uint64_t minRepeat = std::numeric_limits::max(); + uint64_t maxRepeat = 0; + + uint64_t totalRepeatLength = 0; // only needed for strings + if constexpr (nimble::isStringType()) { + totalRepeatLength = data_[0].size(); + } + + T currentValue = data_[0]; + uint64_t currentRepeat = 0; + + for (auto i = 0; i < data_.size(); ++i) { + const auto& value = data_[i]; + + if (value == currentValue) { + ++currentRepeat; + } else { + if constexpr (nimble::isStringType()) { + totalRepeatLength += value.size(); + } + if (currentRepeat > maxRepeat) { + maxRepeat = currentRepeat; + } + + if (currentRepeat < minRepeat) { + minRepeat = currentRepeat; + } + + currentRepeat = 1; + currentValue = value; + ++consecutiveRepeatCount; + } + } + + if (currentRepeat > maxRepeat) { + maxRepeat = currentRepeat; + } + + if (currentRepeat < minRepeat) { + minRepeat = currentRepeat; + } + + ++consecutiveRepeatCount; + minRepeat_ = minRepeat; + maxRepeat_ = maxRepeat; + consecutiveRepeatCount_ = consecutiveRepeatCount; + totalStringsRepeatLength_ = totalRepeatLength; +} + +template +void Statistics::populateMinMax() const { + if constexpr (nimble::isNumericType()) { + const auto [min, max] = std::minmax_element(data_.begin(), data_.end()); + min_ = *min; + max_ = *max; + } else if constexpr (nimble::isStringType()) { + std::string_view minString = data_[0]; + std::string_view maxString = data_[0]; + for (int i = 0; i < data_.size(); ++i) { + const auto& value = data_[i]; + if (value.size() > maxString.size()) { + maxString = value; + } + if (value.size() < minString.size()) { + minString = value; + } + } + min_ = minString; + max_ = maxString; + } +} + +template +void Statistics::populateUniques() const { + MapType uniqueCounts; + if constexpr (nimble::isBoolType()) { + std::array counts{}; + for (int i = 0; i < data_.size(); ++i) { + ++counts[static_cast(data_[i])]; + } + uniqueCounts.reserve(2); + if (counts[0] > 0) { + uniqueCounts[false] = counts[0]; + } + if (counts[1] > 0) { + uniqueCounts[true] = counts[1]; + } + } else { + // Note: There is no science behind the reservation size. Just trying to + // minimize internal allocations... + uniqueCounts.reserve(data_.size() / 3); + for (auto i = 0; i < data_.size(); ++i) { + ++uniqueCounts[data_[i]]; + } + } + uniqueCounts_.emplace(std::move(uniqueCounts)); +} + +template +void Statistics::populateBucketCounts() const { + using UnsignedT = typename std::make_unsigned::type; + // Bucket counts are calculated in two phases. In phase one, we iterate on all + // entries, and (efficiently) count the occurences based on the MSB (most + // significant bit) of the entry. In phase two, we merge the results of phase + // one, for each conscutive 7 bits. + // See benchmarks in + // dwio/nimble/encodings/tests:bucket_benchmark for why this method is used. + std::array::digits + 1> bitCounts{}; + for (auto i = 0; i < data_.size(); ++i) { + ++(bitCounts + [std::numeric_limits::digits - + std::countl_zero(static_cast( + static_cast(data_[i]) - + static_cast(min())))]); + } + + std::vector bucketCounts(sizeof(T) * 8 / 7 + 1, 0); + uint8_t start = 0; + uint8_t end = 8; + uint8_t iteration = 0; + while (start < bitCounts.size()) { + for (auto i = start; i < end; ++i) { + bucketCounts[iteration] += bitCounts[i]; + } + ++iteration; + start = end; + end += 7; + if (bitCounts.size() < end) { + end = bitCounts.size(); + } + } + + bucketCounts_ = std::move(bucketCounts); +} + +template +void Statistics::populateStringLength() const { + uint64_t total = 0; + for (int i = 0; i < data_.size(); ++i) { + total += data_[i].size(); + } + totalStringsLength_ = total; +} + +template +Statistics Statistics::create( + std::span data) { + Statistics statistics; + if (data.size() == 0) { + statistics.consecutiveRepeatCount_ = 0; + statistics.minRepeat_ = 0; + statistics.maxRepeat_ = 0; + statistics.totalStringsLength_ = 0; + statistics.totalStringsRepeatLength_ = 0; + statistics.min_ = T(); + statistics.max_ = T(); + + statistics.bucketCounts_ = {}; + statistics.uniqueCounts_ = {}; + return statistics; + } + + statistics.data_ = data; + return statistics; +} + +template Statistics Statistics::create( + std::span data); +template Statistics Statistics::create( + std::span data); +template Statistics Statistics::create( + std::span data); +template Statistics Statistics::create( + std::span data); +template Statistics Statistics::create( + std::span data); +template Statistics Statistics::create( + std::span data); +template Statistics Statistics::create( + std::span data); +template Statistics Statistics::create( + std::span data); +template Statistics Statistics::create( + std::span data); +template Statistics Statistics::create( + std::span data); +template Statistics Statistics::create(std::span data); +template Statistics Statistics::create( + std::span data); +template Statistics +Statistics::create( + std::span data); + +// populateRepeats works on all types +template void Statistics::populateRepeats() const; +template void Statistics::populateRepeats() const; +template void Statistics::populateRepeats() const; +template void Statistics::populateRepeats() const; +template void Statistics::populateRepeats() const; +template void Statistics::populateRepeats() const; +template void Statistics::populateRepeats() const; +template void Statistics::populateRepeats() const; +template void Statistics::populateRepeats() const; +template void Statistics::populateRepeats() const; +template void Statistics::populateRepeats() const; +template void Statistics::populateRepeats() const; +template void Statistics::populateRepeats() + const; + +// populateUniques works on all types +template void Statistics::populateUniques() const; +template void Statistics::populateUniques() const; +template void Statistics::populateUniques() const; +template void Statistics::populateUniques() const; +template void Statistics::populateUniques() const; +template void Statistics::populateUniques() const; +template void Statistics::populateUniques() const; +template void Statistics::populateUniques() const; +template void Statistics::populateUniques() const; +template void Statistics::populateUniques() const; +template void Statistics::populateUniques() const; +template void Statistics::populateUniques() const; +template void Statistics::populateUniques() + const; + +// populateMinMax works on numeric types only +template void Statistics::populateMinMax() const; +template void Statistics::populateMinMax() const; +template void Statistics::populateMinMax() const; +template void Statistics::populateMinMax() const; +template void Statistics::populateMinMax() const; +template void Statistics::populateMinMax() const; +template void Statistics::populateMinMax() const; +template void Statistics::populateMinMax() const; +template void Statistics::populateMinMax() const; +template void Statistics::populateMinMax() const; +template void Statistics::populateMinMax() const; +template void Statistics::populateMinMax() const; + +// populateBucketCounts works on integral types only +template void Statistics::populateBucketCounts() const; +template void Statistics::populateBucketCounts() const; +template void Statistics::populateBucketCounts() const; +template void Statistics::populateBucketCounts() const; +template void Statistics::populateBucketCounts() const; +template void Statistics::populateBucketCounts() const; +template void Statistics::populateBucketCounts() const; +template void Statistics::populateBucketCounts() const; + +// String functions +template void Statistics::populateStringLength() const; +template void Statistics::populateStringLength() + const; + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/Statistics.h b/dwio/nimble/encodings/Statistics.h new file mode 100644 index 0000000..5ceab6c --- /dev/null +++ b/dwio/nimble/encodings/Statistics.h @@ -0,0 +1,197 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include "dwio/nimble/common/Types.h" + +#include "absl/container/flat_hash_map.h" // @manual=fbsource//third-party/abseil-cpp:container__flat_hash_map + +namespace facebook::nimble { + +template +class UniqueValueCounts { + public: + using MapType = absl::flat_hash_map; + + struct Iterator { + using iterator_category = std::forward_iterator_tag; + using value_type = typename MapType::value_type; + using difference_type = typename MapType::difference_type; + using const_reference = typename MapType::const_reference; + using const_iterator = typename MapType::const_iterator; + + const_reference operator*() const { + return *iterator_; + } + const_iterator operator->() const { + return iterator_; + } + + // Prefix increment + Iterator& operator++() { + ++iterator_; + return *this; + } + + // Postfix increment + Iterator operator++(int) { + Iterator tmp = *this; + ++(*this); + return tmp; + } + + friend bool operator==(const Iterator& a, const Iterator& b) { + return a.iterator_ == b.iterator_; + } + friend bool operator!=(const Iterator& a, const Iterator& b) { + return a.iterator_ != b.iterator_; + } + + private: + explicit Iterator(typename MapType::const_iterator iterator) + : iterator_{iterator} {} + + typename MapType::const_iterator iterator_; + + friend class UniqueValueCounts; + }; + + using const_iterator = Iterator; + + uint64_t at(T key) const noexcept { + auto it = uniqueCounts_.find(key); + if (it == uniqueCounts_.end()) { + return 0; + } + return it->second; + } + + size_t size() const noexcept { + return uniqueCounts_.size(); + } + + const_iterator begin() const noexcept { + return Iterator{uniqueCounts_.cbegin()}; + } + const_iterator cbegin() const noexcept { + return Iterator{uniqueCounts_.cbegin()}; + } + + const_iterator end() const noexcept { + return Iterator{uniqueCounts_.cend()}; + } + const_iterator cend() const noexcept { + return Iterator{uniqueCounts_.cend()}; + } + + UniqueValueCounts() = default; + explicit UniqueValueCounts(MapType&& uniqueCounts) + : uniqueCounts_{std::move(uniqueCounts)} {} + + private: + MapType uniqueCounts_; +}; + +template +class Statistics { + public: + using valueType = T; + + static Statistics create(std::span data); + + uint64_t consecutiveRepeatCount() const noexcept { + if (!consecutiveRepeatCount_.has_value()) { + populateRepeats(); + } + return consecutiveRepeatCount_.value(); + } + + uint64_t minRepeat() const noexcept { + if (!minRepeat_.has_value()) { + populateRepeats(); + } + return minRepeat_.value(); + } + + uint64_t maxRepeat() const noexcept { + if (!maxRepeat_.has_value()) { + populateRepeats(); + } + return maxRepeat_.value(); + } + + uint64_t totalStringsLength() const noexcept { + static_assert(nimble::isStringType()); + if (!totalStringsLength_.has_value()) { + populateStringLength(); + } + return totalStringsLength_.value(); + } + + uint64_t totalStringsRepeatLength() const noexcept { + static_assert(nimble::isStringType()); + if (!totalStringsRepeatLength_.has_value()) { + populateRepeats(); + } + return totalStringsRepeatLength_.value(); + } + + T min() const noexcept { + static_assert(!nimble::isBoolType()); + if (!min_.has_value()) { + populateMinMax(); + } + return min_.value(); + } + + T max() const noexcept { + static_assert(!nimble::isBoolType()); + if (!max_.has_value()) { + populateMinMax(); + } + return max_.value(); + } + + const std::vector& bucketCounts() const noexcept { + static_assert(nimble::isIntegralType()); + if (!bucketCounts_.has_value()) { + populateBucketCounts(); + } + return bucketCounts_.value(); + } + + const UniqueValueCounts& uniqueCounts() const noexcept { + if (!uniqueCounts_.has_value()) { + populateUniques(); + } + return uniqueCounts_.value(); + } + + private: + Statistics() = default; + std::span data_; + + void populateRepeats() const; + void populateUniques() const; + void populateMinMax() const; + void populateBucketCounts() const; + void populateStringLength() const; + + mutable std::optional consecutiveRepeatCount_; + mutable std::optional minRepeat_; + mutable std::optional maxRepeat_; + mutable std::optional totalStringsLength_; + mutable std::optional totalStringsRepeatLength_; + mutable std::optional min_; + mutable std::optional max_; + mutable std::optional> bucketCounts_; + mutable std::optional> uniqueCounts_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/TrivialEncoding.cpp b/dwio/nimble/encodings/TrivialEncoding.cpp new file mode 100644 index 0000000..ae133cb --- /dev/null +++ b/dwio/nimble/encodings/TrivialEncoding.cpp @@ -0,0 +1,231 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/encodings/TrivialEncoding.h" +#include +#include "dwio/nimble/encodings/EncodingFactoryNew.h" + +namespace facebook::nimble { + +TrivialEncoding::TrivialEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data) + : TypedEncoding{memoryPool, data}, + row_{0}, + buffer_{&memoryPool}, + dataUncompressed_{&memoryPool} { + auto pos = data.data() + kDataCompressionOffset; + auto dataCompressionType = + static_cast(encoding::readChar(pos)); + auto lengthsSize = encoding::readUint32(pos); + lengths_ = EncodingFactory::decode(memoryPool, {pos, lengthsSize}); + blob_ = pos + lengthsSize; + + if (dataCompressionType != CompressionType::Uncompressed) { + dataUncompressed_ = Compression::uncompress( + memoryPool, + dataCompressionType, + {blob_, static_cast(data.end() - blob_)}); + blob_ = reinterpret_cast(dataUncompressed_.data()); + } + + pos_ = blob_; +} + +void TrivialEncoding::reset() { + row_ = 0; + pos_ = blob_; + lengths_->reset(); +} + +void TrivialEncoding::skip(uint32_t rowCount) { + buffer_.resize(rowCount); + lengths_->materialize(rowCount, buffer_.data()); + row_ += rowCount; + pos_ += std::accumulate(buffer_.begin(), buffer_.end(), 0U); +} + +void TrivialEncoding::materialize( + uint32_t rowCount, + void* buffer) { + buffer_.resize(rowCount); + lengths_->materialize(rowCount, buffer_.data()); + const char* pos = pos_; + const uint32_t* data = buffer_.data(); + for (int i = 0; i < rowCount; ++i) { + static_cast(buffer)[i] = std::string_view(pos, data[i]); + pos += data[i]; + } + pos_ = pos; + row_ += rowCount; +} + +std::string_view TrivialEncoding::encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer) { + const uint32_t valueCount = values.size(); + std::vector lengths; + lengths.reserve(valueCount); + for (auto value : values) { + lengths.push_back(value.size()); + } + + Buffer tempBuffer{buffer.getMemoryPool()}; + std::string_view serializedLengths = + selection.template encodeNested( + EncodingIdentifiers::Trivial::Lengths, {lengths}, tempBuffer); + + auto dataCompressionPolicy = selection.compressionPolicy(); + auto uncompressedSize = selection.statistics().totalStringsLength(); + + Vector vector{&buffer.getMemoryPool()}; + + CompressionEncoder compressionEncoder{ + buffer.getMemoryPool(), + *dataCompressionPolicy, + DataType::String, + /*bitWidth=*/0, + uncompressedSize, + [&]() { + vector.resize(uncompressedSize); + return std::span{vector.data(), uncompressedSize}; + }, + [&](char*& pos) { + for (auto value : values) { + std::copy(value.cbegin(), value.cend(), pos); + pos += value.size(); + } + }}; + + const uint32_t encodingSize = Encoding::kPrefixSize + + TrivialEncoding::kPrefixSize + + serializedLengths.size() + compressionEncoder.getSize(); + + char* reserved = buffer.reserve(encodingSize); + char* pos = reserved; + Encoding::serializePrefix( + EncodingType::Trivial, DataType::String, valueCount, pos); + encoding::writeChar( + static_cast(compressionEncoder.compressionType()), pos); + encoding::writeUint32(serializedLengths.size(), pos); + encoding::writeBytes(serializedLengths, pos); + compressionEncoder.write(pos); + + NIMBLE_ASSERT(pos - reserved == encodingSize, "Encoding size mismatch."); + return {reserved, encodingSize}; +} + +TrivialEncoding::TrivialEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data) + : TypedEncoding{memoryPool, data}, + row_{0}, + bitmap_{data.data() + kDataOffset}, + uncompressed_{&memoryPool} { + auto compressionType = + static_cast(data[kCompressionTypeOffset]); + if (compressionType != CompressionType::Uncompressed) { + uncompressed_ = Compression::uncompress( + memoryPool, + compressionType, + {bitmap_, static_cast(data.end() - bitmap_)}); + bitmap_ = uncompressed_.data(); + NIMBLE_CHECK( + bitmap_ + FixedBitArray::bufferSize(rowCount(), 1) == + uncompressed_.end(), + "Unexpected trivial encoding end"); + } else { + NIMBLE_CHECK( + bitmap_ + FixedBitArray::bufferSize(rowCount(), 1) == data.end(), + "Unexpected trivial encoding end"); + } +} + +void TrivialEncoding::reset() { + row_ = 0; +} + +void TrivialEncoding::skip(uint32_t rowCount) { + row_ += rowCount; +} + +void TrivialEncoding::materialize(uint32_t rowCount, void* buffer) { + // Align to word boundary, go fast over words, then do remainder. + bool* output = static_cast(buffer); + const uint32_t rowsToWord = (row_ & 63) == 0 ? 0 : 64 - (row_ & 63); + if (rowsToWord >= rowCount) { + for (int i = 0; i < rowCount; ++i) { + *output = bits::getBit(row_, bitmap_); + ++output; + ++row_; + } + return; + } + for (uint32_t i = 0; i < rowsToWord; ++i) { + *output = bits::getBit(row_, bitmap_); + ++output; + ++row_; + } + const uint32_t rowsRemaining = rowCount - rowsToWord; + const uint32_t numWords = rowsRemaining >> 6; + const uint64_t* nextWord = + reinterpret_cast(bitmap_ + (row_ >> 3)); + for (uint32_t i = 0; i < numWords; ++i) { + uint64_t word = nextWord[i]; + for (int j = 0; j < 64; ++j) { + *output = word & 1; + word >>= 1; + ++output; + } + row_ += 64; + } + const uint32_t remainder = rowsRemaining - (numWords << 6); + for (uint32_t i = 0; i < remainder; ++i) { + *output = bits::getBit(row_, bitmap_); + ++output; + ++row_; + } +} + +std::string_view TrivialEncoding::encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer) { + const uint32_t valueCount = values.size(); + const uint32_t bitmapBytes = FixedBitArray::bufferSize(valueCount, 1); + + Vector vector{&buffer.getMemoryPool()}; + + auto dataCompressionPolicy = selection.compressionPolicy(); + CompressionEncoder compressionEncoder{ + buffer.getMemoryPool(), + *dataCompressionPolicy, + DataType::Undefined, + /*bitWidth=*/1, + bitmapBytes, + [&]() { + vector.resize(bitmapBytes); + return std::span{vector}; + }, + [&](char*& pos) { + memset(pos, 0, bitmapBytes); + for (size_t i = 0; i < values.size(); ++i) { + bits::maybeSetBit(i, pos, values[i]); + } + pos += bitmapBytes; + }}; + + const uint32_t encodingSize = kDataOffset + compressionEncoder.getSize(); + char* reserved = buffer.reserve(encodingSize); + char* pos = reserved; + Encoding::serializePrefix( + EncodingType::Trivial, DataType::Bool, valueCount, pos); + encoding::writeChar( + static_cast(compressionEncoder.compressionType()), pos); + compressionEncoder.write(pos); + + NIMBLE_ASSERT(pos - reserved == encodingSize, "Encoding size mismatch."); + return {reserved, encodingSize}; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/TrivialEncoding.h b/dwio/nimble/encodings/TrivialEncoding.h new file mode 100644 index 0000000..cca8620 --- /dev/null +++ b/dwio/nimble/encodings/TrivialEncoding.h @@ -0,0 +1,204 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include + +#include "dwio/nimble/common/Bits.h" +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/EncodingType.h" +#include "dwio/nimble/common/FixedBitArray.h" +#include "dwio/nimble/common/NimbleCompare.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/encodings/Compression.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/encodings/EncodingIdentifier.h" +#include "dwio/nimble/encodings/EncodingSelection.h" +#include "velox/common/memory/Memory.h" + +// Holds data in the the 'trivial' way for each data type. +// Namely as physicalType* for numerics, a packed set of offsets and a blob of +// characters for strings, and bit-packed for bools. + +namespace facebook::nimble { + +// Handles the numeric cases. Bools and strings are specialized below. +// Data layout is: +// Encoding::kPrefixSize bytes: standard Encoding prefix +// 1 byte: lengths compression +// rowCount * sizeof(physicalType) bytes: the data +template +class TrivialEncoding final + : public TypedEncoding::physicalType> { + public: + using cppDataType = T; + using physicalType = typename TypeTraits::physicalType; + + static constexpr int kPrefixSize = 1; + static constexpr int kCompressionTypeOffset = Encoding::kPrefixSize; + static constexpr int kDataOffset = + Encoding::kPrefixSize + TrivialEncoding::kPrefixSize; + + TrivialEncoding(velox::memory::MemoryPool& memoryPool, std::string_view data); + + void reset() final; + void skip(uint32_t rowCount) final; + void materialize(uint32_t rowCount, void* buffer) final; + + static std::string_view encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer); + + private: + uint32_t row_; + const T* values_; + Vector uncompressed_; +}; + +// For the string case the layout is: +// Encoding::kPrefixSize bytes: standard Encoding prefix +// 1 byte: data compression +// 4 bytes: offset to start of blob +// XX bytes: bit packed lengths +// YY bytes: actual characters +template <> +class TrivialEncoding final + : public TypedEncoding { + public: + using cppDataType = std::string_view; + + static constexpr int kPrefixSize = 5; + static constexpr int kDataCompressionOffset = Encoding::kPrefixSize; + static constexpr int kBlobOffsetOffset = Encoding::kPrefixSize + 1; + static constexpr int kLengthOffset = + Encoding::kPrefixSize + TrivialEncoding::kPrefixSize; + + TrivialEncoding(velox::memory::MemoryPool& memoryPool, std::string_view data); + + void reset() final; + void skip(uint32_t rowCount) final; + void materialize(uint32_t rowCount, void* buffer) final; + + static std::string_view encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer); + + private: + uint32_t row_; + const char* blob_; + const char* pos_; + std::unique_ptr lengths_; + Vector buffer_; + Vector dataUncompressed_; +}; + +// For the bool case the layout is: +// Encoding::kPrefixSize bytes: standard Encoding prefix +// 1 byte: compression type +// FBA::BufferSize(rowCount, 1) bytes: the bitmap +template <> +class TrivialEncoding final : public TypedEncoding { + public: + using cppDataType = bool; + + static constexpr int kPrefixSize = 1; + static constexpr int kCompressionTypeOffset = Encoding::kPrefixSize; + static constexpr int kDataOffset = + Encoding::kPrefixSize + TrivialEncoding::kPrefixSize; + + TrivialEncoding(velox::memory::MemoryPool& memoryPool, std::string_view data); + + void reset() final; + void skip(uint32_t rowCount) final; + void materialize(uint32_t rowCount, void* buffer) final; + + static std::string_view encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer); + + private: + uint32_t row_; + const char* bitmap_; + Vector uncompressed_; +}; + +// +// End of public API. Implementations follow. +// + +template +TrivialEncoding::TrivialEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data) + : TypedEncoding{memoryPool, data}, + row_{0}, + values_{reinterpret_cast(data.data() + kDataOffset)}, + uncompressed_{&memoryPool} { + auto compressionType = + static_cast(data[kCompressionTypeOffset]); + if (compressionType != CompressionType::Uncompressed) { + uncompressed_ = Compression::uncompress( + memoryPool, + compressionType, + {data.data() + kDataOffset, data.size() - kDataOffset}); + values_ = reinterpret_cast(uncompressed_.data()); + NIMBLE_CHECK( + reinterpret_cast(values_ + this->rowCount()) == + uncompressed_.end(), + "Unexpected trivial encoding end"); + } else { + NIMBLE_CHECK( + reinterpret_cast(values_ + this->rowCount()) == data.end(), + "Unexpected trivial encoding end"); + } +} + +template +void TrivialEncoding::reset() { + row_ = 0; +} + +template +void TrivialEncoding::skip(uint32_t rowCount) { + row_ += rowCount; +} + +template +void TrivialEncoding::materialize(uint32_t rowCount, void* buffer) { + const auto start = values_ + row_; + std::copy(start, start + rowCount, static_cast(buffer)); + row_ += rowCount; +} + +template +std::string_view TrivialEncoding::encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer) { + const uint32_t rowCount = values.size(); + const uint32_t uncompressedSize = sizeof(T) * rowCount; + + auto compressionPolicy = selection.compressionPolicy(); + CompressionEncoder compressionEncoder{ + buffer.getMemoryPool(), + *compressionPolicy, + TypeTraits::dataType, + {reinterpret_cast(values.data()), uncompressedSize}}; + + const uint32_t encodingSize = kDataOffset + compressionEncoder.getSize(); + char* reserved = buffer.reserve(encodingSize); + char* pos = reserved; + Encoding::serializePrefix( + EncodingType::Trivial, TypeTraits::dataType, rowCount, pos); + encoding::writeChar( + static_cast(compressionEncoder.compressionType()), pos); + compressionEncoder.write(pos); + return {reserved, encodingSize}; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/VarintEncoding.h b/dwio/nimble/encodings/VarintEncoding.h new file mode 100644 index 0000000..ccc5c75 --- /dev/null +++ b/dwio/nimble/encodings/VarintEncoding.h @@ -0,0 +1,153 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include "dwio/nimble/common/Bits.h" +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/EncodingType.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Varint.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/encodings/EncodingSelection.h" +#include "dwio/nimble/encodings/Statistics.h" + +// Stores integer data in a varint encoding. For now we only support encoding +// non-negative values, but we may later add an optional zigzag encoding that +// will let us handle negatives. + +namespace facebook::nimble { + +// Data layout is: +// Encoding::kPrefixSize bytes: standard Encoding prefix +// sizeof(T) bytes: baseline value +// X bytes: the varint encoded bytes +template +class VarintEncoding final + : public TypedEncoding::physicalType> { + public: + using cppDataType = T; + using physicalType = typename TypeTraits::physicalType; + + static const int kBaselineOffset = Encoding::kPrefixSize; + static const int kDataOffset = kBaselineOffset + sizeof(T); + static const int kPrefixSize = kDataOffset - kBaselineOffset; + + explicit VarintEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data); + + void reset() final; + void skip(uint32_t rowCount) final; + void materialize(uint32_t rowCount, void* buffer) final; + + static std::string_view encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer); + + private: + const physicalType baseline_; + uint32_t row_ = 0; + const char* pos_; + Vector buf_; +}; + +// +// End of public API. Implementations follow. +// + +template +VarintEncoding::VarintEncoding( + velox::memory::MemoryPool& memoryPool, + std::string_view data) + : TypedEncoding(memoryPool, data), + baseline_{*reinterpret_cast( + data.data() + kBaselineOffset)}, + buf_{&memoryPool} { + reset(); +} + +template +void VarintEncoding::reset() { + row_ = 0; + pos_ = Encoding::data_.data() + Encoding::kPrefixSize + + VarintEncoding::kPrefixSize; +} + +template +void VarintEncoding::skip(uint32_t rowCount) { + row_ += rowCount; + pos_ = varint::bulkVarintSkip(rowCount, pos_); +} + +template +void VarintEncoding::materialize(uint32_t rowCount, void* buffer) { + static_assert( + sizeof(T) == 4 || sizeof(T) == 8, + "Varint encoding require 4 or 8 bytes data types."); + if constexpr (isFourByteIntegralType()) { + pos_ = varint::bulkVarintDecode32( + rowCount, pos_, static_cast(buffer)); + } else { + pos_ = varint::bulkVarintDecode64( + rowCount, pos_, static_cast(buffer)); + } + + if (baseline_ != 0) { + // Add baseline to values + auto output = reinterpret_cast(buffer); + for (auto i = 0; i < rowCount; ++i) { + output[i] += baseline_; + } + } + row_ += rowCount; +} + +template +std::string_view VarintEncoding::encode( + EncodingSelection& selection, + std::span values, + Buffer& buffer) { + static_assert( + std::is_same_v< + typename std::make_unsigned::type, + physicalType>, + "Physical type must be unsigned."); + static_assert( + sizeof(T) == 4 || sizeof(T) == 8, + "Varint encoding require 4 or 8 bytes data types."); + const uint32_t valueCount = values.size(); + uint8_t index = 0; + uint32_t dataSize = std::accumulate( + selection.statistics().bucketCounts().cbegin(), + selection.statistics().bucketCounts().cend(), + 0, + [&index](uint32_t sum, const uint64_t bucketSize) { + // First (7 bit) bucket is going to consume 1 byte, 2nd bucket is going + // to consume 2 bytes, etc. + return sum + (bucketSize * (++index)); + }); + uint32_t encodingSize = + dataSize + Encoding::kPrefixSize + VarintEncoding::kPrefixSize; + + // Adding 7 bytes, to allow bulk materialization to access the last 64 bit + // word, even if it is not full. + char* reserved = buffer.reserve(encodingSize); + char* pos = reserved; + Encoding::serializePrefix( + EncodingType::Varint, TypeTraits::dataType, valueCount, pos); + encoding::write(selection.statistics().min(), pos); + for (auto value : values) { + varint::writeVarint(value - selection.statistics().min(), &pos); + } + NIMBLE_DASSERT(pos - reserved == encodingSize, "Encoding size mismatch."); + // Adding 7 bytes, to allow bulk materialization to access the last 64 bit + // word, even if it is not full. + return {reserved, encodingSize}; +} +} // namespace facebook::nimble diff --git a/dwio/nimble/encodings/tests/BucketBenchmarks.cpp b/dwio/nimble/encodings/tests/BucketBenchmarks.cpp new file mode 100644 index 0000000..c5d7f6d --- /dev/null +++ b/dwio/nimble/encodings/tests/BucketBenchmarks.cpp @@ -0,0 +1,148 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#include +#include +#include + +#include "common/init/light.h" +#include "folly/Benchmark.h" +#include "folly/Random.h" + +constexpr size_t kDataSize = 20 * 1024 * 1024; // 20M +constexpr size_t kRepeatSize = 5 * 1024; // 5K + +std::vector repeatingData; +std::vector uniqueData; + +BENCHMARK(NaiveBranching, n) { + for (int i = 0; i < n; ++i) { + std::vector bucketCounts(10, 0); + for (auto value : uniqueData) { + if (value < (1 << 7)) { + ++bucketCounts[0]; + } else if (value < (1ULL << 14)) { + ++bucketCounts[1]; + } else if (value < (1ULL << 21)) { + ++bucketCounts[2]; + } else if (value < (1ULL << 28)) { + ++bucketCounts[3]; + } else if (value < (1ULL << 35)) { + ++bucketCounts[4]; + } else if (value < (1ULL << 42)) { + ++bucketCounts[5]; + } else if (value < (1ULL << 49)) { + ++bucketCounts[6]; + } else if (value < (1ULL << 56)) { + ++bucketCounts[7]; + } else if (value < (1ULL << 63)) { + ++bucketCounts[8]; + } else { + ++bucketCounts[9]; + } + } + } +} + +BENCHMARK_RELATIVE(CountLZero, n) { + using T = int64_t; + constexpr int bitSize = sizeof(T) * 8; + for (int i = 0; i < n; ++i) { + std::array bitCounts{}; + for (auto value : uniqueData) { + ++(bitCounts[bitSize - std::countl_zero((uint64_t)value)]); + } + + std::vector bucketCounts(sizeof(uint64_t) * 8 / 7 + 1, 0); + uint8_t start = 0; + uint8_t end = 8; + uint8_t iteration = 0; + while (start < bitCounts.size()) { + if (bitCounts.size() < end) { + end = bitCounts.size(); + } + + for (auto j = start; j < end; ++j) { + bucketCounts[iteration] += bitCounts[j]; + } + ++iteration; + start = end; + end += 7; + } + } +} + +BENCHMARK_DRAW_LINE(); + +BENCHMARK(NaiveBranchingRepeat, n) { + for (int i = 0; i < n; ++i) { + std::vector bucketCounts(10, 0); + for (auto value : repeatingData) { + if (value < (1 << 7)) { + ++bucketCounts[0]; + } else if (value < (1ULL << 14)) { + ++bucketCounts[1]; + + } else if (value < (1ULL << 21)) { + ++bucketCounts[2]; + } else if (value < (1ULL << 28)) { + ++bucketCounts[3]; + } else if (value < (1ULL << 35)) { + ++bucketCounts[4]; + } else if (value < (1ULL << 42)) { + ++bucketCounts[5]; + } else if (value < (1ULL << 49)) { + ++bucketCounts[6]; + } else if (value < (1ULL << 56)) { + ++bucketCounts[7]; + } else if (value < (1ULL << 63)) { + ++bucketCounts[8]; + } else { + ++bucketCounts[9]; + } + } + } +} + +BENCHMARK_RELATIVE(CountLZeroRepeat, n) { + using T = int64_t; + constexpr int bitSize = sizeof(T) * 8; + for (int i = 0; i < n; ++i) { + std::array bitCounts{}; + for (auto value : repeatingData) { + ++(bitCounts[bitSize - std::countl_zero((uint64_t)value)]); + } + + std::vector bucketCounts(sizeof(uint64_t) * 8 / 7 + 1, 0); + uint8_t start = 0; + uint8_t end = 8; + uint8_t iteration = 0; + while (start < bitCounts.size()) { + if (bitCounts.size() < end) { + end = bitCounts.size(); + } + + for (auto j = start; j < end; ++j) { + bucketCounts[iteration] += bitCounts[j]; + } + ++iteration; + start = end; + end += 7; + } + } +} + +int main(int argc, char** argv) { + facebook::init::initFacebookLight(&argc, &argv); + repeatingData.resize(kDataSize); + uniqueData.resize(kDataSize); + for (auto i = 0; i < kDataSize; ++i) { + repeatingData[i] = i % kRepeatSize; + uniqueData[i] = i; + } + + std::random_shuffle(repeatingData.begin(), repeatingData.end()); + std::random_shuffle(uniqueData.begin(), uniqueData.end()); + + folly::runBenchmarks(); + return 0; +} diff --git a/dwio/nimble/encodings/tests/CMakeLists.txt b/dwio/nimble/encodings/tests/CMakeLists.txt new file mode 100644 index 0000000..de7333c --- /dev/null +++ b/dwio/nimble/encodings/tests/CMakeLists.txt @@ -0,0 +1,22 @@ +add_executable( + nimble_encodings_tests + ConstantEncodingTests.cpp + EncodingLayoutTests.cpp + EncodingSelectionTests.cpp + EncodingTestsNew.cpp + MainlyConstantEncodingTests.cpp + NullableEncodingTests.cpp + RleEncodingTests.cpp + SentinelEncodingTests.cpp + StatisticsTests.cpp) + +add_test(nimble_encodings_tests nimble_encodings_tests) + +target_link_libraries( + nimble_encodings_tests + nimble_encodings + nimble_common + nimble_tools_common + gtest + gtest_main + Folly::folly) diff --git a/dwio/nimble/encodings/tests/ConstantEncodingTests.cpp b/dwio/nimble/encodings/tests/ConstantEncodingTests.cpp new file mode 100644 index 0000000..0936a78 --- /dev/null +++ b/dwio/nimble/encodings/tests/ConstantEncodingTests.cpp @@ -0,0 +1,153 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/EncodingType.h" +#include "dwio/nimble/common/NimbleCompare.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/encodings/ConstantEncoding.h" +#include "dwio/nimble/encodings/tests/TestUtils.h" + +#include +#include + +using namespace facebook; + +template +class ConstantEncodingTest : public ::testing::Test { + protected: + void SetUp() override { + pool_ = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + buffer_ = std::make_unique(*pool_); + } + + template + std::vector> prepareValues() { + FAIL() << "unspecialized prepapreValues() should not be called"; + return {}; + } + + template + std::vector> prepareFailureValues() { + FAIL() << "unspecialized prepapreValues() should not be called"; + return {}; + } + + template + using ET = nimble::EncodingPhysicalType; + + double dNaN0 = std::numeric_limits::quiet_NaN(); + double dNaN1 = std::numeric_limits::signaling_NaN(); + double dNaN2 = ET::asEncodingLogicalType( + (ET::asEncodingPhysicalType(dNaN0) | 0x3)); + + float fNaN0 = std::numeric_limits::quiet_NaN(); + float fNaN1 = std::numeric_limits::signaling_NaN(); + float fNaN2 = ET::asEncodingLogicalType( + (ET::asEncodingPhysicalType(fNaN0) | 0x3)); + + template + nimble::Vector toVector(std::initializer_list l) { + nimble::Vector v{pool_.get()}; + v.insert(v.end(), l.begin(), l.end()); + return v; + } + + template <> + std::vector> prepareValues() { + return { + toVector({0.0}), + toVector({0.0, 0.00}), + toVector({-0.0, -0.00}), + toVector({-2.1, -2.1, -2.1, -2.1, -2.1}), + toVector({dNaN0, dNaN0, dNaN0}), + toVector({dNaN1, dNaN1, dNaN1}), + toVector({dNaN2, dNaN2, dNaN2})}; + } + + template <> + std::vector> prepareValues() { + return { + toVector({0.0f}), + toVector({0.0f, 0.00f}), + toVector({-0.0f, -0.00f}), + toVector({-2.1f, -2.1f, -2.1f, -2.1f, -2.1f}), + toVector({fNaN0, fNaN0, fNaN0}), + toVector({fNaN1, fNaN1, fNaN1}), + toVector({fNaN2, fNaN2, fNaN2})}; + } + + template <> + std::vector> prepareValues() { + return {toVector({1}), toVector({3, 3, 3})}; + } + + template <> + std::vector> prepareFailureValues() { + return { + toVector({-0.0, -0.00, -0.0000001}), + toVector({-2.1, -2.1, -2.1, -2.1, -2.2}), + toVector({dNaN0, dNaN0, dNaN1})}; + } + + template <> + std::vector> prepareFailureValues() { + return { + toVector({-0.0f, -0.00f, -0.0000001f}), + toVector({-2.1f, -2.1f, -2.1f, -2.1f, -2.2f}), + toVector({fNaN0, fNaN0, fNaN2})}; + } + + template <> + std::vector> prepareFailureValues() { + return {toVector({3, 2, 3})}; + } + + std::shared_ptr pool_; + std::unique_ptr buffer_; +}; + +#define NUM_TYPES int32_t, double, float + +using TestTypes = ::testing::Types; + +TYPED_TEST_CASE(ConstantEncodingTest, TestTypes); + +TYPED_TEST(ConstantEncodingTest, SerializeThenDeserialize) { + using D = TypeParam; + + auto valueGroups = this->template prepareValues(); + for (const auto& values : valueGroups) { + auto encoding = + nimble::test::Encoder>::createEncoding( + *this->buffer_, values); + + uint32_t rowCount = values.size(); + nimble::Vector result(this->pool_.get(), rowCount); + encoding->materialize(rowCount, result.data()); + + EXPECT_EQ(encoding->encodingType(), nimble::EncodingType::Constant); + EXPECT_EQ(encoding->dataType(), nimble::TypeTraits::dataType); + EXPECT_EQ(encoding->rowCount(), rowCount); + for (uint32_t i = 0; i < rowCount; ++i) { + EXPECT_TRUE(nimble::NimbleCompare::equals(result[i], values[i])); + } + } +} + +TYPED_TEST(ConstantEncodingTest, NonConstantFailure) { + using D = TypeParam; + + auto valueGroups = this->template prepareFailureValues(); + for (const auto& values : valueGroups) { + try { + nimble::test::Encoder>::createEncoding( + *this->buffer_, values); + FAIL() << "ConstantEncodingTest should fail due to non constant data"; + } catch (const nimble::NimbleUserError& e) { + EXPECT_EQ(nimble::error_code::IncompatibleEncoding, e.errorCode()); + EXPECT_EQ("ConstantEncoding requires constant data.", e.errorMessage()); + } + } +} diff --git a/dwio/nimble/encodings/tests/EncodingLayoutTests.cpp b/dwio/nimble/encodings/tests/EncodingLayoutTests.cpp new file mode 100644 index 0000000..4eaf548 --- /dev/null +++ b/dwio/nimble/encodings/tests/EncodingLayoutTests.cpp @@ -0,0 +1,378 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/encodings/EncodingLayout.h" +#include "dwio/nimble/encodings/EncodingLayoutCapture.h" +#include "dwio/nimble/encodings/EncodingSelectionPolicy.h" + +using namespace ::facebook; + +namespace { + +void verifyEncodingLayout( + const std::optional& expected, + const std::optional& actual) { + ASSERT_EQ(expected.has_value(), actual.has_value()); + + if (!expected.has_value()) { + return; + } + + ASSERT_EQ(expected->encodingType(), actual->encodingType()); + ASSERT_EQ(expected->compressionType(), actual->compressionType()); + ASSERT_EQ(expected->childrenCount(), actual->childrenCount()); + + for (auto i = 0; i < expected->childrenCount(); ++i) { + verifyEncodingLayout(expected->child(i), actual->child(i)); + } +} + +void testSerialization(nimble::EncodingLayout expected) { + std::string output; + output.resize(1024); + auto size = expected.serialize(output); + auto actual = nimble::EncodingLayout::create( + {output.data(), static_cast(size)}); + verifyEncodingLayout(expected, actual.first); +} + +template > +void testCapture(nimble::EncodingLayout expected, TCollection data) { + nimble::EncodingSelectionPolicyFactory encodingSelectionPolicyFactory = + [encodingFactory = nimble::ManualEncodingSelectionPolicyFactory{}]( + nimble::DataType dataType) + -> std::unique_ptr { + return encodingFactory.createPolicy(dataType); + }; + + auto defaultPool = velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*defaultPool}; + auto encoding = nimble::EncodingFactory::encode( + std::make_unique>( + expected, + nimble::CompressionOptions{.compressionAcceptRatio = 100}, + encodingSelectionPolicyFactory), + data, + buffer); + + auto actual = nimble::EncodingLayoutCapture::capture(encoding); + verifyEncodingLayout(expected, actual); +} + +} // namespace + +TEST(EncodingLayoutTests, Trivial) { + { + nimble::EncodingLayout expected{ + nimble::EncodingType::Trivial, nimble::CompressionType::Uncompressed}; + + testSerialization(expected); + testCapture(expected, {1, 2, 3}); + } + + { + nimble::EncodingLayout expected{ + nimble::EncodingType::Trivial, nimble::CompressionType::Zstrong}; + + testSerialization(expected); + testCapture(expected, {1, 2, 3}); + } +} + +TEST(EncodingLayoutTests, TrivialString) { + { + nimble::EncodingLayout expected{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Uncompressed, + { + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Uncompressed}, + }}; + + testSerialization(expected); + testCapture(expected, {"a", "b", "c"}); + } +} + +TEST(EncodingLayoutTests, FixedBitWidth) { + { + nimble::EncodingLayout expected{ + nimble::EncodingType::FixedBitWidth, + nimble::CompressionType::Uncompressed, + }; + + testSerialization(expected); + testCapture(expected, {1, 2, 3}); + } + + { + nimble::EncodingLayout expected{ + nimble::EncodingType::FixedBitWidth, + nimble::CompressionType::Zstd, + }; + + testSerialization(expected); + // NOTE: We need this artitifical long input data, because if Zstd + // compressed buffer is bigger than the uncompressed buffer, it is not + // picked up, which then leads to the captured encloding layout to be + // uncompressed. + testCapture( + expected, {0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, + 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, + 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, + 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, + 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, + 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, + 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, + 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, + 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, + 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, + 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, + 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF}); + } +} + +TEST(EncodingLayoutTests, Varint) { + nimble::EncodingLayout expected{ + nimble::EncodingType::Varint, + nimble::CompressionType::Uncompressed, + }; + + testSerialization(expected); + testCapture(expected, {1, 2, 3}); +} + +TEST(EncodingLayoutTests, Constant) { + nimble::EncodingLayout expected{ + nimble::EncodingType::Constant, + nimble::CompressionType::Uncompressed, + }; + + testSerialization(expected); + testCapture(expected, {1, 1, 1}); +} + +TEST(EncodingLayoutTests, SparseBool) { + nimble::EncodingLayout expected{ + nimble::EncodingType::SparseBool, + nimble::CompressionType::Uncompressed, + { + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, nimble::CompressionType::Zstrong}, + }}; + + testSerialization(expected); + testCapture( + expected, std::array{false, false, false, true, false}); +} + +TEST(EncodingLayoutTests, MainlyConst) { + nimble::EncodingLayout expected{ + nimble::EncodingType::MainlyConstant, + nimble::CompressionType::Uncompressed, + { + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, nimble::CompressionType::Zstrong}, + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Uncompressed}, + }}; + + testSerialization(expected); + testCapture(expected, {1, 1, 1, 1, 5, 1}); +} + +TEST(EncodingLayoutTests, Dictionary) { + nimble::EncodingLayout expected{ + nimble::EncodingType::Dictionary, + nimble::CompressionType::Uncompressed, + { + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Uncompressed}, + nimble::EncodingLayout{ + nimble::EncodingType::FixedBitWidth, + nimble::CompressionType::Uncompressed}, + }}; + + testSerialization(expected); + testCapture(expected, {1, 1, 1, 1, 5, 1}); +} + +TEST(EncodingLayoutTests, Rle) { + nimble::EncodingLayout expected{ + nimble::EncodingType::RLE, + nimble::CompressionType::Uncompressed, + { + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Uncompressed}, + nimble::EncodingLayout{ + nimble::EncodingType::FixedBitWidth, + nimble::CompressionType::Uncompressed}, + }}; + + testSerialization(expected); + testCapture(expected, {1, 1, 1, 1, 5, 1}); +} + +TEST(EncodingLayoutTests, RleBool) { + nimble::EncodingLayout expected{ + nimble::EncodingType::RLE, + nimble::CompressionType::Uncompressed, + { + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Uncompressed}, + }}; + + testSerialization(expected); + testCapture(expected, std::array{false, false, true, true}); +} + +TEST(EncodingLayoutTests, Nullable) { + nimble::EncodingLayout expected{ + nimble::EncodingType::Nullable, + nimble::CompressionType::Uncompressed, + {nimble::EncodingLayout{ + nimble::EncodingType::FixedBitWidth, + nimble::CompressionType::Uncompressed}, + nimble::EncodingLayout{ + nimble::EncodingType::SparseBool, + nimble::CompressionType::Uncompressed, + { + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Zstrong}, + }}}}; + + testSerialization(expected); + + nimble::EncodingSelectionPolicyFactory encodingSelectionPolicyFactory = + [encodingFactory = nimble::ManualEncodingSelectionPolicyFactory{}]( + nimble::DataType dataType) + -> std::unique_ptr { + return encodingFactory.createPolicy(dataType); + }; + + auto defaultPool = velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*defaultPool}; + auto encoding = nimble::EncodingFactory::encodeNullable( + std::make_unique>( + expected.child(nimble::EncodingIdentifiers::Nullable::Data).value(), + nimble::CompressionOptions{}, + encodingSelectionPolicyFactory), + std::vector{1, 1, 1, 1, 5, 1}, + std::array{false, false, true, false, false, false}, + buffer); + + std::string output; + output.resize(1024); + auto captured = nimble::EncodingLayoutCapture::capture(encoding); + auto size = captured.serialize(output); + + auto actual = nimble::EncodingLayout::create( + {output.data(), static_cast(size)}); + + // For nullable, captured encoding layout strips out the nullable node and + // just captures the data node. + verifyEncodingLayout( + expected.child(nimble::EncodingIdentifiers::Nullable::Data), + actual.first); +} + +TEST(EncodingLayoutTests, SizeTooSmall) { + { + nimble::EncodingLayout expected{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Uncompressed, + }; + + std::string output; + // Encoding needs minimum of 5 bytes. 4 is not enough. + output.resize(4); + EXPECT_THROW(expected.serialize(output), facebook::nimble::NimbleUserError); + } + { + nimble::EncodingLayout expected{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Uncompressed, + }; + + std::string output; + // Encoding needs minimum of 5 bytes. Should not throw. + output.resize(5); + EXPECT_EQ(5, expected.serialize(output)); + } + { + nimble::EncodingLayout expected{ + nimble::EncodingType::MainlyConstant, + nimble::CompressionType::Uncompressed, + { + std::nullopt, + std::nullopt, + }}; + + std::string output; + // 5 bytes for the top level encoding, plus 2 "exists" bytes. + // Total of 7 bytes. 6 bytes is not enough. + output.resize(6); + EXPECT_THROW(expected.serialize(output), facebook::nimble::NimbleUserError); + } + { + nimble::EncodingLayout expected{ + nimble::EncodingType::MainlyConstant, + nimble::CompressionType::Uncompressed, + { + std::nullopt, + std::nullopt, + }}; + + std::string output; + // 5 bytes for the top level encoding, plus 2 "exists" bytes. + // Total of 7 bytes. 7 bytes is enough. + output.resize(7); + EXPECT_EQ(7, expected.serialize(output)); + } + { + nimble::EncodingLayout expected{ + nimble::EncodingType::MainlyConstant, + nimble::CompressionType::Uncompressed, + { + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Zstrong}, + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Uncompressed}, + }}; + + std::string output; + // Each sub-encoding is 5 bytes (total of 10), plus 5 for the top level one. + // Plus 2 "exists" bytes. Total of 17 bytes. 16 bytes is not enough. + output.resize(16); + EXPECT_THROW(expected.serialize(output), facebook::nimble::NimbleUserError); + } + + { + nimble::EncodingLayout expected{ + nimble::EncodingType::MainlyConstant, + nimble::CompressionType::Uncompressed, + { + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Zstrong}, + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Uncompressed}, + }}; + + std::string output; + // Each sub-encoding is 5 bytes (total of 10), plus 5 for the top level one. + // Plus 2 "exists" bytes. Total of 17 bytes. 17 bytes is enough. + output.resize(17); + EXPECT_EQ(17, expected.serialize(output)); + } +} diff --git a/dwio/nimble/encodings/tests/EncodingSelectionTests.cpp b/dwio/nimble/encodings/tests/EncodingSelectionTests.cpp new file mode 100644 index 0000000..f79c65f --- /dev/null +++ b/dwio/nimble/encodings/tests/EncodingSelectionTests.cpp @@ -0,0 +1,869 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#define NIMBLE_ENCODING_SELECTION_DEBUG + +#include +#include +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" +#include "dwio/nimble/encodings/EncodingSelectionPolicy.h" +#include "dwio/nimble/encodings/NullableEncoding.h" +#include "dwio/nimble/tools/EncodingUtilities.h" + +using namespace ::facebook; + +namespace { +template +std::unique_ptr> +getRootManualSelectionPolicy() { + return std::make_unique< + nimble::ManualEncodingSelectionPolicy>( + std::vector>{ + {nimble::EncodingType::Constant, 1.0}, + {nimble::EncodingType::Trivial, 0.7}, + {nimble::EncodingType::FixedBitWidth, 1.05}, + {nimble::EncodingType::MainlyConstant, 1.05}, + {nimble::EncodingType::SparseBool, 1.05}, + {nimble::EncodingType::Dictionary, 1.05}, + {nimble::EncodingType::RLE, 1.05}, + {nimble::EncodingType::Varint, 1.1}, + }, + nimble::CompressionOptions{ + .compressionAcceptRatio = 0.9, + .zstrongCompressionLevel = 9, + .zstrongDecompressionLevel = 2, + .useVariableBitWidthCompressor = false, + }, + std::nullopt); +} +} // namespace + +struct EncodingDetails { + nimble::EncodingType encodingType; + nimble::DataType dataType; + uint32_t level; + std::string nestedEncodingName; +}; + +void verifyEncodingTree( + std::string_view stream, + std::vector expected) { + ASSERT_GT(expected.size(), 0); + std::vector actual; + nimble::tools::traverseEncodings( + stream, + [&](auto encodingType, + auto dataType, + auto level, + auto /* index */, + auto nestedEncodingName, + std::unordered_map< + nimble::tools::EncodingPropertyType, + nimble::tools::EncodingProperty> /* properties */) { + actual.push_back( + {.encodingType = encodingType, + .dataType = dataType, + .level = level, + .nestedEncodingName = nestedEncodingName}); + + return true; + }); + + ASSERT_EQ(expected.size(), actual.size()); + for (auto i = 0; i < expected.size(); ++i) { + LOG(INFO) << "Expected: " << expected[i].encodingType << "<" + << expected[i].dataType << ">[" << expected[i].nestedEncodingName + << ":" << expected[i].level << "]"; + LOG(INFO) << "Actual: " << actual[i].encodingType << "<" + << actual[i].dataType << ">[" << actual[i].nestedEncodingName + << ":" << actual[i].level << "]"; + EXPECT_EQ(expected[i].encodingType, actual[i].encodingType); + EXPECT_EQ(expected[i].dataType, actual[i].dataType); + EXPECT_EQ(expected[i].level, actual[i].level); + EXPECT_EQ(expected[i].nestedEncodingName, actual[i].nestedEncodingName); + } +} + +template +typename nimble::TypeTraits::physicalType asPhysicalType(T value) { + return *reinterpret_cast::physicalType*>( + &value); +} + +template +void test(std::span values, std::vector expected) { + auto pool = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + auto policy = getRootManualSelectionPolicy(); + nimble::Buffer buffer{*pool}; + + auto serialized = + nimble::EncodingFactory::encode(std::move(policy), values, buffer); + + LOG(INFO) << "Final size: " << serialized.size(); + + ASSERT_GT(expected.size(), 0); + std::vector actual; + nimble::tools::traverseEncodings( + serialized, + [&](auto encodingType, + auto dataType, + auto level, + auto /* index */, + auto nestedEncodingName, + std::unordered_map< + nimble::tools::EncodingPropertyType, + nimble::tools::EncodingProperty> /* properties */) { + actual.push_back( + {.encodingType = encodingType, + .dataType = dataType, + .level = level, + .nestedEncodingName = nestedEncodingName}); + + return true; + }); + + ASSERT_EQ(expected.size(), actual.size()); + for (auto i = 0; i < expected.size(); ++i) { + LOG(INFO) << "Expected: " << expected[i].encodingType << "<" + << expected[i].dataType << ">[" << expected[i].nestedEncodingName + << ":" << expected[i].level << "]"; + LOG(INFO) << "Actual: " << actual[i].encodingType << "<" + << actual[i].dataType << ">[" << actual[i].nestedEncodingName + << ":" << actual[i].level << "]"; + EXPECT_EQ(expected[i].encodingType, actual[i].encodingType); + EXPECT_EQ(expected[i].dataType, actual[i].dataType); + EXPECT_EQ(expected[i].level, actual[i].level); + EXPECT_EQ(expected[i].nestedEncodingName, actual[i].nestedEncodingName); + } + + nimble::Vector materialized{pool.get()}; + auto encoding = nimble::EncodingFactory::decode(*pool, serialized); + nimble::Vector result{pool.get()}; + result.resize(values.size()); + encoding->materialize(values.size(), result.data()); + + for (auto i = 0; i < values.size(); ++i) { + ASSERT_EQ(asPhysicalType(values[i]), asPhysicalType(result[i])) << i; + } +} + +using NumericTypes = ::testing::Types< + int8_t, + uint8_t, + int16_t, + uint16_t, + int32_t, + uint32_t, + int64_t, + uint64_t, + float, + double>; + +TYPED_TEST_CASE(EncodingSelectionNumericTests, NumericTypes); + +template +class EncodingSelectionNumericTests : public ::testing::Test {}; + +TYPED_TEST(EncodingSelectionNumericTests, SelectConst) { + using T = TypeParam; + + for (const T value : + {std::numeric_limits::min(), + static_cast(0), + std::numeric_limits::max()}) { + std::vector values; + values.resize(1000); + for (auto i = 0; i < values.size(); ++i) { + values[i] = value; + } + + test( + values, + { + {.encodingType = nimble::EncodingType::Constant, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + }); + } +} + +TYPED_TEST(EncodingSelectionNumericTests, SelectMainlyConst) { + using T = TypeParam; + + for (const T value : { + std::numeric_limits::min(), + static_cast(0), + std::numeric_limits::max(), + }) { + LOG(INFO) << "Verifying type " << folly::demangle(typeid(T)) + << " with data: " << value; + + std::vector values; + values.resize(1000); + for (auto i = 0; i < values.size(); ++i) { + values[i] = value; + } + + for (auto i = 0; i < values.size(); i += 20) { + if constexpr (nimble::isFloatingPointType()) { + values[i] = i; + } else { + values[i] = i % std::numeric_limits::max(); + } + } + + if constexpr (nimble::isFloatingPointType() || sizeof(T) < 4) { + // Floating point types and small types use Trivial encoding to encode + // the exception values. + test( + values, + { + {.encodingType = nimble::EncodingType::MainlyConstant, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + {.encodingType = nimble::EncodingType::SparseBool, + .dataType = nimble::DataType::Bool, + .level = 1, + .nestedEncodingName = "IsCommon"}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::DataType::Uint32, + .level = 2, + .nestedEncodingName = "Indices"}, + {.encodingType = nimble::EncodingType::Trivial, + .dataType = nimble::TypeTraits< + typename nimble::EncodingPhysicalType::type>::dataType, + .level = 1, + .nestedEncodingName = "OtherValues"}, + }); + } else { + // All other numeric types use FixedBitWidth encoding to encode the + // exception values. + test( + values, + { + {.encodingType = nimble::EncodingType::MainlyConstant, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + {.encodingType = nimble::EncodingType::SparseBool, + .dataType = nimble::DataType::Bool, + .level = 1, + .nestedEncodingName = "IsCommon"}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::DataType::Uint32, + .level = 2, + .nestedEncodingName = "Indices"}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::TypeTraits< + typename nimble::EncodingPhysicalType::type>::dataType, + .level = 1, + .nestedEncodingName = "OtherValues"}, + }); + } + } +} + +TYPED_TEST(EncodingSelectionNumericTests, SelectTrivial) { + using T = TypeParam; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + std::vector values; + values.resize(10000); + for (auto i = 0; i < values.size(); ++i) { + auto random = folly::Random::rand64(rng); + values[i] = *reinterpret_cast(&random); + } + + test( + values, + { + {.encodingType = nimble::EncodingType::Trivial, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + }); +} + +TYPED_TEST(EncodingSelectionNumericTests, SelectFixedBitWidth) { + using T = TypeParam; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (const auto bitWidth : {3, 4, 5}) { + LOG(INFO) << "Testing with bit width: " << bitWidth; + uint32_t maxValue = 1 << bitWidth; + + std::vector values; + values.resize(10000); + for (auto i = 0; i < values.size(); ++i) { + auto random = folly::Random::rand64(rng) % maxValue; + values[i] = *reinterpret_cast(&random); + } + + test( + values, + { + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + }); + } +} + +TYPED_TEST(EncodingSelectionNumericTests, SelectDictionary) { + using T = TypeParam; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + std::array uniqueValues{ + std::numeric_limits::min(), 0, std::numeric_limits::max()}; + + std::vector values; + values.resize(10000); + for (auto i = 0; i < values.size(); ++i) { + values[i] = uniqueValues[folly::Random::rand32(rng) % uniqueValues.size()]; + } + + test( + values, + { + {.encodingType = nimble::EncodingType::Dictionary, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + {.encodingType = nimble::EncodingType::Trivial, + .dataType = nimble::TypeTraits< + typename nimble::EncodingPhysicalType::type>::dataType, + .level = 1, + .nestedEncodingName = "Alphabet"}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::DataType::Uint32, + .level = 1, + .nestedEncodingName = "Indices"}, + }); +} + +TYPED_TEST(EncodingSelectionNumericTests, SelectRunLength) { + using T = TypeParam; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + std::array uniqueValues{ + std::numeric_limits::min(), 0, std::numeric_limits::max()}; + std::vector runLengths; + runLengths.resize(20); + uint32_t valueCount = 0; + for (auto i = 0; i < runLengths.size(); ++i) { + runLengths[i] = (folly::Random::rand32(rng) % 700) + 20; + valueCount += runLengths[i]; + } + + std::vector values; + values.reserve(valueCount); + auto index = 0; + for (const auto length : runLengths) { + for (auto i = 0; i < length; ++i) { + values.push_back(uniqueValues[index % uniqueValues.size()]); + } + ++index; + } + + if constexpr ( + nimble::isFloatingPointType() || std::is_same_v || + sizeof(T) > 4) { + // Floating point types and big types prefer storing the run values as + // dictionary + test( + values, + { + {.encodingType = nimble::EncodingType::RLE, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::DataType::Uint32, + .level = 1, + .nestedEncodingName = "Lengths"}, + {.encodingType = nimble::EncodingType::Dictionary, + .dataType = nimble::TypeTraits< + typename nimble::EncodingPhysicalType::type>::dataType, + .level = 1, + .nestedEncodingName = "Values"}, + {.encodingType = nimble::EncodingType::Trivial, + .dataType = nimble::TypeTraits< + typename nimble::EncodingPhysicalType::type>::dataType, + .level = 2, + .nestedEncodingName = "Alphabet"}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::DataType::Uint32, + .level = 2, + .nestedEncodingName = "Indices"}, + }); + } else { + test( + values, + { + {.encodingType = nimble::EncodingType::RLE, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::DataType::Uint32, + .level = 1, + .nestedEncodingName = "Lengths"}, + {.encodingType = nimble::EncodingType::Trivial, + .dataType = nimble::TypeTraits< + typename nimble::EncodingPhysicalType::type>::dataType, + .level = 1, + .nestedEncodingName = "Values"}, + }); + } +} + +TYPED_TEST(EncodingSelectionNumericTests, SelectVarint) { + using T = TypeParam; + + if constexpr (sizeof(T) < 4) { + return; + } + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + std::vector values; + values.resize(700); + for (auto i = 0; i < values.size(); ++i) { + auto value = (std::numeric_limits::max() / 2); + auto random = + *reinterpret_cast::physicalType*>( + &value) + + (folly::Random::rand32(rng) % 128); + values[i] = *reinterpret_cast(&random); + } + + for (auto i = 0; i < values.size(); i += 30) { + values[i] = std::numeric_limits::max() - folly::Random::rand32(128, rng); + } + + test( + values, + { + {.encodingType = nimble::EncodingType::Varint, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + }); +} + +TEST(EncodingSelectionBoolTests, SelectConst) { + using T = bool; + + for (const T value : {true, false}) { + std::array values; + for (auto i = 0; i < values.size(); ++i) { + values[i] = value; + } + + test( + values, + { + {.encodingType = nimble::EncodingType::Constant, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + }); + } +} + +TEST(EncodingSelectionBoolTests, SelectSparseBool) { + using T = bool; + + for (const T value : {false, true}) { + std::array values; + for (auto i = 0; i < values.size(); ++i) { + values[i] = value; + } + + for (auto i = 0; i < values.size(); i += 200) { + values[i] = !value; + } + + test( + values, + { + {.encodingType = nimble::EncodingType::SparseBool, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::DataType::Uint32, + .level = 1, + .nestedEncodingName = "Indices"}, + }); + } +} + +TEST(EncodingSelectionBoolTests, SelectTrivial) { + using T = bool; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + std::vector>> policies; + policies.push_back(getRootManualSelectionPolicy()); + // use ml policy + policies.push_back( + std::make_unique>()); + + auto pool = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*pool}; + + for (auto& policy : policies) { + std::array values; + for (auto i = 0; i < values.size(); ++i) { + values[i] = folly::Random::oneIn(2, rng); + } + + auto serialized = + nimble::EncodingFactory::encode(std::move(policy), values, buffer); + + LOG(INFO) << "Final size: " << serialized.size(); + + verifyEncodingTree( + serialized, + { + {.encodingType = nimble::EncodingType::Trivial, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + }); + } +} + +TEST(EncodingSelectionBoolTests, SelectRunLength) { + using T = bool; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + auto pool = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*pool}; + + std::vector runLengths; + runLengths.resize(20); + uint32_t valueCount = 0; + for (auto i = 0; i < runLengths.size(); ++i) { + runLengths[i] = (folly::Random::rand32(rng) % 700) + 20; + valueCount += runLengths[i]; + } + + nimble::Vector values{pool.get()}; + values.reserve(valueCount); + auto index = 0; + for (const auto length : runLengths) { + for (auto i = 0; i < length; ++i) { + values.push_back(index % 2 == 0); + } + ++index; + } + + auto policy = getRootManualSelectionPolicy(); + auto serialized = + nimble::EncodingFactory::encode(std::move(policy), values, buffer); + + LOG(INFO) << "Final size: " << serialized.size(); + + verifyEncodingTree( + serialized, + { + {.encodingType = nimble::EncodingType::RLE, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::DataType::Uint32, + .level = 1, + .nestedEncodingName = "Lengths"}, + }); +} + +TEST(EncodingSelectionStringTests, SelectConst) { + using T = std::string_view; + + auto pool = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*pool}; + + for (const T value : + {std::string(""), std::string("aaaaa"), std::string(5000, '\0')}) { + LOG(INFO) << "Testing string with value: " << value; + auto policy = getRootManualSelectionPolicy(); + + std::vector values; + values.resize(1000); + for (auto i = 0; i < values.size(); ++i) { + values[i] = value; + } + + auto serialized = + nimble::EncodingFactory::encode(std::move(policy), values, buffer); + + LOG(INFO) << "Final size: " << serialized.size(); + + verifyEncodingTree( + serialized, + { + {.encodingType = nimble::EncodingType::Constant, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + }); + } +} + +TEST(EncodingSelectionStringTests, SelectMainlyConst) { + using T = std::string_view; + + auto pool = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*pool}; + + for (const T value : { + std::string("aaaaa"), + std::string(5000, '\0'), + }) { + std::vector values; + values.resize(1000); + for (auto i = 0; i < values.size(); ++i) { + values[i] = value; + } + + std::vector uncommonValues; + for (auto i = 0; i < values.size() / 20; ++i) { + uncommonValues.emplace_back(i, 'b'); + } + + for (auto i = 0; i < uncommonValues.size(); ++i) { + values[i * 20] = uncommonValues[i]; + } + + auto policy = getRootManualSelectionPolicy(); + auto serialized = + nimble::EncodingFactory::encode(std::move(policy), values, buffer); + + LOG(INFO) << "Final size: " << serialized.size(); + + verifyEncodingTree( + serialized, + { + {.encodingType = nimble::EncodingType::MainlyConstant, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + {.encodingType = nimble::EncodingType::SparseBool, + .dataType = nimble::DataType::Bool, + .level = 1, + .nestedEncodingName = "IsCommon"}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::DataType::Uint32, + .level = 2, + .nestedEncodingName = "Indices"}, + {.encodingType = nimble::EncodingType::Trivial, + .dataType = nimble::TypeTraits< + typename nimble::EncodingPhysicalType::type>::dataType, + .level = 1, + .nestedEncodingName = "OtherValues"}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::DataType::Uint32, + .level = 2, + .nestedEncodingName = "Lengths"}, + }); + } +} + +TEST(EncodingSelectionStringTests, SelectTrivial) { + using T = std::string_view; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + auto pool = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*pool}; + + auto policy = getRootManualSelectionPolicy(); + + std::vector cache; + cache.resize(10000); + for (auto i = 0; i < cache.size(); ++i) { + std::string value(folly::Random::rand32(128, rng), ' '); + for (auto j = 0; j < value.size(); ++j) { + value[j] = static_cast(folly::Random::rand32(256, rng)); + } + cache[i] = std::move(value); + } + + std::vector values; + values.resize(cache.size()); + for (auto i = 0; i < cache.size(); ++i) { + values[i] = cache[i]; + } + + auto serialized = + nimble::EncodingFactory::encode(std::move(policy), values, buffer); + + LOG(INFO) << "Final size: " << serialized.size(); + + verifyEncodingTree( + serialized, + { + {.encodingType = nimble::EncodingType::Trivial, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::DataType::Uint32, + .level = 1, + .nestedEncodingName = "Lengths"}, + }); +} + +TEST(EncodingSelectionStringTests, SelectDictionary) { + using T = std::string_view; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + std::vector uniqueValues{"", "abcdef", std::string(5000, '\0')}; + + auto pool = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*pool}; + + auto policy = getRootManualSelectionPolicy(); + + std::vector values; + values.resize(10000); + for (auto i = 0; i < values.size(); ++i) { + values[i] = uniqueValues[folly::Random::rand32(rng) % uniqueValues.size()]; + } + + auto serialized = + nimble::EncodingFactory::encode(std::move(policy), values, buffer); + + LOG(INFO) << "Final size: " << serialized.size(); + + verifyEncodingTree( + serialized, + { + {.encodingType = nimble::EncodingType::Dictionary, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + {.encodingType = nimble::EncodingType::Trivial, + .dataType = nimble::TypeTraits< + typename nimble::EncodingPhysicalType::type>::dataType, + .level = 1, + .nestedEncodingName = "Alphabet"}, + {.encodingType = nimble::EncodingType::Varint, + .dataType = nimble::DataType::Uint32, + .level = 2, + .nestedEncodingName = "Lengths"}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::DataType::Uint32, + .level = 1, + .nestedEncodingName = "Indices"}, + }); +} + +TEST(EncodingSelectionStringTests, SelectRunLength) { + using T = std::string_view; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + auto pool = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*pool}; + + std::vector runLengths; + runLengths.resize(20); + uint32_t valueCount = 0; + for (auto i = 0; i < runLengths.size(); ++i) { + runLengths[i] = (folly::Random::rand32(rng) % 700) + 20; + valueCount += runLengths[i]; + } + + std::vector values; + values.reserve(valueCount); + auto index = 0; + for (const auto length : runLengths) { + for (auto i = 0; i < length; ++i) { + values.emplace_back( + index % 2 == 0 ? "abcdefghijklmnopqrstuvwxyz" : "1234567890"); + } + ++index; + } + + auto policy = getRootManualSelectionPolicy(); + auto serialized = + nimble::EncodingFactory::encode(std::move(policy), values, buffer); + + LOG(INFO) << "Final size: " << serialized.size(); + + verifyEncodingTree( + serialized, + { + {.encodingType = nimble::EncodingType::RLE, + .dataType = nimble::TypeTraits::dataType, + .level = 0, + .nestedEncodingName = ""}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::DataType::Uint32, + .level = 1, + .nestedEncodingName = "Lengths"}, + {.encodingType = nimble::EncodingType::Dictionary, + .dataType = nimble::TypeTraits::dataType, + .level = 1, + .nestedEncodingName = "Values"}, + {.encodingType = nimble::EncodingType::Trivial, + .dataType = nimble::TypeTraits::dataType, + .level = 2, + .nestedEncodingName = "Alphabet"}, + {.encodingType = nimble::EncodingType::Varint, + .dataType = nimble::DataType::Uint32, + .level = 3, + .nestedEncodingName = "Lengths"}, + {.encodingType = nimble::EncodingType::FixedBitWidth, + .dataType = nimble::DataType::Uint32, + .level = 2, + .nestedEncodingName = "Indices"}, + }); +} + +TEST(EncodingSelectionTests, TestNullable) { + using T = std::string_view; + auto pool = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*pool}; + auto policy = getRootManualSelectionPolicy(); + std::vector data{"abcd", "efg", "hijk", "lmno"}; + std::array nulls{ + true, false, true, true, false, false, true, true, true, false}; + + auto serialized = nimble::EncodingFactory::encodeNullable( + std::move(policy), data, nulls, buffer); + LOG(INFO) << "Final size: " << serialized.size(); +} diff --git a/dwio/nimble/encodings/tests/EncodingTestsNew.cpp b/dwio/nimble/encodings/tests/EncodingTestsNew.cpp new file mode 100644 index 0000000..213e1e2 --- /dev/null +++ b/dwio/nimble/encodings/tests/EncodingTestsNew.cpp @@ -0,0 +1,555 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include +#include +#include +#include +#include "dwio/nimble/common/Bits.h" +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/FixedBitArray.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/common/tests/TestUtils.h" +#include "dwio/nimble/encodings/ConstantEncoding.h" +#include "dwio/nimble/encodings/DeltaEncoding.h" +#include "dwio/nimble/encodings/DictionaryEncoding.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" +#include "dwio/nimble/encodings/EncodingSelectionPolicy.h" +#include "dwio/nimble/encodings/FixedBitWidthEncoding.h" +#include "dwio/nimble/encodings/MainlyConstantEncoding.h" +#include "dwio/nimble/encodings/RleEncoding.h" +#include "dwio/nimble/encodings/SparseBoolEncoding.h" +#include "dwio/nimble/encodings/TrivialEncoding.h" +#include "dwio/nimble/encodings/VarintEncoding.h" +#include "folly/Random.h" +#include "folly/container/F14Set.h" +#include "velox/common/memory/Memory.h" + +// Tests the Encoding API for all basic Encoding implementations and data types. +// +// Nullable data will be tested via a separate file, as will structured data, +// and any data-type-specific functions (e.g. for there are some Encoding calls +// only valid for bools, and those are testing in BoolEncodingTests.cpp). +// +// The tests are templated so as to cover all data types and encoding types with +// a single code path. + +using namespace ::facebook; + +namespace { + +class TestCompressPolicy : public nimble::CompressionPolicy { + public: + explicit TestCompressPolicy(bool compress, bool useVariableBitWidthCompressor) + : compress_{compress}, + useVariableBitWidthCompressor_{useVariableBitWidthCompressor} {} + + nimble::CompressionInformation compression() const override { + if (!compress_) { + return {.compressionType = nimble::CompressionType::Uncompressed}; + } + + nimble::CompressionInformation information{ + .compressionType = nimble::CompressionType::Zstrong}; + information.parameters.zstrong.compressionLevel = 9; + information.parameters.zstrong.decompressionLevel = 2; + information.parameters.zstrong.useVariableBitWidthCompressor = + useVariableBitWidthCompressor_; + return information; + } + + virtual bool shouldAccept( + nimble::CompressionType /* compressionType */, + uint64_t /* uncompressedSize */, + uint64_t /* compressedSize */) const override { + return true; + } + + private: + bool compress_; + bool useVariableBitWidthCompressor_; +}; + +template +class TestTrivialEncodingSelectionPolicy + : public nimble::EncodingSelectionPolicy { + using physicalType = typename nimble::TypeTraits::physicalType; + + public: + explicit TestTrivialEncodingSelectionPolicy( + bool shouldCompress, + bool useVariableBitWidthCompressor) + : shouldCompress_{shouldCompress}, + useVariableBitWidthCompressor_{useVariableBitWidthCompressor} {} + + nimble::EncodingSelectionResult select( + std::span /* values */, + const nimble::Statistics& /* statistics */) override { + return { + .encodingType = nimble::EncodingType::Trivial, + .compressionPolicyFactory = [this]() { + return std::make_unique( + shouldCompress_, useVariableBitWidthCompressor_); + }}; + } + + nimble::EncodingSelectionResult selectNullable( + std::span /* values */, + std::span /* nulls */, + const nimble::Statistics& /* statistics */) override { + return { + .encodingType = nimble::EncodingType::Nullable, + }; + } + + std::unique_ptr createImpl( + nimble::EncodingType /* encodingType */, + nimble::NestedEncodingIdentifier /* identifier */, + nimble::DataType type) override { + UNIQUE_PTR_FACTORY( + type, + TestTrivialEncodingSelectionPolicy, + shouldCompress_, + useVariableBitWidthCompressor_); + } + + private: + bool shouldCompress_; + bool useVariableBitWidthCompressor_; +}; +} // namespace +// C is the encoding type. +template +class EncodingTests : public ::testing::Test { + protected: + using E = typename C::cppDataType; + + void SetUp() override { + pool_ = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + buffer_ = std::make_unique(*this->pool_); + util_ = std::make_unique(*this->pool_); + } + + template + struct EncodingTypeTraits {}; + + template <> + struct EncodingTypeTraits> { + static inline nimble::EncodingType encodingType = + nimble::EncodingType::Constant; + }; + + template <> + struct EncodingTypeTraits> { + static inline nimble::EncodingType encodingType = + nimble::EncodingType::Dictionary; + }; + + template <> + struct EncodingTypeTraits> { + static inline nimble::EncodingType encodingType = + nimble::EncodingType::FixedBitWidth; + }; + + template <> + struct EncodingTypeTraits> { + static inline nimble::EncodingType encodingType = + nimble::EncodingType::MainlyConstant; + }; + + template <> + struct EncodingTypeTraits> { + static inline nimble::EncodingType encodingType = nimble::EncodingType::RLE; + }; + + template <> + struct EncodingTypeTraits { + static inline nimble::EncodingType encodingType = + nimble::EncodingType::SparseBool; + }; + + template <> + struct EncodingTypeTraits> { + static inline nimble::EncodingType encodingType = + nimble::EncodingType::Trivial; + }; + + template <> + struct EncodingTypeTraits> { + static inline nimble::EncodingType encodingType = + nimble::EncodingType::Varint; + }; + + std::unique_ptr createEncoding( + const nimble::Vector& values, + bool compress, + bool useVariableBitWidthCompressor) { + using physicalType = typename nimble::TypeTraits::physicalType; + auto physicalValues = std::span( + reinterpret_cast(values.data()), values.size()); + nimble::EncodingSelection selection{ + {.encodingType = EncodingTypeTraits::encodingType, + .compressionPolicyFactory = + [compress, useVariableBitWidthCompressor]() { + return std::make_unique( + compress, useVariableBitWidthCompressor); + }}, + nimble::Statistics::create(physicalValues), + std::make_unique>( + compress, useVariableBitWidthCompressor)}; + + auto encoded = C::encode(selection, physicalValues, *buffer_); + return std::make_unique(*this->pool_, encoded); + } + + // Each unit test runs on randomized data this many times before + // we conclude the unit test passed. + static constexpr int32_t kNumRandomRuns = 10; + + // We want the number of row tested to potentially be large compared to a + // skip block. When we actually generate data we pick a random length between + // 1 and this size. + static constexpr int32_t kMaxRows = 2000; + + std::shared_ptr pool_; + std::unique_ptr buffer_; + std::unique_ptr util_; +}; + +#define VARINT_TYPES(EncodingName) \ + EncodingName, EncodingName, EncodingName, \ + EncodingName, EncodingName, EncodingName + +#define NUMERIC_TYPES(EncodingName) \ + VARINT_TYPES(EncodingName), EncodingName, EncodingName, \ + EncodingName, EncodingName + +#define NON_BOOL_TYPES(EncodingName) \ + NUMERIC_TYPES(EncodingName), EncodingName + +#define ALL_TYPES(EncodingName) NON_BOOL_TYPES(EncodingName), EncodingName + +using TestTypes = ::testing::Types< + nimble::SparseBoolEncoding, + VARINT_TYPES(nimble::VarintEncoding), + NUMERIC_TYPES(nimble::FixedBitWidthEncoding), + NON_BOOL_TYPES(nimble::DictionaryEncoding), + NON_BOOL_TYPES(nimble::MainlyConstantEncoding), + ALL_TYPES(nimble::TrivialEncoding), + ALL_TYPES(nimble::RLEEncoding), + ALL_TYPES(nimble::ConstantEncoding)>; + +TYPED_TEST_CASE(EncodingTests, TestTypes); + +TYPED_TEST(EncodingTests, Materialize) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + using E = typename TypeParam::cppDataType; + + for (int run = 0; run < this->kNumRandomRuns; ++run) { + const std::vector> dataPatterns = + this->util_->template makeDataPatterns( + rng, this->kMaxRows, this->buffer_.get()); + for (const auto& data : dataPatterns) { + for (auto compress : {false, true}) { + for (auto useVariableBitWidthCompressor : {false, true}) { + const int rowCount = data.size(); + ASSERT_GT(rowCount, 0); + std::unique_ptr encoding; + try { + encoding = this->createEncoding( + data, compress, useVariableBitWidthCompressor); + } catch (const nimble::NimbleUserError& e) { + if (e.errorCode() == nimble::error_code::IncompatibleEncoding) { + continue; + } + throw; + } + ASSERT_EQ(encoding->dataType(), nimble::TypeTraits::dataType); + nimble::Vector buffer(this->pool_.get(), rowCount); + + encoding->materialize(rowCount, buffer.data()); + for (int i = 0; i < rowCount; ++i) { + ASSERT_EQ(buffer[i], data[i]) << "i: " << i; + } + + encoding->reset(); + const int firstBlock = rowCount / 2; + encoding->materialize(firstBlock, buffer.data()); + for (int i = 0; i < firstBlock; ++i) { + ASSERT_EQ(buffer[i], data[i]); + } + const int secondBlock = rowCount - firstBlock; + encoding->materialize(secondBlock, buffer.data()); + for (int i = 0; i < secondBlock; ++i) { + ASSERT_EQ(buffer[i], data[firstBlock + i]); + } + + encoding->reset(); + for (int i = 0; i < rowCount; ++i) { + encoding->materialize(1, buffer.data()); + ASSERT_EQ(buffer[0], data[i]); + } + + encoding->reset(); + int start = 0; + int len = 0; + for (int i = 0; i < rowCount; ++i) { + start += len; + len += 1; + if (start + len > rowCount) { + break; + } + encoding->materialize(len, buffer.data()); + for (int j = 0; j < len; ++j) { + ASSERT_EQ(data[start + j], buffer[j]); + } + } + + const uint32_t offset = folly::Random::rand32(rng) % data.size(); + const uint32_t length = + 1 + folly::Random::rand32(rng) % (data.size() - offset); + encoding->reset(); + encoding->skip(offset); + encoding->materialize(length, buffer.data()); + for (uint32_t i = 0; i < length; ++i) { + ASSERT_EQ(buffer[i], data[offset + i]); + } + } + } + } + } +} + +template +void checkScatteredOutput( + bool hasNulls, + size_t index, + const bool* scatter, + const T* actual, + const T* expected, + const char* nulls, + uint32_t& expectedRow) { + if (scatter[index]) { + ASSERT_EQ(actual[index], expected[expectedRow]); + ++expectedRow; + } + + if (hasNulls) { + ASSERT_EQ(scatter[index], nimble::bits::getBit(index, nulls)); + } +} + +TYPED_TEST(EncodingTests, ScatteredMaterialize) { + using E = typename TypeParam::cppDataType; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int run = 0; run < this->kNumRandomRuns; ++run) { + const std::vector> dataPatterns = + this->util_->template makeDataPatterns( + rng, this->kMaxRows, this->buffer_.get()); + for (const auto& data : dataPatterns) { + for (auto compress : {false, true}) { + for (auto useVariableBitWidthCompressor : {false, true}) { + const int rowCount = data.size(); + ASSERT_GT(rowCount, 0); + std::unique_ptr encoding; + try { + encoding = this->createEncoding( + data, compress, useVariableBitWidthCompressor); + } catch (const nimble::NimbleUserError& e) { + if (e.errorCode() == nimble::error_code::IncompatibleEncoding) { + continue; + } + throw; + } + ASSERT_EQ(encoding->dataType(), nimble::TypeTraits::dataType); + + int setBits = 0; + std::vector scatterSizes(rowCount + 1); + scatterSizes[0] = 0; + nimble::Vector scatter(this->pool_.get()); + while (setBits < rowCount) { + scatter.push_back(folly::Random::oneIn(2, rng)); + if (scatter.back()) { + scatterSizes[++setBits] = scatter.size(); + } + } + + auto newRowCount = scatter.size(); + auto requiredBytes = nimble::bits::bytesRequired(newRowCount); + // Note: Internally, some bit implementations use word boundaries to + // efficiently iterate on bitmaps. If the buffer doesn't end on a word + // boundary, this leads to ASAN buffer overflow (debug builds). So for + // now, we are allocating extra 7 bytes to make sure the buffer ends + // or exceeds a word boundary. + nimble::Buffer scatterBuffer{*this->pool_, requiredBytes + 7}; + nimble::Buffer nullsBuffer{*this->pool_, requiredBytes + 7}; + auto scatterPtr = scatterBuffer.reserve(requiredBytes); + auto nullsPtr = nullsBuffer.reserve(requiredBytes); + memset(scatterPtr, 0, requiredBytes); + nimble::bits::packBitmap(scatter, scatterPtr); + + nimble::Vector buffer(this->pool_.get(), newRowCount); + + uint32_t expectedRow = 0; + uint32_t actualRows = 0; + { + nimble::bits::Bitmap scatterBitmap(scatterPtr, newRowCount); + actualRows = encoding->materializeNullable( + rowCount, + buffer.data(), + [&]() { return nullsPtr; }, + &scatterBitmap); + } + ASSERT_EQ(actualRows, rowCount); + auto hasNulls = actualRows != newRowCount; + for (int i = 0; i < newRowCount; ++i) { + checkScatteredOutput( + hasNulls, + i, + scatter.data(), + buffer.data(), + data.data(), + nullsPtr, + expectedRow); + } + EXPECT_EQ(rowCount, expectedRow); + + encoding->reset(); + const int firstBlock = rowCount / 2; + const int firstScatterSize = scatterSizes[firstBlock]; + expectedRow = 0; + { + nimble::bits::Bitmap scatterBitmap(scatterPtr, firstScatterSize); + actualRows = encoding->materializeNullable( + firstBlock, + buffer.data(), + [&]() { return nullsPtr; }, + &scatterBitmap); + } + ASSERT_EQ(actualRows, firstBlock); + hasNulls = actualRows != firstScatterSize; + for (int i = 0; i < firstScatterSize; ++i) { + checkScatteredOutput( + hasNulls, + i, + scatter.data(), + buffer.data(), + data.data(), + nullsPtr, + expectedRow); + } + ASSERT_EQ(actualRows, expectedRow); + + const int secondBlock = rowCount - firstBlock; + const int secondScatterSize = scatter.size() - firstScatterSize; + expectedRow = actualRows; + { + nimble::bits::Bitmap scatterBitmap(scatterPtr, newRowCount); + actualRows = encoding->materializeNullable( + secondBlock, + buffer.data(), + [&]() { return nullsPtr; }, + &scatterBitmap, + firstScatterSize); + } + ASSERT_EQ(actualRows, secondBlock); + hasNulls = actualRows != secondScatterSize; + auto previousRows = expectedRow; + for (int i = firstScatterSize; i < newRowCount; ++i) { + checkScatteredOutput( + hasNulls, + i, + scatter.data(), + buffer.data(), + data.data(), + nullsPtr, + expectedRow); + } + ASSERT_EQ(actualRows, expectedRow - previousRows); + ASSERT_EQ(rowCount, expectedRow); + + encoding->reset(); + expectedRow = 0; + for (int i = 0; i < rowCount; ++i) { + // Note: Internally, some bit implementations use word boundaries to + // efficiently iterate on bitmaps. If the buffer doesn't end on a + // word boundary, this leads to ASAN buffer overflow (debug builds). + // So for now, we are using uint64_t as the bitmap to make sure the + // buffer ends on a word boundary. + auto scatterStart = scatterSizes[i]; + auto scatterEnd = scatterSizes[i + 1]; + { + nimble::bits::Bitmap scatterBitmap(scatterPtr, scatterEnd); + actualRows = encoding->materializeNullable( + 1, + buffer.data(), + [&]() { return nullsPtr; }, + &scatterBitmap, + scatterStart); + } + ASSERT_EQ(actualRows, 1); + previousRows = expectedRow; + for (int j = scatterStart; j < scatterEnd; ++j) { + checkScatteredOutput( + scatterEnd - scatterStart > 1, + j, + scatter.data(), + buffer.data(), + data.data(), + nullsPtr, + expectedRow); + } + ASSERT_EQ(actualRows, expectedRow - previousRows); + } + ASSERT_EQ(rowCount, expectedRow); + + encoding->reset(); + expectedRow = 0; + int start = 0; + int len = 0; + for (int i = 0; i < rowCount; ++i) { + start += len; + len += 1; + if (start + len > rowCount) { + break; + } + auto scatterStart = scatterSizes[start]; + auto scatterEnd = scatterSizes[start + len]; + { + nimble::bits::Bitmap scatterBitmap(scatterPtr, scatterEnd); + actualRows = encoding->materializeNullable( + len, + buffer.data(), + [&]() { return nullsPtr; }, + &scatterBitmap, + scatterStart); + } + ASSERT_EQ(actualRows, len); + hasNulls = actualRows != scatterEnd - scatterStart; + previousRows = expectedRow; + for (int j = scatterStart; j < scatterEnd; ++j) { + checkScatteredOutput( + hasNulls, + j, + scatter.data(), + buffer.data(), + data.data(), + nullsPtr, + expectedRow); + } + ASSERT_EQ(actualRows, expectedRow - previousRows); + } + } + } + } + } +} diff --git a/dwio/nimble/encodings/tests/MainlyConstantEncodingTests.cpp b/dwio/nimble/encodings/tests/MainlyConstantEncodingTests.cpp new file mode 100644 index 0000000..2e24def --- /dev/null +++ b/dwio/nimble/encodings/tests/MainlyConstantEncodingTests.cpp @@ -0,0 +1,103 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/EncodingType.h" +#include "dwio/nimble/common/NimbleCompare.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/encodings/MainlyConstantEncoding.h" +#include "dwio/nimble/encodings/tests/TestUtils.h" + +#include +#include + +using namespace facebook; + +template +class MainlyConstantEncodingTest : public ::testing::Test { + protected: + void SetUp() override { + pool_ = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + buffer_ = std::make_unique(*pool_); + } + + template + std::vector> prepareValues() { + FAIL() << "unspecialized prepapreValues() should not be called"; + return {}; + } + + template + nimble::Vector toVector(std::initializer_list l) { + nimble::Vector v{pool_.get()}; + v.insert(v.end(), l.begin(), l.end()); + return v; + } + + template + using ET = nimble::EncodingPhysicalType; + + double dNaN0 = std::numeric_limits::quiet_NaN(); + double dNaN1 = std::numeric_limits::signaling_NaN(); + double dNaN2 = ET::asEncodingLogicalType( + (ET::asEncodingPhysicalType(dNaN0) | 0x3)); + + template <> + std::vector> prepareValues() { + return { + toVector({0.0}), + toVector({0.0, 0.00, 0.12}), + toVector({-2.1, -2.1, -2.3, -2.1, -2.1}), + toVector({dNaN0, dNaN0, dNaN1, dNaN2, dNaN0})}; + } + + float fNaN0 = std::numeric_limits::quiet_NaN(); + float fNaN1 = std::numeric_limits::signaling_NaN(); + float fNaN2 = ET::asEncodingLogicalType( + (ET::asEncodingPhysicalType(fNaN0) | 0x3)); + + template <> + std::vector> prepareValues() { + return { + toVector({0.0f}), + toVector({0.0f, 0.00f, 0.12f}), + toVector({-2.1f, -2.1f, -2.3f, -2.1f, -2.1f}), + toVector({fNaN0, fNaN0, fNaN1, fNaN2, fNaN2})}; + } + + template <> + std::vector> prepareValues() { + return {toVector({3, 3, 3, 1, 3})}; + } + + std::shared_ptr pool_; + std::unique_ptr buffer_; +}; + +#define NUM_TYPES int32_t, double, float + +using TestTypes = ::testing::Types; + +TYPED_TEST_CASE(MainlyConstantEncodingTest, TestTypes); + +TYPED_TEST(MainlyConstantEncodingTest, SerializeThenDeserialize) { + using D = TypeParam; + + auto valueGroups = this->template prepareValues(); + for (const auto& values : valueGroups) { + auto encoding = nimble::test::Encoder>:: + createEncoding(*this->buffer_, values); + + uint32_t rowCount = values.size(); + nimble::Vector result(this->pool_.get(), rowCount); + encoding->materialize(rowCount, result.data()); + + EXPECT_EQ(encoding->encodingType(), nimble::EncodingType::MainlyConstant); + EXPECT_EQ(encoding->dataType(), nimble::TypeTraits::dataType); + EXPECT_EQ(encoding->rowCount(), rowCount); + for (uint32_t i = 0; i < rowCount; ++i) { + EXPECT_TRUE(nimble::NimbleCompare::equals(result[i], values[i])); + } + } +} diff --git a/dwio/nimble/encodings/tests/MapBenchmarks.cpp b/dwio/nimble/encodings/tests/MapBenchmarks.cpp new file mode 100644 index 0000000..973a424 --- /dev/null +++ b/dwio/nimble/encodings/tests/MapBenchmarks.cpp @@ -0,0 +1,463 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#include + +#include +#include "absl/container/flat_hash_map.h" +#include "common/init/light.h" +#include "folly/Benchmark.h" +#include "folly/Demangle.h" +#include "folly/container/F14Map.h" + +constexpr size_t kDataSize = 20 * 1024 * 1024; // 20M +constexpr size_t kRepeatSize = 5 * 1024; // 5K + +std::vector repeatingData; +std::vector uniqueData; + +template +struct TestParameters { + size_t elementCount; + bool preAllocate; + const std::vector& data; +}; + +template +struct is_robin_map : std::integral_constant< + bool, + std::is_same_v< + Map, + tsl::robin_map< + typename Map::key_type, + typename Map::mapped_type>> || + std::is_same_v< + Map, + tsl::robin_pg_map< + typename Map::key_type, + typename Map::mapped_type>>> {}; + +template ::value, bool> = true> +void incrementValue(Map& map, typename Map::key_type key) { + ++(map.try_emplace(key, 0).first->second); +} + +template ::value, bool> = true> +void incrementValue(Map& map, typename Map::key_type key) { + ++(map.try_emplace(key, 0).first.value()); +} + +template +void incrementCounts(size_t iters, const TestParameters& parameters) { + for (int i = 0; i < iters; ++i) { + Map map; + if (parameters.preAllocate) { + map.reserve(10000); + } + for (auto j = 0; j < parameters.elementCount; ++j) { + incrementValue(map, parameters.data[i]); + } + } +} + +BENCHMARK_NAMED_PARAM( + incrementCounts, + UnorderedMapInt64Repeating20K, + TestParameters>{ + .elementCount = 20000, + .preAllocate = true, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + UnorderedMapInt64Repeating20KNoReserve, + TestParameters>{ + .elementCount = 20000, + .preAllocate = false, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinMapInt64Repeating20K, + TestParameters>{ + .elementCount = 20000, + .preAllocate = true, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinMapInt64Repeating20KNoReserve, + TestParameters>{ + .elementCount = 20000, + .preAllocate = false, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinPGMapInt64Repeating20K, + TestParameters>{ + .elementCount = 20000, + .preAllocate = true, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinPGMapInt64Repeating20KNoReserve, + TestParameters>{ + .elementCount = 20000, + .preAllocate = false, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14ValueMapInt64Repeating20K, + TestParameters>{ + .elementCount = 20000, + .preAllocate = true, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14ValueMapInt64Repeating20KNoReserve, + TestParameters>{ + .elementCount = 20000, + .preAllocate = false, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14NodeMapInt64Repeating20K, + TestParameters>{ + .elementCount = 20000, + .preAllocate = true, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14NodeMapInt64Repeating20KNoReserve, + TestParameters>{ + .elementCount = 20000, + .preAllocate = false, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + AbslMapInt64Repeating20K, + TestParameters>{ + .elementCount = 20000, + .preAllocate = true, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + AbslMapInt64Repeating20KNoReserve, + TestParameters>{ + .elementCount = 20000, + .preAllocate = false, + .data = repeatingData}) + +BENCHMARK_DRAW_LINE(); + +BENCHMARK_NAMED_PARAM( + incrementCounts, + UnorderedMapInt64Repeating20M, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = true, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + UnorderedMapInt64Repeating20MNoReserve, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = false, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinMapInt64Repeating20M, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = true, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinMapInt64Repeating20MNoReserve, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = false, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinPGMapInt64Repeating20M, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = true, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinPGMapInt64Repeating20MNoReserve, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = false, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14ValueMapInt64Repeating20M, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = true, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14ValueMapInt64Repeating20MNoReserve, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = false, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14NodeMapInt64Repeating20M, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = true, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14NodeMapInt64Repeating20MNoReserve, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = false, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + AbslMapInt64Repeating20M, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = true, + .data = repeatingData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + AbslMapInt64Repeating20MNoReserve, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = false, + .data = repeatingData}) + +BENCHMARK_DRAW_LINE(); + +BENCHMARK_NAMED_PARAM( + incrementCounts, + UnorderedMapInt64Unique20K, + TestParameters>{ + .elementCount = 20000, + .preAllocate = true, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + UnorderedMapInt64Unique20KNoReserve, + TestParameters>{ + .elementCount = 20000, + .preAllocate = false, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinMapInt64Unique20K, + TestParameters>{ + .elementCount = 20000, + .preAllocate = true, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinMapInt64Unique20KNoReserve, + TestParameters>{ + .elementCount = 20000, + .preAllocate = false, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinPGMapInt64Unique20K, + TestParameters>{ + .elementCount = 20000, + .preAllocate = true, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinPGMapInt64Unique20KNoReserve, + TestParameters>{ + .elementCount = 20000, + .preAllocate = false, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14ValueMapInt64Unique20K, + TestParameters>{ + .elementCount = 20000, + .preAllocate = true, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14ValueMapInt64Unique20KNoReserve, + TestParameters>{ + .elementCount = 20000, + .preAllocate = false, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14NodeMapInt64Unique20K, + TestParameters>{ + .elementCount = 20000, + .preAllocate = true, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14NodeMapInt64Unique20KNoReserve, + TestParameters>{ + .elementCount = 20000, + .preAllocate = false, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + AbslMapInt64Unique20K, + TestParameters>{ + .elementCount = 20000, + .preAllocate = true, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + AbslMapInt64Unique20KNoReserve, + TestParameters>{ + .elementCount = 20000, + .preAllocate = false, + .data = uniqueData}) + +BENCHMARK_DRAW_LINE(); + +BENCHMARK_NAMED_PARAM( + incrementCounts, + UnorderedMapInt64Unique20M, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = true, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + UnorderedMapInt64Unique20MNoReserve, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = false, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinMapInt64Unique20M, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = true, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinMapInt64Unique20MNoReserve, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = false, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinPGMapInt64Unique20M, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = true, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + RobinPGMapInt64Unique20MNoReserve, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = false, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14ValueMapInt64Unique20M, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = true, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14ValueMapInt64Unique20MNoReserve, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = false, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14NodeMapInt64Unique20M, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = true, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + F14NodeMapInt64Unique20MNoReserve, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = false, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + AbslMapInt64Unique20M, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = true, + .data = uniqueData}) + +BENCHMARK_RELATIVE_NAMED_PARAM( + incrementCounts, + AbslMapInt64Unique20MNoReserve, + TestParameters>{ + .elementCount = 2000000, + .preAllocate = false, + .data = uniqueData}) + +int main(int argc, char** argv) { + facebook::init::initFacebookLight(&argc, &argv); + repeatingData.resize(kDataSize); + uniqueData.resize(kDataSize); + for (auto i = 0; i < kDataSize; ++i) { + repeatingData[i] = i % kRepeatSize - (kRepeatSize / 2); + uniqueData[i] = i - kRepeatSize; + } + + folly::runBenchmarks(); + return 0; +} diff --git a/dwio/nimble/encodings/tests/NullableEncodingTests.cpp b/dwio/nimble/encodings/tests/NullableEncodingTests.cpp new file mode 100644 index 0000000..46b6fa1 --- /dev/null +++ b/dwio/nimble/encodings/tests/NullableEncodingTests.cpp @@ -0,0 +1,528 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include +#include +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/common/tests/TestUtils.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" +#include "dwio/nimble/encodings/NullableEncoding.h" +#include "dwio/nimble/encodings/SentinelEncoding.h" +#include "dwio/nimble/encodings/TrivialEncoding.h" +#include "dwio/nimble/encodings/tests/TestUtils.h" +#include "folly/Random.h" +#include "velox/common/memory/Memory.h" + +// Tests the Encoding API for all nullable Encoding implementations + data +// types. +// +// These encodings generally use the factory themselves to encode their non-null +// values as another encoding. We assume that the other encodings are thoroughly +// tested and conform to the API, so we can use just a single underlying +// encoding implementation for the non-null values (namely, the TrivialEncoding) +// rather than having to test all the numNullableEncodings X numNormalEncodings +// combinations. + +using namespace ::facebook; + +namespace { +enum class NullsPattern { + None, + All, + Random, +}; +} + +// C is the encoding type. +template +class NullableEncodingTest : public ::testing::Test { + protected: + using E = typename C::cppDataType; + + void SetUp() override { + pool_ = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + buffer_ = std::make_unique(*pool_); + util_ = std::make_unique(*pool_); + } + + // Makes a random-length nulls vector with num_nonNulls values set to true + // (and at least one value set to false) scattered randomly throughout. + template + nimble::Vector + makeRandomNulls(RNG&& rng, uint32_t dataSize, NullsPattern pattern) { + if (pattern == NullsPattern::None) { + return nimble::Vector{pool_.get(), dataSize, false}; + } else if (pattern == NullsPattern::All) { + return nimble::Vector{pool_.get(), dataSize, true}; + } + const uint32_t rowCount = + dataSize + folly::Random::rand32(3 * kMaxRows, std::forward(rng)); + nimble::Vector nulls(pool_.get(), rowCount, false); + for (uint32_t i = 0; i < dataSize; ++i) { + nulls[i] = true; + } + std::random_shuffle(nulls.begin(), nulls.end()); + return nulls; + } + + // Each unit test runs on randomized data this many times before + // we conclude the unit test passed. + static constexpr int kNumRandomRuns = 20; + // We want the number of row tested to potentially be large compared to a + // skip block. When we actually generate data we pick a random length between + // 1 and this size. + static constexpr int kMaxRows = 2000; + + std::shared_ptr pool_; + std::unique_ptr buffer_; + std::unique_ptr util_; +}; + +#define ALL_TYPES(EncodingName) \ + EncodingName, EncodingName, EncodingName, \ + EncodingName, EncodingName, EncodingName, \ + EncodingName, EncodingName + +#define NON_BOOL_TYPES(EncodingName) \ + EncodingName, EncodingName, EncodingName, \ + EncodingName, EncodingName, EncodingName, \ + EncodingName + +using TestTypes = ::testing::Types< + ALL_TYPES(nimble::NullableEncoding), + NON_BOOL_TYPES(nimble::SentinelEncoding)>; + +TYPED_TEST_CASE(NullableEncodingTest, TestTypes); + +//.Spreads the nonNulls out into a vector of length |nulls|, with a non-null +// placed at each true value in |nulls|. Equivalent to Encoding::Materialize. +template +nimble::Vector spreadNullsIntoData( + velox::memory::MemoryPool& memoryPool, + std::span nonNulls, + std::span nulls) { + nimble::Vector result(&memoryPool); + auto nonNullsIt = nonNulls.begin(); + for (auto nulls_it = nulls.begin(); nulls_it < nulls.end(); ++nulls_it) { + if (*nulls_it) { + result.push_back(*nonNullsIt++); + } else { + result.push_back(E()); + } + } + + return result; +} + +TYPED_TEST(NullableEncodingTest, Materialize) { + using E = typename TypeParam::cppDataType; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int run = 0; run < this->kNumRandomRuns; ++run) { + const std::vector> dataPatterns = + this->util_->template makeDataPatterns( + rng, this->kMaxRows, this->buffer_.get()); + for (const auto& data : dataPatterns) { + for (auto nullPattern : + {NullsPattern::None, NullsPattern::All, NullsPattern::Random}) { + const nimble::Vector nulls = this->makeRandomNulls( + rng, folly::to(data.size()), nullPattern); + // Spreading the data out will help us check correctness more easily. + const nimble::Vector spreadData = + spreadNullsIntoData(*this->pool_, data, nulls); + + auto encoding = nimble::test::Encoder>:: + createNullableEncoding(*this->buffer_, data, nulls); + ASSERT_EQ(encoding->dataType(), nimble::TypeTraits::dataType); + ASSERT_TRUE(encoding->isNullable()); + const uint32_t rowCount = encoding->rowCount(); + + nimble::Vector buffer(this->pool_.get(), rowCount); + encoding->materialize(rowCount, buffer.data()); + for (int i = 0; i < rowCount; ++i) { + ASSERT_EQ(buffer[i], spreadData[i]); + } + + encoding->reset(); + const int firstBlock = folly::to(rowCount / 2); + encoding->materialize(firstBlock, buffer.data()); + for (int i = 0; i < firstBlock; ++i) { + ASSERT_EQ(buffer[i], spreadData[i]); + } + const int secondBlock = rowCount - firstBlock; + encoding->materialize(secondBlock, buffer.data()); + for (int i = 0; i < secondBlock; ++i) { + ASSERT_EQ(buffer[i], spreadData[firstBlock + i]); + } + + encoding->reset(); + for (int i = 0; i < rowCount; ++i) { + encoding->materialize(1, buffer.data()); + ASSERT_EQ(buffer[0], spreadData[i]); + } + + encoding->reset(); + int start = 0; + int len = 0; + for (int i = 0; i < rowCount; ++i) { + start += len; + len += 1; + if (start + len > rowCount) { + break; + } + encoding->materialize(len, buffer.data()); + for (int j = 0; j < len; ++j) { + ASSERT_EQ(spreadData[start + j], buffer[j]); + } + } + + const uint32_t offset = + folly::to(folly::Random::rand32(rng) % data.size()); + const uint32_t length = folly::to( + 1 + folly::Random::rand32(rng) % (data.size() - offset)); + encoding->reset(); + encoding->skip(offset); + encoding->materialize(length, buffer.data()); + for (uint32_t i = 0; i < length; ++i) { + ASSERT_EQ(buffer[i], spreadData[offset + i]); + } + } + } + } +} + +template +void checkOutput( + size_t index, + const bool* nulls, + const T* data, + const char* actualNulls, + const T* actualData, + bool hasNulls) { + if (nulls[index]) { + ASSERT_EQ(data[index], actualData[index]) << index; + } + if (hasNulls) { + ASSERT_EQ(nimble::bits::getBit(index, actualNulls), nulls[index]) << index; + } +} + +TYPED_TEST(NullableEncodingTest, ScatteredMaterialize) { + using E = typename TypeParam::cppDataType; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (int run = 0; run < this->kNumRandomRuns; ++run) { + const std::vector> dataPatterns = + this->util_->template makeDataPatterns( + rng, this->kMaxRows, this->buffer_.get()); + for (const auto& data : dataPatterns) { + for (auto nullPattern : + {NullsPattern::None, NullsPattern::All, NullsPattern::Random}) { + const nimble::Vector nulls = + this->makeRandomNulls(rng, data.size(), nullPattern); + // Spreading the data out will help us check correctness more easily. + const nimble::Vector spreadData = + spreadNullsIntoData(*this->pool_, data, nulls); + + auto encoding = nimble::test::Encoder>:: + createNullableEncoding(*this->buffer_, data, nulls); + ASSERT_EQ(encoding->dataType(), nimble::TypeTraits::dataType); + ASSERT_TRUE(encoding->isNullable()); + const uint32_t rowCount = encoding->rowCount(); + ASSERT_EQ(rowCount, nulls.size()); + + int setBits = 0; + std::vector scatterSizes(rowCount + 1); + scatterSizes[0] = 0; + nimble::Vector scatter(this->pool_.get()); + while (setBits < rowCount) { + scatter.push_back(folly::Random::rand32(2, rng) ? true : false); + if (scatter.back()) { + scatterSizes[++setBits] = scatter.size(); + } + } + + auto newRowCount = scatter.size(); + auto requiredBytes = nimble::bits::bytesRequired(newRowCount); + // Note: Internally, some bit implementations use word boundaries to + // efficiently iterate on bitmaps. If the buffer doesn't end on a word + // boundary, this leads to ASAN buffer overflow (debug builds). So for + // now, we are allocating extra 7 bytes to make sure the buffer ends or + // exceeds a word boundary. + nimble::Buffer scatterBuffer{*this->pool_, requiredBytes + 7}; + nimble::Buffer nullsBuffer{*this->pool_, requiredBytes + 7}; + auto scatterPtr = scatterBuffer.reserve(requiredBytes); + auto nullsPtr = nullsBuffer.reserve(requiredBytes); + memset(scatterPtr, 0, requiredBytes); + nimble::bits::packBitmap(scatter, scatterPtr); + + nimble::Vector buffer(this->pool_.get(), newRowCount); + + auto test = [&encoding, &scatter, &nulls, &spreadData]( + uint32_t rowCount, + E* buffer, + void* nullsBitmap, + uint32_t scatterCount, + void* scatterBitmap, + uint32_t scatterOffset = 0, + uint32_t expectedOffset = 0) { + uint32_t expectedRow = 0; + nimble::bits::Bitmap bitmap{ + scatterBitmap, scatterOffset + scatterCount}; + auto nonNullCount = encoding->materializeNullable( + rowCount, + buffer, + [&]() { return nullsBitmap; }, + &bitmap, + scatterOffset); + for (int i = 0; i < scatterCount; ++i) { + auto isSet = false; + if (scatter[i + scatterOffset]) { + if (nulls[expectedRow + expectedOffset]) { + ASSERT_EQ( + buffer[i + scatterOffset], + spreadData[expectedRow + expectedOffset]); + isSet = true; + } + ++expectedRow; + } + if (nonNullCount != scatterCount) { + ASSERT_EQ( + isSet, + nimble::bits::getBit( + i + scatterOffset, + reinterpret_cast(nullsBitmap))); + } + } + + ASSERT_EQ(rowCount, expectedRow); + }; + + // Test reading all data + test(rowCount, buffer.data(), nullsPtr, newRowCount, scatterPtr); + + encoding->reset(); + const int firstBlock = newRowCount / 2; + + auto firstBlockSetBits = + nimble::bits::countSetBits(0, firstBlock, scatterPtr); + + // Test reading first half of the data + test( + firstBlockSetBits, buffer.data(), nullsPtr, firstBlock, scatterPtr); + + const int secondBlock = newRowCount - firstBlock; + + // Test reading second half of the data + test( + nimble::bits::countSetBits(firstBlock, secondBlock, scatterPtr), + buffer.data(), + nullsPtr, + secondBlock, + scatterPtr, + /* scatterOffset */ firstBlock, + /* expectedOffset */ firstBlockSetBits); + + encoding->reset(); + uint32_t expectedRow = 0; + for (int i = 0; i < rowCount; ++i) { + // Note: Internally, some bit implementations use word boundaries to + // efficiently iterate on bitmaps. If the buffer doesn't end on a word + // boundary, this leads to ASAN buffer overflow (debug builds). So for + // now, we are using uint64_t as the bitmap to make sure the buffer + // ends on a word boundary. + auto scatterStart = scatterSizes[i]; + auto scatterSize = scatterSizes[i + 1] - scatterStart; + + // Test reading one item at a time + test( + 1, + buffer.data(), + nullsPtr, + scatterSize, + scatterPtr, + /* scatterOffset */ scatterStart, + /* expectedOffset */ expectedRow); + + ++expectedRow; + } + + encoding->reset(); + expectedRow = 0; + int start = 0; + int len = 0; + while (true) { + start += len; + len += 1; + if (start + len > rowCount) { + break; + } + auto scatterStart = scatterSizes[start]; + auto scatterSize = scatterSizes[start + len] - scatterStart; + + // Test reading different ranges of data + test( + len, + buffer.data(), + nullsPtr, + scatterSize, + scatterPtr, + /* scatterOffset */ scatterStart, + /* expectedOffset */ expectedRow); + + expectedRow += len; + } + } + } + } +} + +TYPED_TEST(NullableEncodingTest, MaterializeNullable) { + using E = typename TypeParam::cppDataType; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (auto nullPattern : + {NullsPattern::None, NullsPattern::All, NullsPattern::Random}) { + for (int run = 0; run < this->kNumRandomRuns; ++run) { + const std::vector> dataPatterns = + this->util_->template makeDataPatterns( + rng, this->kMaxRows, this->buffer_.get()); + for (const auto& data : dataPatterns) { + const nimble::Vector nulls = + this->makeRandomNulls(rng, data.size(), nullPattern); + const nimble::Vector spreadData = + spreadNullsIntoData(*this->pool_, data, nulls); + + auto encoding = nimble::test::Encoder>:: + createNullableEncoding(*this->buffer_, data, nulls); + ASSERT_TRUE(encoding->isNullable()); + const uint32_t rowCount = encoding->rowCount(); + nimble::Vector buffer(this->pool_.get(), rowCount); + nimble::Vector bitmap(this->pool_.get(), rowCount); + + auto nonNullCount = encoding->materializeNullable( + rowCount, buffer.data(), [&]() { return bitmap.data(); }); + EXPECT_EQ( + std::accumulate(nulls.data(), nulls.data() + nulls.size(), 0), + nonNullCount); + + for (int i = 0; i < rowCount; ++i) { + checkOutput( + i, + nulls.data(), + spreadData.data(), + bitmap.data(), + buffer.data(), + nonNullCount != rowCount); + } + + encoding->reset(); + const int firstBlock = rowCount / 2; + nonNullCount = encoding->materializeNullable( + firstBlock, buffer.data(), [&]() { return bitmap.data(); }); + EXPECT_EQ( + std::accumulate(nulls.data(), nulls.data() + firstBlock, 0), + nonNullCount); + + for (int i = 0; i < firstBlock; ++i) { + checkOutput( + i, + nulls.data(), + spreadData.data(), + bitmap.data(), + buffer.data(), + nonNullCount != firstBlock); + } + const int secondBlock = rowCount - firstBlock; + nonNullCount = encoding->materializeNullable( + secondBlock, buffer.data(), [&]() { return bitmap.data(); }); + EXPECT_EQ( + std::accumulate( + nulls.data() + firstBlock, nulls.data() + rowCount, 0), + nonNullCount); + + for (int i = 0; i < secondBlock; ++i) { + checkOutput( + i, + nulls.data() + firstBlock, + spreadData.data() + firstBlock, + bitmap.data(), + buffer.data(), + nonNullCount != secondBlock); + } + + encoding->reset(); + for (int i = 0; i < rowCount; ++i) { + nonNullCount = encoding->materializeNullable( + 1, buffer.data(), [&]() { return bitmap.data(); }); + checkOutput( + 0, + nulls.data() + i, + spreadData.data() + i, + bitmap.data(), + buffer.data(), + nonNullCount == 0); + } + + encoding->reset(); + int start = 0; + int len = 0; + for (int i = 0; i < rowCount; ++i) { + start += len; + len += 1; + if (start + len > rowCount) { + break; + } + nonNullCount = encoding->materializeNullable( + len, buffer.data(), [&]() { return bitmap.data(); }); + EXPECT_EQ( + std::accumulate( + nulls.data() + start, nulls.data() + start + len, 0), + nonNullCount); + for (int j = 0; j < len; ++j) { + checkOutput( + j, + nulls.data() + start, + spreadData.data() + start, + bitmap.data(), + buffer.data(), + nonNullCount != len); + } + } + + const uint32_t offset = folly::Random::rand32(rng) % data.size(); + const uint32_t length = + 1 + folly::Random::rand32(rng) % (data.size() - offset); + encoding->reset(); + encoding->skip(offset); + nonNullCount = encoding->materializeNullable( + length, buffer.data(), [&]() { return bitmap.data(); }); + EXPECT_EQ( + std::accumulate( + nulls.data() + offset, nulls.data() + offset + length, 0), + nonNullCount); + for (uint32_t i = 0; i < length; ++i) { + checkOutput( + i, + nulls.data() + offset, + spreadData.data() + offset, + bitmap.data(), + buffer.data(), + nonNullCount != length); + } + } + } + } +} diff --git a/dwio/nimble/encodings/tests/RleEncodingTests.cpp b/dwio/nimble/encodings/tests/RleEncodingTests.cpp new file mode 100644 index 0000000..3f14a42 --- /dev/null +++ b/dwio/nimble/encodings/tests/RleEncodingTests.cpp @@ -0,0 +1,152 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/EncodingType.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/encodings/RleEncoding.h" +#include "dwio/nimble/encodings/tests/TestUtils.h" + +#include +#include + +using namespace facebook; + +template +class RleEncodingTest : public ::testing::Test { + protected: + void SetUp() override { + pool_ = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + buffer_ = std::make_unique(*pool_); + } + + template + std::vector> prepareValues() { + FAIL() << "unspecialized prepapreValues() should not be called"; + return {}; + } + + template + nimble::Vector toVector(std::initializer_list l) { + nimble::Vector v{pool_.get()}; + v.insert(v.end(), l.begin(), l.end()); + return v; + } + + template + using ET = nimble::EncodingPhysicalType; + + double dNaN0 = std::numeric_limits::quiet_NaN(); + double dNaN1 = std::numeric_limits::signaling_NaN(); + double dNaN2 = ET::asEncodingLogicalType( + (ET::asEncodingPhysicalType(dNaN0) | 0x3)); + double dNaN3 = ET::asEncodingLogicalType( + (ET::asEncodingPhysicalType(dNaN0) | 0x5)); + + template <> + std::vector> prepareValues() { + return { + toVector({0.0, -0.0}), + toVector({-0.0, -0.0, -0.0, -0.0, -0.0, -0.0, -0.0, -0.0, -0.0, -0.0}), + toVector({-0.0, -0.0, -0.0, +0.0, -0.0, +0.0, -0.0, -0.0, -0.0, -0.0}), + toVector({-2.1, -0.0, -0.0, 3.54, 9.87, -0.0, -0.0, -0.0, -0.0, 10.6}), + toVector({0.00, 1.11, 2.22, 3.33, 4.44, 5.55, 6.66, 7.77, 8.88, 9.99}), + toVector( + {dNaN0, dNaN0, dNaN0, dNaN1, dNaN1, dNaN2, dNaN3, dNaN3, dNaN0})}; + } + + float fNaN0 = std::numeric_limits::quiet_NaN(); + float fNaN1 = std::numeric_limits::signaling_NaN(); + float fNaN2 = ET::asEncodingLogicalType( + (ET::asEncodingPhysicalType(fNaN0) | 0x3)); + float fNaN3 = ET::asEncodingLogicalType( + (ET::asEncodingPhysicalType(fNaN0) | 0x5)); + + template <> + std::vector> prepareValues() { + return { + toVector({0.0f, -0.0f}), + toVector( + {-0.0f, + -0.0f, + -0.0f, + -0.0f, + -0.0f, + -0.0f, + -0.0f, + -0.0f, + -0.0f, + -0.0f}), + toVector( + {-0.0f, + -0.0f, + -0.0f, + +0.0f, + -0.0f, + +0.0f, + -0.0f, + -0.0f, + -0.0f, + -0.0f}), + toVector( + {-2.1f, + -0.0f, + -0.0f, + 3.54f, + 9.87f, + -0.0f, + -0.0f, + -0.0f, + -0.0f, + 10.6f}), + toVector( + {0.00f, + 1.11f, + 2.22f, + 3.33f, + 4.44f, + 5.55f, + 6.66f, + 7.77f, + 8.88f, + 9.99f}), + toVector( + {fNaN0, fNaN0, fNaN0, fNaN1, fNaN1, fNaN2, fNaN3, fNaN3, fNaN0})}; + } + + template <> + std::vector> prepareValues() { + return {toVector({2, 3, 3}), toVector({1, 2, 2, 3, 3, 3, 4, 4, 4, 4})}; + } + + std::shared_ptr pool_; + std::unique_ptr buffer_; +}; + +#define NUM_TYPES int32_t, double, float + +using TestTypes = ::testing::Types; + +TYPED_TEST_CASE(RleEncodingTest, TestTypes); + +TYPED_TEST(RleEncodingTest, SerializeThenDeserialize) { + using D = TypeParam; + + auto valueGroups = this->template prepareValues(); + for (const auto& values : valueGroups) { + auto encoding = + nimble::test::Encoder>::createEncoding( + *this->buffer_, values); + uint32_t rowCount = values.size(); + nimble::Vector result(this->pool_.get(), rowCount); + encoding->materialize(rowCount, result.data()); + + EXPECT_EQ(encoding->encodingType(), nimble::EncodingType::RLE); + EXPECT_EQ(encoding->dataType(), nimble::TypeTraits::dataType); + EXPECT_EQ(encoding->rowCount(), rowCount); + for (uint32_t i = 0; i < rowCount; ++i) { + EXPECT_TRUE(nimble::NimbleCompare::equals(result[i], values[i])); + } + } +} diff --git a/dwio/nimble/encodings/tests/SentinelEncodingTests.cpp b/dwio/nimble/encodings/tests/SentinelEncodingTests.cpp new file mode 100644 index 0000000..f18928f --- /dev/null +++ b/dwio/nimble/encodings/tests/SentinelEncodingTests.cpp @@ -0,0 +1,142 @@ +// // Copyright 2004-present Facebook. All Rights Reserved. + +// #include +// #include +// #include "dwio/nimble/common/Buffer.h" +// #include "dwio/nimble/common/Vector.h" +// #include "dwio/nimble/encodings/SentinelEncoding.h" + +// #include + +// using namespace facebook; + +// class SentinelEncodingTest : public ::testing::Test { +// protected: +// void SetUp() override { +// pool_ = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); +// buffer_ = std::make_unique(*pool_); +// } + +// std::shared_ptr pool_; +// std::unique_ptr buffer_; +// }; + +// namespace { +// template +// void serializeAndDeserializeNullableValues( +// velox::memory::MemoryPool* memoryPool, +// nimble::Buffer* buffer, +// uint32_t count, +// uint32_t minNonNullCount, +// std::function valueFunc, +// std::function nullFunc, +// bool sentinelFound) { +// nimble::Vector values(memoryPool); +// nimble::Vector valuesNullRemoved(memoryPool); +// nimble::Vector nulls(memoryPool); + +// for (uint32_t i = 0; i < count; ++i) { +// auto notNull = nullFunc(count, i); +// nulls.push_back(notNull); +// auto v = valueFunc(count, i); +// values.push_back(v); +// if (notNull) { +// valuesNullRemoved.push_back(v); +// } +// } + +// std::string sentinelString; +// auto sentinelValue = +// nimble::findSentinelValue(valuesNullRemoved, &sentinelString); +// EXPECT_EQ(sentinelValue.has_value(), sentinelFound) +// << (sentinelFound ? "Expect to find sentinel value" +// : "Expect not able to find sentinel value"); + +// ASSERT_GE(valuesNullRemoved.size(), minNonNullCount) +// << "Too few non-null values"; + +// auto serializedData = +// nimble::SentinelEncoding::serialize(valuesNullRemoved, nulls, +// buffer); +// auto sentinelEncoding = +// std::make_unique>(*memoryPool, +// serializedData); +// nimble::Vector valuesResult(memoryPool, count); +// auto requiredBytes = nimble::bits::bytesRequired(count); +// nimble::Buffer nullsBuffer{*memoryPool, requiredBytes + 7}; +// auto nullsPtr = nullsBuffer.reserve(requiredBytes); +// sentinelEncoding->materializeNullable(count, valuesResult.data(), +// nullsPtr); EXPECT_EQ(sentinelEncoding->encodingType(), +// nimble::EncodingType::Sentinel); EXPECT_EQ(sentinelEncoding->dataType(), +// nimble::TypeTraits::dataType); EXPECT_EQ(sentinelEncoding->rowCount(), +// count); for (uint32_t i = 0; i < count; ++i) { +// EXPECT_EQ( +// nulls[i], +// nimble::bits::getBit(i, reinterpret_cast(nullsPtr))) +// << "Wrong null value at index " << i; +// if (nulls[i]) { +// EXPECT_EQ(values[i], valuesResult[i]) << "Wrong value at i " << i; +// } +// } +// } +// } // namespace + +// TEST_F(SentinelEncodingTest, SerializeAndDeserialize) { +// uint32_t count = 300; +// uint32_t minNonNullCount = 256; +// auto valueFunc = [](uint32_t /* count */, uint32_t i) -> uint8_t { +// uint8_t v = i % 256; +// return v == 100 ? 99 : v; +// }; + +// auto nullsFunc = [](uint32_t /* count */, uint32_t i) -> bool { +// if (i >= 256) { +// return i % 5 != 0; +// } + +// return true; +// }; + +// serializeAndDeserializeNullableValues( +// pool_.get(), +// this->buffer_.get(), +// count, +// minNonNullCount, +// valueFunc, +// nullsFunc, +// true); +// } + +// TEST_F(SentinelEncodingTest, CannotFoundSentinel) { +// uint32_t count = 300; +// uint32_t minNonNullCount = 256; +// auto valueFunc = [](uint32_t /* count */, uint32_t i) -> uint8_t { +// return i % 256; +// }; + +// auto nullsFunc = [](uint32_t /* count */, uint32_t i) -> bool { +// if (i >= 256) { +// return i % 5 != 0; +// } + +// return true; +// }; + +// try { +// serializeAndDeserializeNullableValues( +// pool_.get(), +// this->buffer_.get(), +// count, +// minNonNullCount, +// valueFunc, +// nullsFunc, +// false); +// } catch (const nimble::NimbleUserError& e) { +// EXPECT_EQ( +// "Cannot use SentinelEncoding when no value is left for sentinel.", +// e.errorMessage()); +// return; +// } + +// FAIL() << "Expect nimble::NimbleUserError"; +// } diff --git a/dwio/nimble/encodings/tests/StatisticsTests.cpp b/dwio/nimble/encodings/tests/StatisticsTests.cpp new file mode 100644 index 0000000..e540e08 --- /dev/null +++ b/dwio/nimble/encodings/tests/StatisticsTests.cpp @@ -0,0 +1,383 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include +#include +#include +#include +#include + +#include "dwio/nimble/common/Bits.h" +#include "dwio/nimble/common/Varint.h" +#include "dwio/nimble/encodings/Statistics.h" + +using namespace ::facebook; + +#define INTEGRAL_TYPES(StatisticType) \ + StatisticType, StatisticType, StatisticType, \ + StatisticType, StatisticType, \ + StatisticType, StatisticType, StatisticType + +#define NUMERIC_TYPES(StatisticType) \ + INTEGRAL_TYPES(StatisticType), StatisticType, StatisticType + +using IntegerTypes = ::testing::Types; +using NumericTypes = ::testing::Types; +using BoolTypes = ::testing::Types>; +using StringTypes = + ::testing::Types>; + +TYPED_TEST_CASE(StatisticsNumericTests, NumericTypes); +TYPED_TEST_CASE(StatisticsIntegerTests, IntegerTypes); +TYPED_TEST_CASE(StatisticsBoolTests, BoolTypes); +TYPED_TEST_CASE(StatisticsStringTests, StringTypes); + +template +class StatisticsNumericTests : public ::testing::Test {}; + +template +class StatisticsIntegerTests : public ::testing::Test {}; + +template +class StatisticsBoolTests : public ::testing::Test {}; + +template +class StatisticsStringTests : public ::testing::Test {}; + +TYPED_TEST(StatisticsNumericTests, Create) { + using T = TypeParam; + using ValueType = typename T::valueType; + + constexpr auto dataSize = 100; + constexpr auto repeatSize = 25; + constexpr auto offset = 10; + static_assert(dataSize % repeatSize == 0); + + ValueType minValue = + std::is_signed() ? (ValueType)-offset : (ValueType)offset; + + std::vector data; + data.resize(dataSize); + + // Repeating data (non-consecutive) + for (auto i = 0; i < data.size(); ++i) { + data[i] = minValue + (i % repeatSize); + } + + auto statistics = T::create({data}); + EXPECT_EQ(repeatSize, statistics.uniqueCounts().size()); + EXPECT_EQ(minValue, statistics.min()); + EXPECT_EQ(minValue + repeatSize - 1, statistics.max()); + EXPECT_EQ(1, statistics.minRepeat()); + EXPECT_EQ(1, statistics.maxRepeat()); + EXPECT_EQ(dataSize, statistics.consecutiveRepeatCount()); + + for (const auto& pair : statistics.uniqueCounts()) { + EXPECT_EQ(dataSize / repeatSize, pair.second); + } + + // Repeating data (consecutive) + std::sort(data.begin(), data.end()); + statistics = T::create({data}); + EXPECT_EQ(repeatSize, statistics.uniqueCounts().size()); + EXPECT_EQ(minValue, statistics.min()); + EXPECT_EQ(minValue + repeatSize - 1, statistics.max()); + EXPECT_EQ(dataSize / repeatSize, statistics.minRepeat()); + EXPECT_EQ(dataSize / repeatSize, statistics.maxRepeat()); + EXPECT_EQ(repeatSize, statistics.consecutiveRepeatCount()); + + for (const auto& pair : statistics.uniqueCounts()) { + EXPECT_EQ(dataSize / repeatSize, pair.second); + } + + // Unique data + for (auto i = 0; i < data.size(); ++i) { + data[i] = minValue + i; + } + + statistics = T::create({data}); + EXPECT_EQ(dataSize, statistics.uniqueCounts().size()); + EXPECT_EQ(minValue, statistics.min()); + EXPECT_EQ(minValue + dataSize - 1, statistics.max()); + EXPECT_EQ(1, statistics.minRepeat()); + EXPECT_EQ(1, statistics.maxRepeat()); + EXPECT_EQ(dataSize, statistics.consecutiveRepeatCount()); + + for (const auto& pair : statistics.uniqueCounts()) { + EXPECT_EQ(1, pair.second); + } + + // Limits + if (nimble::isFloatingPointType()) { + data = { + std::numeric_limits::min(), + 1, + std::numeric_limits::lowest(), + std::numeric_limits::max()}; + statistics = T::create({data}); + EXPECT_EQ(4, statistics.uniqueCounts().size()); + EXPECT_EQ(std::numeric_limits::lowest(), statistics.min()); + EXPECT_EQ(std::numeric_limits::max(), statistics.max()); + EXPECT_EQ(1, statistics.minRepeat()); + EXPECT_EQ(1, statistics.maxRepeat()); + EXPECT_EQ(4, statistics.consecutiveRepeatCount()); + + for (const auto& pair : statistics.uniqueCounts()) { + EXPECT_EQ(1, pair.second); + } + } else { + data = { + std::numeric_limits::min(), + 1, + std::numeric_limits::max()}; + statistics = T::create({data}); + EXPECT_EQ(3, statistics.uniqueCounts().size()); + EXPECT_EQ(std::numeric_limits::min(), statistics.min()); + EXPECT_EQ(std::numeric_limits::max(), statistics.max()); + EXPECT_EQ(1, statistics.minRepeat()); + EXPECT_EQ(1, statistics.maxRepeat()); + EXPECT_EQ(3, statistics.consecutiveRepeatCount()); + + for (const auto& pair : statistics.uniqueCounts()) { + EXPECT_EQ(1, pair.second); + } + } +} + +TYPED_TEST(StatisticsBoolTests, Create) { + using T = TypeParam; + constexpr auto trueCount = 100; + constexpr auto falseCount = 230; + + auto data = std::make_unique(trueCount + falseCount); + for (auto i = 0; i < trueCount; ++i) { + data.get()[i] = true; + } + for (auto i = 0; i < falseCount; ++i) { + data.get()[i + trueCount] = false; + } + + auto statistics = + T::create(std::span(data.get(), trueCount + falseCount)); + uint64_t expectedDistinctValuesCount = 0; + expectedDistinctValuesCount += trueCount > 0 ? 1 : 0; + expectedDistinctValuesCount += falseCount > 0 ? 1 : 0; + EXPECT_EQ(expectedDistinctValuesCount, statistics.uniqueCounts().size()); + EXPECT_EQ(trueCount, statistics.uniqueCounts().at(true)); + EXPECT_EQ(falseCount, statistics.uniqueCounts().at(false)); + EXPECT_EQ(std::min(trueCount, falseCount), statistics.minRepeat()); + EXPECT_EQ(std::max(trueCount, falseCount), statistics.maxRepeat()); + EXPECT_EQ(2, statistics.consecutiveRepeatCount()); + + std::random_shuffle(data.get(), data.get() + trueCount + falseCount); + + statistics = + T::create(std::span(data.get(), trueCount + falseCount)); + EXPECT_EQ(expectedDistinctValuesCount, statistics.uniqueCounts().size()); + EXPECT_EQ(trueCount, statistics.uniqueCounts().at(true)); + EXPECT_EQ(falseCount, statistics.uniqueCounts().at(false)); +} + +template +void verifyString( + std::function data)> genStatisticsType) { + constexpr auto uniqueStrings = 10; + constexpr auto maxRepeat = 20; + + std::vector data; + data.reserve(uniqueStrings * maxRepeat); + + uint64_t totalLength = 0; + uint64_t totalRepeatLength = 0; + auto currentRepeat = maxRepeat; + for (auto i = 0; i < uniqueStrings; ++i) { + for (auto j = 0; j < currentRepeat; ++j) { + data.emplace_back(i + 1, 'a' + i); + totalLength += data.back().size(); + } + totalRepeatLength += i + 1; + --currentRepeat; + } + + T statistics = genStatisticsType(data); + + EXPECT_EQ(uniqueStrings, statistics.uniqueCounts().size()); + EXPECT_EQ(maxRepeat - uniqueStrings + 1, statistics.minRepeat()); + EXPECT_EQ(maxRepeat, statistics.maxRepeat()); + EXPECT_EQ(uniqueStrings, statistics.consecutiveRepeatCount()); + EXPECT_EQ(totalLength, statistics.totalStringsLength()); + EXPECT_EQ(totalRepeatLength, statistics.totalStringsRepeatLength()); + EXPECT_EQ( + std::string(uniqueStrings, 'a' + uniqueStrings - 1), statistics.max()); + EXPECT_EQ(std::string(1, 'a'), statistics.min()); + + currentRepeat = maxRepeat; + for (auto i = 0; i < statistics.uniqueCounts().size(); ++i) { + EXPECT_EQ( + statistics.uniqueCounts().at(std::string(i + 1, 'a' + i)), + currentRepeat--); + } + + std::random_shuffle(data.begin(), data.end()); + + statistics = genStatisticsType(data); + + EXPECT_EQ(uniqueStrings, statistics.uniqueCounts().size()); + EXPECT_EQ(totalLength, statistics.totalStringsLength()); + EXPECT_EQ( + std::string(uniqueStrings, 'a' + uniqueStrings - 1), statistics.max()); + EXPECT_EQ(std::string(1, 'a'), statistics.min()); + + currentRepeat = maxRepeat; + for (auto i = 0; i < statistics.uniqueCounts().size(); ++i) { + EXPECT_EQ( + statistics.uniqueCounts().at(std::string(i + 1, 'a' + i)), + currentRepeat--); + } +} + +TYPED_TEST(StatisticsStringTests, Create) { + using T = TypeParam; + + constexpr auto uniqueStrings = 10; + constexpr auto maxRepeat = 20; + + std::vector data; + data.reserve(uniqueStrings * maxRepeat); + + uint64_t totalLength = 0; + uint64_t totalRepeatLength = 0; + auto currentRepeat = maxRepeat; + for (auto i = 0; i < uniqueStrings; ++i) { + for (auto j = 0; j < currentRepeat; ++j) { + data.emplace_back(i + 1, 'a' + i); + totalLength += data.back().size(); + } + totalRepeatLength += i + 1; + --currentRepeat; + } + + T statistics = + nimble::Statistics::create(data); + + EXPECT_EQ(uniqueStrings, statistics.uniqueCounts().size()); + EXPECT_EQ(maxRepeat - uniqueStrings + 1, statistics.minRepeat()); + EXPECT_EQ(maxRepeat, statistics.maxRepeat()); + EXPECT_EQ(uniqueStrings, statistics.consecutiveRepeatCount()); + EXPECT_EQ(totalLength, statistics.totalStringsLength()); + EXPECT_EQ(totalRepeatLength, statistics.totalStringsRepeatLength()); + EXPECT_EQ( + std::string(uniqueStrings, 'a' + uniqueStrings - 1), statistics.max()); + EXPECT_EQ(std::string(1, 'a'), statistics.min()); + + currentRepeat = maxRepeat; + for (auto i = 0; i < statistics.uniqueCounts().size(); ++i) { + EXPECT_EQ( + statistics.uniqueCounts().at(std::string(i + 1, 'a' + i)), + currentRepeat--); + } + + std::random_shuffle(data.begin(), data.end()); + + statistics = + nimble::Statistics::create({data}); + + EXPECT_EQ(uniqueStrings, statistics.uniqueCounts().size()); + EXPECT_EQ(totalLength, statistics.totalStringsLength()); + EXPECT_EQ( + std::string(uniqueStrings, 'a' + uniqueStrings - 1), statistics.max()); + EXPECT_EQ(std::string(1, 'a'), statistics.min()); + + currentRepeat = maxRepeat; + for (auto i = 0; i < statistics.uniqueCounts().size(); ++i) { + EXPECT_EQ( + statistics.uniqueCounts().at(std::string(i + 1, 'a' + i)), + currentRepeat--); + } +} + +TYPED_TEST(StatisticsNumericTests, Repeat) { + using T = TypeParam; + using ValueType = typename T::valueType; + + struct Test { + std::vector data; + uint64_t expectedConsecutiveRepeatCount; + uint64_t expectedMinRepeat; + uint64_t expectedMaxRepeat; + ValueType expectedMin; + ValueType expectedMax; + }; + + std::vector tests{ + {{}, 0, 0, 0, 0, 0}, + {{1}, 1, 1, 1, 1, 1}, + {{1, 1}, 1, 2, 2, 1, 1}, + {{1, 2}, 2, 1, 1, 1, 2}, + {{1, 1, 2}, 2, 1, 2, 1, 2}, + {{1, 2, 2}, 2, 1, 2, 1, 2}, + {{1, 2, 1}, 3, 1, 1, 1, 2}, + {{1, 1, 2, 1}, 3, 1, 2, 1, 2}, + {{1, 2, 2, 1}, 3, 1, 2, 1, 2}, + {{1, 2, 1, 1}, 3, 1, 2, 1, 2}, + {{1, 1, 1, 2, 2, 2, 1, 1}, 3, 2, 3, 1, 2}, + }; + + for (const auto& test : tests) { + auto statistics = T::create({test.data}); + EXPECT_EQ( + test.expectedConsecutiveRepeatCount, + statistics.consecutiveRepeatCount()); + EXPECT_EQ(test.expectedMinRepeat, statistics.minRepeat()); + EXPECT_EQ(test.expectedMaxRepeat, statistics.maxRepeat()); + EXPECT_EQ(test.expectedMin, statistics.min()); + EXPECT_EQ(test.expectedMax, statistics.max()); + } +} + +TYPED_TEST(StatisticsIntegerTests, Buckets) { + using T = TypeParam; + using ValueType = typename T::valueType; + using UnsignedValueType = typename std::make_unsigned::type; + + constexpr auto offset = 10; + constexpr auto repeatSize = 25; + constexpr auto dataSize = 100; + ValueType minValue = + std::is_signed() ? (ValueType)-offset : (ValueType)offset; + + std::vector data; + data.reserve(dataSize + (sizeof(ValueType) * 8)); + + std::array expectedBuckets{}; + + // Repeating data (non-consecutive) + for (auto i = 0; i < dataSize; ++i) { + data.push_back(minValue + (i % repeatSize)); + } + + for (auto i = 0; i < sizeof(ValueType) * 8; ++i) { + auto value = std::numeric_limits::max() >> i; + if (value >= offset) { + data.push_back(value); + } + } + + for (auto i = 0; i < data.size(); ++i) { + char buffer[10]; + auto pos = buffer; + nimble::varint::writeVarint( + static_cast( + static_cast(data[i]) - + static_cast(minValue)), + &pos); + ++expectedBuckets[pos - buffer - 1]; + } + + auto statistics = T::create({data}); + auto& buckets = statistics.bucketCounts(); + EXPECT_LE(buckets.size(), expectedBuckets.size()); + EXPECT_GT(buckets.size(), 0); + for (auto i = 0; i < buckets.size(); ++i) { + EXPECT_EQ(expectedBuckets[i], buckets[i]) << "index: " << i; + } +} diff --git a/dwio/nimble/encodings/tests/TestGenerator.cpp b/dwio/nimble/encodings/tests/TestGenerator.cpp new file mode 100644 index 0000000..63755cc --- /dev/null +++ b/dwio/nimble/encodings/tests/TestGenerator.cpp @@ -0,0 +1,393 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include +#include +#include +#include +#include +#include + +#include "common/init/light.h" +#include "dwio/nimble/encodings/tests/TestUtils.h" +#include "folly/experimental/io/FsUtil.h" + +using namespace ::facebook; + +DEFINE_string(output_dir, "", "Output directory to write test artifacts to."); +DEFINE_string( + encoding_file, + "", + "If porivided, prints the content of the encoding file."); + +namespace { +template +T randomValue(RNG&& rng, nimble::Buffer& buffer) { + if constexpr (std::is_same_v) { + return static_cast(folly::Random::rand32(2, std::forward(rng))); + } else if constexpr (std::is_same_v) { + auto size = folly::Random::rand32(32, std::forward(rng)); + auto data = buffer.reserve(size); + for (auto i = 0; i < size; ++i) { + data[i] = static_cast(folly::Random::rand32( + std::numeric_limits::max(), std::forward(rng))); + } + return std::string_view{data, size}; + } else { + uint64_t rand = SmallNumbers + ? folly::Random::rand64(15, std::forward(rng)) + : folly::Random::rand64(std::forward(rng)); + return *reinterpret_cast(&rand); + } +} + +template +inline nimble::Vector +generateRandomData(RNG&& rng, nimble::Buffer& buffer, int count) { + nimble::Vector data{&buffer.getMemoryPool()}; + data.reserve(count); + + for (int i = 0; i < count; ++i) { + data.push_back( + randomValue(std::forward(rng), buffer)); + } + + return data; +} + +template +inline nimble::Vector +generateConstData(RNG&& rng, nimble::Buffer& buffer, int count) { + nimble::Vector data{&buffer.getMemoryPool()}; + data.reserve(count); + + T value = randomValue(std::forward(rng), buffer); + for (int i = 0; i < count; ++i) { + data.push_back(value); + } + + return data; +} + +template +inline nimble::Vector +generateFixedBitWidthData(RNG&& rng, nimble::Buffer& buffer, int count) { + nimble::Vector data{&buffer.getMemoryPool()}; + data.reserve(count); + + using physicalType = typename nimble::TypeTraits::physicalType; + + auto bits = 4 * sizeof(T) - 1; + physicalType mask = (1 << bits) - 1u; + LOG(INFO) << nimble::bits::printBits(mask); + physicalType val = + (randomValue(std::forward(rng), buffer)) & mask; + for (int i = 0; i < count; ++i) { + physicalType value = + randomValue(std::forward(rng), buffer) & mask; + value = value ? value : 1; + data.push_back(value << bits); + } + + return data; +} + +template +inline void randomZero(RNG&& rng, std::vector& data) { + for (int i = 0; i < data.size(); ++i) { + if (folly::Random::rand64(std::forward(rng)) % 2) { + data[i] = 0; + } + } +} + +} // namespace + +template +void writeFile( + RNG&& rng, + const std::string& path, + uint32_t rowCount, + std::function(RNG&&, nimble::Buffer&, int)> dataGenerator = + [](RNG&& rng, nimble::Buffer& buffer, int count) { + return generateRandomData( + std::forward(rng), buffer, count); + }, + std::string suffix = "") { + auto identifier = fmt::format( + "{}_{}_{}", + nimble::test::Encoder::encodingType(), + toString(nimble::TypeTraits::dataType), + rowCount); + + if (!suffix.empty()) { + identifier += "_" + suffix; + } + + LOG(INFO) << "Writing " << identifier; + + auto pool = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*pool}; + auto data = dataGenerator(rng, buffer, rowCount); + + { + std::ofstream file{ + fmt::format("{}/{}.data", path, identifier), + std::ios::out | std::ios::binary | std::ios::trunc}; + auto count = data.size(); + for (const auto& value : data) { + if constexpr (std::is_same_v) { + auto size = value.size(); + file.write(reinterpret_cast(&size), sizeof(size_t)); + file.write(value.data(), value.size()); + } else { + file.write( + reinterpret_cast(&value), + sizeof(typename E::cppDataType)); + } + } + } + + std::optional> nulls; + + if (nimble::test::Encoder::encodingType() == + nimble::EncodingType::Nullable) { + nulls = generateRandomData(std::forward(rng), buffer, rowCount); + } + + if (nulls) { + std::ofstream file{ + fmt::format("{}/{}.nulls", path, identifier), + std::ios::out | std::ios::binary | std::ios::trunc}; + auto count = data.size(); + for (const auto& value : nulls.value()) { + file.write(reinterpret_cast(&value), sizeof(uint8_t)); + } + } + + for (auto compressionType : + {nimble::CompressionType::Uncompressed, + nimble::CompressionType::Zstd, + nimble::CompressionType::Zstrong}) { + std::string_view encoded; + if constexpr ( + nimble::test::Encoder::encodingType() == + nimble::EncodingType::Nullable) { + encoded = nimble::test::Encoder::encodeNullable( + buffer, data, nulls.value(), compressionType); + } else { + encoded = nimble::test::Encoder::encode(buffer, data, compressionType); + } + { + std::ofstream file{ + fmt::format( + "{}/{}_{}.encoding", path, identifier, toString(compressionType)), + std::ios::out | std::ios::binary | std::ios::trunc}; + file << encoded; + } + } +} + +template +void writeFileSmallNumbers( + RNG&& rng, + const std::string& path, + uint32_t rowCount) { + writeFile( + std::forward(rng), + path, + rowCount, + [](RNG&& rng, nimble::Buffer& buffer, int count) { + return generateRandomData< + typename E::cppDataType, + RNG, + /* SmallNumbers */ true>(std::forward(rng), buffer, count); + }, + "small-numbers"); +} + +template +void writeFileConstant(RNG&& rng, const std::string& path, uint32_t rowCount) { + writeFile( + std::forward(rng), + path, + rowCount, + [](RNG&& rng, nimble::Buffer& buffer, int count) { + return generateConstData( + std::forward(rng), buffer, count); + }, + "constant"); +} + +template +void writeFileFixedBitWidth( + RNG&& rng, + const std::string& path, + uint32_t rowCount) { + writeFile( + std::forward(rng), + path, + rowCount, + [](RNG&& rng, nimble::Buffer& buffer, int count) { + return generateFixedBitWidthData( + std::forward(rng), buffer, count); + }, + "bits"); +} + +template +void printScalarData( + std::ostream& ostream, + velox::memory::MemoryPool& pool, + nimble::Encoding& stream, + uint32_t rowCount) { + nimble::Vector buffer(&pool); + nimble::Vector nulls(&pool); + buffer.resize(rowCount); + nulls.resize((nimble::FixedBitArray::bufferSize(rowCount, 1))); + nulls.zero_out(); + if (stream.isNullable()) { + stream.materializeNullable( + rowCount, buffer.data(), [&]() { return nulls.data(); }); + } else { + stream.materialize(rowCount, buffer.data()); + nulls.fill(0xff); + } + for (uint32_t i = 0; i < rowCount; ++i) { + if (stream.isNullable() && nimble::bits::getBit(i, nulls.data()) == 0) { + ostream << "NULL" << std::endl; + } else { + ostream << folly::to(buffer[i]) + << std::endl; // Have to use folly::to as Int8 was getting + // converted to char. + } + } +} + +void printScalarType( + std::ostream& ostream, + velox::memory::MemoryPool& pool, + nimble::Encoding& stream, + uint32_t rowCount) { + switch (stream.dataType()) { +#define CASE(KIND, cppType) \ + case nimble::DataType::KIND: { \ + printScalarData(ostream, pool, stream, rowCount); \ + break; \ + } + CASE(Int8, int8_t); + CASE(Uint8, uint8_t); + CASE(Int16, int16_t); + CASE(Uint16, uint16_t); + CASE(Int32, int32_t); + CASE(Uint32, uint32_t); + CASE(Int64, int64_t); + CASE(Uint64, uint64_t); + CASE(Float, float); + CASE(Double, double); + CASE(Bool, bool); + CASE(String, std::string_view); +#undef CASE + case nimble::DataType::Undefined: { + NIMBLE_UNREACHABLE( + fmt::format("Undefined type for stream: {}", stream.dataType())); + } + } +} + +#define _WRITE_NUMERIC_FILES_(encoding, rowCount, type) \ + writeFile##type>(rng, FLAGS_output_dir, rowCount); \ + writeFile##type>(rng, FLAGS_output_dir, rowCount); \ + writeFile##type>(rng, FLAGS_output_dir, rowCount); \ + writeFile##type>( \ + rng, FLAGS_output_dir, rowCount); \ + writeFile##type>(rng, FLAGS_output_dir, rowCount); \ + writeFile##type>( \ + rng, FLAGS_output_dir, rowCount); \ + writeFile##type>(rng, FLAGS_output_dir, rowCount); \ + writeFile##type>( \ + rng, FLAGS_output_dir, rowCount); \ + writeFile##type>(rng, FLAGS_output_dir, rowCount); \ + writeFile##type>(rng, FLAGS_output_dir, rowCount); + +#define _WRITE_FILES_(encoding, rowCount, type) \ + _WRITE_NUMERIC_FILES_(encoding, rowCount, type) \ + writeFile##type>(rng, FLAGS_output_dir, rowCount); \ + writeFile##type>( \ + rng, FLAGS_output_dir, rowCount); + +#define WRITE_FILES(encoding, rowCount) _WRITE_FILES_(encoding, rowCount, ) + +#define WRITE_NUMERIC_FILES(encoding, rowCount) \ + _WRITE_NUMERIC_FILES_(encoding, rowCount, ) + +#define WRITE_CONSTANT_FILES(rowCount) \ + _WRITE_FILES_(ConstantEncoding, rowCount, Constant); \ + _WRITE_FILES_(RLEEncoding, rowCount, Constant) + +#define WRITE_FBW_FILES(rowCount) \ + _WRITE_NUMERIC_FILES_(FixedBitWidthEncoding, rowCount, FixedBitWidth); \ + _WRITE_NUMERIC_FILES_(FixedBitWidthEncoding, rowCount, SmallNumbers) + +#define WRITE_NON_BOOL_FILES(encoding, rowCount) \ + _WRITE_NUMERIC_FILES_(encoding, rowCount, ) \ + writeFile>( \ + rng, FLAGS_output_dir, rowCount); + +int main(int argc, char* argv[]) { + init::initFacebookLight(&argc, &argv); + + if (!FLAGS_encoding_file.empty()) { + std::ifstream stream{FLAGS_encoding_file, std::ios::binary}; + stream.seekg(0, std::ios::end); + auto size = stream.tellg(); + stream.seekg(0, std::ios::beg); + std::vector buffer(size); + stream.read(buffer.data(), buffer.size()); + + auto pool = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + auto encoding = + nimble::EncodingFactory::decode(*pool, {buffer.data(), buffer.size()}); + auto rowCount = encoding->rowCount(); + printScalarType(std::cout, *pool, *encoding, rowCount); + return 0; + } + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + if (FLAGS_output_dir.empty()) { + LOG(ERROR) << "Output directory is empty"; + return 1; + } + + if (folly::fs::exists(FLAGS_output_dir)) { + folly::fs::remove_all(FLAGS_output_dir); + } + + folly::fs::create_directory(FLAGS_output_dir); + + WRITE_CONSTANT_FILES(256); + WRITE_FBW_FILES(256); + WRITE_NON_BOOL_FILES(MainlyConstantEncoding, 256); + + for (auto rowCount : {0, 256}) { + WRITE_FILES(TrivialEncoding, rowCount); + WRITE_FILES(DictionaryEncoding, rowCount); + WRITE_FILES(RLEEncoding, rowCount); + WRITE_FILES(NullableEncoding, rowCount); + WRITE_NUMERIC_FILES(FixedBitWidthEncoding, rowCount); + writeFile(rng, FLAGS_output_dir, rowCount); + writeFile>(rng, FLAGS_output_dir, rowCount); + writeFile>( + rng, FLAGS_output_dir, rowCount); + writeFile>(rng, FLAGS_output_dir, rowCount); + writeFile>( + rng, FLAGS_output_dir, rowCount); + writeFile>(rng, FLAGS_output_dir, rowCount); + writeFile>(rng, FLAGS_output_dir, rowCount); + } + + return 0; +} diff --git a/dwio/nimble/encodings/tests/TestUtils.h b/dwio/nimble/encodings/tests/TestUtils.h new file mode 100644 index 0000000..28a838e --- /dev/null +++ b/dwio/nimble/encodings/tests/TestUtils.h @@ -0,0 +1,226 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/encodings/ConstantEncoding.h" +#include "dwio/nimble/encodings/DeltaEncoding.h" +#include "dwio/nimble/encodings/DictionaryEncoding.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" +#include "dwio/nimble/encodings/EncodingSelectionPolicy.h" +#include "dwio/nimble/encodings/FixedBitWidthEncoding.h" +#include "dwio/nimble/encodings/MainlyConstantEncoding.h" +#include "dwio/nimble/encodings/NullableEncoding.h" +#include "dwio/nimble/encodings/RleEncoding.h" +#include "dwio/nimble/encodings/SparseBoolEncoding.h" +#include "dwio/nimble/encodings/TrivialEncoding.h" +#include "dwio/nimble/encodings/VarintEncoding.h" + +namespace facebook::nimble::test { + +template +class Encoder { + using T = typename E::cppDataType; + + private: + class TestCompressPolicy : public nimble::CompressionPolicy { + public: + explicit TestCompressPolicy(CompressionType compressionType) + : compressionType_{compressionType} {} + + nimble::CompressionInformation compression() const override { + if (compressionType_ == CompressionType::Uncompressed) { + return {.compressionType = CompressionType::Uncompressed}; + } + + if (compressionType_ == CompressionType::Zstd) { + nimble::CompressionInformation information{ + .compressionType = nimble::CompressionType::Zstd}; + information.parameters.zstd.compressionLevel = 3; + return information; + } + + nimble::CompressionInformation information{ + .compressionType = nimble::CompressionType::Zstrong}; + information.parameters.zstrong.compressionLevel = 9; + information.parameters.zstrong.decompressionLevel = 2; + return information; + } + + virtual bool shouldAccept( + nimble::CompressionType /* compressionType */, + uint64_t /* uncompressedSize */, + uint64_t /* compressedSize */) const override { + return true; + } + + private: + CompressionType compressionType_; + }; + + template + class TestTrivialEncodingSelectionPolicy + : public nimble::EncodingSelectionPolicy { + using physicalType = typename nimble::TypeTraits::physicalType; + + public: + explicit TestTrivialEncodingSelectionPolicy(CompressionType compressionType) + : compressionType_{compressionType} {} + + nimble::EncodingSelectionResult select( + std::span /* values */, + const nimble::Statistics& /* statistics */) override { + return { + .encodingType = nimble::EncodingType::Trivial, + .compressionPolicyFactory = [this]() { + return std::make_unique(compressionType_); + }}; + } + + EncodingSelectionResult selectNullable( + std::span /* values */, + std::span /* nulls */, + const Statistics& /* statistics */) override { + return { + .encodingType = EncodingType::Nullable, + }; + } + + std::unique_ptr createImpl( + nimble::EncodingType /* encodingType */, + nimble::NestedEncodingIdentifier /* identifier */, + nimble::DataType type) override { + UNIQUE_PTR_FACTORY( + type, TestTrivialEncodingSelectionPolicy, compressionType_); + } + + private: + CompressionType compressionType_; + }; + + template + struct EncodingTypeTraits {}; + + template <> + struct EncodingTypeTraits> { + static constexpr inline nimble::EncodingType encodingType = + nimble::EncodingType::Constant; + }; + + template <> + struct EncodingTypeTraits> { + static constexpr inline nimble::EncodingType encodingType = + nimble::EncodingType::Dictionary; + }; + + template <> + struct EncodingTypeTraits> { + static constexpr inline nimble::EncodingType encodingType = + nimble::EncodingType::FixedBitWidth; + }; + + template <> + struct EncodingTypeTraits> { + static constexpr inline nimble::EncodingType encodingType = + nimble::EncodingType::MainlyConstant; + }; + + template <> + struct EncodingTypeTraits> { + static constexpr inline nimble::EncodingType encodingType = + nimble::EncodingType::RLE; + }; + + template <> + struct EncodingTypeTraits { + static constexpr inline nimble::EncodingType encodingType = + nimble::EncodingType::SparseBool; + }; + + template <> + struct EncodingTypeTraits> { + static constexpr inline nimble::EncodingType encodingType = + nimble::EncodingType::Trivial; + }; + + template <> + struct EncodingTypeTraits> { + static constexpr inline nimble::EncodingType encodingType = + nimble::EncodingType::Varint; + }; + + template <> + struct EncodingTypeTraits> { + static constexpr inline nimble::EncodingType encodingType = + nimble::EncodingType::Nullable; + }; + + public: + static constexpr EncodingType encodingType() { + return EncodingTypeTraits::encodingType; + } + + static std::string_view encode( + nimble::Buffer& buffer, + const nimble::Vector& values, + CompressionType compressionType = CompressionType::Uncompressed) { + using physicalType = typename nimble::TypeTraits::physicalType; + + auto physicalValues = std::span( + reinterpret_cast(values.data()), values.size()); + nimble::EncodingSelection selection{ + {.encodingType = EncodingTypeTraits::encodingType, + .compressionPolicyFactory = + [compressionType]() { + return std::make_unique(compressionType); + }}, + nimble::Statistics::create(physicalValues), + std::make_unique>( + compressionType)}; + + return E::encode(selection, physicalValues, buffer); + } + + static std::string_view encodeNullable( + nimble::Buffer& buffer, + const nimble::Vector& values, + const nimble::Vector& nulls, + nimble::CompressionType compressionType = + nimble::CompressionType::Uncompressed) { + using physicalType = typename nimble::TypeTraits::physicalType; + + auto physicalValues = std::span( + reinterpret_cast(values.data()), values.size()); + nimble::EncodingSelection selection{ + {.encodingType = EncodingTypeTraits::encodingType, + .compressionPolicyFactory = + [compressionType]() { + return std::make_unique(compressionType); + }}, + nimble::Statistics::create(physicalValues), + std::make_unique>( + compressionType)}; + + return NullableEncoding::encodeNullable( + selection, physicalValues, nulls, buffer); + } + + static std::unique_ptr createEncoding( + nimble::Buffer& buffer, + const nimble::Vector& values, + CompressionType compressionType = CompressionType::Uncompressed) { + return std::make_unique( + buffer.getMemoryPool(), encode(buffer, values, compressionType)); + } + + static std::unique_ptr createNullableEncoding( + nimble::Buffer& buffer, + const nimble::Vector& values, + const nimble::Vector& nulls, + CompressionType compressionType = CompressionType::Uncompressed) { + return std::make_unique( + buffer.getMemoryPool(), + encodeNullable(buffer, values, nulls, compressionType)); + } +}; +} // namespace facebook::nimble::test diff --git a/dwio/nimble/tablet/CMakeLists.txt b/dwio/nimble/tablet/CMakeLists.txt new file mode 100644 index 0000000..1fb8e43 --- /dev/null +++ b/dwio/nimble/tablet/CMakeLists.txt @@ -0,0 +1,23 @@ +add_subdirectory(tests) + +# Nimble code expects an upper case suffix to the generated file. +list(PREPEND FLATBUFFERS_FLATC_SCHEMA_EXTRA_ARGS "--filename-suffix" + "Generated") + +build_flatbuffers( + "${CMAKE_CURRENT_SOURCE_DIR}/Footer.fbs" + "" + nimble_footer_schema_fb + "" + "${CMAKE_CURRENT_BINARY_DIR}" + "" + "") +add_library(nimble_footer_fb INTERFACE) +target_include_directories( + nimble_footer_fb + INTERFACE ${FLATBUFFERS_INCLUDE_DIR}) +add_dependencies(nimble_footer_fb nimble_footer_schema_fb) + +add_library(nimble_tablet Compression.cpp StreamInput.cpp Tablet.cpp) + +target_link_libraries(nimble_tablet nimble_footer_fb Folly::folly) diff --git a/dwio/nimble/tablet/Compression.cpp b/dwio/nimble/tablet/Compression.cpp new file mode 100644 index 0000000..187b0a7 --- /dev/null +++ b/dwio/nimble/tablet/Compression.cpp @@ -0,0 +1,52 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include +#include +#include + +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/tablet/Compression.h" + +namespace facebook::nimble { + +std::optional> ZstdCompression::compress( + velox::memory::MemoryPool& memoryPool, + std::string_view source, + int32_t level) { + Vector buffer{&memoryPool, source.size() + sizeof(uint32_t)}; + auto pos = buffer.data(); + encoding::writeUint32(source.size(), pos); + auto ret = ZSTD_compress( + pos, + source.size() - sizeof(uint32_t), + source.data(), + source.size(), + level); + if (ZSTD_isError(ret)) { + NIMBLE_ASSERT( + ZSTD_getErrorCode(ret) == ZSTD_ErrorCode::ZSTD_error_dstSize_tooSmall, + fmt::format( + "Error while compressing data: {}", ZSTD_getErrorName(ret))); + return std::nullopt; + } + + buffer.resize(ret + sizeof(uint32_t)); + return {buffer}; +} + +Vector ZstdCompression::uncompress( + velox::memory::MemoryPool& memoryPool, + std::string_view source) { + auto pos = source.data(); + const uint32_t uncompressedSize = encoding::readUint32(pos); + Vector buffer{&memoryPool, uncompressedSize}; + auto ret = ZSTD_decompress( + buffer.data(), buffer.size(), pos, source.size() - sizeof(uint32_t)); + NIMBLE_CHECK( + !ZSTD_isError(ret), + fmt::format("Error uncompressing data: {}", ZSTD_getErrorName(ret))); + return buffer; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/tablet/Compression.h b/dwio/nimble/tablet/Compression.h new file mode 100644 index 0000000..c832bde --- /dev/null +++ b/dwio/nimble/tablet/Compression.h @@ -0,0 +1,21 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/common/Vector.h" + +namespace facebook::nimble { + +class ZstdCompression { + public: + static std::optional> compress( + velox::memory::MemoryPool& memoryPool, + std::string_view source, + int32_t level = 1); + + static Vector uncompress( + velox::memory::MemoryPool& memoryPool, + std::string_view source); +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/tablet/Footer.fbs b/dwio/nimble/tablet/Footer.fbs new file mode 100644 index 0000000..8acf1f1 --- /dev/null +++ b/dwio/nimble/tablet/Footer.fbs @@ -0,0 +1,44 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +namespace facebook.nimble.serialization; + +enum CompressionType:uint8 { + Uncompressed = 0, + Zstd = 1, + Zstrong = 2, +} + +table Stripes { + row_counts:[uint32]; + offsets:[uint64]; + sizes:[uint32]; + group_indices:[uint32]; +} + +table StripeGroup { + stripe_count:uint32; + stream_offsets:[uint32]; + stream_sizes:[uint32]; +} + +table MetadataSection { + offset:uint64; + size:uint32; + compression_type:CompressionType; +} + +table OptionalMetadataSections { + names:[string]; + offsets:[uint64]; + sizes:[uint32]; + compression_types:[CompressionType]; +} + +table Footer { + row_count:uint64; + stripes:MetadataSection; + stripe_groups:[MetadataSection]; + optional_sections:OptionalMetadataSections; +} + +root_type Footer; diff --git a/dwio/nimble/tablet/Tablet.cpp b/dwio/nimble/tablet/Tablet.cpp new file mode 100644 index 0000000..0b0680d --- /dev/null +++ b/dwio/nimble/tablet/Tablet.cpp @@ -0,0 +1,918 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include "dwio/nimble/tablet/Tablet.h" +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/tablet/Compression.h" +#include "dwio/nimble/tablet/FooterGenerated.h" +#include "folly/compression/Compression.h" +#include "folly/io/Cursor.h" + +#include +#include +#include +#include +#include +#include + +namespace facebook::nimble { + +DEFINE_bool( + nimble_disable_coalesce, + false, + "Disable read coalescing in Nimble reader."); + +DEFINE_uint64( + nimble_coalesce_max_distance, + 1024 * 1024 * 1.25, + "Maximum read coalescing distance in Nimble reader. And gap smaller than this value will be coalesced."); + +// Here's the layout of the tablet: +// +// stripe 1 streams +// stripe 2 streams +// ... +// stripe k streams +// footer +// +// where the footer is a flatbuffer payload, as described here: +// dwio/nimble/tablet/footer.fbs +// followed by fixed payload: +// 4 bytes footer size + 1 byte footer compression type + +// 1 byte checksum type + 8 bytes checksum + +// 2 bytes major version + 2 bytes minor version + +// 4 bytes magic number. +namespace { + +constexpr uint16_t kMagicNumber = 0xA1FA; +constexpr uint64_t kInitialFooterSize = 8 * 1024 * 1024; // 8Mb +constexpr uint16_t kVersionMajor = 0; +constexpr uint16_t kVersionMinor = 1; + +// Total size of the fields after the flatbuffer. +constexpr uint32_t kPostscriptSize = 20; + +// The following fields in postscript are included in checksum calculation. +// 4 bytes footer size + 1 byte compression type +constexpr uint32_t kPostscriptChecksumedSize = 5; + +template +flatbuffers::Offset> createFlattenedVector( + flatbuffers::FlatBufferBuilder& builder, + size_t streamCount, + const std::vector>& source, + Target defaultValue) { + // This method is performing the following: + // 1. Converts the source vector into FlatBuffers representation. The flat + // buffer representation is a single dimension array, so this method flattens + // the source vectors into this single dimension array. + // 2. Source vector contains a child vector per stripe. Each one of those + // child vectors might contain different entry count than the final number of + // streams in the tablet. This is because each stripe might have + // encountered just subset of the streams. Therefore, this method fills up all + // missing offsets in the target vector with the provided default value. + + auto stripeCount = source.size(); + std::vector target(stripeCount * streamCount); + size_t targetIndex = 0; + + for (auto& items : source) { + NIMBLE_ASSERT( + items.size() <= streamCount, + "Corrupted footer. Stream count exceeds expected total number of streams."); + + for (auto i = 0; i < streamCount; ++i) { + target[targetIndex++] = i < items.size() + ? static_cast(items[i]) + : static_cast(defaultValue); + } + } + + return builder.CreateVector(target); +} + +std::string_view asView(const flatbuffers::FlatBufferBuilder& builder) { + return { + reinterpret_cast(builder.GetBufferPointer()), + builder.GetSize()}; +} + +template +const T* asFlatBuffersRoot(std::string_view content) { + return flatbuffers::GetRoot(content.data()); +} + +size_t copyTo(const folly::IOBuf& source, void* target, size_t size) { + NIMBLE_DASSERT( + source.computeChainDataLength() <= size, "Target buffer too small."); + size_t offset = 0; + for (const auto& chunk : source) { + std::copy(chunk.begin(), chunk.end(), static_cast(target) + offset); + offset += chunk.size(); + } + + return offset; +} + +folly::IOBuf +cloneAndCoalesce(const folly::IOBuf& src, size_t offset, size_t size) { + folly::io::Cursor cursor(&src); + NIMBLE_ASSERT(cursor.totalLength() >= offset, "Offset out of range"); + cursor.skip(offset); + NIMBLE_ASSERT(cursor.totalLength() >= size, "Size out of range"); + folly::IOBuf result; + cursor.clone(result, size); + result.coalesceWithHeadroomTailroom(0, 0); + return result; +} + +std::string_view toStringView(const folly::IOBuf& buf) { + return {reinterpret_cast(buf.data()), buf.length()}; +} + +} // namespace + +void TabletWriter::close() { + auto stripeCount = stripeOffsets_.size(); + NIMBLE_ASSERT( + stripeCount == stripeSizes_.size() && + stripeCount == stripeGroupIndices_.size() && + stripeCount == rowCounts_.size(), + "Stripe count mismatch."); + + const uint64_t totalRows = + std::accumulate(rowCounts_.begin(), rowCounts_.end(), uint64_t{0}); + + flatbuffers::FlatBufferBuilder builder(kInitialFooterSize); + + // write remaining stripe groups + tryWriteStripeGroup(true); + + // write stripes + MetadataSection stripes; + if (stripeCount > 0) { + flatbuffers::FlatBufferBuilder stripesBuilder(kInitialFooterSize); + stripesBuilder.Finish(serialization::CreateStripes( + stripesBuilder, + stripesBuilder.CreateVector(rowCounts_), + stripesBuilder.CreateVector(stripeOffsets_), + stripesBuilder.CreateVector(stripeSizes_), + stripesBuilder.CreateVector(stripeGroupIndices_))); + stripes = createMetadataSection(asView(stripesBuilder)); + } + + auto createOptionalMetadataSection = + [](flatbuffers::FlatBufferBuilder& builder, + const std::vector< + std::pair>& + optionalSections) { + return serialization::CreateOptionalMetadataSections( + builder, + builder.CreateVector>( + optionalSections.size(), + [&builder, &optionalSections](size_t i) { + return builder.CreateString(optionalSections[i].first); + }), + builder.CreateVector( + optionalSections.size(), + [&optionalSections](size_t i) { + return optionalSections[i].second.offset; + }), + builder.CreateVector( + optionalSections.size(), + [&optionalSections](size_t i) { + return optionalSections[i].second.size; + }), + builder.CreateVector( + optionalSections.size(), [&optionalSections](size_t i) { + return static_cast( + optionalSections[i].second.compressionType); + })); + }; + + // write footer + builder.Finish(serialization::CreateFooter( + builder, + totalRows, + stripeCount > 0 ? serialization::CreateMetadataSection( + builder, + stripes.offset, + stripes.size, + static_cast( + stripes.compressionType)) + : 0, + !stripeGroups_.empty() + ? builder.CreateVector< + flatbuffers::Offset>( + stripeGroups_.size(), + [this, &builder](size_t i) { + return serialization::CreateMetadataSection( + builder, + stripeGroups_[i].offset, + stripeGroups_[i].size, + static_cast( + stripeGroups_[i].compressionType)); + }) + : 0, + !optionalSections_.empty() + ? createOptionalMetadataSection( + builder, {optionalSections_.begin(), optionalSections_.end()}) + : 0)); + + auto footerStart = file_->size(); + auto footerCompressionType = writeMetadata(asView(builder)); + + // End with the fixed length constants. + const uint64_t footerSize64Bit = (file_->size() - footerStart); + NIMBLE_ASSERT( + footerSize64Bit <= std::numeric_limits::max(), + fmt::format("Footer size too big: {}.", footerSize64Bit)); + const uint32_t footerSize = footerSize64Bit; + writeWithChecksum({reinterpret_cast(&footerSize), 4}); + writeWithChecksum({reinterpret_cast(&footerCompressionType), 1}); + uint64_t checksum = checksum_->getChecksum(); + file_->append({reinterpret_cast(&options_.checksumType), 1}); + file_->append({reinterpret_cast(&checksum), 8}); + file_->append({reinterpret_cast(&kVersionMajor), 2}); + file_->append({reinterpret_cast(&kVersionMinor), 2}); + file_->append({reinterpret_cast(&kMagicNumber), 2}); +} + +void TabletWriter::writeStripe(uint32_t rowCount, std::vector streams) { + if (UNLIKELY(rowCount == 0)) { + return; + } + + rowCounts_.push_back(rowCount); + stripeOffsets_.push_back(file_->size()); + + auto streamCount = streams.empty() + ? 0 + : std::max_element( + streams.begin(), + streams.end(), + [](const auto& a, const auto& b) { + return a.offset < b.offset; + })->offset + + 1; + + auto& stripeStreamOffsets = streamOffsets_.emplace_back(streamCount, 0); + auto& stripeStreamSizes = streamSizes_.emplace_back(streamCount, 0); + + if (options_.layoutPlanner) { + streams = options_.layoutPlanner->getLayout(std::move(streams)); + } + + for (const auto& stream : streams) { + const uint32_t index = stream.offset; + + // @lint-ignore CLANGTIDY facebook-hte-LocalUncheckedArrayBounds + stripeStreamOffsets[index] = file_->size() - stripeOffsets_.back(); + + for (auto output : stream.content) { + writeWithChecksum(output); + } + + // @lint-ignore CLANGTIDY facebook-hte-LocalUncheckedArrayBounds + stripeStreamSizes[index] = + file_->size() - (stripeStreamOffsets[index] + stripeOffsets_.back()); + } + + stripeSizes_.push_back(file_->size() - stripeOffsets_.back()); + stripeGroupIndices_.push_back(stripeGroupIndex_); + + // Write stripe group if size of column offsets/sizes is too large. + tryWriteStripeGroup(); +} + +CompressionType TabletWriter::writeMetadata(std::string_view metadata) { + auto size = metadata.size(); + bool shouldCompress = size > options_.metadataCompressionThreshold; + CompressionType compressionType = CompressionType::Uncompressed; + std::optional> compressed; + if (shouldCompress) { + compressed = ZstdCompression::compress(memoryPool_, metadata); + if (compressed.has_value()) { + compressionType = CompressionType::Zstd; + metadata = {compressed->data(), compressed->size()}; + } + } + writeWithChecksum(metadata); + return compressionType; +} + +TabletWriter::MetadataSection TabletWriter::createMetadataSection( + std::string_view metadata) { + auto offset = file_->size(); + auto compressionType = writeMetadata(metadata); + auto size = static_cast(file_->size() - offset); + return MetadataSection{ + .offset = offset, .size = size, .compressionType = compressionType}; +} + +void TabletWriter::writeOptionalSection( + std::string name, + std::string_view content) { + NIMBLE_CHECK(!name.empty(), "Optional section name cannot be empty."); + NIMBLE_CHECK( + optionalSections_.find(name) == optionalSections_.end(), + fmt::format("Optional section '{}' already exists.", name)); + optionalSections_.try_emplace( + std::move(name), createMetadataSection(content)); +} + +void TabletWriter::tryWriteStripeGroup(bool force) { + auto stripeCount = streamOffsets_.size(); + NIMBLE_ASSERT(stripeCount == streamSizes_.size(), "Stripe count mismatch."); + if (stripeCount == 0) { + return; + } + + // Estimate size + // 8 bytes for offsets, 4 for size, 1 for compression type, so 13. + size_t estimatedSize = 4 + stripeCount * streamOffsets_.back().size() * 13; + if (!force && (estimatedSize < options_.metadataFlushThreshold)) { + return; + } + + auto maxStreamCountIt = std::max_element( + streamOffsets_.begin(), + streamOffsets_.end(), + [](const auto& a, const auto& b) { return a.size() < b.size(); }); + + auto streamCount = + maxStreamCountIt == streamOffsets_.end() ? 0 : maxStreamCountIt->size(); + + // Each stripe may have different stream count recorded. + // We need to pad shorter stripes to the full length of stream count. All + // these is handled by the |createFlattenedVector| function. + flatbuffers::FlatBufferBuilder builder(kInitialFooterSize); + auto streamOffsets = createFlattenedVector( + builder, streamCount, streamOffsets_, 0); + auto streamSizes = createFlattenedVector( + builder, streamCount, streamSizes_, 0); + + builder.Finish(serialization::CreateStripeGroup( + builder, stripeCount, streamOffsets, streamSizes)); + + stripeGroups_.push_back(createMetadataSection(asView(builder))); + ++stripeGroupIndex_; + + streamOffsets_.clear(); + streamSizes_.clear(); +} + +void TabletWriter::writeWithChecksum(std::string_view data) { + file_->append(data); + checksum_->update(data); +} + +void TabletWriter::writeWithChecksum(const folly::IOBuf& buf) { + for (auto buffer : buf) { + writeWithChecksum( + {reinterpret_cast(buffer.data()), buffer.size()}); + } +} + +MetadataBuffer::MetadataBuffer( + velox::memory::MemoryPool& memoryPool, + std::string_view ref, + CompressionType type) + : buffer_{&memoryPool} { + switch (type) { + case CompressionType::Uncompressed: { + buffer_.resize(ref.size()); + std::copy(ref.cbegin(), ref.cend(), buffer_.begin()); + break; + } + case CompressionType::Zstd: { + buffer_ = ZstdCompression::uncompress(memoryPool, ref); + break; + } + default: + NIMBLE_UNREACHABLE(fmt::format( + "Unexpected stream compression type: {}", toString(type))); + } +} + +MetadataBuffer::MetadataBuffer( + velox::memory::MemoryPool& memoryPool, + const folly::IOBuf& iobuf, + size_t offset, + size_t length, + CompressionType type) + : buffer_{&memoryPool} { + switch (type) { + case CompressionType::Uncompressed: { + buffer_.resize(length); + folly::io::Cursor cursor(&iobuf); + cursor.skip(offset); + cursor.pull(buffer_.data(), length); + break; + } + case CompressionType::Zstd: { + auto compressed = cloneAndCoalesce(iobuf, offset, length); + buffer_ = + ZstdCompression::uncompress(memoryPool, toStringView(compressed)); + break; + } + default: + NIMBLE_UNREACHABLE(fmt::format( + "Unexpected stream compression type: {}", toString(type))); + } +} + +MetadataBuffer::MetadataBuffer( + velox::memory::MemoryPool& memoryPool, + const folly::IOBuf& iobuf, + CompressionType type) + : MetadataBuffer{ + memoryPool, + iobuf, + 0, + iobuf.computeChainDataLength(), + type} {} + +void Tablet::StripeGroup::reset( + uint32_t stripeGroupIndex, + const MetadataBuffer& stripes, + uint32_t stripeIndex, + std::unique_ptr stripeGroup) { + index_ = stripeGroupIndex; + metadata_ = std::move(stripeGroup); + auto metadataRoot = + asFlatBuffersRoot(metadata_->content()); + auto stripesRoot = + asFlatBuffersRoot(stripes.content()); + + auto streamCount = metadataRoot->stream_offsets()->size(); + NIMBLE_ASSERT( + streamCount == metadataRoot->stream_sizes()->size(), + "Unexpected stream metadata"); + + auto stripeCount = metadataRoot->stripe_count(); + NIMBLE_ASSERT(stripeCount > 0, "Unexpected stripe count"); + streamCount_ = streamCount / stripeCount; + + streamOffsets_ = metadataRoot->stream_offsets()->data(); + streamSizes_ = metadataRoot->stream_sizes()->data(); + + // Find the first stripe that use this stripe group + auto groupIndices = stripesRoot->group_indices()->data(); + while (stripeIndex > 0) { + if (groupIndices[stripeIndex] != groupIndices[stripeIndex - 1]) { + break; + } + --stripeIndex; + } + firstStripe_ = stripeIndex; +} + +std::span Tablet::StripeGroup::streamOffsets( + uint32_t stripe) const { + return { + streamOffsets_ + (stripe - firstStripe_) * streamCount_, streamCount_}; +} + +std::span Tablet::StripeGroup::streamSizes( + uint32_t stripe) const { + return {streamSizes_ + (stripe - firstStripe_) * streamCount_, streamCount_}; +} + +Postscript Postscript::parse(std::string_view data) { + NIMBLE_CHECK(data.size() >= kPostscriptSize, "Invalid postscript length"); + + Postscript ps; + // Read and validate magic + auto pos = data.data() + data.size() - 2; + const uint16_t magicNumber = *reinterpret_cast(pos); + + NIMBLE_CHECK( + magicNumber == kMagicNumber, + "Magic number mismatch. Not an nimble file!"); + + // Read and validate versions + pos -= 4; + ps.majorVersion_ = *reinterpret_cast(pos); + ps.minorVersion_ = *reinterpret_cast(pos + 2); + + NIMBLE_CHECK( + ps.majorVersion_ <= kVersionMajor, + fmt::format( + "Unsupported file version. Reader version: {}, file version: {}", + kVersionMajor, + ps.majorVersion_)); + + pos -= 14; + ps.footerSize_ = *reinterpret_cast(pos); + + // How CompressionType is written into and read from postscript requires + // its size must be 1 byte. + static_assert(sizeof(CompressionType) == 1); + ps.footerCompressionType_ = + *reinterpret_cast(pos + 4); + + // How ChecksumType is written into and read from postscript requires + // its size must be 1 byte. + static_assert(sizeof(ChecksumType) == 1); + ps.checksumType_ = *reinterpret_cast(pos + 5); + ps.checksum_ = *reinterpret_cast(pos + 6); + + return ps; +} + +Tablet::Tablet( + MemoryPool& memoryPool, + std::shared_ptr readFile, + Postscript postscript, + std::string_view footer, + std::string_view stripes, + std::string_view stripeGroup, + std::unordered_map optionalSections) + : memoryPool_{memoryPool}, + file_{readFile.get()}, + ownedFile_{std::move(readFile)}, + ps_{std::move(postscript)}, + footer_{std::make_unique(memoryPool, footer)}, + stripes_{std::make_unique(memoryPool, stripes)} { + stripeGroup_.reset( + /* stripeGroupIndex */ 0, + *stripes_, + /* stripeIndex */ 0, + std::make_unique(memoryPool, stripeGroup)); + initStripes(); + for (auto& pair : optionalSections) { + optionalSectionsCache_.insert( + {pair.first, + std::make_unique(memoryPool, pair.second)}); + } +} + +Tablet::Tablet( + MemoryPool& memoryPool, + std::shared_ptr readFile, + const std::vector& preloadOptionalSections) + : Tablet{memoryPool, readFile.get(), preloadOptionalSections} { + ownedFile_ = std::move(readFile); +} + +Tablet::Tablet( + MemoryPool& memoryPool, + velox::ReadFile* readFile, + const std::vector& preloadOptionalSections) + : memoryPool_{memoryPool}, file_{readFile} { + // We make an initial read of the last piece of the file, and then do + // another read if our first one didn't cover the whole footer. We could + // make this a parameter to the constructor later. + const auto fileSize = file_->size(); + const uint64_t readSize = std::min(kInitialFooterSize, fileSize); + + NIMBLE_CHECK_FILE( + readSize >= kPostscriptSize, "Corrupted file. Footer is too small."); + + const uint64_t offset = fileSize - readSize; + velox::common::Region footerRegion{offset, readSize, "footer"}; + folly::IOBuf footerIOBuf; + file_->preadv({&footerRegion, 1}, {&footerIOBuf, 1}); + + { + folly::IOBuf psIOBuf = cloneAndCoalesce( + footerIOBuf, + footerIOBuf.computeChainDataLength() - kPostscriptSize, + kPostscriptSize); + ps_ = Postscript::parse(toStringView(psIOBuf)); + } + + NIMBLE_CHECK( + ps_.footerSize() + kPostscriptSize <= readSize, "Unexpected footer size"); + footer_ = std::make_unique( + memoryPool_, + footerIOBuf, + footerIOBuf.computeChainDataLength() - kPostscriptSize - ps_.footerSize(), + ps_.footerSize(), + ps_.footerCompressionType()); + + auto footerRoot = + asFlatBuffersRoot(footer_->content()); + + auto stripes = footerRoot->stripes(); + if (stripes) { + // For now, assume stripes section will always be within the initial fetch + NIMBLE_CHECK( + stripes->offset() + readSize >= fileSize, + "Incomplete stripes metadata."); + stripes_ = std::make_unique( + memoryPool_, + footerIOBuf, + stripes->offset() + readSize - fileSize, + stripes->size(), + static_cast(stripes->compression_type())); + auto stripesRoot = + asFlatBuffersRoot(stripes_->content()); + + auto stripeGroups = footerRoot->stripe_groups(); + NIMBLE_CHECK( + stripeGroups && + (stripeGroups->size() == + *stripesRoot->group_indices()->rbegin() + 1), + "Unexpected stripe group count"); + + // Always eagerly load if it's the only stripe group and is already + // fetched + auto stripeGroup = stripeGroups->Get(0); + if (stripeGroups->size() == 1 && + stripeGroup->offset() + readSize >= fileSize) { + stripeGroup_.reset( + /* stripeGroupIndex */ 0, + *stripes_, + /* stripeIndex */ 0, + std::make_unique( + memoryPool_, + footerIOBuf, + stripeGroup->offset() + readSize - fileSize, + stripeGroup->size(), + static_cast(stripeGroup->compression_type()))); + } + } + + initStripes(); + + auto optionalSections = footerRoot->optional_sections(); + if (optionalSections) { + NIMBLE_CHECK( + optionalSections->names() && optionalSections->offsets() && + optionalSections->sizes() && + optionalSections->compression_types() && + optionalSections->names()->size() == + optionalSections->offsets()->size() && + optionalSections->names()->size() == + optionalSections->sizes()->size() && + optionalSections->names()->size() == + optionalSections->compression_types()->size(), + "Invalid optional sections metadata."); + + optionalSections_.reserve(optionalSections->names()->size()); + + for (auto i = 0; i < optionalSections->names()->size(); ++i) { + optionalSections_.insert(std::make_pair( + optionalSections->names()->GetAsString(i)->str(), + std::make_tuple( + optionalSections->offsets()->Get(i), + optionalSections->sizes()->Get(i), + static_cast( + optionalSections->compression_types()->Get(i))))); + } + } + + std::vector mustRead; + + for (const auto& preload : preloadOptionalSections) { + auto it = optionalSections_.find(preload); + if (it == optionalSections_.end()) { + continue; + } + + const auto sectionOffset = std::get<0>(it->second); + const auto sectionSize = std::get<1>(it->second); + const auto sectionCompressionType = std::get<2>(it->second); + + if (sectionOffset < offset) { + // Section was not read yet. Need to read from file. + mustRead.emplace_back(sectionOffset, sectionSize, preload); + } else { + // Section already loaded from file + auto metadata = std::make_unique( + memoryPool_, + footerIOBuf, + sectionOffset - offset, + sectionSize, + sectionCompressionType); + optionalSectionsCache_.insert({preload, std::move(metadata)}); + } + } + if (!mustRead.empty()) { + std::vector result(mustRead.size()); + file_->preadv(mustRead, {result.data(), result.size()}); + NIMBLE_ASSERT( + result.size() == mustRead.size(), + "Region and IOBuf vector sizes don't match"); + for (size_t i = 0; i < result.size(); ++i) { + auto iobuf = std::move(result[i]); + const std::string preload{mustRead[i].label}; + auto metadata = std::make_unique( + memoryPool_, iobuf, std::get<2>(optionalSections_[preload])); + optionalSectionsCache_.insert({preload, std::move(metadata)}); + } + } +} + +uint64_t Tablet::calculateChecksum( + velox::memory::MemoryPool& memoryPool, + velox::ReadFile* readFile, + uint64_t chunkSize) { + auto postscriptStart = readFile->size() - kPostscriptSize; + Vector postscript(&memoryPool, kPostscriptSize); + readFile->pread(postscriptStart, kPostscriptSize, postscript.data()); + ChecksumType checksumType = *reinterpret_cast( + postscript.data() + kPostscriptChecksumedSize); + + auto checksum = ChecksumFactory::create(checksumType); + Vector buffer(&memoryPool); + uint64_t sizeToRead = + readFile->size() - kPostscriptSize + kPostscriptChecksumedSize; + uint64_t offset = 0; + while (sizeToRead > 0) { + auto sizeOneRead = std::min(chunkSize, sizeToRead); + buffer.resize(sizeOneRead); + std::string_view bufferRead = + readFile->pread(offset, sizeOneRead, buffer.data()); + checksum->update(bufferRead); + sizeToRead -= sizeOneRead; + offset += sizeOneRead; + } + + return checksum->getChecksum(); +} + +namespace { + +// LoadTask describes the task of PReading a region and splitting it into the +// streams within (described by StreamTask). + +struct StreamTask { + // The stream index from the input indices + uint32_t index; + + // Byte offset for the stream, relative to the beginning of the stripe + uint32_t offset; + + // Size of the stream + uint32_t length; +}; + +struct LoadTask { + // Relative to first byte in file. + uint64_t readStart; + uint32_t readLength; + std::vector streamTasks; +}; + +class PreloadedStreamLoader : public StreamLoader { + public: + explicit PreloadedStreamLoader(Vector&& stream) + : stream_{std::move(stream)} {} + const std::string_view getStream() const override { + return {stream_.data(), stream_.size()}; + } + + private: + const Vector stream_; +}; + +} // namespace + +void Tablet::initStripes() { + auto footerRoot = + asFlatBuffersRoot(footer_->content()); + tabletRowCount_ = footerRoot->row_count(); + + if (stripes_) { + auto stripesRoot = + asFlatBuffersRoot(stripes_->content()); + + stripeCount_ = stripesRoot->row_counts()->size(); + NIMBLE_CHECK(stripeCount_ > 0, "Unexpected stripe count"); + NIMBLE_CHECK( + stripeCount_ == stripesRoot->offsets()->size() && + stripeCount_ == stripesRoot->sizes()->size() && + stripeCount_ == stripesRoot->group_indices()->size(), + "Unexpected stripe count"); + + stripeRowCounts_ = stripesRoot->row_counts()->data(); + stripeOffsets_ = stripesRoot->offsets()->data(); + } +} + +void Tablet::ensureStripeGroup(uint32_t stripe) const { + auto footerRoot = + asFlatBuffersRoot(footer_->content()); + auto stripesRoot = + asFlatBuffersRoot(stripes_->content()); + auto targetIndex = stripesRoot->group_indices()->Get(stripe); + if (targetIndex != stripeGroup_.index()) { + auto stripeGroup = footerRoot->stripe_groups()->Get(targetIndex); + velox::common::Region stripeGroupRegion{ + stripeGroup->offset(), stripeGroup->size(), "StripeGroup"}; + folly::IOBuf result; + file_->preadv({&stripeGroupRegion, 1}, {&result, 1}); + + stripeGroup_.reset( + targetIndex, + *stripes_, + stripe, + std::make_unique( + memoryPool_, + result, + static_cast(stripeGroup->compression_type()))); + } +} + +std::span Tablet::streamOffsets(uint32_t stripe) const { + ensureStripeGroup(stripe); + return stripeGroup_.streamOffsets(stripe); +} + +std::span Tablet::streamSizes(uint32_t stripe) const { + ensureStripeGroup(stripe); + return stripeGroup_.streamSizes(stripe); +} + +uint32_t Tablet::streamCount(uint32_t stripe) const { + ensureStripeGroup(stripe); + return stripeGroup_.streamCount(); +} + +std::vector> Tablet::load( + uint32_t stripe, + std::span streamIdentifiers, + std::function streamLabel) const { + NIMBLE_CHECK(stripe < stripeCount_, "Stripe is out of range."); + + const uint64_t stripeOffset = this->stripeOffset(stripe); + ensureStripeGroup(stripe); + const auto stripeStreamOffsets = stripeGroup_.streamOffsets(stripe); + const auto stripeStreamSizes = stripeGroup_.streamSizes(stripe); + const uint32_t streamsToLoad = streamIdentifiers.size(); + + std::vector> streams(streamsToLoad); + std::vector regions; + std::vector streamIdx; + regions.reserve(streamsToLoad); + streamIdx.reserve(streamsToLoad); + + for (uint32_t i = 0; i < streamsToLoad; ++i) { + const uint32_t streamIdentifier = streamIdentifiers[i]; + if (streamIdentifier >= stripeGroup_.streamCount()) { + streams[i] = nullptr; + continue; + } + + const uint32_t streamSize = stripeStreamSizes[streamIdentifier]; + if (streamSize == 0) { + streams[i] = nullptr; + continue; + } + + const auto streamStart = + stripeOffset + stripeStreamOffsets[streamIdentifier]; + regions.emplace_back( + streamStart, streamSize, streamLabel(streamIdentifier)); + streamIdx.push_back(i); + } + if (!regions.empty()) { + std::vector iobufs(regions.size()); + file_->preadv(regions, {iobufs.data(), iobufs.size()}); + NIMBLE_DASSERT(iobufs.size() == streamIdx.size(), "Buffer size mismatch."); + for (uint32_t i = 0; i < streamIdx.size(); ++i) { + const auto size = iobufs[i].computeChainDataLength(); + Vector vector{&memoryPool_, size}; + copyTo(iobufs[i], vector.data(), vector.size()); + streams[streamIdx[i]] = + std::make_unique(std::move(vector)); + } + } + + return streams; +} + +std::optional
Tablet::loadOptionalSection( + const std::string& name, + bool keepCache) const { + NIMBLE_CHECK(!name.empty(), "Optional section name cannot be empty."); + auto itCache = optionalSectionsCache_.find(name); + if (itCache != optionalSectionsCache_.end()) { + if (keepCache) { + return Section{MetadataBuffer{*itCache->second}}; + } else { + auto metadata = std::move(itCache->second); + optionalSectionsCache_.erase(itCache); + return Section{std::move(*metadata)}; + } + } + + auto it = optionalSections_.find(name); + if (it == optionalSections_.end()) { + return std::nullopt; + } + + const auto offset = std::get<0>(it->second); + const auto size = std::get<1>(it->second); + const auto compressionType = std::get<2>(it->second); + + velox::common::Region region{offset, size, name}; + folly::IOBuf iobuf; + file_->preadv({®ion, 1}, {&iobuf, 1}); + return Section{MetadataBuffer{memoryPool_, iobuf, compressionType}}; +} +} // namespace facebook::nimble diff --git a/dwio/nimble/tablet/Tablet.h b/dwio/nimble/tablet/Tablet.h new file mode 100644 index 0000000..864c8e7 --- /dev/null +++ b/dwio/nimble/tablet/Tablet.h @@ -0,0 +1,391 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once +#include + +#include "dwio/nimble/common/Checksum.h" +#include "dwio/nimble/common/Vector.h" +#include "folly/Range.h" +#include "folly/io/IOBuf.h" +#include "velox/common/file/File.h" +#include "velox/common/memory/Memory.h" + +// The Tablet class is the on-disk layout for nimble. +// +// As data is streamed into a tablet, we buffer it until the total amount +// of memory used for buffering hits a chosen limit. Then we convert the +// buffered memory to streams and write them out to disk in a stripe, recording +// their byte ranges. This continues until all data for the file has been +// streamed in, at which point we write out any remaining buffered data and +// write out the byte ranges + some other metadata in the footer. +// +// The general recommendation for the buffering limit is to make it as large +// as the amount of memory you've allocated to a single processing task. The +// rationale being that the highest memory read case (select *) loads all the +// encoded stream, and in the worst case (totally random data) the encoded data +// will be the same size as the raw data. + +namespace facebook::nimble { + +using MemoryPool = facebook::velox::memory::MemoryPool; + +class MetadataBuffer { + public: + MetadataBuffer( + velox::memory::MemoryPool& memoryPool, + std::string_view ref, + CompressionType type = CompressionType::Uncompressed); + + MetadataBuffer( + velox::memory::MemoryPool& memoryPool, + const folly::IOBuf& iobuf, + size_t offset, + size_t length, + CompressionType type = CompressionType::Uncompressed); + + MetadataBuffer( + velox::memory::MemoryPool& memoryPool, + const folly::IOBuf& iobuf, + CompressionType type = CompressionType::Uncompressed); + + std::string_view content() const { + return {buffer_.data(), buffer_.size()}; + } + + private: + Vector buffer_; +}; + +class Section { + public: + explicit Section(MetadataBuffer&& buffer) : buffer_{std::move(buffer)} {} + + std::string_view content() const { + return buffer_.content(); + } + explicit operator std::string_view() const { + return content(); + } + + private: + MetadataBuffer buffer_; +}; + +class Postscript { + public: + uint32_t footerSize() const { + return footerSize_; + } + + CompressionType footerCompressionType() const { + return footerCompressionType_; + } + + uint64_t checksum() const { + return checksum_; + } + + ChecksumType checksumType() const { + return checksumType_; + } + + uint32_t majorVersion() const { + return majorVersion_; + } + + uint32_t minorVersion() const { + return minorVersion_; + } + + static Postscript parse(std::string_view data); + + private: + uint32_t footerSize_; + CompressionType footerCompressionType_; + uint64_t checksum_; + ChecksumType checksumType_; + uint32_t majorVersion_; + uint32_t minorVersion_; +}; + +// Stream loader abstraction. +// This is the returned object when loading streams from a tablet. +class StreamLoader { + public: + virtual ~StreamLoader() = default; + virtual const std::string_view getStream() const = 0; +}; + +// Provides read access to a Tablet written by a TabletWriter. +// Example usage to read all streams from stripe 0 in a file: +// auto readFile = std::make_unique("/tmp/myfile"); +// Tablet tablet(std::move(readFile)); +// auto serializedStreams = tablet.load(0, std::vector{1, 2}); +// |serializedStreams[i]| now contains the stream corresponding to +// the stream identifier provided in the input vector. +class Tablet { + public: + // Compute checksum from the beginning of the file all the way to footer + // size and footer compression type field in postscript. + // chunkSize means each time reads up to chunkSize, until all data are done. + static uint64_t calculateChecksum( + MemoryPool& memoryPool, + velox::ReadFile* readFile, + uint64_t chunkSize = 256 * 1024 * 1024); + + Tablet( + MemoryPool& memoryPool, + velox::ReadFile* readFile, + const std::vector& preloadOptionalSections = {}); + Tablet( + MemoryPool& memoryPool, + std::shared_ptr readFile, + const std::vector& preloadOptionalSections = {}); + + // Returns a collection of stream loaders for the given stripe. The stream + // loaders are returned in the same order as the input stream identifiers + // span. If a stream was not present in the given stripe a nullptr is returned + // in its slot. + std::vector> load( + uint32_t stripe, + std::span streamIdentifiers, + std::function streamLabel = [](uint32_t) { + return std::string_view{}; + }) const; + + std::optional
loadOptionalSection( + const std::string& name, + bool keepCache = false) const; + + uint64_t fileSize() const { + return file_->size(); + } + + uint32_t footerSize() const { + return ps_.footerSize(); + } + + CompressionType footerCompressionType() const { + return ps_.footerCompressionType(); + } + + uint64_t checksum() const { + return ps_.checksum(); + } + + ChecksumType checksumType() const { + return ps_.checksumType(); + } + + uint32_t majorVersion() const { + return ps_.majorVersion(); + } + + uint32_t minorVersion() const { + return ps_.minorVersion(); + } + + // Number of rows in the whole tablet. + uint64_t tabletRowCount() const { + return tabletRowCount_; + } + + // The number of rows in the given stripe. These sum to tabletRowCount(). + uint32_t stripeRowCount(uint32_t stripe) const { + return stripeRowCounts_[stripe]; + } + + // The number of stripes in the tablet. + uint32_t stripeCount() const { + return stripeCount_; + } + + uint64_t stripeOffset(uint32_t stripe) const { + return stripeOffsets_[stripe]; + } + + // Returns stream offsets for the specified stripe. Number of streams is + // determined by schema node count at the time when stripe is written, so it + // may have fewer number of items than the final schema node count + std::span streamOffsets(uint32_t stripe) const; + + // Returns stream sizes for the specified stripe. Has same constraint as + // `streamOffsets()`. + std::span streamSizes(uint32_t stripe) const; + + // Returns stream count for the specified stripe. Has same constraint as + // `streamOffsets()`. + uint32_t streamCount(uint32_t stripe) const; + + private: + struct StripeGroup { + uint32_t index() const { + return index_; + } + + uint32_t streamCount() const { + return streamCount_; + } + + void reset( + uint32_t stripeGroupIndex, + const MetadataBuffer& stripes, + uint32_t stripeIndex, + std::unique_ptr metadata); + + std::span streamOffsets(uint32_t stripe) const; + std::span streamSizes(uint32_t stripe) const; + + private: + std::unique_ptr metadata_; + uint32_t index_{std::numeric_limits::max()}; + uint32_t streamCount_{0}; + uint32_t firstStripe_{0}; + const uint32_t* streamOffsets_{nullptr}; + const uint32_t* streamSizes_{nullptr}; + }; + + void initStripes(); + + void ensureStripeGroup(uint32_t stripe) const; + + // For testing use + Tablet( + MemoryPool& memoryPool, + std::shared_ptr readFile, + Postscript postscript, + std::string_view footer, + std::string_view stripes, + std::string_view stripeGroup, + std::unordered_map optionalSections = {}); + + MemoryPool& memoryPool_; + velox::ReadFile* file_; + std::shared_ptr ownedFile_; + + Postscript ps_; + std::unique_ptr footer_; + std::unique_ptr stripes_; + mutable StripeGroup stripeGroup_; + + uint64_t tabletRowCount_; + uint32_t stripeCount_{0}; + const uint32_t* stripeRowCounts_{nullptr}; + const uint64_t* stripeOffsets_{nullptr}; + std::unordered_map< + std::string, + std::tuple> + optionalSections_; + mutable std::unordered_map> + optionalSectionsCache_; + + friend class TabletHelper; +}; + +struct Stream { + uint32_t offset; + std::vector content; +}; + +class LayoutPlanner { + public: + virtual std::vector getLayout(std::vector&& streams) = 0; + + virtual ~LayoutPlanner() = default; +}; + +struct TabletWriterOptions { + std::unique_ptr layoutPlanner{nullptr}; + uint32_t metadataFlushThreshold{8 * 1024 * 1024}; // 8Mb + uint32_t metadataCompressionThreshold{4 * 1024 * 1024}; // 4Mb + ChecksumType checksumType{ChecksumType::XXH3_64}; +}; + +// Writes a new nimble file. +class TabletWriter { + public: + TabletWriter( + velox::memory::MemoryPool& memoryPool, + velox::WriteFile* writeFile, + TabletWriterOptions options = {}) + : file_{writeFile}, + memoryPool_(memoryPool), + options_(std::move(options)), + checksum_{ChecksumFactory::create(options_.checksumType)} {} + + // Writes out the footer. Remember that the underlying file is not readable + // until the write file itself is closed. + void close(); + + // TODO: do we want a stripe header? E.g. per stream min/max (at least for + // key cols), that sort of thing? We can add later. We'll also want to + // be able to control the stream order on disk, presumably via a options + // param to the constructor. + // + // The first argument in the map gives the stream name. The second's first + // gives part gives the uncompressed string returned from a Serialize + // function, and the second part the compression type to apply when writing to + // disk. Later we may want to generalize that compression type to include a + // level or other params. + // + // A stream's type must be the same across all stripes. + void writeStripe(uint32_t rowCount, std::vector streams); + + void writeOptionalSection(std::string name, std::string_view content); + + // The number of bytes written so far. + uint64_t size() const { + return file_->size(); + } + + // For testing purpose + uint32_t stripeGroupCount() const { + return stripeGroups_.size(); + } + + private: + struct MetadataSection { + uint64_t offset; + uint32_t size; + CompressionType compressionType; + }; + + // Write metadata entry to file + CompressionType writeMetadata(std::string_view metadata); + // Write stripe group metadata entry and also add that to footer sections if + // exceeds metadata flush size. + void tryWriteStripeGroup(bool force = false); + // Create metadata section in the file + MetadataSection createMetadataSection(std::string_view metadata); + + void writeWithChecksum(std::string_view data); + void writeWithChecksum(const folly::IOBuf& buf); + + velox::WriteFile* file_; + velox::memory::MemoryPool& memoryPool_; + TabletWriterOptions options_; + std::unique_ptr checksum_; + + // Number of rows in each stripe. + std::vector rowCounts_; + // Offsets from start of file to first byte in each stripe. + std::vector stripeOffsets_; + // Sizes of the stripes + std::vector stripeSizes_; + // Stripe group indices + std::vector stripeGroupIndices_; + // Accumulated offsets within each stripe relative to start of the stripe, + // with one value for each seen stream in the current OR PREVIOUS stripes. + std::vector> streamOffsets_; + // Accumulated stream sizes within each stripe. Same behavior as + // streamOffsets_. + std::vector> streamSizes_; + // Current stripe group index + uint32_t stripeGroupIndex_{0}; + // Stripe groups + std::vector stripeGroups_; + // Optional sections + std::unordered_map optionalSections_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/tablet/footer_flatc.sh b/dwio/nimble/tablet/footer_flatc.sh new file mode 100755 index 0000000..a67ec01 --- /dev/null +++ b/dwio/nimble/tablet/footer_flatc.sh @@ -0,0 +1,4 @@ +#!/bin/bash +# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +$FLATC -j -o "${OUT}" --filename-suffix Generated ./Footer.fbs diff --git a/dwio/nimble/tablet/tests/CMakeLists.txt b/dwio/nimble/tablet/tests/CMakeLists.txt new file mode 100644 index 0000000..e379ded --- /dev/null +++ b/dwio/nimble/tablet/tests/CMakeLists.txt @@ -0,0 +1,14 @@ +add_executable(nimble_tablet_tests TabletTests.cpp) + +add_test(nimble_tablet_tests nimble_tablet_tests) + +target_link_libraries( + nimble_tablet_tests + nimble_tablet + nimble_common + velox_memory + velox_file + gtest + gtest_main + glog::glog + Folly::folly) diff --git a/dwio/nimble/tablet/tests/TabletTests.cpp b/dwio/nimble/tablet/tests/TabletTests.cpp new file mode 100644 index 0000000..69f04eb --- /dev/null +++ b/dwio/nimble/tablet/tests/TabletTests.cpp @@ -0,0 +1,724 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include +#include +#include +#include + +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/Checksum.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/tests/TestUtils.h" +#include "dwio/nimble/tablet/Tablet.h" +#include "folly/FileUtil.h" +#include "folly/Random.h" +#include "folly/experimental/coro/Generator.h" +#include "velox/common/file/File.h" +#include "velox/common/memory/Memory.h" + +using namespace ::facebook; + +namespace { + +// Total size of the fields after the flatbuffer. +constexpr uint32_t kPostscriptSize = 20; + +struct StripeSpecifications { + uint32_t rowCount; + std::vector streamOffsets; +}; + +struct StripeData { + uint32_t rowCount; + std::vector streams; +}; + +void printData(std::string prefix, std::string_view data) { + std::string output; + for (auto i = 0; i < data.size(); ++i) { + output += folly::to((uint8_t)data[i]) + " "; + } + + LOG(INFO) << prefix << " (" << (void*)data.data() << "): " << output; +} + +std::vector createStripesData( + std::mt19937& rng, + const std::vector& stripes, + nimble::Buffer& buffer) { + std::vector stripesData; + stripesData.reserve(stripes.size()); + + // Each generator iteration returns a single stripe to write + for (auto& stripe : stripes) { + std::vector streams; + streams.reserve(stripe.streamOffsets.size()); + std::transform( + stripe.streamOffsets.cbegin(), + stripe.streamOffsets.cend(), + std::back_inserter(streams), + [&rng, &buffer](auto offset) { + const auto size = folly::Random::rand32(32, rng) + 2; + auto pos = buffer.reserve(size); + for (auto i = 0; i < size; ++i) { + pos[i] = folly::Random::rand32(256, rng); + } + printData(folly::to("Stream ", offset), {pos, size}); + + return nimble::Stream{.offset = offset, .content = {{pos, size}}}; + }); + stripesData.push_back({ + .rowCount = stripe.rowCount, + .streams = std::move(streams), + }); + } + + return stripesData; +} + +// Runs a single write/read test using input parameters +void parameterizedTest( + std::mt19937& rng, + velox::memory::MemoryPool& memoryPool, + uint32_t metadataFlushThreshold, + uint32_t metadataCompressionThreshold, + std::vector stripes, + const std::optional>& + errorVerifier = std::nullopt) { + try { + std::string file; + velox::InMemoryWriteFile writeFile(&file); + nimble::TabletWriter tabletWriter{ + memoryPool, + &writeFile, + {nullptr, metadataFlushThreshold, metadataCompressionThreshold}}; + + EXPECT_EQ(0, tabletWriter.size()); + + struct StripeData { + uint32_t rowCount; + std::vector streams; + }; + + nimble::Buffer buffer{memoryPool}; + auto stripesData = createStripesData(rng, stripes, buffer); + for (auto& stripe : stripesData) { + tabletWriter.writeStripe(stripe.rowCount, stripe.streams); + } + + tabletWriter.close(); + EXPECT_LT(0, tabletWriter.size()); + writeFile.close(); + EXPECT_EQ(writeFile.size(), tabletWriter.size()); + + auto stripeGroupCount = tabletWriter.stripeGroupCount(); + + folly::writeFile(file, "/tmp/test.nimble"); + + for (auto useChainedBuffers : {false, true}) { + nimble::testing::InMemoryTrackableReadFile readFile( + file, useChainedBuffers); + nimble::Tablet tablet{memoryPool, &readFile}; + EXPECT_EQ(stripesData.size(), tablet.stripeCount()); + EXPECT_EQ( + std::accumulate( + stripesData.begin(), + stripesData.end(), + uint64_t{0}, + [](uint64_t r, const auto& s) { return r + s.rowCount; }), + tablet.tabletRowCount()); + + VLOG(1) << "Output Tablet -> StripeCount: " << tablet.stripeCount() + << ", RowCount: " << tablet.tabletRowCount(); + + // Now, read all stripes and verify results + size_t extraReads = 0; + for (auto stripe = 0; stripe < stripesData.size(); ++stripe) { + EXPECT_EQ(stripesData[stripe].rowCount, tablet.stripeRowCount(stripe)); + + readFile.resetChunks(); + std::vector identifiers(tablet.streamCount(stripe)); + std::iota(identifiers.begin(), identifiers.end(), 0); + auto serializedStreams = + tablet.load(stripe, {identifiers.cbegin(), identifiers.cend()}); + auto chunks = readFile.chunks(); + auto expectedReads = stripesData[stripe].streams.size(); + auto diff = chunks.size() - expectedReads; + EXPECT_LE(diff, 1); + extraReads += diff; + + for (const auto& chunk : chunks) { + VLOG(1) << "Chunk Offset: " << chunk.offset + << ", Size: " << chunk.size << ", Stripe: " << stripe; + } + + for (auto i = 0; i < serializedStreams.size(); ++i) { + // Verify streams content. If stream wasn't written in this stripe, it + // should return nullopt optional. + auto found = false; + for (const auto& stream : stripesData[stripe].streams) { + if (stream.offset == i) { + found = true; + EXPECT_TRUE(serializedStreams[i]); + printData( + folly::to("Expected Stream ", stream.offset), + stream.content.front()); + const auto& actual = serializedStreams[i]; + std::string_view actualData = actual->getStream(); + printData( + folly::to("Actual Stream ", stream.offset), + actualData); + EXPECT_EQ(stream.content.front(), actualData); + } + } + if (!found) { + EXPECT_FALSE(serializedStreams[i]); + } + } + } + + EXPECT_EQ(extraReads, (stripeGroupCount == 1 ? 0 : stripeGroupCount)); + + if (errorVerifier.has_value()) { + FAIL() << "Error verifier is provided, but no exception was thrown."; + } + } + } catch (const std::exception& e) { + if (!errorVerifier.has_value()) { + FAIL() << "Unexpected exception: " << e.what(); + } + + errorVerifier.value()(e); + + // If errorVerifier detected an error, log the exception + if (testing::Test::HasFatalFailure()) { + FAIL() << "Failed verifying exception: " << e.what(); + } else if (testing::Test::HasNonfatalFailure()) { + LOG(WARNING) << "Failed verifying exception: " << e.what(); + } + } +} + +// Run all permutations of a test using all test parameters +void test( + velox::memory::MemoryPool& memoryPool, + std::vector stripes, + std::optional> errorVerifier = + std::nullopt) { + std::vector metadataCompressionThresholds{ + // use size 0 here so it will always force a footer compression + 0, + // use a large number here so it will not do a footer compression + 1024 * 1024 * 1024}; + std::vector metadataFlushThresholds{ + 0, // force flush + 100, // flush a few + 1024 * 1024 * 1024, // never flush + }; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (auto flushThreshold : metadataFlushThresholds) { + for (auto compressionThreshold : metadataCompressionThresholds) { + LOG(INFO) << "FlushThreshold: " << flushThreshold + << ", CompressionThreshold: " << compressionThreshold; + parameterizedTest( + rng, + memoryPool, + flushThreshold, + compressionThreshold, + stripes, + errorVerifier); + } + } +} + +} // namespace + +class TabletTestSuite : public ::testing::Test { + protected: + void SetUp() override { + pool_ = velox::memory::deprecatedAddDefaultLeafMemoryPool(); + } + + std::shared_ptr pool_; +}; + +TEST_F(TabletTestSuite, EmptyWrite) { + // Creating an Nimble file without writing any stripes + test( + *this->pool_, + /* stripes */ {}); +} + +TEST_F(TabletTestSuite, WriteDifferentStreamsPerStripe) { + // Write different subset of streams in each stripe + test( + *this->pool_, + /* stripes */ + { + {.rowCount = 20, .streamOffsets = {3, 1}}, + {.rowCount = 30, .streamOffsets = {2}}, + {.rowCount = 20, .streamOffsets = {4}}, + }); +} + +namespace { +void checksumTest( + std::mt19937& rng, + velox::memory::MemoryPool& memoryPool, + uint32_t metadataCompressionThreshold, + nimble::ChecksumType checksumType, + bool checksumChunked, + std::vector stripes) { + std::string file; + velox::InMemoryWriteFile writeFile(&file); + nimble::TabletWriter tabletWriter{ + memoryPool, + &writeFile, + {.layoutPlanner = nullptr, + .metadataCompressionThreshold = metadataCompressionThreshold, + .checksumType = checksumType}}; + EXPECT_EQ(0, tabletWriter.size()); + + struct StripeData { + uint32_t rowCount; + std::vector streams; + }; + + nimble::Buffer buffer{memoryPool}; + auto stripesData = createStripesData(rng, stripes, buffer); + + for (auto& stripe : stripesData) { + tabletWriter.writeStripe(stripe.rowCount, stripe.streams); + } + + tabletWriter.close(); + EXPECT_LT(0, tabletWriter.size()); + writeFile.close(); + EXPECT_EQ(writeFile.size(), tabletWriter.size()); + + for (auto useChaniedBuffers : {false, true}) { + // Velidate checksum on a good file + nimble::testing::InMemoryTrackableReadFile readFile( + file, useChaniedBuffers); + nimble::Tablet tablet{memoryPool, &readFile}; + auto storedChecksum = tablet.checksum(); + EXPECT_EQ( + storedChecksum, + nimble::Tablet::calculateChecksum( + memoryPool, + &readFile, + checksumChunked ? writeFile.size() / 3 : writeFile.size())) + << "metadataCompressionThreshold: " << metadataCompressionThreshold + << ", checksumType: " << nimble::toString(checksumType) + << ", checksumChunked: " << checksumChunked; + + // Flip a bit in the stream and verify that checksum can catch the error + { + // First, make sure we are working on a clean file + nimble::testing::InMemoryTrackableReadFile readFileUnchanged( + file, useChaniedBuffers); + EXPECT_EQ( + storedChecksum, + nimble::Tablet::calculateChecksum(memoryPool, &readFileUnchanged)); + + char& c = file[10]; + c ^= 0x80; + nimble::testing::InMemoryTrackableReadFile readFileChanged( + file, useChaniedBuffers); + EXPECT_NE( + storedChecksum, + nimble::Tablet::calculateChecksum( + memoryPool, + &readFileChanged, + checksumChunked ? writeFile.size() / 3 : writeFile.size())) + << "Checksum didn't find corruption when stream content is changed. " + << "metadataCompressionThreshold: " << metadataCompressionThreshold + << ", checksumType: " << nimble::toString(checksumType) + << ", checksumChunked: " << checksumChunked; + // revert the file back. + c ^= 0x80; + } + + // Flip a bit in the flatbuffer footer and verify that checksum can catch + // the error + { + // First, make sure we are working on a clean file + nimble::testing::InMemoryTrackableReadFile readFileUnchanged( + file, useChaniedBuffers); + EXPECT_EQ( + storedChecksum, + nimble::Tablet::calculateChecksum(memoryPool, &readFileUnchanged)); + + auto posInFooter = + tablet.fileSize() - kPostscriptSize - tablet.footerSize() / 2; + uint8_t& byteInFooter = + *reinterpret_cast(file.data() + posInFooter); + byteInFooter ^= 0x1; + nimble::testing::InMemoryTrackableReadFile readFileChanged( + file, useChaniedBuffers); + EXPECT_NE( + storedChecksum, + nimble::Tablet::calculateChecksum( + memoryPool, + &readFileChanged, + checksumChunked ? writeFile.size() / 3 : writeFile.size())) + << "Checksum didn't find corruption when footer content is changed. " + << "metadataCompressionThreshold: " << metadataCompressionThreshold + << ", checksumType: " << nimble::toString(checksumType) + << ", checksumChunked: " << checksumChunked; + // revert the file back. + byteInFooter ^= 0x1; + } + + // Flip a bit in the footer size field and verify that checksum can catch + // the error + { + // First, make sure we are working on a clean file + nimble::testing::InMemoryTrackableReadFile readFileUnchanged( + file, useChaniedBuffers); + EXPECT_EQ( + storedChecksum, + nimble::Tablet::calculateChecksum(memoryPool, &readFileUnchanged)); + + auto footerSizePos = tablet.fileSize() - kPostscriptSize; + uint32_t& footerSize = + *reinterpret_cast(file.data() + footerSizePos); + ASSERT_EQ(footerSize, tablet.footerSize()); + footerSize ^= 0x1; + nimble::testing::InMemoryTrackableReadFile readFileChanged( + file, useChaniedBuffers); + EXPECT_NE( + storedChecksum, + nimble::Tablet::calculateChecksum( + memoryPool, + &readFileChanged, + checksumChunked ? writeFile.size() / 3 : writeFile.size())) + << "Checksum didn't find corruption when footer size field is changed. " + << "metadataCompressionThreshold: " << metadataCompressionThreshold + << ", checksumType: " << nimble::toString(checksumType) + << ", checksumChunked: " << checksumChunked; + // revert the file back. + footerSize ^= 0x1; + } + + // Flip a bit in the footer compression type field and verify that checksum + // can catch the error + { + // First, make sure we are working on a clean file + nimble::testing::InMemoryTrackableReadFile readFileUnchanged( + file, useChaniedBuffers); + EXPECT_EQ( + storedChecksum, + nimble::Tablet::calculateChecksum(memoryPool, &readFileUnchanged)); + + auto footerCompressionTypePos = tablet.fileSize() - kPostscriptSize + 4; + nimble::CompressionType& footerCompressionType = + *reinterpret_cast( + file.data() + footerCompressionTypePos); + ASSERT_EQ(footerCompressionType, tablet.footerCompressionType()); + // Cannot do bit operation on enums, so cast it to integer type. + uint8_t& typeAsInt = *reinterpret_cast(&footerCompressionType); + typeAsInt ^= 0x1; + nimble::testing::InMemoryTrackableReadFile readFileChanged( + file, useChaniedBuffers); + EXPECT_NE( + storedChecksum, + nimble::Tablet::calculateChecksum( + memoryPool, + &readFileChanged, + checksumChunked ? writeFile.size() / 3 : writeFile.size())) + << "Checksum didn't find corruption when compression type field is changed. " + << "metadataCompressionThreshold: " << metadataCompressionThreshold + << ", checksumType: " << nimble::toString(checksumType) + << ", checksumChunked: " << checksumChunked; + // revert the file back. + typeAsInt ^= 0x1; + } + } +} +} // namespace + +TEST_F(TabletTestSuite, ChecksumValidation) { + std::vector metadataCompressionThresholds{ + // use size 0 here so it will always force a footer compression + 0, + // use a large number here so it will not do a footer compression + 1024 * 1024 * 1024}; + + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + for (auto metadataCompressionThreshold : metadataCompressionThresholds) { + for (auto algorithm : {nimble::ChecksumType::XXH3_64}) { + for (auto checksumChunked : {true, false}) { + checksumTest( + rng, + *this->pool_, + metadataCompressionThreshold, + algorithm, + checksumChunked, + // Write different subset of streams in each stripe + /* stripes */ + { + {.rowCount = 20, .streamOffsets = {1}}, + {.rowCount = 30, .streamOffsets = {2}}, + }); + } + } + } +} + +TEST(TabletTests, OptionalSections) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng{seed}; + + auto pool = velox::memory::deprecatedAddDefaultLeafMemoryPool(); + std::string file; + velox::InMemoryWriteFile writeFile(&file); + nimble::TabletWriter tabletWriter{*pool, &writeFile}; + + auto randomSize = folly::Random::rand32(20, 2000000, rng); + std::string random; + random.resize(folly::Random::rand32(20, 2000000, rng)); + for (auto i = 0; i < random.size(); ++i) { + random[i] = folly::Random::rand32(256); + } + { + const std::string& content = random; + tabletWriter.writeOptionalSection("section1", content); + } + { + std::string content; + content.resize(randomSize); + for (auto i = 0; i < content.size(); ++i) { + content[i] = '\0'; + } + + tabletWriter.writeOptionalSection("section2", content); + } + { + std::string content; + tabletWriter.writeOptionalSection("section3", content); + } + + tabletWriter.close(); + + for (auto useChaniedBuffers : {false, true}) { + nimble::testing::InMemoryTrackableReadFile readFile( + file, useChaniedBuffers); + nimble::Tablet tablet{*pool, &readFile}; + + auto section = tablet.loadOptionalSection("section1"); + ASSERT_TRUE(section.has_value()); + ASSERT_EQ(random, section->content()); + + std::string expectedContent; + expectedContent.resize(randomSize); + for (auto i = 0; i < expectedContent.size(); ++i) { + expectedContent[i] = '\0'; + } + section = tablet.loadOptionalSection("section2"); + ASSERT_TRUE(section.has_value()); + ASSERT_EQ(expectedContent, section->content()); + + section = tablet.loadOptionalSection("section3"); + ASSERT_TRUE(section.has_value()); + ASSERT_EQ(std::string(), section->content()); + + section = tablet.loadOptionalSection("section4"); + ASSERT_FALSE(section.has_value()); + } +} + +TEST(TabletTests, OptionalSectionsEmpty) { + auto pool = velox::memory::deprecatedAddDefaultLeafMemoryPool(); + std::string file; + velox::InMemoryWriteFile writeFile(&file); + nimble::TabletWriter tabletWriter{*pool, &writeFile}; + + tabletWriter.close(); + + for (auto useChaniedBuffers : {false, true}) { + nimble::testing::InMemoryTrackableReadFile readFile( + file, useChaniedBuffers); + nimble::Tablet tablet{*pool, &readFile}; + + auto section = tablet.loadOptionalSection("section1"); + ASSERT_FALSE(section.has_value()); + } +} + +TEST(TabletTests, OptionalSectionsPreload) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng{seed}; + + for (const auto footerCompressionThreshold : + {0U, std::numeric_limits::max()}) { + auto pool = velox::memory::deprecatedAddDefaultLeafMemoryPool(); + std::string file; + velox::InMemoryWriteFile writeFile(&file); + nimble::TabletWriter tabletWriter{*pool, &writeFile}; + + // Using random string to make sure compression can't compress it well + std::string random; + random.resize(20 * 1024 * 1024); + for (auto i = 0; i < random.size(); ++i) { + random[i] = folly::Random::rand32(256); + } + + tabletWriter.writeOptionalSection("section1", "aaaa"); + tabletWriter.writeOptionalSection("section2", "bbbb"); + tabletWriter.writeOptionalSection("section3", random); + tabletWriter.writeOptionalSection("section4", "dddd"); + tabletWriter.writeOptionalSection("section5", "eeee"); + tabletWriter.close(); + + auto verify = [&](std::vector preload, + size_t expectedInitialReads, + std::vector> + expected) { + for (auto useChaniedBuffers : {false, true}) { + nimble::testing::InMemoryTrackableReadFile readFile( + file, useChaniedBuffers); + nimble::Tablet tablet{*pool, &readFile, preload}; + + // Expecting only the initial footer read. + ASSERT_EQ(expectedInitialReads, readFile.chunks().size()); + + for (const auto& e : expected) { + auto expectedSection = std::get<0>(e); + auto expectedReads = std::get<1>(e); + auto expectedContent = std::get<2>(e); + auto section = tablet.loadOptionalSection(expectedSection); + ASSERT_TRUE(section.has_value()); + ASSERT_EQ(expectedContent, section->content()); + ASSERT_EQ(expectedReads, readFile.chunks().size()); + } + } + }; + + // Not preloading anything. + // Expecting one initial read, to read the footer. + // Each section load should increment read count + verify( + /* preload */ {}, + /* expectedInitialReads */ 1, + { + {"section1", 2, "aaaa"}, + {"section2", 3, "bbbb"}, + {"section3", 4, random}, + {"section4", 5, "dddd"}, + {"section5", 6, "eeee"}, + {"section1", 7, "aaaa"}, + }); + + // Preloading small section covered by footer. + // Expecting one initial read, to read only the footer. + verify( + /* preload */ {"section4"}, + /* expectedInitialReads */ 1, + { + {"section1", 2, "aaaa"}, + {"section2", 3, "bbbb"}, + {"section3", 4, random}, + // This is the preloaded section, so it should not trigger a read + {"section4", 4, "dddd"}, + {"section5", 5, "eeee"}, + // This is the preloaded section, but the section was already + // consumed, so it should now trigger a read. + {"section4", 6, "dddd"}, + }); + + // Preloading section partially covered by footer. + // Expecting one initial read for the footer, and an additional read to read + // the full preloaded section. + verify( + /* preload */ {"section3"}, + /* expectedInitialReads */ 2, + { + {"section1", 3, "aaaa"}, + {"section2", 4, "bbbb"}, + // This is the preloaded section, so it should not trigger a read + {"section3", 4, random}, + {"section4", 5, "dddd"}, + {"section5", 6, "eeee"}, + // This is the preloaded section, but the section was already + // consumed, so it should now trigger a read. + {"section3", 7, random}, + }); + + // Preloading section completely not covered by footer. + // Expecting one initial read for the footer, and an additional read to read + // the uncovered preloaded section. + verify( + /* preload */ {"section1"}, + /* expectedInitialReads */ 2, + { + // This is the preloaded section, so it should not trigger a read + {"section1", 2, "aaaa"}, + {"section2", 3, "bbbb"}, + {"section3", 4, random}, + {"section4", 5, "dddd"}, + {"section5", 6, "eeee"}, + // This is the preloaded section, but the section was already + // consumed, so it should now trigger a read. + {"section1", 7, "aaaa"}, + }); + + // Preloading multiple sections. One covered, and the other is not. + // Expecting one initial read for the footer, and an additional read to read + // the uncovered preloaded section. + verify( + /* preload */ {"section2", "section4"}, + /* expectedInitialReads */ 2, + { + {"section1", 3, "aaaa"}, + // This is one of the preloaded sections, so it should not trigger a + // read + {"section2", 3, "bbbb"}, + {"section3", 4, random}, + // This is the other preloaded sections, so it should not trigger a + // read + {"section4", 4, "dddd"}, + {"section5", 5, "eeee"}, + // This is one of the preloaded section, but the section was already + // consumed, so it should now trigger a read. + {"section2", 6, "bbbb"}, + // This is the other preloaded section, but the section was already + // consumed, so it should now trigger a read. + {"section4", 7, "dddd"}, + }); + + // Preloading all sections. Two are fully covered, and three are + // partially or not covered. + // Expecting one initial read for the footer (and the covered sections), and + // additional reads to read the partial/uncovered preloaded sections. + verify( + /* preload */ + {"section2", "section4", "section1", "section3", "section5"}, + /* expectedInitialReads */ 4, + { + // All sections are preloaded, so no addtional reads + {"section1", 4, "aaaa"}, + {"section2", 4, "bbbb"}, + {"section3", 4, random}, + {"section4", 4, "dddd"}, + {"section5", 4, "eeee"}, + // Now all sections are consumed so all should trigger reads + {"section1", 5, "aaaa"}, + {"section2", 6, "bbbb"}, + {"section3", 7, random}, + {"section4", 8, "dddd"}, + {"section5", 9, "eeee"}, + }); + } +} diff --git a/dwio/nimble/tools/CMakeLists.txt b/dwio/nimble/tools/CMakeLists.txt new file mode 100644 index 0000000..0995dc9 --- /dev/null +++ b/dwio/nimble/tools/CMakeLists.txt @@ -0,0 +1,3 @@ +add_library(nimble_tools_common EncodingUtilities.cpp) + +target_link_libraries(nimble_tools_common nimble_common nimble_tablet) diff --git a/dwio/nimble/tools/EncodingLayoutTrainer.cpp b/dwio/nimble/tools/EncodingLayoutTrainer.cpp new file mode 100644 index 0000000..51401c1 --- /dev/null +++ b/dwio/nimble/tools/EncodingLayoutTrainer.cpp @@ -0,0 +1,328 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/tools/EncodingLayoutTrainer.h" +#include +#include +#include +#include "common/strings/Zstd.h" +#include "dwio/api/NimbleWriterOptionBuilder.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" +#include "dwio/nimble/encodings/EncodingLayoutCapture.h" +#include "dwio/nimble/velox/ChunkedStream.h" +#include "dwio/nimble/velox/VeloxReader.h" +#include "fbjava/datainfra-metastore/api/if/gen-cpp2/hive_metastore_types.h" +#include "thrift/lib/cpp/protocol/TBase64Utils.h" +#include "thrift/lib/cpp2/protocol/CompactProtocol.h" +#include "velox/dwio/common/ExecutorBarrier.h" + +namespace facebook::nimble::tools { + +namespace { + +template +class TrainingNode { + public: + TrainingNode( + std::string name, + std::unique_ptr&& state, + std::vector>>&& children) + : name_{std::move(name)}, + state_{std::move(state)}, + children_{std::move(children)} {} + + const std::string& name() const { + return name_; + } + + T& state() const { + return *state_; + } + + const std::vector>>& children() const { + return children_; + } + + private: + std::string name_; + std::unique_ptr state_; + std::vector>> children_; +}; + +template +T deserialize(const std::string& source) { + T result; + auto compressed = apache::thrift::protocol::base64Decode(source); + std::string uncompressed; + if (!strings::zstdDecompress(compressed->moveToFbString(), &uncompressed)) { + throw std::runtime_error( + fmt::format("Unable to decompress data: {}", source)); + } + apache::thrift::CompactSerializer::deserialize(uncompressed, result); + return result; +} + +struct State { + explicit State(const Type& type) : type{type} {} + + const Type& type; + std::unordered_map + encodingLayouts; + std::mutex mutex; +}; + +std::unique_ptr> createTrainingTree( + const Type& type, + const std::function< + void(const StreamDescriptor&, std::function)>& + train, + const std::string& name = "") { + auto state = std::make_unique(type); + std::vector>> children; + +#define _ASYNC_TRAIN(descriptor, identifier) \ + train(descriptor, [state = state.get()](EncodingLayout&& layout) { \ + std::lock_guard lock{state->mutex}; \ + state->encodingLayouts.insert({ \ + EncodingLayoutTree::StreamIdentifiers::identifier, \ + std::move(layout), \ + }); \ + }); + + switch (type.kind()) { + case Kind::Scalar: { + _ASYNC_TRAIN(type.asScalar().scalarDescriptor(), Scalar::ScalarStream); + break; + } + case Kind::Array: { + auto& array = type.asArray(); + _ASYNC_TRAIN(array.lengthsDescriptor(), Array::LengthsStream); + children.reserve(1); + children.emplace_back(createTrainingTree(*array.elements(), train)); + break; + } + case Kind::Map: { + auto& map = type.asMap(); + _ASYNC_TRAIN(map.lengthsDescriptor(), Map::LengthsStream); + children.reserve(2); + children.emplace_back(createTrainingTree(*map.keys(), train)); + children.emplace_back(createTrainingTree(*map.values(), train)); + break; + } + case Kind::Row: { + auto& row = type.asRow(); + _ASYNC_TRAIN(row.nullsDescriptor(), Row::NullsStream); + children.reserve(row.childrenCount()); + for (auto i = 0; i < row.childrenCount(); ++i) { + children.emplace_back(createTrainingTree(*row.childAt(i), train)); + } + break; + } + case Kind::FlatMap: { + auto& map = type.asFlatMap(); + _ASYNC_TRAIN(map.nullsDescriptor(), FlatMap::NullsStream); + children.reserve(map.childrenCount()); + for (auto i = 0; i < map.childrenCount(); ++i) { + children.emplace_back( + createTrainingTree(*map.childAt(i), train, map.nameAt(i))); + } + break; + } + case Kind::ArrayWithOffsets: { + auto& array = type.asArrayWithOffsets(); + _ASYNC_TRAIN(array.offsetsDescriptor(), ArrayWithOffsets::OffsetsStream); + _ASYNC_TRAIN(array.lengthsDescriptor(), ArrayWithOffsets::LengthsStream); + children.reserve(1); + children.emplace_back(createTrainingTree(*array.elements(), train)); + break; + } + } + + return std::make_unique>( + name, std::move(state), std::move(children)); +} + +class PreloadedStreamLoader : public StreamLoader { + public: + explicit PreloadedStreamLoader(std::string_view stream) : stream_{stream} {} + const std::string_view getStream() const override { + return {stream_.data(), stream_.size()}; + } + + private: + const std::string_view stream_; +}; + +template +EncodingLayout trainEncoding( + velox::memory::MemoryPool& memoryPool, + const VeloxWriterOptions& options, + const std::vector& streams) { + // Train on a single schema node. Load all data from all stripes and perform + // basic encoding selection + + std::vector> chunks; + std::vector> encodings; + uint64_t rowCount = 0; + for (const auto& stream : streams) { + InMemoryChunkedStream chunkedStream{ + memoryPool, std::make_unique(stream)}; + while (chunkedStream.hasNext()) { + auto chunk = chunkedStream.nextChunk(); + auto encoding = nimble::EncodingFactory::decode(memoryPool, chunk); + Vector data{&memoryPool}; + data.resize(encoding->rowCount()); + encoding->materialize(encoding->rowCount(), data.data()); + rowCount += encoding->rowCount(); + chunks.push_back(std::move(data)); + encodings.push_back(std::move(encoding)); + } + } + + Vector data{&memoryPool}; + data.reserve(rowCount); + for (const auto& chunk : chunks) { + for (const auto& item : chunk) { + data.push_back(item); + } + } + + auto policy = std::unique_ptr>( + static_cast*>( + options.encodingSelectionPolicyFactory(TypeTraits::dataType) + .release())); + Buffer buffer{memoryPool}; + std::string_view encoding; + encoding = + nimble::EncodingFactory::encode(std::move(policy), data, buffer); + return EncodingLayoutCapture::capture(encoding); +} + +EncodingLayoutTree toEncodingLayoutTree(const TrainingNode& node) { + auto& state = node.state(); + std::vector children; + children.reserve(node.children().size()); + for (const auto& child : node.children()) { + children.push_back(toEncodingLayoutTree(*child)); + } + return { + state.type.kind(), + std::move(state.encodingLayouts), + node.name(), + std::move(children), + }; +} + +} // namespace + +EncodingLayoutTrainer::EncodingLayoutTrainer( + velox::memory::MemoryPool& memoryPool, + std::vector files, + std::string serializedSerde) + : memoryPool_{memoryPool}, + files_{std::move(files)}, + serializedSerde_{std::move(serializedSerde)} { + NIMBLE_CHECK(!files_.empty(), "No files provided to train on"); +} + +EncodingLayoutTree EncodingLayoutTrainer::train(folly::Executor& executor) { + // Initial "training" implementation is very basic. + // It loads a single file, and for each schema node (stream), it loads all + // data from the file and performs encoding selection on it. + // + // Future versions will: + // * Support multiple files + // * Verify encoding selection stability across files/stripes. + // * Perform better encoding selection (brute forcing, etc.) + // * Measure read/write performance + // * Support different cost functions + + // One file for now + NIMBLE_CHECK(files_.size() == 1, "Only supporting single file training."); + auto& file = files_.front(); + + LOG(INFO) << "Opening file " << file; + + std::unique_ptr readFile = + std::make_unique(file); + + std::shared_ptr tablet = + std::make_shared(memoryPool_, std::move(readFile)); + std::unique_ptr reader = + std::make_unique(memoryPool_, tablet); + + std::vector>> stripeStreams; + stripeStreams.reserve(tablet->stripeCount()); + for (auto i = 0; i < tablet->stripeCount(); ++i) { + std::vector identifiers; + identifiers.resize(tablet->streamCount(i)); + std::iota(identifiers.begin(), identifiers.end(), 0); + stripeStreams.push_back(tablet->load(i, identifiers)); + } + + dwio::api::NimbleWriterOptionBuilder optionBuilder; + if (!serializedSerde_.empty()) { + optionBuilder.withSerdeParams( + reader->getType(), + deserialize(serializedSerde_) + .get_parameters()); + } + auto options = optionBuilder.build(); + + LOG(INFO) << "Training parameters: CompressionAcceptRatio = " + << options.compressionOptions.compressionAcceptRatio + << ", Zstrong.CompressionLevel = " + << options.compressionOptions.zstrongCompressionLevel + << ", Zstrong.DecompressionLevel = " + << options.compressionOptions.zstrongDecompressionLevel; + + velox::dwio::common::ExecutorBarrier barrier{executor}; + // Traverse schema. For each node, load all data and capture basic encoding + // selection on data. + auto taskTree = createTrainingTree( + *reader->schema(), + [&](const StreamDescriptor& descriptor, + std::function setLayout) { + barrier.add([&, setLayout = std::move(setLayout)]() { + std::vector streams; + for (auto& stripeStream : stripeStreams) { + const auto offset = descriptor.offset(); + if (stripeStream.size() > offset && stripeStream[offset]) { + streams.push_back(stripeStream[offset]->getStream()); + } + } + +#define _SCALAR_CASE(kind, dataType) \ + case ScalarKind::kind: \ + setLayout(trainEncoding(memoryPool_, options, streams)); \ + break; + + switch (descriptor.scalarKind()) { + _SCALAR_CASE(Int8, int8_t); + _SCALAR_CASE(UInt8, uint8_t); + _SCALAR_CASE(Int16, int16_t); + _SCALAR_CASE(UInt16, uint16_t); + _SCALAR_CASE(Int32, int32_t); + _SCALAR_CASE(UInt32, uint32_t); + _SCALAR_CASE(Int64, int64_t); + _SCALAR_CASE(UInt64, uint64_t); + _SCALAR_CASE(Float, float); + _SCALAR_CASE(Double, double); + _SCALAR_CASE(Bool, bool); + _SCALAR_CASE(String, std::string_view); + _SCALAR_CASE(Binary, std::string_view); + case ScalarKind::Undefined: + NIMBLE_UNREACHABLE("Scalar kind cannot be undefined."); + } + +#undef _SCALAR_KIND + }); + }); + + barrier.waitAll(); + + return toEncodingLayoutTree(*taskTree); +} + +} // namespace facebook::nimble::tools diff --git a/dwio/nimble/tools/EncodingLayoutTrainer.h b/dwio/nimble/tools/EncodingLayoutTrainer.h new file mode 100644 index 0000000..a9a08eb --- /dev/null +++ b/dwio/nimble/tools/EncodingLayoutTrainer.h @@ -0,0 +1,30 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/velox/EncodingLayoutTree.h" +#include "folly/Executor.h" +#include "velox/common/memory/Memory.h" + +namespace facebook::nimble::tools { + +class EncodingLayoutTrainer { + public: + EncodingLayoutTrainer( + velox::memory::MemoryPool& memoryPool, + std::vector files, + std::string serializedSerde = ""); + + EncodingLayoutTree train(folly::Executor& executor); + + private: + velox::memory::MemoryPool& memoryPool_; + const std::vector files_; + const std::string serializedSerde_; +}; + +} // namespace facebook::nimble::tools diff --git a/dwio/nimble/tools/EncodingSelectionLogger.cpp b/dwio/nimble/tools/EncodingSelectionLogger.cpp new file mode 100644 index 0000000..4f2ca29 --- /dev/null +++ b/dwio/nimble/tools/EncodingSelectionLogger.cpp @@ -0,0 +1,125 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#define NIMBLE_ENCODING_SELECTION_DEBUG + +#include +#include +#include "common/init/light.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" +#include "dwio/nimble/encodings/EncodingSelectionPolicy.h" +#include "dwio/nimble/tools/EncodingUtilities.h" + +DEFINE_string( + data_type, + "", + "If provided, used as the data type for the encoded values. " + "If not provided, first input line is parsed to try and find the data type."); + +DEFINE_string( + file, + "", + "If provided, used as the input file to read from. Otherwise, cin is used."); + +DEFINE_string(read_factors, "", "Read factors to use for encoding selection."); +DEFINE_double(compression_acceptance_ratio, 0.9, ""); + +using namespace ::facebook; + +template +void logEncodingSelection(const std::vector& source) { + auto pool = facebook::velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Vector values{pool.get()}; + values.reserve(source.size()); + for (const auto& value : source) { + values.push_back(folly::to(value)); + } + + auto policy = std::unique_ptr>( + static_cast*>( + facebook::nimble::ManualEncodingSelectionPolicyFactory{ + FLAGS_read_factors.empty() + ? nimble::ManualEncodingSelectionPolicyFactory:: + defaultReadFactors() + : nimble::ManualEncodingSelectionPolicyFactory:: + parseReadFactors(FLAGS_read_factors), + nimble::CompressionOptions{ + .compressionAcceptRatio = + folly::to(FLAGS_compression_acceptance_ratio)}} + .createPolicy(nimble::TypeTraits::dataType) + .release())); + nimble::Buffer buffer{*pool}; + + auto serialized = + nimble::EncodingFactory::encode(std::move(policy), values, buffer); + + LOG(INFO) << "Encoding: " << GREEN + << nimble::tools::getEncodingLabel(serialized) << RESET_COLOR; +} + +int main(int argc, char* argv[]) { + auto init = init::InitFacebookLight{&argc, &argv}; + + std::istream* stream{&std::cin}; + std::ifstream file; + if (!FLAGS_file.empty()) { + file = std::ifstream{FLAGS_file}; + stream = &file; + } + + std::string dataTypeString = FLAGS_data_type; + if (dataTypeString.empty()) { + std::string line; + std::getline(*stream, line); + + const std::string kDataTypePrefix = "DataType "; + auto index = line.find(kDataTypePrefix); + if (index == std::string::npos) { + LOG(FATAL) + << "Unable to find data type prefix '" << kDataTypePrefix + << "' in first input row. Consider providing the data type using " + "the 'data_type' command line argument."; + } + + line = line.substr(index + kDataTypePrefix.size()); + auto end = line.find(','); + if (end == std::string::npos) { + dataTypeString = line; + } else { + dataTypeString = line.substr(0, end); + } + } + +#define DATA_TYPE(type) \ + {toString(nimble::TypeTraits::dataType), logEncodingSelection} + std::unordered_map< + std::string, + std::function&)>> + dataTypes{ + DATA_TYPE(int8_t), + DATA_TYPE(uint8_t), + DATA_TYPE(int16_t), + DATA_TYPE(uint16_t), + DATA_TYPE(int32_t), + DATA_TYPE(uint32_t), + DATA_TYPE(int64_t), + DATA_TYPE(uint64_t), + DATA_TYPE(float), + DATA_TYPE(double), + DATA_TYPE(bool), + DATA_TYPE(std::string_view), + }; +#undef DATA_TYPE + + auto it = dataTypes.find(dataTypeString); + if (it == dataTypes.end()) { + LOG(FATAL) << "Unknown data type: " << dataTypeString; + } + + std::vector values; + std::string value; + while (*stream >> value) { + values.push_back(value); + } + + it->second(values); +} diff --git a/dwio/nimble/tools/EncodingUtilities.cpp b/dwio/nimble/tools/EncodingUtilities.cpp new file mode 100644 index 0000000..15e8795 --- /dev/null +++ b/dwio/nimble/tools/EncodingUtilities.cpp @@ -0,0 +1,292 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/tools/EncodingUtilities.h" +#include "dwio/nimble/common/Exceptions.h" + +namespace facebook::nimble::tools { +namespace { +constexpr uint32_t kEncodingPrefixSize = 6; + +void extractCompressionType( + EncodingType encodingType, + DataType /* dataType */, + std::string_view stream, + std::unordered_map& properties) { + switch (encodingType) { + // Compression type is the byte right after the encoding header for both + // encodings. + case EncodingType::Trivial: + case EncodingType::FixedBitWidth: { + auto pos = stream.data() + kEncodingPrefixSize; + properties.insert( + {EncodingPropertyType::Compression, + EncodingProperty{ + .value = toString( + static_cast(encoding::readChar(pos))), + }}); + break; + } + case EncodingType::RLE: + case EncodingType::Dictionary: + case EncodingType::Sentinel: + case EncodingType::Nullable: + case EncodingType::SparseBool: + case EncodingType::Varint: + case EncodingType::Delta: + case EncodingType::Constant: + case EncodingType::MainlyConstant: + break; + } +} + +std::unordered_map +extractEncodingProperties( + EncodingType encodingType, + DataType dataType, + std::string_view stream) { + std::unordered_map properties{}; + extractCompressionType(encodingType, dataType, stream, properties); + properties.insert( + {EncodingPropertyType::EncodedSize, + {.value = folly::to(stream.size())}}); + return properties; +} + +void traverseEncodings( + std::string_view stream, + uint32_t level, + uint32_t index, + const std::string& nestedEncodingName, + std::function /* properties */)> visitor) { + NIMBLE_CHECK( + stream.size() >= kEncodingPrefixSize, "Unexpected end of stream."); + + const EncodingType encodingType = static_cast(stream[0]); + auto dataType = static_cast(stream[1]); + bool continueTraversal = visitor( + encodingType, + dataType, + level, + index, + nestedEncodingName, + extractEncodingProperties(encodingType, dataType, stream)); + + if (!continueTraversal) { + return; + } + + switch (encodingType) { + case EncodingType::FixedBitWidth: + case EncodingType::Varint: + case EncodingType::Constant: { + // don't have any nested encoding + break; + } + case EncodingType::Trivial: { + if (dataType == DataType::String) { + const char* pos = stream.data() + kEncodingPrefixSize + 1; + const uint32_t lengthsBytes = encoding::readUint32(pos); + traverseEncodings( + {pos, lengthsBytes}, level + 1, 0, "Lengths", visitor); + } + break; + } + case EncodingType::SparseBool: { + const char* pos = stream.data() + kEncodingPrefixSize + 1; + traverseEncodings( + {pos, stream.size() - (pos - stream.data())}, + level + 1, + 0, + "Indices", + visitor); + break; + } + case EncodingType::MainlyConstant: { + const char* pos = stream.data() + kEncodingPrefixSize; + const uint32_t isCommonBytes = encoding::readUint32(pos); + traverseEncodings( + {pos, isCommonBytes}, level + 1, 0, "IsCommon", visitor); + pos += isCommonBytes; + const uint32_t otherValueBytes = encoding::readUint32(pos); + traverseEncodings( + {pos, otherValueBytes}, level + 1, 1, "OtherValues", visitor); + break; + } + case EncodingType::Dictionary: { + const char* pos = stream.data() + kEncodingPrefixSize; + const uint32_t alphabetBytes = encoding::readUint32(pos); + traverseEncodings( + {pos, alphabetBytes}, level + 1, 0, "Alphabet", visitor); + pos += alphabetBytes; + traverseEncodings( + {pos, stream.size() - (pos - stream.data())}, + level + 1, + 1, + "Indices", + visitor); + break; + } + case EncodingType::RLE: { + const char* pos = stream.data() + kEncodingPrefixSize; + const uint32_t runLengthBytes = encoding::readUint32(pos); + traverseEncodings( + {pos, runLengthBytes}, level + 1, 0, "Lengths", visitor); + if (dataType != DataType::Bool) { + pos += runLengthBytes; + traverseEncodings( + {pos, stream.size() - (pos - stream.data())}, + level + 1, + 1, + "Values", + visitor); + } + break; + } + case EncodingType::Delta: { + const char* pos = stream.data() + kEncodingPrefixSize; + const uint32_t deltaBytes = encoding::readUint32(pos); + const uint32_t restatementBytes = encoding::readUint32(pos); + traverseEncodings({pos, deltaBytes}, level + 1, 0, "Deltas", visitor); + pos += deltaBytes; + traverseEncodings( + {pos, restatementBytes}, level + 1, 1, "Restatements", visitor); + pos += restatementBytes; + traverseEncodings( + {pos, stream.size() - (pos - stream.data())}, + level + 1, + 2, + "IsRestatements", + visitor); + break; + } + case EncodingType::Nullable: { + const char* pos = stream.data() + kEncodingPrefixSize; + const uint32_t dataBytes = encoding::readUint32(pos); + traverseEncodings({pos, dataBytes}, level + 1, 0, "Data", visitor); + pos += dataBytes; + traverseEncodings( + {pos, stream.size() - (pos - stream.data())}, + level + 1, + 1, + "Nulls", + visitor); + break; + } + case EncodingType::Sentinel: { + const char* pos = stream.data() + kEncodingPrefixSize + 8; + traverseEncodings( + {pos, stream.size() - (pos - stream.data())}, + level + 1, + 0, + "Sentinels", + visitor); + break; + } + } +} + +} // namespace + +std::ostream& operator<<(std::ostream& out, EncodingPropertyType propertyType) { + switch (propertyType) { + case EncodingPropertyType::Compression: + return out << "Compression"; + default: + return out << "Unknown"; + } +} + +void traverseEncodings( + std::string_view stream, + std::function /* properties */)> visitor) { + traverseEncodings(stream, 0, 0, "", visitor); +} + +std::string getStreamInputLabel(nimble::ChunkedStream& stream) { + std::string label; + uint32_t chunkId = 0; + while (stream.hasNext()) { + label += folly::to(chunkId++); + auto compression = stream.peekCompressionType(); + if (compression != CompressionType::Uncompressed) { + label += "{" + toString(compression) + "}"; + } + label += ":" + getEncodingLabel(stream.nextChunk()); + + if (stream.hasNext()) { + label += ";"; + } + } + return label; +} + +std::string getEncodingLabel(std::string_view stream) { + std::string label; + uint32_t currentLevel = 0; + + traverseEncodings( + stream, + [&](EncodingType encodingType, + DataType dataType, + uint32_t level, + uint32_t index, + std::string nestedEncodingName, + std::unordered_map properties) + -> bool { + if (level > currentLevel) { + label += "[" + nestedEncodingName + ":"; + } else if (level < currentLevel) { + label += "]"; + } + + if (index > 0) { + label += "," + nestedEncodingName + ":"; + } + + currentLevel = level; + + label += toString(encodingType) + "<" + toString(dataType); + const auto& encodedSize = + properties.find(EncodingPropertyType::EncodedSize); + if (encodedSize != properties.end()) { + label += "," + encodedSize->second.value; + } + label += ">"; + + const auto& compression = + properties.find(EncodingPropertyType::Compression); + if (compression != properties.end() + // If the "Uncompressed" label clutters the output, uncomment the + // next line. + // && compression->second.value != + // toString(CompressionType::Uncompressed) + ) { + label += "{" + compression->second.value + "}"; + } + + return true; + }); + while (currentLevel-- > 0) { + label += "]"; + } + + return label; +} + +} // namespace facebook::nimble::tools diff --git a/dwio/nimble/tools/EncodingUtilities.h b/dwio/nimble/tools/EncodingUtilities.h new file mode 100644 index 0000000..57fe936 --- /dev/null +++ b/dwio/nimble/tools/EncodingUtilities.h @@ -0,0 +1,42 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include +#include +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/velox/ChunkedStream.h" + +namespace facebook::nimble::tools { + +enum class EncodingPropertyType { + Compression, + EncodedSize, +}; + +std::ostream& operator<<(std::ostream& out, EncodingPropertyType propertyType); + +struct EncodingProperty { + std::string value; + std::string_view data; +}; + +void traverseEncodings( + std::string_view stream, + std::function /* properties */)> visitor); + +std::string getStreamInputLabel(nimble::ChunkedStream& stream); +std::string getEncodingLabel(std::string_view stream); + +} // namespace facebook::nimble::tools diff --git a/dwio/nimble/tools/NimbleDump.cpp b/dwio/nimble/tools/NimbleDump.cpp new file mode 100644 index 0000000..9919b50 --- /dev/null +++ b/dwio/nimble/tools/NimbleDump.cpp @@ -0,0 +1,304 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include +#include +#include +#include + +#include "common/base/BuildInfo.h" +#include "common/init/light.h" +#include "dwio/nimble/tools/NimbleDumpLib.h" +#include "folly/Singleton.h" +#include "folly/experimental/NestedCommandLineApp.h" +#include "velox/common/base/StatsReporter.h" + +using namespace ::facebook; +namespace po = ::boost::program_options; + +template +std::optional getOptional(const po::variable_value& val) { + return val.empty() ? std::nullopt : std::optional(val.as()); +} + +int main(int argc, char* argv[]) { + // Disable GLOG + FLAGS_minloglevel = 5; + + auto init = init::InitFacebookLight{ + &argc, &argv, folly::InitOptions().useGFlags(false)}; + + std::string version{BuildInfo::getBuildTimeISO8601()}; + if (!version.empty()) { + auto buildRevision = BuildInfo::getBuildRevision(); + if (buildRevision && buildRevision[0] != '\0') { + version += folly::to(" [", buildRevision, "]"); + } + } + + folly::NestedCommandLineApp app{"", version}; + int style = po::command_line_style::default_style; + style &= ~po::command_line_style::allow_guessing; + app.setOptionStyle(static_cast(style)); + + po::positional_options_description positionalArgs; + positionalArgs.add("file", /*max_count*/ 1); + + app.addCommand( + "info", + "", + "Print file information", + "Prints file information from the file footer.", + [](const po::variables_map& options, + const std::vector& /*args*/) { + nimble::tools::NimbleDumpLib{ + std::cout, options["file"].as()} + .emitInfo(); + }, + positionalArgs) + .add_options()( + "file", + po::value()->required(), + "Nimble file path. Can be a local path or a Warm Storage path."); + + app.addCommand( + "schema", + "", + "Print file schema", + "Prints the file schema tree.", + [](const po::variables_map& options, + const std::vector& /*args*/) { + nimble::tools::NimbleDumpLib{ + std::cout, options["file"].as()} + .emitSchema(!options["full"].as()); + }, + positionalArgs) + // clang-format off + .add_options() + ( + "file", + po::value()->required(), + "Nimble file path. Can be a local path or a Warm Storage path." + )( + "full,f", + po::bool_switch()->default_value(false), + "Emit full flat map schemas. Default is to collapse flat map schemas." + ); + // clang-format on + + app.addCommand( + "stripes", + "", + "Print stripe information", + "Prints detailed stripe information.", + [](const po::variables_map& options, + const std::vector& /*args*/) { + nimble::tools::NimbleDumpLib{ + std::cout, options["file"].as()} + .emitStripes(options["no_header"].as()); + }, + positionalArgs) + // clang-format off + .add_options() + ( + "file", + po::value()->required(), + "Nimble file path. Can be a local path or a Warm Storage path." + )( + "no_header,n", + po::bool_switch()->default_value(false), + "Don't print column names. Default is to include column names." + ); + // clang-format on + + app.addCommand( + "streams", + "", + "Print stream information", + "Prints detailed stream information.", + [](const po::variables_map& options, + const std::vector& /*args*/) { + nimble::tools::NimbleDumpLib{ + std::cout, options["file"].as()} + .emitStreams( + options["no_header"].as(), + options["labels"].as(), + getOptional(options["stripe"])); + }, + positionalArgs) + // clang-format off + .add_options() + ( + "file", + po::value()->required(), + "Nimble file path. Can be a local path or a Warm Storage path." + )( + "stripe,s", + po::value(), + "Limit output to a single stripe with the provided stripe id. " + "Default is to print streams for all stripes." + )( + "no_header,n", + po::bool_switch()->default_value(false), + "Don't print column names. Default is to include column names." + )( + "labels,l", + po::bool_switch()->default_value(false), + "Include stream labels. Lables provide a readable path from the " + "root node to the stream, as they appear in the schema tree." + ); + // clang-format on + + app.addCommand( + "histogram", + "", + "Print encoding histogram", + "Prints encoding histogram, counting how many times each encoding " + "appears in the file.", + [](const po::variables_map& options, + const std::vector& /*args*/) { + nimble::tools::NimbleDumpLib{ + std::cout, options["file"].as()} + .emitHistogram( + options["root_only"].as(), + options["no_header"].as(), + getOptional(options["stripe"])); + }, + positionalArgs) + // clang-format off + .add_options() + ( + "file", + po::value()->required(), + "Nimble file path. Can be a local path or a Warm Storage path." + )( + "stripe,s", + po::value(), + "Limit analysis to a single stripe with the provided stripe id. " + "Default is to analyze encodings in all stripes." + )( + "root_only,r", + po::bool_switch()->default_value(false), + "Include only root (top level) encodings in histogram. " + "Default is to analyze full encoding trees." + )( + "no_header,n", + po::bool_switch()->default_value(false), + "Don't print column names. Default is to include column names." + ); + // clang-format on + + app.addCommand( + "content", + "", + "Print the content of a stream", + "Prints the materialized content (actual values) of a stream.", + [](const po::variables_map& options, + const std::vector& /*args*/) { + nimble::tools::NimbleDumpLib{ + std::cout, options["file"].as()} + .emitContent( + options["stream"].as(), + getOptional(options["stripe"])); + }, + positionalArgs) + // clang-format off + .add_options() + ( + "file", + po::value()->required(), + "Nimble file path. Can be a local path or a Warm Storage path." + )( + "stream", + po::value()->required(), + "The content of this stream id will be emitted." + )( + "stripe", + po::value(), + "Limit output to a single stripe with the provided stripe id. " + "Default is to output stream content across in all stripes." + ); + // clang-format on + + app.addCommand( + "binary", + "", + "Dumps stream binary content", + "Dumps stream binary content to a file.", + [](const po::variables_map& options, + const std::vector& /*args*/) { + nimble::tools::NimbleDumpLib{ + std::cout, options["file"].as()} + .emitBinary( + [path = options["output"].as()]() { + return std::make_unique( + path, + std::ios::out | std::ios::binary | std::ios::trunc); + }, + options["stream"].as(), + options["stripe"].as()); + }, + positionalArgs) + // clang-format off + .add_options() + ( + "file", + po::value()->required(), + "Nimble file path. Can be a local path or a Warm Storage path." + )( + "output,o", + po::value()->required(), + "Output file path." + )( + "stream", + po::value()->required(), + "The content of this stream id will be dumped to the output file." + )( + "stripe", + po::value()->required(), + "Dumps the stream from this stripe id." + ); + // clang-format on + + app.addCommand( + "layout", + "", + "Dumps layout file", + "Dumps captured layout file content.", + [](const po::variables_map& options, + const std::vector& /*args*/) { + nimble::tools::NimbleDumpLib{ + std::cout, options["file"].as()} + .emitLayout( + options["no_header"].as(), + !options["uncompressed"].as()); + }, + positionalArgs) + // clang-format off + .add_options() + ( + "file", + po::value()->required(), + "Encoding layout file path." + )( + "no_header,n", + po::bool_switch()->default_value(false), + "Don't print column names. Default is to include column names." + )( + "uncompressed,u", + po::bool_switch()->default_value(false), + "Is the layout file uncompressed. Default is false, which means " + "the layout file is compressed." + ); + // clang-format on + + app.addAlias("i", "info"); + app.addAlias("b", "binary"); + app.addAlias("c", "content"); + + return app.run(argc, argv); +} + +// Initialize dummy Velox stats reporter +folly::Singleton reporter([]() { + return new facebook::velox::DummyStatsReporter(); +}); diff --git a/dwio/nimble/tools/NimbleDumpLib.cpp b/dwio/nimble/tools/NimbleDumpLib.cpp new file mode 100644 index 0000000..6768cbf --- /dev/null +++ b/dwio/nimble/tools/NimbleDumpLib.cpp @@ -0,0 +1,687 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. +#include +#include +#include +#include +#include +#include +#include + +#include "common/strings/Zstd.h" +#include "dwio/common/filesystem/FileSystem.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/FixedBitArray.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" +#include "dwio/nimble/encodings/EncodingLayout.h" +#include "dwio/nimble/tools/EncodingUtilities.h" +#include "dwio/nimble/tools/NimbleDumpLib.h" +#include "dwio/nimble/velox/EncodingLayoutTree.h" +#include "dwio/nimble/velox/VeloxReader.h" +#include "folly/experimental/NestedCommandLineApp.h" + +namespace facebook::nimble::tools { +#define RED "\033[31m" +#define GREEN "\033[32m" +#define YELLOW "\033[33m" +#define BLUE "\033[34m" +#define PURPLE "\033[35m" +#define CYAN "\033[36m" +#define RESET_COLOR "\033[0m" + +namespace { + +constexpr uint32_t kBufferSize = 1000; + +struct GroupingKey { + EncodingType encodingType; + DataType dataType; + std::optional compressinType; +}; + +struct GroupingKeyCompare { + size_t operator()(const GroupingKey& lhs, const GroupingKey& rhs) const { + if (lhs.encodingType != rhs.encodingType) { + return lhs.encodingType < rhs.encodingType; + } else if (lhs.dataType != rhs.dataType) { + return lhs.dataType < rhs.dataType; + } else { + return lhs.compressinType < rhs.compressinType; + } + } +}; + +class TableFormatter { + public: + TableFormatter( + std::ostream& ostream, + std::vector> fields, + bool noHeader = false) + : ostream_{ostream}, fields_{std::move(fields)} { + if (!noHeader) { + ostream << YELLOW; + for (const auto& field : fields_) { + ostream << std::left << std::setw(std::get<1>(field) + 2) + << std::get<0>(field); + } + ostream << RESET_COLOR << std::endl; + } + } + + void writeRow(const std::vector& values) { + assert(values.size() == fields_.size()); + for (auto i = 0; i < values.size(); ++i) { + ostream_ << std::left << std::setw(std::get<1>(fields_[i]) + 2) + << values[i]; + } + ostream_ << std::endl; + } + + private: + std::ostream& ostream_; + std::vector> + fields_; +}; + +void traverseTablet( + velox::memory::MemoryPool& memoryPool, + const Tablet& tablet, + std::optional stripeIndex, + std::function stripeVisitor = nullptr, + std::function streamVisitor = nullptr) { + uint32_t startStripe = stripeIndex ? *stripeIndex : 0; + uint32_t endStripe = stripeIndex ? *stripeIndex : tablet.stripeCount() - 1; + for (uint32_t i = startStripe; i <= endStripe; ++i) { + if (stripeVisitor) { + stripeVisitor(i); + } + if (streamVisitor) { + std::vector identifiers(tablet.streamCount(i)); + std::iota(identifiers.begin(), identifiers.end(), 0); + auto streams = tablet.load(i, {identifiers.cbegin(), identifiers.cend()}); + for (uint32_t j = 0; j < streams.size(); ++j) { + auto& stream = streams[j]; + if (stream) { + InMemoryChunkedStream chunkedStream{memoryPool, std::move(stream)}; + streamVisitor(chunkedStream, i, j); + } + } + } + } +} + +template +void printScalarData( + std::ostream& ostream, + velox::memory::MemoryPool& pool, + Encoding& stream, + uint32_t rowCount) { + nimble::Vector buffer(&pool); + nimble::Vector nulls(&pool); + buffer.resize(rowCount); + nulls.resize((nimble::FixedBitArray::bufferSize(rowCount, 1))); + nulls.zero_out(); + if (stream.isNullable()) { + stream.materializeNullable( + rowCount, buffer.data(), [&]() { return nulls.data(); }); + } else { + stream.materialize(rowCount, buffer.data()); + nulls.fill(0xff); + } + for (uint32_t i = 0; i < rowCount; ++i) { + if (stream.isNullable() && nimble::bits::getBit(i, nulls.data()) == 0) { + ostream << "NULL" << std::endl; + } else { + ostream << folly::to(buffer[i]) + << std::endl; // Have to use folly::to as Int8 was getting + // converted to char. + } + } +} + +void printScalarType( + std::ostream& ostream, + velox::memory::MemoryPool& pool, + Encoding& stream, + uint32_t rowCount) { + switch (stream.dataType()) { +#define CASE(KIND, cppType) \ + case DataType::KIND: { \ + printScalarData(ostream, pool, stream, rowCount); \ + break; \ + } + CASE(Int8, int8_t); + CASE(Uint8, uint8_t); + CASE(Int16, int16_t); + CASE(Uint16, uint16_t); + CASE(Int32, int32_t); + CASE(Uint32, uint32_t); + CASE(Int64, int64_t); + CASE(Uint64, uint64_t); + CASE(Float, float); + CASE(Double, double); + CASE(Bool, bool); + CASE(String, std::string_view); +#undef CASE + case DataType::Undefined: { + NIMBLE_UNREACHABLE( + fmt::format("Undefined type for stream: {}", stream.dataType())); + } + } +} + +template +auto commaSeparated(T value) { + return fmt::format(std::locale("en_US.UTF-8"), "{:L}", value); +} + +} // namespace + +NimbleDumpLib::NimbleDumpLib(std::ostream& ostream, const std::string& file) + : pool_{velox::memory::deprecatedAddDefaultLeafMemoryPool()}, + file_{dwio::file_system::FileSystem::openForRead( + file, + dwio::common::request::AccessDescriptorBuilder() + .withClientId("nimble_dump") + .build())}, + ostream_{ostream} {} + +void NimbleDumpLib::emitInfo() { + auto tablet = std::make_shared(*pool_, file_.get()); + ostream_ << CYAN << "Nimble File " << RESET_COLOR << "Version " + << tablet->majorVersion() << "." << tablet->minorVersion() + << std::endl; + ostream_ << "File Size: " << commaSeparated(tablet->fileSize()) << std::endl; + ostream_ << "Checksum: " << tablet->checksum() << " [" + << nimble::toString(tablet->checksumType()) << "]" << std::endl; + ostream_ << "Footer Compression: " << tablet->footerCompressionType() + << std::endl; + ostream_ << "Footer Size: " << commaSeparated(tablet->footerSize()) + << std::endl; + ostream_ << "Stripe Count: " << commaSeparated(tablet->stripeCount()) + << std::endl; + ostream_ << "Row Count: " << commaSeparated(tablet->tabletRowCount()) + << std::endl; + + VeloxReader reader{*pool_, tablet}; + auto& metadata = reader.metadata(); + if (!metadata.empty()) { + ostream_ << "Metadata:"; + for (const auto& pair : metadata) { + ostream_ << std::endl << " " << pair.first << ": " << pair.second; + } + } + ostream_ << std::endl; +} + +void NimbleDumpLib::emitSchema(bool collapseFlatMap) { + auto tablet = std::make_shared(*pool_, file_.get()); + VeloxReader reader{*pool_, tablet}; + + auto emitOffsets = [](const Type& type) { + std::string offsets; + switch (type.kind()) { + case Kind::Scalar: { + offsets = + folly::to(type.asScalar().scalarDescriptor().offset()); + break; + } + case Kind::Array: { + offsets = + folly::to(type.asArray().lengthsDescriptor().offset()); + break; + } + case Kind::Map: { + offsets = + folly::to(type.asMap().lengthsDescriptor().offset()); + break; + } + case Kind::Row: { + offsets = + folly::to(type.asRow().nullsDescriptor().offset()); + break; + } + case Kind::FlatMap: { + offsets = + folly::to(type.asFlatMap().nullsDescriptor().offset()); + break; + } + case Kind::ArrayWithOffsets: { + offsets = "o:" + + folly::to( + type.asArrayWithOffsets().offsetsDescriptor().offset()) + + ",l:" + + folly::to( + type.asArrayWithOffsets().lengthsDescriptor().offset()); + break; + } + } + + return offsets; + }; + + bool skipping = false; + SchemaReader::traverseSchema( + reader.schema(), + [&](uint32_t level, + const Type& type, + const SchemaReader::NodeInfo& info) { + auto parentType = info.parentType; + if (parentType != nullptr && parentType->isFlatMap()) { + auto childrenCount = parentType->asFlatMap().childrenCount(); + if (childrenCount > 2 && collapseFlatMap) { + if (info.placeInSibling == 1) { + ostream_ << std::string( + (std::basic_string::size_type)level * 2, + ' ') + << "..." << std::endl; + skipping = true; + } else if (info.placeInSibling == childrenCount - 1) { + skipping = false; + } + } + } + if (!skipping) { + ostream_ << std::string( + (std::basic_string::size_type)level * 2, ' ') + << "[" << emitOffsets(type) << "] " << info.name << " : "; + if (type.isScalar()) { + ostream_ << toString(type.kind()) << "<" + << toString( + type.asScalar().scalarDescriptor().scalarKind()) + << ">" << std::endl; + } else { + ostream_ << toString(type.kind()) << std::endl; + } + } + }); +} + +void NimbleDumpLib::emitStripes(bool noHeader) { + Tablet tablet{*pool_, file_.get()}; + TableFormatter formatter( + ostream_, + {{"Stripe Id", 11}, + {"Stripe Offset", 15}, + {"Stripe Size", 15}, + {"Row Count", 15}}, + noHeader); + traverseTablet(*pool_, tablet, std::nullopt, [&](uint32_t stripeIndex) { + auto sizes = tablet.streamSizes(stripeIndex); + auto stripeSize = std::accumulate(sizes.begin(), sizes.end(), 0UL); + formatter.writeRow({ + folly::to(stripeIndex), + folly::to(tablet.stripeOffset(stripeIndex)), + folly::to(stripeSize), + folly::to(tablet.stripeRowCount(stripeIndex)), + }); + }); +} + +void NimbleDumpLib::emitStreams( + bool noHeader, + bool streamLabels, + std::optional stripeId) { + auto tablet = std::make_shared(*pool_, file_.get()); + + std::vector> fields; + fields.push_back({"Stripe Id", 11}); + fields.push_back({"Stream Id", 11}); + fields.push_back({"Stream Offset", 13}); + fields.push_back({"Stream Size", 13}); + fields.push_back({"Item Count", 13}); + if (streamLabels) { + fields.push_back({"Stream Label", 16}); + } + fields.push_back({"Type", 30}); + + TableFormatter formatter(ostream_, fields, noHeader); + + std::optional labels{}; + if (streamLabels) { + VeloxReader reader{*pool_, tablet}; + labels.emplace(reader.schema()); + } + + traverseTablet( + *pool_, + *tablet, + stripeId, + nullptr /* stripeVisitor */, + [&](ChunkedStream& stream, uint32_t stripeId, uint32_t streamId) { + uint32_t itemCount = 0; + while (stream.hasNext()) { + auto chunk = stream.nextChunk(); + itemCount += *reinterpret_cast(chunk.data() + 2); + } + stream.reset(); + std::vector values; + values.push_back(folly::to(stripeId)); + values.push_back(folly::to(streamId)); + values.push_back( + folly::to(tablet->streamOffsets(stripeId)[streamId])); + values.push_back( + folly::to(tablet->streamSizes(stripeId)[streamId])); + values.push_back(folly::to(itemCount)); + if (streamLabels) { + auto it = values.emplace_back(labels->streamLabel(streamId)); + } + values.push_back(getStreamInputLabel(stream)); + formatter.writeRow(values); + }); +} + +void NimbleDumpLib::emitHistogram( + bool topLevel, + bool noHeader, + std::optional stripeId) { + Tablet tablet{*pool_, file_.get()}; + std::map encodingHistogram; + const std::unordered_map compressionMap{ + {toString(CompressionType::Uncompressed), CompressionType::Uncompressed}, + {toString(CompressionType::Zstd), CompressionType::Zstd}, + {toString(CompressionType::Zstrong), CompressionType::Zstrong}, + }; + traverseTablet( + *pool_, + tablet, + stripeId, + nullptr, + [&](ChunkedStream& stream, auto /*stripeIndex*/, auto /*streamIndex*/) { + while (stream.hasNext()) { + traverseEncodings( + stream.nextChunk(), + [&](EncodingType encodingType, + DataType dataType, + uint32_t level, + uint32_t /* index */, + std::string /*nestedEncodingName*/, + std::unordered_map + properties) { + GroupingKey key{ + .encodingType = encodingType, .dataType = dataType}; + const auto& compression = + properties.find(EncodingPropertyType::Compression); + if (compression != properties.end()) { + key.compressinType = + compressionMap.at(compression->second.value); + } + ++encodingHistogram[key]; + return !(topLevel && level == 1); + }); + } + }); + + TableFormatter formatter( + ostream_, + {{"Encoding Type", 17}, + {"Data Type", 13}, + {"Compression", 15}, + {"Count", 15}}, + noHeader); + + for (auto& [key, value] : encodingHistogram) { + formatter.writeRow({ + toString(key.encodingType), + toString(key.dataType), + key.compressinType ? toString(*key.compressinType) : "", + folly::to(value), + }); + } +} + +void NimbleDumpLib::emitContent( + uint32_t streamId, + std::optional stripeId) { + Tablet tablet{*pool_, file_.get()}; + + uint32_t maxStreamCount; + bool found = false; + traverseTablet(*pool_, tablet, stripeId, [&](uint32_t stripeId) { + maxStreamCount = std::max(maxStreamCount, tablet.streamCount(stripeId)); + if (streamId >= tablet.streamCount(stripeId)) { + return; + } + + found = true; + + auto streams = tablet.load(stripeId, std::vector{streamId}); + + if (auto& stream = streams[0]) { + InMemoryChunkedStream chunkedStream{*pool_, std::move(stream)}; + while (chunkedStream.hasNext()) { + auto encoding = + EncodingFactory::decode(*pool_, chunkedStream.nextChunk()); + uint32_t totalRows = encoding->rowCount(); + while (totalRows > 0) { + auto currentReadSize = std::min(kBufferSize, totalRows); + printScalarType(ostream_, *pool_, *encoding, currentReadSize); + totalRows -= currentReadSize; + } + } + } + }); + + if (!found) { + throw folly::ProgramExit( + -1, + fmt::format( + "Stream identifier {} is out of bound. Must be between 0 and {}\n", + streamId, + maxStreamCount)); + } +} + +void NimbleDumpLib::emitBinary( + std::function()> outputFactory, + uint32_t streamId, + uint32_t stripeId) { + Tablet tablet{*pool_, file_.get()}; + if (streamId >= tablet.streamCount(stripeId)) { + throw folly::ProgramExit( + -1, + fmt::format( + "Stream identifier {} is out of bound. Must be between 0 and {}\n", + streamId, + tablet.streamCount(stripeId))); + } + + auto streams = tablet.load(stripeId, std::vector{streamId}); + + if (auto& stream = streams[0]) { + auto output = outputFactory(); + output->write(stream->getStream().data(), stream->getStream().size()); + output->flush(); + } +} + +void traverseEncodingLayout( + const std::optional& node, + const std::optional& parentNode, + uint32_t& nodeId, + uint32_t parentId, + uint32_t level, + uint8_t childIndex, + const std::function&, + const std::optional&, + uint32_t, + uint32_t, + uint32_t, + uint8_t)>& visitor) { + auto currentNodeId = nodeId; + visitor(node, parentNode, currentNodeId, parentId, level, childIndex); + + if (node.has_value()) { + for (int i = 0; i < node->childrenCount(); ++i) { + traverseEncodingLayout( + node->child(i), node, ++nodeId, currentNodeId, level + 1, i, visitor); + } + } +} + +void traverseEncodingLayoutTree( + const EncodingLayoutTree& node, + const EncodingLayoutTree& parentNode, + uint32_t& nodeId, + uint32_t parentId, + uint32_t level, + uint8_t childIndex, + const std::function& visitor) { + auto currentNodeId = nodeId; + visitor(node, parentNode, currentNodeId, parentId, level, childIndex); + + for (int i = 0; i < node.childrenCount(); ++i) { + traverseEncodingLayoutTree( + node.child(i), node, ++nodeId, currentNodeId, level + 1, i, visitor); + } +} + +std::string getEncodingLayoutLabel( + const std::optional& root) { + std::string label; + uint32_t currentLevel = 0; + std::unordered_map> + identifierNames{ + {nimble::EncodingType::Dictionary, {"Alphabet", "Indices"}}, + {nimble::EncodingType::MainlyConstant, {"IsCommon", "OtherValues"}}, + {nimble::EncodingType::Nullable, {"Data", "Nulls"}}, + {nimble::EncodingType::RLE, {"RunLengths", "RunValues"}}, + {nimble::EncodingType::SparseBool, {"Indices"}}, + {nimble::EncodingType::Trivial, {"Lengths"}}, + }; + + auto getIdentifierName = [&](nimble::EncodingType encodingType, + uint8_t identifier) { + auto it = identifierNames.find(encodingType); + LOG(INFO) << (it == identifierNames.end()) << ", " + << (it != identifierNames.end() + ? (int)(identifier >= it->second.size()) + : -1); + return it == identifierNames.end() || identifier >= it->second.size() + ? "Unknown" + : it->second[identifier]; + }; + + uint32_t id = 0; + traverseEncodingLayout( + root, + root, + id, + id, + 0, + (uint8_t)0, + [&](const std::optional& node, + const std::optional& parentNode, + uint32_t /* nodeId */, + uint32_t /* parentId */, + uint32_t level, + uint8_t identifier) { + if (!node.has_value()) { + label += "N/A"; + return true; + } + + if (level > currentLevel) { + label += "[" + + getIdentifierName(parentNode->encodingType(), identifier) + ":"; + + } else if (level < currentLevel) { + label += "]"; + } + + if (identifier > 0) { + label += "," + + getIdentifierName(parentNode->encodingType(), identifier) + ":"; + } + + currentLevel = level; + + label += toString(node->encodingType()) + "{" + + toString(node->compressionType()) + "}"; + + return true; + }); + + while (currentLevel-- > 0) { + label += "]"; + } + + return label; +} + +void NimbleDumpLib::emitLayout(bool noHeader, bool compressed) { + auto size = file_->size(); + std::string buffer; + buffer.resize(size); + file_->pread(0, size, buffer.data()); + if (compressed) { + std::string uncompressed; + strings::zstdDecompress(buffer, &uncompressed); + buffer = std::move(uncompressed); + } + + auto layout = nimble::EncodingLayoutTree::create(buffer); + + TableFormatter formatter( + ostream_, + { + {"Node Id", 11}, + {"Parent Id", 11}, + {"Node Type", 15}, + {"Node Name", 17}, + {"Encoding Layout", 20}, + }, + noHeader); + + uint32_t id = 0; + traverseEncodingLayoutTree( + layout, + layout, + id, + id, + 0, + 0, + [&](const EncodingLayoutTree& node, + const EncodingLayoutTree& /* parentNode */, + uint32_t nodeId, + uint32_t parentId, + uint32_t /* level */, + uint8_t /* identifier */) { + auto identifiers = node.encodingLayoutIdentifiers(); + std::sort(identifiers.begin(), identifiers.end()); + + std::string encodingLayout; + for (auto identifier : identifiers) { + if (!encodingLayout.empty()) { + encodingLayout += "|"; + } + encodingLayout += folly::to(identifier) + ":" + + getEncodingLayoutLabel(*node.encodingLayout(identifier)); + } + + formatter.writeRow( + {folly::to(nodeId), + folly::to(parentId), + toString(node.schemaKind()), + std::string(node.name()), + encodingLayout}); + }); +} + +} // namespace facebook::nimble::tools diff --git a/dwio/nimble/tools/NimbleDumpLib.h b/dwio/nimble/tools/NimbleDumpLib.h new file mode 100644 index 0000000..1bc6213 --- /dev/null +++ b/dwio/nimble/tools/NimbleDumpLib.h @@ -0,0 +1,37 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once +#include +#include + +#include "dwio/nimble/encodings/Encoding.h" +#include "velox/common/file/File.h" + +namespace facebook::nimble::tools { + +class NimbleDumpLib { + public: + NimbleDumpLib(std::ostream& ostream, const std::string& file); + + void emitInfo(); + void emitSchema(bool collapseFlatMap = true); + void emitStripes(bool noHeader); + void emitStreams( + bool noHeader, + bool flatmapKeys, + std::optional stripeId); + void + emitHistogram(bool topLevel, bool noHeader, std::optional stripeId); + void emitContent(uint32_t streamId, std::optional stripeId); + void emitBinary( + std::function()> outputFactory, + uint32_t streamId, + uint32_t stripeId); + void emitLayout(bool noHeader, bool compressed); + + private: + std::shared_ptr pool_; + std::shared_ptr file_; + std::ostream& ostream_; +}; +} // namespace facebook::nimble::tools diff --git a/dwio/nimble/tools/ParallelReader.cpp b/dwio/nimble/tools/ParallelReader.cpp new file mode 100644 index 0000000..465ef64 --- /dev/null +++ b/dwio/nimble/tools/ParallelReader.cpp @@ -0,0 +1,228 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include +#include "common/files/FileUtil.h" +#include "common/init/light.h" +#include "dwio/api/DwioConfig.h" +#include "dwio/common/filesystem/FileSystem.h" +#include "dwio/nimble/velox/VeloxReader.h" +#include "dwio/utils/FeatureFlatteningUtils.h" +#include "folly/executors/CPUThreadPoolExecutor.h" +#include "folly/experimental/FunctionScheduler.h" +#include "velox/dwio/dwrf/reader/DwrfReader.h" + +DEFINE_string(file, "", "Input file to read"); +DEFINE_bool( + memory_mode, + false, + "Should file be loaded and served from memory?"); +DEFINE_bool( + shared_memory_pool, + false, + "Should we use single memory pool for all readers?"); +DEFINE_uint32(concurrency, 16, "Number of reader to use in parallel"); +DEFINE_uint32(read_count, 500, "Total number of file reads to perform"); +DEFINE_uint32(batch_size, 32, "Batch size to use when reading from file"); +DEFINE_uint32(feature_shuffle_seed, 140185, "Seed used to shuffle features"); +DEFINE_uint32( + feature_percentage, + 20, + "Percentage of features to use (out of 100)"); + +using namespace ::facebook; + +namespace { +class Reader { + public: + virtual ~Reader() = default; + + virtual bool next(uint32_t count, velox::VectorPtr& vector) = 0; +}; + +class ReaderFactory { + public: + explicit ReaderFactory(std::shared_ptr readFile) + : readFile_{std::move(readFile)} { + NIMBLE_ASSERT(readFile_->size() > 2, "Invalid read file size."); + uint16_t magicNumber; + readFile_->pread(readFile_->size() - 2, 2, &magicNumber); + format_ = (magicNumber == 0xA1FA) ? velox::dwio::common::FileFormat::ALPHA + : velox::dwio::common::FileFormat::DWRF; + auto columnFeatures = + dwio::utils::feature_flattening::extractFeatureNames(readFile_); + features_.reserve(columnFeatures.size()); + for (const auto& pair : columnFeatures) { + std::vector features{pair.second.cbegin(), pair.second.cend()}; + std::sort(features.begin(), features.end()); + std::mt19937 gen; + gen.seed(FLAGS_feature_shuffle_seed); + std::shuffle(features.begin(), features.end(), gen); + features.resize(features.size() * FLAGS_feature_percentage / 100); + LOG(INFO) << "Loading " << features.size() << " features for column " + << pair.first; + features_.insert({pair.first, std::move(features)}); + } + } + + std::unique_ptr create(velox::memory::MemoryPool& memoryPool); + + private: + std::shared_ptr readFile_; + velox::dwio::common::FileFormat format_; + folly::F14FastMap> features_; +}; + +class NimbleReader : public Reader { + public: + NimbleReader( + velox::memory::MemoryPool& memoryPool, + std::shared_ptr readFile, + const folly::F14FastMap>& features) { + nimble::VeloxReadParams params; + for (const auto& feature : features) { + if (feature.second.size() == 0) { + continue; + } + + params.readFlatMapFieldAsStruct.insert(feature.first); + auto& target = params.flatMapFeatureSelector[feature.first]; + + std::transform( + feature.second.begin(), + feature.second.end(), + std::inserter(target.features, target.features.end()), + [](auto f) { return folly::to(f); }); + } + reader_ = std::make_unique( + memoryPool, readFile, /* selector */ nullptr, std::move(params)); + } + + bool next(uint32_t count, velox::VectorPtr& vector) override { + return reader_->next(count, vector); + } + + private: + std::shared_ptr file_; + std::unique_ptr reader_; +}; + +class DwrfReader : public Reader { + public: + DwrfReader( + velox::memory::MemoryPool& memoryPool, + std::shared_ptr readFile, + const folly::F14FastMap>& features) { + velox::dwio::common::ReaderOptions readerOptions{&memoryPool}; + reader_ = std::make_unique( + readerOptions, + std::make_unique( + std::move(readFile), readerOptions.getMemoryPool())); + + auto& schema = reader_->rowType(); + velox::dwio::common::RowReaderOptions rowReaderOptions; + std::vector projection; + std::vector mapAsStruct; + for (auto& col : schema->names()) { + if (features.count(col) == 0 || features.at(col).size() == 0) { + projection.push_back(col); + } else { + std::string joinedFeatures; + folly::join(",", features.at(col), joinedFeatures); + projection.push_back(fmt::format("{}#[{}]", col, joinedFeatures)); + mapAsStruct.push_back(fmt::format("{}#[{}]", col, joinedFeatures)); + } + } + auto cs = std::make_shared( + schema, projection); + rowReaderOptions.select(cs); + rowReaderOptions.setFlatmapNodeIdsAsStruct( + facebook::dwio::api::getNodeIdToSelectedKeysMap(*cs, mapAsStruct)); + rowReaderOptions.setReturnFlatVector(true); + + rowReader_ = reader_->createRowReader(rowReaderOptions); + } + + bool next(uint32_t count, velox::VectorPtr& vector) override { + return rowReader_->next(count, vector); + } + + private: + std::unique_ptr reader_; + std::unique_ptr rowReader_; +}; + +std::unique_ptr ReaderFactory::create( + velox::memory::MemoryPool& memoryPool) { + if (format_ == velox::dwio::common::FileFormat::ALPHA) { + return std::make_unique(memoryPool, readFile_, features_); + } else { + return std::make_unique(memoryPool, readFile_, features_); + } +} +} // namespace + +int main(int argc, char* argv[]) { + auto init = facebook::init::InitFacebookLight{&argc, &argv}; + + std::shared_ptr readFile; + if (FLAGS_memory_mode) { + std::string content; + files::FileUtil::readFileToString(FLAGS_file, &content); + readFile = std::make_shared(content); + } else { + readFile = std::make_shared(FLAGS_file); + } + + folly::FunctionScheduler scheduler; + std::atomic rowCount = 0; + std::atomic completedCount = 0; + std::atomic inflightCount = 0; + scheduler.addFunction( + [&]() { + LOG(INFO) << "Rows per second: " << rowCount + << ", Completed Reader Count: " << completedCount + << ", In Flight Readers: " << inflightCount; + rowCount = 0; + }, + std::chrono::seconds(1), + "Counts"); + + auto executor = std::make_shared( + FLAGS_concurrency, + std::make_unique>(FLAGS_concurrency), + std::make_shared("reader")); + + std::shared_ptr sharedMemoryPool = + velox::memory::deprecatedAddDefaultLeafMemoryPool(); + + ReaderFactory readerFactory{readFile}; + + scheduler.start(); + + for (auto i = 0; i < FLAGS_read_count; ++i) { + executor->add([&]() { + std::shared_ptr localMemoryPool = nullptr; + if (FLAGS_shared_memory_pool) { + localMemoryPool = sharedMemoryPool; + } else { + localMemoryPool = velox::memory::deprecatedAddDefaultLeafMemoryPool(); + } + + auto reader = readerFactory.create(*localMemoryPool); + ++inflightCount; + velox::VectorPtr result; + while (reader->next(FLAGS_batch_size, result)) { + rowCount += result->size(); + } + + --inflightCount; + ++completedCount; + }); + } + + executor->join(); + + return 0; +} diff --git a/dwio/nimble/tools/ParallelWriter.cpp b/dwio/nimble/tools/ParallelWriter.cpp new file mode 100644 index 0000000..a046f09 --- /dev/null +++ b/dwio/nimble/tools/ParallelWriter.cpp @@ -0,0 +1,189 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "common/init/light.h" +#include "dwio/common/filesystem/FileSystem.h" +#include "dwio/nimble/velox/VeloxWriter.h" +#include "dwio/nimble/velox/VeloxWriterOptions.h" +#include "dwio/utils/BufferedWriteFile.h" +#include "dwio/utils/FileSink.h" +#include "dwio/utils/InputStreamFactory.h" +#include "fb_velox/nimble/reader/NimbleReader.h" +#include "velox/dwio/common/Options.h" +#include "velox/dwio/dwrf/writer/Writer.h" + +DEFINE_string(input_file, "", "Input file to read"); +DEFINE_string( + output_file, + "", + "Output file to write. If not specified, writes to memory"); +DEFINE_int32( + threads, + -1, + "Defines the # of threads to use in parallel writes. Default is std::hardware_concurrency()"); +DEFINE_int32( + buffered_write_size, + 72 * 1024 * 1024, + "Set the buffered write size to use when writing to memory. Default is 72MB"); +DEFINE_bool(use_dwrf, false, "Use the DWRF writer instead of Nimble"); +DEFINE_uint32(batch_size, 32, "Batch size to use when reading from file"); +DEFINE_uint32(num_batches, 20, "Number of batches to read/write"); +DEFINE_string( + dictionary_array_names, + "id_list_features", + "Column names to write as dictionary array, semicolon separaated list"); +DEFINE_bool( + nimble_parallelize, + true, + "Use inter-writer parallelization. On by default"); + +using namespace ::facebook; + +namespace { + +class Writer { + public: + virtual ~Writer() = default; + + virtual void write(velox::VectorPtr& vector) = 0; + virtual void close() = 0; +}; + +class DWRFWriter : public Writer { + public: + DWRFWriter( + velox::memory::MemoryPool& rootPool, + std::shared_ptr leafPool, + // std::unique_ptr fileSink, + velox::TypePtr schema) + : leafPool_(leafPool) { + std::unique_ptr writeSink; + if (FLAGS_output_file == "") { + writeSink = std::make_unique( + 1000 * 1024 * 1024, + facebook::velox::dwio::common::FileSink::Options{ + .pool = leafPool_.get()}); + } else { + auto writeFile = + std::make_unique(FLAGS_output_file); + writeSink = std::make_unique( + std::move(writeFile), FLAGS_output_file); + } + velox::dwrf::WriterOptions options; + options.schema = schema; + writer_ = std::make_unique( + std::move(options), std::move(writeSink), rootPool); + } + + void write(velox::VectorPtr& vector) override { + writer_->write(vector); + } + + void close() override { + writer_->close(); + } + + private: + std::shared_ptr leafPool_; + std::unique_ptr writer_; +}; // DWRFWriter + +class NimbleWriter : public Writer { + public: + NimbleWriter( + velox::memory::MemoryPool& memoryPool, + std::shared_ptr leafPool, + velox::TypePtr schema) + : leafPool_(leafPool) { + std::unique_ptr writeFile; + if (FLAGS_output_file == "") { + buffer_.reserve(512 * 1024 * 1024); + auto inMemoryFile = std::make_unique(&buffer_); + writeFile = std::make_unique( + leafPool_, FLAGS_buffered_write_size, std::move(inMemoryFile)); + } else { + writeFile = std::make_unique(FLAGS_output_file); + } + nimble::VeloxWriterOptions options; + std::vector dictionaryArrays; + folly::split(';', FLAGS_dictionary_array_names, dictionaryArrays); + for (const auto& columnName : dictionaryArrays) { + options.dictionaryArrayColumns.insert(columnName); + } + if (FLAGS_nimble_parallelize) { + auto numThreads = std::thread::hardware_concurrency(); + if (FLAGS_threads > 0) { + numThreads = FLAGS_threads; + } + options.encodingExecutor = + std::make_shared(numThreads); + } + writer_ = std::make_unique( + memoryPool, schema, std::move(writeFile), std::move(options)); + } + + void write(velox::VectorPtr& vector) override { + writer_->write(vector); + } + + void close() override { + writer_->close(); + } + + private: + std::string buffer_; + std::shared_ptr leafPool_; + std::unique_ptr writer_; +}; // NimbleWriter +} // namespace + +int main(int argc, char* argv[]) { + auto init = facebook::init::InitFacebookLight{&argc, &argv}; + facebook::velox::memory::initializeMemoryManager( + facebook::velox::memory::MemoryManagerOptions{}); + auto rootPool = facebook::velox::memory::memoryManager()->addRootPool( + "velox_parallel_writer"); + auto leafPool = rootPool->addLeafChild("leaf_pool"); + auto accessDescriptor = + facebook::dwio::common::request::AccessDescriptorBuilder{} + .withClientId("parallel_writer") + .build(); + auto inputStream = facebook::dwio::utils::InputStreamFactory::create( + FLAGS_input_file, accessDescriptor); + facebook::velox::dwio::common::ReaderOptions readerOpts{leafPool.get()}; + auto readerFactory = facebook::velox::nimble::NimbleReaderFactory(); + auto reader = readerFactory.createReader( + std::make_unique( + std::move(inputStream), *leafPool), + readerOpts); + auto rowReader = reader->createRowReader(); + velox::VectorPtr schemaVector; + rowReader->next(1, schemaVector); + auto schema = schemaVector->type(); + std::vector vectorsToWrite; + vectorsToWrite.push_back(schemaVector); + for (int i = 0; i < FLAGS_num_batches - 1; ++i) { + velox::VectorPtr resultVector; + auto read = rowReader->next(FLAGS_batch_size, resultVector); + if (read == 0) { + break; + } + vectorsToWrite.push_back(resultVector); + } + LOG(INFO) << "Read " << vectorsToWrite.size() << " batches"; + std::unique_ptr writer; + if (FLAGS_output_file == "") { + LOG(INFO) << "No output file provided, writing to memory"; + } + if (FLAGS_use_dwrf) { + writer = std::make_unique(*rootPool, leafPool, schema); + } else { + writer = std::make_unique(*rootPool, leafPool, schema); + } + for (auto& vector : vectorsToWrite) { + writer->write(vector); + } + writer->close(); + LOG(INFO) << "Wrote " << vectorsToWrite.size() << " batches while writing"; + + return 0; +} diff --git a/dwio/nimble/velox/BufferGrowthPolicy.cpp b/dwio/nimble/velox/BufferGrowthPolicy.cpp new file mode 100644 index 0000000..bf624f0 --- /dev/null +++ b/dwio/nimble/velox/BufferGrowthPolicy.cpp @@ -0,0 +1,38 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/BufferGrowthPolicy.h" + +#include + +namespace facebook::nimble { + +uint64_t DefaultInputBufferGrowthPolicy::getExtendedCapacity( + uint64_t newSize, + uint64_t capacity) { + // Short circuit when we don't need to grow further. + if (newSize <= capacity) { + return capacity; + } + + auto iter = rangeConfigs_.upper_bound(newSize); + if (iter == rangeConfigs_.begin()) { + return minCapacity_; + } + + // Move to the range that actually contains the newSize. + --iter; + // The sizes are item counts, hence we should really not run into overflow or + // precision loss. + // TODO: We determine the growth factor only once and grow the + // capacity until it suffices. This doesn't matter that much in practice when + // capacities start from the min capacity and the range boundaries are + // aligned. We could start the capacity at the range boundaries instead, after + // some testing. + double extendedCapacity = folly::to(std::max(capacity, minCapacity_)); + while (extendedCapacity < newSize) { + extendedCapacity *= iter->second; + } + return folly::to(std::floor(extendedCapacity)); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/BufferGrowthPolicy.h b/dwio/nimble/velox/BufferGrowthPolicy.h new file mode 100644 index 0000000..1bc87db --- /dev/null +++ b/dwio/nimble/velox/BufferGrowthPolicy.h @@ -0,0 +1,60 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include + +#include "dwio/nimble/common/Exceptions.h" + +namespace facebook::nimble { +class InputBufferGrowthPolicy { + public: + virtual ~InputBufferGrowthPolicy() = default; + + virtual uint64_t getExtendedCapacity(uint64_t newSize, uint64_t capacity) = 0; +}; + +// Default growth policy for input buffering is a step function across +// the ranges. +// The default growth policy is based on item count in the vectors because +// it tend to be more uniform (tied together by row count) across types. +// However, bytes based policies are worth evaluation as well. +class DefaultInputBufferGrowthPolicy : public InputBufferGrowthPolicy { + public: + explicit DefaultInputBufferGrowthPolicy( + std::map rangeConfigs) + : rangeConfigs_{std::move(rangeConfigs)} { + NIMBLE_CHECK( + !rangeConfigs_.empty() && rangeConfigs_.begin()->first > 0, + "Invalid range config supplied for default buffer input growth policy."); + minCapacity_ = rangeConfigs_.begin()->first; + } + + ~DefaultInputBufferGrowthPolicy() override = default; + + uint64_t getExtendedCapacity(uint64_t newSize, uint64_t capacity) override; + + static std::unique_ptr withDefaultRanges() { + return std::make_unique( + std::map{ + {32UL, 4.0}, {512UL, 1.414}, {4096UL, 1.189}}); + } + + private: + // Map of range lowerbounds and the growth factor for the range. + // The first lowerbound is the smallest unit of allocation and the last range + // extends uint64_t max. + std::map rangeConfigs_; + uint64_t minCapacity_; +}; + +class ExactGrowthPolicy : public InputBufferGrowthPolicy { + public: + uint64_t getExtendedCapacity(uint64_t newSize, uint64_t /* capacity */) + override { + return newSize; + } +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/CMakeLists.txt b/dwio/nimble/velox/CMakeLists.txt new file mode 100644 index 0000000..00f0e46 --- /dev/null +++ b/dwio/nimble/velox/CMakeLists.txt @@ -0,0 +1,97 @@ +add_subdirectory(tests) + +add_library(nimble_velox_common SchemaUtils.cpp) +target_link_libraries(nimble_velox_common nimble_common velox_type) + +add_library(nimble_velox_schema SchemaTypes.cpp) +target_link_libraries(nimble_velox_schema nimble_common Folly::folly) + +add_library(nimble_velox_schema_reader SchemaReader.cpp) +target_link_libraries(nimble_velox_schema_reader nimble_velox_schema + nimble_common Folly::folly) + +add_library(nimble_velox_schema_builder SchemaBuilder.cpp) +target_link_libraries(nimble_velox_schema_builder nimble_velox_schema_reader + nimble_velox_schema nimble_common Folly::folly) + +add_library(nimble_velox_config Config.cpp) +target_link_libraries(nimble_velox_config nimble_encodings nimble_common + Folly::folly) + +add_library(nimble_velox_field_reader FieldReader.cpp) +target_link_libraries(nimble_velox_field_reader nimble_velox_schema_reader + nimble_common Folly::folly) + +add_library(nimble_velox_flatmap_layout_planner FlatMapLayoutPlanner.cpp) +target_link_libraries(nimble_velox_flatmap_layout_planner + nimble_velox_schema_reader nimble_tablet velox_file) + +add_library(nimble_velox_field_writer BufferGrowthPolicy.cpp FieldWriter.cpp) +target_link_libraries(nimble_velox_field_writer nimble_velox_schema + nimble_velox_schema_builder Folly::folly) + +# Nimble code expects an upper case suffix to the generated file. +list(PREPEND FLATBUFFERS_FLATC_SCHEMA_EXTRA_ARGS "--filename-suffix" + "Generated") + +build_flatbuffers( + "${CMAKE_CURRENT_SOURCE_DIR}/Schema.fbs" + "" + nimble_velox_schema_schema_fb + "" + "${CMAKE_CURRENT_BINARY_DIR}" + "" + "") +add_library(nimble_velox_schema_fb INTERFACE) +target_include_directories(nimble_velox_schema_fb + INTERFACE ${FLATBUFFERS_INCLUDE_DIR}) +add_dependencies(nimble_velox_schema_fb nimble_velox_schema_schema_fb) + +build_flatbuffers( + "${CMAKE_CURRENT_SOURCE_DIR}/Metadata.fbs" + "" + nimble_velox_metadata_schema_fb + "" + "${CMAKE_CURRENT_BINARY_DIR}" + "" + "") +add_library(nimble_velox_metadata_fb INTERFACE) +target_include_directories(nimble_velox_metadata_fb + INTERFACE ${FLATBUFFERS_INCLUDE_DIR}) +add_dependencies(nimble_velox_metadata_fb nimble_velox_metadata_schema_fb) + +add_library(nimble_velox_schema_serialization SchemaSerialization.cpp) +target_link_libraries( + nimble_velox_schema_serialization nimble_velox_schema_reader + nimble_velox_schema_builder nimble_velox_schema_fb) + +add_library(nimble_velox_deserializer Deserializer.cpp) +target_link_libraries( + nimble_velox_deserializer + nimble_common + nimble_velox_common + nimble_velox_field_reader + velox_vector + velox_memory + velox_dwio_common) + +add_library(nimble_velox_serializer Serializer.cpp) +target_link_libraries(nimble_velox_serializer nimble_common + nimble_velox_schema_builder velox_vector) + +add_library(nimble_velox_reader StreamInputDecoder.cpp StreamLabels.cpp + VeloxReader.cpp) +target_link_libraries( + nimble_velox_reader + nimble_velox_schema + nimble_velox_schema_serialization + nimble_velox_schema_fb + nimble_velox_metadata_fb + nimble_common + Folly::folly) + +add_library( + nimble_velox_writer EncodingLayoutTree.cpp FlushPolicy.cpp VeloxWriter.cpp + VeloxWriterDefaultMetadataOSS.cpp) +target_link_libraries(nimble_velox_writer nimble_encodings nimble_common + nimble_velox_metadata_fb Folly::folly) diff --git a/dwio/nimble/velox/ChunkedStream.cpp b/dwio/nimble/velox/ChunkedStream.cpp new file mode 100644 index 0000000..4889ce4 --- /dev/null +++ b/dwio/nimble/velox/ChunkedStream.cpp @@ -0,0 +1,71 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/ChunkedStream.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/tablet/Compression.h" +#include "folly/io/Cursor.h" + +namespace facebook::nimble { + +void InMemoryChunkedStream::ensureLoaded() { + if (!pos_) { + stream_ = streamLoader_->getStream(); + pos_ = stream_.data(); + } +} + +bool InMemoryChunkedStream::hasNext() { + ensureLoaded(); + return pos_ - stream_.data() < stream_.size(); +} + +std::string_view InMemoryChunkedStream::nextChunk() { + ensureLoaded(); + uncompressed_.clear(); + NIMBLE_ASSERT( + sizeof(uint32_t) + sizeof(char) <= + stream_.size() - (pos_ - stream_.data()), + "Read beyond end of stream"); + auto length = encoding::readUint32(pos_); + auto compressionType = static_cast(encoding::readChar(pos_)); + NIMBLE_ASSERT( + length <= stream_.size() - (pos_ - stream_.data()), + "Read beyond end of stream"); + std::string_view chunk; + switch (compressionType) { + case CompressionType::Uncompressed: { + chunk = {pos_, length}; + break; + } + case CompressionType::Zstd: { + uncompressed_ = ZstdCompression::uncompress( + *uncompressed_.memoryPool(), {pos_, length}); + chunk = {uncompressed_.data(), uncompressed_.size()}; + break; + } + default: { + NIMBLE_UNREACHABLE(fmt::format( + "Unexpected stream compression type: ", toString(compressionType))); + } + } + pos_ += length; + return chunk; +} + +CompressionType InMemoryChunkedStream::peekCompressionType() { + ensureLoaded(); + NIMBLE_ASSERT( + sizeof(uint32_t) + sizeof(char) <= + stream_.size() - (pos_ - stream_.data()), + "Read beyond end of stream"); + auto pos = pos_ + sizeof(uint32_t); + return static_cast(encoding::readChar(pos)); +} + +void InMemoryChunkedStream::reset() { + uncompressed_.clear(); + pos_ = stream_.data(); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/ChunkedStream.h b/dwio/nimble/velox/ChunkedStream.h new file mode 100644 index 0000000..10882f2 --- /dev/null +++ b/dwio/nimble/velox/ChunkedStream.h @@ -0,0 +1,55 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include + +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/tablet/Tablet.h" +#include "folly/io/IOBuf.h" + +namespace facebook::nimble { + +class ChunkedStream { + public: + virtual ~ChunkedStream() = default; + + virtual bool hasNext() = 0; + + virtual std::string_view nextChunk() = 0; + + virtual CompressionType peekCompressionType() = 0; + + virtual void reset() = 0; +}; + +class InMemoryChunkedStream : public ChunkedStream { + public: + InMemoryChunkedStream( + velox::memory::MemoryPool& memoryPool, + std::unique_ptr streamLoader) + : streamLoader_{std::move(streamLoader)}, + pos_{nullptr}, + uncompressed_{&memoryPool} {} + + bool hasNext() override; + + std::string_view nextChunk() override; + + CompressionType peekCompressionType() override; + + void reset() override; + + private: + void ensureLoaded(); + + std::unique_ptr streamLoader_; + std::string_view stream_; + const char* pos_; + Vector uncompressed_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/ChunkedStreamDecoder.cpp b/dwio/nimble/velox/ChunkedStreamDecoder.cpp new file mode 100644 index 0000000..29f4878 --- /dev/null +++ b/dwio/nimble/velox/ChunkedStreamDecoder.cpp @@ -0,0 +1,172 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/ChunkedStreamDecoder.h" +#include "dwio/nimble/encodings/EncodingFactoryNew.h" + +namespace facebook::nimble { + +namespace { +void bufferStringContent( + Vector& buffer, + void* values, + void* FOLLY_NULLABLE notNulls, + uint32_t offset, + uint32_t count) { + auto* source = reinterpret_cast(values) + offset; + uint64_t size = 0; + if (notNulls) { + for (auto i = 0; i < count; ++i) { + if (bits::getBit(i + offset, reinterpret_cast(notNulls))) { + size += (source + i)->size(); + } + } + } else { + for (auto i = 0; i < count; ++i) { + size += (source + i)->size(); + } + } + + uint64_t targetOffset = 0; + buffer.resize(size); + + if (notNulls) { + for (auto i = 0; i < count; ++i) { + if (bits::getBit(i + offset, reinterpret_cast(notNulls))) { + auto* value = source + i; + auto* target = buffer.data() + targetOffset; + if (!value->empty()) { + std::copy(value->cbegin(), value->cend(), target); + targetOffset += value->size(); + } + *value = std::string_view{target, value->size()}; + } + } + } else { + for (auto i = 0; i < count; ++i) { + auto* value = source + i; + auto* target = buffer.data() + targetOffset; + if (!value->empty()) { + std::copy(value->cbegin(), value->cend(), target); + targetOffset += value->size(); + } + *value = std::string_view{target, value->size()}; + } + } +} + +} // namespace + +uint32_t ChunkedStreamDecoder::next( + uint32_t count, + void* output, + std::function nulls, + const bits::Bitmap* scatterBitmap) { + stringBuffers_.clear(); + + if (count == 0) { + if (nulls && scatterBitmap) { + auto nullsPtr = nulls(); + // @lint-ignore CLANGTIDY facebook-hte-BadMemset + memset(nullsPtr, 0, bits::bytesRequired(scatterBitmap->size())); + } + return 0; + } + + LoggingScope scope{logger_}; + + uint32_t nonNullCount = 0; + bool hasNulls = false; + uint32_t offset = 0; + void* nullsPtr = nullptr; + std::function initNulls = [&]() { + if (!nullsPtr) { + nullsPtr = nulls(); + } + return nullsPtr; + }; + + while (count > 0) { + ensureLoaded(); + + auto rowsToRead = std::min(count, remaining_); + uint32_t chunkNonNullCount = 0; + uint32_t endOffset = 0; + + if (!nulls || !scatterBitmap) { + NIMBLE_CHECK(!scatterBitmap, "unexpected scatter bitmap"); + chunkNonNullCount = encoding_->materializeNullable( + rowsToRead, output, initNulls, nullptr, offset); + endOffset = offset + rowsToRead; + } else { + endOffset = bits::findSetBit( + static_cast(scatterBitmap->bits()), + offset, + scatterBitmap->size(), + rowsToRead + 1); + bits::Bitmap localBitmap{scatterBitmap->bits(), endOffset}; + chunkNonNullCount = encoding_->materializeNullable( + rowsToRead, output, initNulls, &localBitmap, offset); + } + + auto chunkHasNulls = chunkNonNullCount != (endOffset - offset); + if (chunkHasNulls && !hasNulls) { + // back fill the nulls bitmap to all non-nulls + bits::BitmapBuilder builder{nullsPtr, offset}; + builder.set(0, offset); + } + hasNulls = hasNulls || chunkHasNulls; + if (hasNulls && !chunkHasNulls) { + // fill nulls bitmap to reflect that all values are non-null + bits::BitmapBuilder builder{nullsPtr, endOffset}; + builder.set(offset, endOffset); + } + + if (encoding_->dataType() == DataType::String && rowsToRead != count) { + // We are going to load a new chunk. + // For string values, this means that the memory pointed by the + // string_views is going to be freed. Before we do so, we copy all + // strings to a temporary buffer and fix the string_views to point to the + // new location. + // NOTE1: Instead of copying the data, we can just hold on to + // the previous chunk(s) for a while. However, this means holding on to + // more memory, which is undesirable. + // NOTE2: We perform an additional copy of the strings later on, into the + // string buffers of the Velox Vector. Later diff will change this logic + // to directly copy the strings into the Velox string buffers directly. + auto& buffer = stringBuffers_.emplace_back(&pool_); + bufferStringContent(buffer, output, nullsPtr, offset, rowsToRead); + } + + offset = endOffset; + remaining_ -= rowsToRead; + count -= rowsToRead; + nonNullCount += chunkNonNullCount; + } + + return nonNullCount; +} + +void ChunkedStreamDecoder::skip(uint32_t count) { + while (count > 0) { + ensureLoaded(); + auto toSkip = std::min(count, remaining_); + encoding_->skip(toSkip); + count -= toSkip; + remaining_ -= toSkip; + } +} + +void ChunkedStreamDecoder::reset() { + stream_->reset(); + remaining_ = 0; +} + +void ChunkedStreamDecoder::ensureLoaded() { + if (UNLIKELY(remaining_ == 0)) { + encoding_ = EncodingFactory::decode(pool_, stream_->nextChunk()); + remaining_ = encoding_->rowCount(); + NIMBLE_ASSERT(remaining_ > 0, "Empty chunk"); + } +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/ChunkedStreamDecoder.h b/dwio/nimble/velox/ChunkedStreamDecoder.h new file mode 100644 index 0000000..9e1d8cd --- /dev/null +++ b/dwio/nimble/velox/ChunkedStreamDecoder.h @@ -0,0 +1,41 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/common/MetricsLogger.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/velox/ChunkedStream.h" +#include "dwio/nimble/velox/Decoder.h" + +namespace facebook::nimble { + +class ChunkedStreamDecoder : public Decoder { + public: + ChunkedStreamDecoder( + velox::memory::MemoryPool& pool, + std::unique_ptr stream, + const MetricsLogger& logger) + : pool_{pool}, stream_{std::move(stream)}, logger_{logger} {} + + uint32_t next( + uint32_t count, + void* output, + std::function nulls = nullptr, + const bits::Bitmap* scatterBitmap = nullptr) override; + + void skip(uint32_t count) override; + + void reset() override; + + private: + void ensureLoaded(); + + velox::memory::MemoryPool& pool_; + std::unique_ptr stream_; + std::unique_ptr encoding_; + uint32_t remaining_{0}; + const MetricsLogger& logger_; + std::vector> stringBuffers_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/ChunkedStreamWriter.cpp b/dwio/nimble/velox/ChunkedStreamWriter.cpp new file mode 100644 index 0000000..f414b51 --- /dev/null +++ b/dwio/nimble/velox/ChunkedStreamWriter.cpp @@ -0,0 +1,45 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/ChunkedStreamWriter.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/tablet/Compression.h" + +namespace facebook::nimble { + +ChunkedStreamWriter::ChunkedStreamWriter( + Buffer& buffer, + CompressionParams compressionParams) + : buffer_{buffer}, compressionParams_{compressionParams} { + NIMBLE_CHECK( + compressionParams_.type == CompressionType::Uncompressed || + compressionParams_.type == CompressionType::Zstd, + fmt::format( + "Unsupported chunked stream compression type: {}", + toString(compressionParams_.type))); +} + +std::vector ChunkedStreamWriter::encode( + std::string_view chunk) { + constexpr uint8_t headerSize = sizeof(uint32_t) + sizeof(CompressionType); + auto* header = buffer_.reserve(headerSize); + auto* pos = header; + + if (compressionParams_.type == CompressionType::Zstd) { + auto compressed = ZstdCompression::compress( + buffer_.getMemoryPool(), chunk, compressionParams_.zstdLevel); + if (compressed.has_value()) { + encoding::writeUint32(compressed->size(), pos); + encoding::write(CompressionType::Zstd, pos); + return { + {header, headerSize}, + buffer_.takeOwnership(compressed->releaseOwnership())}; + } + } + + encoding::writeUint32(chunk.size(), pos); + encoding::write(CompressionType::Uncompressed, pos); + return {{header, headerSize}, chunk}; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/ChunkedStreamWriter.h b/dwio/nimble/velox/ChunkedStreamWriter.h new file mode 100644 index 0000000..fd828b7 --- /dev/null +++ b/dwio/nimble/velox/ChunkedStreamWriter.h @@ -0,0 +1,24 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/Types.h" + +namespace facebook::nimble { + +class ChunkedStreamWriter { + public: + explicit ChunkedStreamWriter( + Buffer& buffer, + CompressionParams compressionParams = { + .type = CompressionType::Uncompressed}); + + std::vector encode(std::string_view chunk); + + private: + Buffer& buffer_; + CompressionParams compressionParams_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/Config.cpp b/dwio/nimble/velox/Config.cpp new file mode 100644 index 0000000..ca6e081 --- /dev/null +++ b/dwio/nimble/velox/Config.cpp @@ -0,0 +1,139 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include +#include + +#include "dwio/nimble/encodings/EncodingSelectionPolicy.h" +#include "dwio/nimble/velox/Config.h" + +DEFINE_string( + nimble_selection_read_factors, + "Constant=1.0;Trivial=0.5;FixedBitWidth=0.9;MainlyConstant=1.0;SparseBool=1.0;Dictionary=1.0;RLE=1.0;Varint=1.0", + "Encoding selection read factors, in the format: " + "=;=;..."); + +DEFINE_double( + nimble_selection_compression_accept_ratio, + 0.97, + "Encoding selection compression accept ratio."); + +DEFINE_bool( + nimble_zstrong_enable_variable_bit_width_compressor, + false, + "Enable zstrong variable bit width compressor at write time. Transparent at read time."); + +DEFINE_string( + nimble_writer_input_buffer_default_growth_config, + "{\"32\":4.0,\"512\":1.414,\"4096\":1.189}", + "Default growth config for writer input buffers, each entry in the format of {range_start,growth_factor}"); + +namespace facebook::nimble { +namespace { +template +std::vector parseVector(const std::string& str) { + std::vector result; + if (!str.empty()) { + std::vector pieces; + folly::split(',', str, pieces, true); + for (auto& p : pieces) { + const auto& trimmedCol = folly::trimWhitespace(p); + if (!trimmedCol.empty()) { + result.push_back(folly::to(trimmedCol)); + } + } + } + return result; +} + +std::map parseGrowthConfigMap(const std::string& str) { + std::map ret; + NIMBLE_CHECK(!str.empty(), "Can't supply an empty growth config."); + folly::dynamic json = folly::parseJson(str); + for (const auto& pair : json.items()) { + auto [_, inserted] = ret.emplace( + folly::to(pair.first.asString()), pair.second.asDouble()); + NIMBLE_CHECK( + inserted, fmt::format("Duplicate key: {}.", pair.first.asString())); + } + return ret; +} +} // namespace + +/* static */ Config::Entry Config::FLATTEN_MAP("orc.flatten.map", false); + +/* static */ Config::Entry> Config::MAP_FLAT_COLS( + "orc.map.flat.cols", + {}, + [](const std::vector& val) { return folly::join(",", val); }, + [](const std::string& /* key */, const std::string& val) { + return parseVector(val); + }); + +/* static */ Config::Entry> + Config::BATCH_REUSE_COLS( + "alpha.dictionaryarray.cols", + {}, + [](const std::vector& val) { return folly::join(",", val); }, + [](const std::string& /* key */, const std::string& val) { + return parseVector(val); + }); + +/* static */ Config::Entry Config::RAW_STRIPE_SIZE( + "alpha.raw.stripe.size", + 512L * 1024L * 1024L); + +/* static */ Config::Entry>> + Config::MANUAL_ENCODING_SELECTION_READ_FACTORS( + "alpha.encodingselection.read.factors", + ManualEncodingSelectionPolicyFactory::parseReadFactors( + FLAGS_nimble_selection_read_factors), + [](const std::vector>& val) { + std::vector encodingFactorStrings; + std::transform( + val.cbegin(), + val.cend(), + std::back_inserter(encodingFactorStrings), + [](const auto& readFactor) { + return fmt::format( + "{}={}", toString(readFactor.first), readFactor.second); + }); + return folly::join(";", encodingFactorStrings); + }, + [](const std::string& /* key */, const std::string& val) { + return ManualEncodingSelectionPolicyFactory::parseReadFactors(val); + }); + +/* static */ Config::Entry + Config::ENCODING_SELECTION_COMPRESSION_ACCEPT_RATIO( + "alpha.encodingselection.compression.accept.ratio", + FLAGS_nimble_selection_compression_accept_ratio); + +/* static */ Config::Entry Config::ZSTRONG_COMPRESSION_LEVEL( + "alpha.zstrong.compression.level", + 4); + +/* static */ Config::Entry Config::ZSTRONG_DECOMPRESSION_LEVEL( + "alpha.zstrong.decompression.level", + 2); + +/* static */ Config::Entry + Config::ENABLE_ZSTRONG_VARIABLE_BITWIDTH_COMPRESSOR( + "alpha.zstrong.enable.variable.bit.width.compressor", + FLAGS_nimble_zstrong_enable_variable_bit_width_compressor); + +/* static */ Config::Entry> + Config::INPUT_BUFFER_DEFAULT_GROWTH_CONFIGS( + "alpha.writer.input.buffer.default.growth.configs", + parseGrowthConfigMap( + FLAGS_nimble_writer_input_buffer_default_growth_config), + [](const std::map& val) { + folly::dynamic obj = folly::dynamic::object; + for (const auto& [rangeStart, growthFactor] : val) { + obj[folly::to(rangeStart)] = growthFactor; + } + return folly::toJson(obj); + }, + [](const std::string& /* key */, const std::string& val) { + return parseGrowthConfigMap(val); + }); +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/Config.h b/dwio/nimble/velox/Config.h new file mode 100644 index 0000000..e41429f --- /dev/null +++ b/dwio/nimble/velox/Config.h @@ -0,0 +1,36 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/common/Types.h" +#include "velox/common/config/Config.h" + +namespace facebook::nimble { + +class Config : public velox::common::ConfigBase { + public: + template + using Entry = velox::common::ConfigBase::Entry; + + static Entry FLATTEN_MAP; + static Entry> MAP_FLAT_COLS; + static Entry> BATCH_REUSE_COLS; + static Entry RAW_STRIPE_SIZE; + static Entry>> + MANUAL_ENCODING_SELECTION_READ_FACTORS; + static Entry ENCODING_SELECTION_COMPRESSION_ACCEPT_RATIO; + static Entry ZSTRONG_COMPRESSION_LEVEL; + static Entry ZSTRONG_DECOMPRESSION_LEVEL; + static Entry ENABLE_ZSTRONG_VARIABLE_BITWIDTH_COMPRESSOR; + static Entry> + INPUT_BUFFER_DEFAULT_GROWTH_CONFIGS; + + static std::shared_ptr fromMap( + const std::map& map) { + auto ret = std::make_shared(); + ret->configs_.insert(map.cbegin(), map.cend()); + return ret; + } +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/Decoder.h b/dwio/nimble/velox/Decoder.h new file mode 100644 index 0000000..6508623 --- /dev/null +++ b/dwio/nimble/velox/Decoder.h @@ -0,0 +1,26 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include "dwio/nimble/common/Bits.h" + +namespace facebook::nimble { + +class Decoder { + public: + virtual ~Decoder() = default; + + virtual uint32_t next( + uint32_t count, + void* output, + std::function nulls = nullptr, + const bits::Bitmap* scatterBitmap = nullptr) = 0; + + virtual void skip(uint32_t count) = 0; + + virtual void reset() = 0; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/Deserializer.cpp b/dwio/nimble/velox/Deserializer.cpp new file mode 100644 index 0000000..6fdf07a --- /dev/null +++ b/dwio/nimble/velox/Deserializer.cpp @@ -0,0 +1,186 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/Deserializer.h" +#include +#include +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/velox/Decoder.h" +#include "dwio/nimble/velox/SchemaUtils.h" +#include "velox/dwio/common/TypeWithId.h" + +namespace facebook::nimble { + +namespace { + +uint32_t getTypeStorageWidth(const Type& type) { + switch (type.kind()) { + case Kind::Scalar: { + auto scalarKind = type.asScalar().scalarDescriptor().scalarKind(); + switch (scalarKind) { + case ScalarKind::Bool: + case ScalarKind::Int8: + case ScalarKind::UInt8: + return 1; + case ScalarKind::Int16: + case ScalarKind::UInt16: + return 2; + case ScalarKind::Int32: + case ScalarKind::Float: + case ScalarKind::UInt32: + return 4; + case ScalarKind::Int64: + case ScalarKind::UInt64: + case ScalarKind::Double: + return 8; + case ScalarKind::String: + case ScalarKind::Binary: + case ScalarKind::Undefined: + return 0; + } + break; + } + case Kind::Row: + case Kind::FlatMap: + return 1; + case Kind::Array: + case Kind::ArrayWithOffsets: + case Kind::Map: + return 4; + } +} + +class DeserializerImpl : public Decoder { + public: + explicit DeserializerImpl(const Type& type) : type_{type} {} + + uint32_t next( + uint32_t count, + void* output, + std::function /* nulls */, + const bits::Bitmap* /* scatterBitmap */) override { + if (count == 0) { + return count; + } + + if (data_.empty()) { + // TODO: This is less ideal. Need a way to avoid sending back the values. + NIMBLE_CHECK( + type_.isRow() || type_.isFlatMap(), + fmt::format("missing input data for {}", toString(type_.kind()))); + auto bools = static_cast(output); + std::fill(bools, bools + count, true); + return count; + } + + auto width = getTypeStorageWidth(type_); + if (width > 0) { + auto expectedSize = count * width; + auto pos = data_.begin(); + auto compression = static_cast(encoding::readChar(pos)); + switch (compression) { + case CompressionType::Uncompressed: { + NIMBLE_CHECK(expectedSize == data_.end() - pos, "unexpected size"); + std::copy(pos, data_.end(), static_cast(output)); + break; + } + case CompressionType::Zstd: { + // TODO: share compression implementation + auto ret = + ZSTD_decompress(output, expectedSize, pos, data_.size() - 1); + NIMBLE_CHECK( + !ZSTD_isError(ret), + fmt::format( + "Error uncompressing data: {}", ZSTD_getErrorName(ret))); + NIMBLE_CHECK(ret == expectedSize, "unexpected size"); + break; + } + default: + NIMBLE_NOT_SUPPORTED( + fmt::format("Unsupported compression {}", toString(compression))); + } + return count; + } + + // TODO: handle string compression. One option is to share implementations + // with trivial encoding. That way, we also get bit packed booleans for + // free. + auto scalarKind = type_.asScalar().scalarDescriptor().scalarKind(); + NIMBLE_CHECK( + scalarKind == ScalarKind::String || scalarKind == ScalarKind::Binary, + fmt::format("Unexpected scalar kind {}", toString(scalarKind))); + auto sv = static_cast(output); + auto pos = data_.data(); + for (auto i = 0; i < count; ++i) { + sv[i] = encoding::readString(pos); + } + return count; + } + + void skip(uint32_t /* count */) override { + NIMBLE_UNREACHABLE("unexpected call"); + } + + void reset() override { + NIMBLE_UNREACHABLE("unexpected call"); + } + + void reset(std::string_view data) { + data_ = data; + } + + private: + const Type& type_; + std::string_view data_; +}; + +const StreamDescriptor& getMainDescriptor(const Type& type) { + switch (type.kind()) { + case Kind::Scalar: + return type.asScalar().scalarDescriptor(); + case Kind::Array: + return type.asArray().lengthsDescriptor(); + case Kind::Map: + return type.asMap().lengthsDescriptor(); + case Kind::Row: + return type.asRow().nullsDescriptor(); + default: + NIMBLE_NOT_SUPPORTED(fmt::format( + "Schema type {} is not supported.", toString(type.kind()))); + } +} + +} // namespace + +Deserializer::Deserializer( + velox::memory::MemoryPool& pool, + std::shared_ptr schema) + : pool_{pool}, schema_{std::move(schema)} { + FieldReaderParams params; + std::shared_ptr schemaWithId = + velox::dwio::common::TypeWithId::create(convertToVeloxType(*schema_)); + std::vector offsets; + rootFactory_ = + FieldReaderFactory::create(params, pool_, schema_, schemaWithId, offsets); + SchemaReader::traverseSchema(schema_, [this](auto, auto& type, auto&) { + deserializers_[getMainDescriptor(type).offset()] = + std::make_unique(type); + }); + rootReader_ = rootFactory_->createReader(deserializers_); +} + +void Deserializer::deserialize( + std::string_view data, + velox::VectorPtr& vector) { + auto pos = data.data(); + auto rows = encoding::readUint32(pos); + for (uint32_t stream = 0; stream < deserializers_.size(); ++stream) { + static_cast(deserializers_[stream].get()) + ->reset(encoding::readString(pos)); + } + NIMBLE_CHECK(pos == data.end(), "unexpected end"); + rootReader_->next(rows, vector, nullptr); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/Deserializer.h b/dwio/nimble/velox/Deserializer.h new file mode 100644 index 0000000..6d5412d --- /dev/null +++ b/dwio/nimble/velox/Deserializer.h @@ -0,0 +1,28 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/velox/FieldReader.h" +#include "folly/container/F14Map.h" +#include "velox/common/memory/Memory.h" +#include "velox/vector/BaseVector.h" + +namespace facebook::nimble { + +class Deserializer { + public: + Deserializer( + velox::memory::MemoryPool& pool, + std::shared_ptr schema); + + void deserialize(std::string_view data, velox::VectorPtr& vector); + + private: + velox::memory::MemoryPool& pool_; + std::shared_ptr schema_; + folly::F14FastMap> deserializers_; + std::unique_ptr rootFactory_; + std::unique_ptr rootReader_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/EncodingLayoutTree.cpp b/dwio/nimble/velox/EncodingLayoutTree.cpp new file mode 100644 index 0000000..82e6962 --- /dev/null +++ b/dwio/nimble/velox/EncodingLayoutTree.cpp @@ -0,0 +1,177 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/EncodingLayoutTree.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/Exceptions.h" + +namespace facebook::nimble { + +namespace { + +constexpr uint32_t kMinBufferSize = 8; + +std::pair createInternal(std::string_view tree) { + // Layout: + // 1 byte: Schema Kind + // 2 byte: Name length + // X bytes: Name bytes + // 1 byte: Stream encoding layout count + // Repeat next for "Stream encoding layout count" times: + // 1 byte: Stream identifier + // 2 byte: Encoding layout length + // Y bytes: Encoding layout bytes + // End repeat + // 4 byte: Children count + // Z bytes: Children + + NIMBLE_CHECK( + tree.size() >= kMinBufferSize, + "Invalid captured encoding tree. Buffer too small."); + + auto pos = tree.data(); + auto schemaKind = encoding::read(pos); + auto nameLength = encoding::read(pos); + + NIMBLE_CHECK( + tree.size() >= nameLength + kMinBufferSize, + "Invalid captured encoding tree. Buffer too small."); + + std::string_view name{pos, nameLength}; + pos += nameLength; + + auto encodingLayoutCount = encoding::read(pos); + std::unordered_map + encodingLayouts; + encodingLayouts.reserve(encodingLayoutCount); + for (auto i = 0; i < encodingLayoutCount; ++i) { + NIMBLE_CHECK( + tree.size() - (pos - tree.data()) >= 3, + "Invalid captured encoding tree. Buffer too small."); + auto streamIdentifier = encoding::read(pos); + auto encodingLength = encoding::read(pos); + + NIMBLE_CHECK( + tree.size() - (pos - tree.data()) >= encodingLength, + "Invalid captured encoding tree. Buffer too small."); + + auto layout = EncodingLayout::create({pos, encodingLength}); + encodingLayouts.insert({streamIdentifier, std::move(layout.first)}); + pos += layout.second; + + NIMBLE_CHECK( + layout.second == encodingLength, + "Invalid captured encoding tree. Encoding size mismatch."); + } + + auto childrenCount = encoding::read(pos); + uint32_t offset = pos - tree.data(); + std::vector children; + children.reserve(childrenCount); + for (auto i = 0; i < childrenCount; ++i) { + auto encodingLayoutTree = createInternal(tree.substr(offset)); + offset += encodingLayoutTree.second; + children.push_back(std::move(encodingLayoutTree.first)); + } + + return { + {schemaKind, + std::move(encodingLayouts), + std::string{name}, + std::move(children)}, + offset}; +} + +} // namespace + +EncodingLayoutTree::EncodingLayoutTree( + Kind schemaKind, + std::unordered_map encodingLayouts, + std::string name, + std::vector children) + : schemaKind_{schemaKind}, + encodingLayouts_{std::move(encodingLayouts)}, + name_{std::move(name)}, + children_{std::move(children)} { + NIMBLE_CHECK( + encodingLayouts_.size() < std::numeric_limits::max(), + "Too many encoding layout streams."); +} + +uint32_t EncodingLayoutTree::serialize(std::span output) const { + NIMBLE_CHECK( + output.size() >= kMinBufferSize + name_.size(), + "Captured encoding buffer too small."); + + auto pos = output.data(); + nimble::encoding::write(schemaKind_, pos); + nimble::encoding::write(name_.size(), pos); + if (!name_.empty()) { + nimble::encoding::writeBytes(name_, pos); + } + + nimble::encoding::write(encodingLayouts_.size(), pos); + for (const auto& pair : encodingLayouts_) { + uint32_t encodingSize = 0; + nimble::encoding::write(pair.first, pos); + encodingSize = pair.second.serialize( + output.subspan(pos - output.data() + sizeof(uint16_t))); + nimble::encoding::write(encodingSize, pos); + pos += encodingSize; + } + + nimble::encoding::write(children_.size(), pos); + + for (auto i = 0; i < children_.size(); ++i) { + pos += children_[i].serialize(output.subspan(pos - output.data())); + } + + return pos - output.data(); +} + +EncodingLayoutTree EncodingLayoutTree::create(std::string_view tree) { + return std::move(createInternal(tree).first); +} + +Kind EncodingLayoutTree::schemaKind() const { + return schemaKind_; +} + +const EncodingLayout* FOLLY_NULLABLE EncodingLayoutTree::encodingLayout( + EncodingLayoutTree::StreamIdentifier identifier) const { + auto it = encodingLayouts_.find(identifier); + if (it == encodingLayouts_.end()) { + return nullptr; + } + return &it->second; +} + +const std::string& EncodingLayoutTree::name() const { + return name_; +} + +uint32_t EncodingLayoutTree::childrenCount() const { + return children_.size(); +} + +const EncodingLayoutTree& EncodingLayoutTree::child(uint32_t index) const { + NIMBLE_DCHECK( + index < childrenCount(), + "Encoding layout tree child index is out of range."); + + return children_[index]; +} + +std::vector +EncodingLayoutTree::encodingLayoutIdentifiers() const { + std::vector identifiers; + identifiers.reserve(encodingLayouts_.size()); + std::transform( + encodingLayouts_.cbegin(), + encodingLayouts_.cend(), + std::back_inserter(identifiers), + [](const auto& pair) { return pair.first; }); + + return identifiers; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/EncodingLayoutTree.h b/dwio/nimble/velox/EncodingLayoutTree.h new file mode 100644 index 0000000..960e8d9 --- /dev/null +++ b/dwio/nimble/velox/EncodingLayoutTree.h @@ -0,0 +1,62 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include "dwio/nimble/encodings/EncodingLayout.h" +#include "dwio/nimble/velox/SchemaTypes.h" +#include "folly/CppAttributes.h" + +namespace facebook::nimble { + +class EncodingLayoutTree { + public: + using StreamIdentifier = uint8_t; + + struct StreamIdentifiers { + struct Scalar { + constexpr static StreamIdentifier ScalarStream = 0; + }; + struct Array { + constexpr static StreamIdentifier LengthsStream = 0; + }; + struct Map { + constexpr static StreamIdentifier LengthsStream = 0; + }; + struct Row { + constexpr static StreamIdentifier NullsStream = 0; + }; + struct FlatMap { + constexpr static StreamIdentifier NullsStream = 0; + }; + struct ArrayWithOffsets { + constexpr static StreamIdentifier OffsetsStream = 0; + constexpr static StreamIdentifier LengthsStream = 1; + }; + }; + + EncodingLayoutTree( + Kind schemaKind, + std::unordered_map encodingLayouts, + std::string name, + std::vector children = {}); + + Kind schemaKind() const; + const EncodingLayout* FOLLY_NULLABLE encodingLayout(StreamIdentifier) const; + const std::string& name() const; + uint32_t childrenCount() const; + const EncodingLayoutTree& child(uint32_t index) const; + + uint32_t serialize(std::span output) const; + static EncodingLayoutTree create(std::string_view tree); + + std::vector encodingLayoutIdentifiers() const; + + private: + const Kind schemaKind_; + const std::unordered_map encodingLayouts_; + const std::string name_; + const std::vector children_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/FieldReader.cpp b/dwio/nimble/velox/FieldReader.cpp new file mode 100644 index 0000000..fa5a8d3 --- /dev/null +++ b/dwio/nimble/velox/FieldReader.cpp @@ -0,0 +1,2457 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/FieldReader.h" +#include "dwio/nimble/common/Bits.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" +#include "velox/dwio/common/FlatMapHelper.h" +#include "velox/vector/ComplexVector.h" +#include "velox/vector/DictionaryVector.h" +#include "velox/vector/FlatVector.h" + +namespace facebook::nimble { + +namespace { + +constexpr uint32_t kSkipBatchSize = 1024; + +uint32_t scatterCount(uint32_t count, const bits::Bitmap* scatterBitmap) { + return scatterBitmap ? scatterBitmap->size() : count; +} + +// Bytes needed for a packed null bitvector in velox. +constexpr uint64_t nullBytes(uint32_t rowCount) { + return velox::bits::nbytes(rowCount) + velox::simd::kPadding; +} + +// Bits needed for a packed null bitvector in velox. +constexpr uint64_t nullBits(uint32_t rowCount) { + return rowCount + 8 * velox::simd::kPadding; +} + +// Returns the nulls for a vector properly padded to hold |rowCount|. +char* paddedNulls(velox::BaseVector* vec, uint32_t rowCount) { + return vec->mutableNulls(nullBits(rowCount))->template asMutable(); +} + +// Zeroes vec's null vector (aka make it 'all null'). +void zeroNulls(velox::BaseVector* vec, uint32_t rowCount) { + memset(paddedNulls(vec, rowCount), 0, nullBytes(rowCount)); +} + +template +T* FOLLY_NULLABLE verifyVectorState(velox::VectorPtr& vector) { + // we want vector to not be referenced by anyone else (e.g. ref count of 1) + if (vector) { + auto casted = vector->as(); + if (casted && vector.use_count() == 1) { + return casted; + } + vector.reset(); + } + return nullptr; +} + +// Ensure the internal buffer to vector are refCounted to one +template +inline void resetIfNotWritable(velox::VectorPtr& vector, T&... buffer) { + // The result vector and the buffer both hold reference, so refCount is at + // least 2 + auto resetIfShared = [](auto& buffer) { + if (!buffer) { + return false; + } + const bool reset = buffer->refCount() > 2; + if (reset) { + buffer.reset(); + } + return reset; + }; + + if ((... || resetIfShared(buffer))) { + vector.reset(); + } +} + +template +struct VectorInitializer {}; + +template +struct VectorInitializer> { + static velox::FlatVector* initialize( + velox::memory::MemoryPool* pool, + velox::VectorPtr& output, + const velox::TypePtr& veloxType, + uint64_t rowCount, + velox::BufferPtr values = nullptr) { + auto vector = verifyVectorState>(output); + velox::BufferPtr nulls; + if (vector) { + nulls = vector->nulls(); + values = vector->mutableValues(rowCount); + resetIfNotWritable(output, nulls, values); + } + if (!values) { + values = velox::AlignedBuffer::allocate(rowCount, pool); + } + if (!output) { + output = std::make_shared>( + pool, + veloxType, + nulls, + rowCount, + values, + std::vector()); + } + return static_cast*>(output.get()); + } +}; + +template <> +struct VectorInitializer { + static velox::ArrayVector* initialize( + velox::memory::MemoryPool* pool, + velox::VectorPtr& output, + const velox::TypePtr& veloxType, + uint64_t rowCount) { + auto vector = verifyVectorState(output); + velox::BufferPtr nulls, sizes, offsets; + velox::VectorPtr elements; + if (vector) { + nulls = vector->nulls(); + sizes = vector->mutableSizes(rowCount); + offsets = vector->mutableOffsets(rowCount); + elements = vector->elements(); + resetIfNotWritable(output, nulls, sizes, offsets); + } + if (!offsets) { + offsets = + velox::AlignedBuffer::allocate(rowCount, pool); + } + if (!sizes) { + sizes = + velox::AlignedBuffer::allocate(rowCount, pool); + } + if (!output) { + output = std::make_shared( + pool, + veloxType, + nulls, + rowCount, + std::move(offsets), + std::move(sizes), + /* elements */ elements, + 0 /*nullCount*/); + } + return static_cast(output.get()); + } +}; + +template <> +struct VectorInitializer { + static velox::MapVector* initialize( + velox::memory::MemoryPool* pool, + velox::VectorPtr& output, + const velox::TypePtr& veloxType, + uint64_t rowCount) { + auto vector = verifyVectorState(output); + velox::BufferPtr nulls, sizes, offsets; + velox::VectorPtr mapKeys, mapValues; + if (vector) { + nulls = vector->nulls(); + sizes = vector->mutableSizes(rowCount); + offsets = vector->mutableOffsets(rowCount); + mapKeys = vector->mapKeys(); + mapValues = vector->mapValues(); + resetIfNotWritable(output, nulls, sizes, offsets); + } + if (!offsets) { + offsets = + velox::AlignedBuffer::allocate(rowCount, pool); + } + if (!sizes) { + sizes = + velox::AlignedBuffer::allocate(rowCount, pool); + } + if (!output) { + output = std::make_shared( + pool, + veloxType, + nulls, + rowCount, + std::move(offsets), + std::move(sizes), + /* keys*/ mapKeys, + /*values*/ mapValues, + 0 /*nullCount*/); + } + return static_cast(output.get()); + } +}; + +template <> +struct VectorInitializer { + static velox::RowVector* initialize( + velox::memory::MemoryPool* pool, + velox::VectorPtr& output, + const velox::TypePtr& veloxType, + uint64_t rowCount) { + auto vector = verifyVectorState(output); + velox::BufferPtr nulls; + std::vector childrenVectors; + if (vector) { + nulls = vector->nulls(); + childrenVectors = vector->children(); + resetIfNotWritable(output, nulls); + } else { + childrenVectors.resize(veloxType->size()); + } + if (!output) { + output = std::make_shared( + pool, + veloxType, + nulls, + rowCount, + std::move(childrenVectors), + 0 /*nullCount*/); + } + return static_cast(output.get()); + } +}; + +class NullColumnReader final : public FieldReader { + public: + NullColumnReader(velox::memory::MemoryPool& pool, velox::TypePtr type) + : FieldReader{pool, std::move(type), nullptr} {} + + void next( + uint32_t count, + velox::VectorPtr& output, + const bits::Bitmap* scatterBitmap) final { + ensureNullConstant(scatterCount(count, scatterBitmap), output, type_); + } + + void skip(uint32_t /* count */) final {} +}; + +class NullFieldReaderFactory final : public FieldReaderFactory { + public: + NullFieldReaderFactory( + velox::memory::MemoryPool& pool, + velox::TypePtr veloxType) + : FieldReaderFactory{pool, std::move(veloxType), nullptr} {} + + std::unique_ptr createReader( + const folly::F14FastMap< + offset_size, + std::unique_ptr>& /* decoders */) final { + return createNullColumnReader(); + } +}; + +template +static auto wrap(T& t) { + return [&]() -> T& { return t; }; +} + +template +struct IsBool : std::false_type {}; + +template <> +struct IsBool : std::true_type {}; + +template +struct ScalarFieldReaderBase; + +template +struct ScalarFieldReaderBase< + TRequested, + TData, + std::enable_if_t::value>> { + explicit ScalarFieldReaderBase(velox::memory::MemoryPool& pool) + : buf_{&pool} {} + + bool* ensureBuffer(uint32_t rowCount) { + buf_.reserve(rowCount); + auto data = buf_.data(); + std::fill(data, data + rowCount, false); + return data; + } + + Vector buf_; +}; + +template +struct ScalarFieldReaderBase< + TRequested, + TData, + std::enable_if_t::value>> { + explicit ScalarFieldReaderBase(velox::memory::MemoryPool& /* pool */) {} +}; + +// TRequested is the requested data type from the reader, TData is the +// data type as stored in the file's schema +template +class ScalarFieldReader final + : public FieldReader, + private ScalarFieldReaderBase { + public: + ScalarFieldReader( + velox::memory::MemoryPool& pool, + velox::TypePtr type, + Decoder* decoder) + : FieldReader(pool, std::move(type), decoder), + ScalarFieldReaderBase{pool} { + if constexpr ( + (isSignedIntegralType() && !isSignedIntegralType() && + !isBoolType()) || + (isUnsignedIntegralType() && + !isUnsignedIntegralType()) || + (isFloatingPointType() && !isFloatingPointType()) || + sizeof(TRequested) < sizeof(TData)) { + NIMBLE_ASSERT(false, "Incompatabile data type and requested type"); + } + } + + using FieldReader::FieldReader; + + void next( + uint32_t count, + velox::VectorPtr& output, + const bits::Bitmap* scatterBitmap) final { + auto rowCount = scatterCount(count, scatterBitmap); + auto vector = VectorInitializer>::initialize( + &pool_, output, type_, rowCount); + vector->resize(rowCount); + + auto upcastNoNulls = [&vector]() { + auto vecRowCount = vector->size(); + if (vecRowCount == 0) { + return; + } + auto to = vector->mutableRawValues(); + const auto from = vector->template rawValues(); + // we can't use for (uint32_t i = vecRowCount - 1; i >= 0; --i) + // for the loop control, because for unsigned int, i >= 0 is always true, + // it becomes an infinite loop + for (uint32_t i = 0; i < vecRowCount; ++i) { + to[vecRowCount - i - 1] = + static_cast(from[vecRowCount - i - 1]); + } + }; + + auto upcastWithNulls = [&vector]() { + auto vecRowCount = vector->size(); + if (vecRowCount == 0) { + return; + } + auto to = vector->mutableRawValues(); + const auto from = vector->template rawValues(); + for (uint32_t i = 0; i < vecRowCount; ++i) { + if (vector->isNullAt(vecRowCount - i - 1)) { + to[vecRowCount - i - 1] = TRequested(); + } else { + to[vecRowCount - i - 1] = + static_cast(from[vecRowCount - i - 1]); + } + } + }; + + uint32_t nonNullCount = 0; + if constexpr (IsBool::value) { + // TODO: implement method for bitpacked bool + auto buf = this->ensureBuffer(rowCount); + nonNullCount = decoder_->next( + count, + buf, + [&]() { return paddedNulls(vector, rowCount); }, + scatterBitmap); + + auto target = vector->mutableValues(rowCount)->template asMutable(); + std::fill(target, target + bits::bytesRequired(rowCount), 0); + for (uint32_t i = 0; i < rowCount; ++i) { + bits::maybeSetBit(i, target, buf[i]); + } + } else { + nonNullCount = decoder_->next( + count, + vector->mutableValues(rowCount)->template asMutable(), + [&]() { return paddedNulls(vector, rowCount); }, + scatterBitmap); + } + + if (nonNullCount == rowCount) { + vector->resetNulls(); + if constexpr (sizeof(TRequested) > sizeof(TData)) { + upcastNoNulls(); + } + } else { + vector->setNullCount(rowCount - nonNullCount); + if constexpr (sizeof(TRequested) > sizeof(TData)) { + upcastWithNulls(); + } + } + } + + void skip(uint32_t count) final { + decoder_->skip(count); + } +}; + +template +class ScalarFieldReaderFactory final : public FieldReaderFactory { + public: + ScalarFieldReaderFactory( + velox::memory::MemoryPool& pool, + velox::TypePtr veloxType, + const Type* type) + : FieldReaderFactory{pool, std::move(veloxType), type} {} + + std::unique_ptr createReader( + const folly::F14FastMap>& decoders) + final { + const auto& descriptor = nimbleType_->asScalar().scalarDescriptor(); + switch (descriptor.scalarKind()) { + case ScalarKind::Bool: { + return createReaderImpl>( + decoders, descriptor); + } + case ScalarKind::Int8: { + return createReaderImpl>( + decoders, descriptor); + } + case ScalarKind::Int16: { + return createReaderImpl>( + decoders, descriptor); + } + case ScalarKind::Int32: { + return createReaderImpl>( + decoders, descriptor); + } + case ScalarKind::Int64: { + return createReaderImpl>( + decoders, descriptor); + } + case ScalarKind::Float: { + return createReaderImpl>( + decoders, descriptor); + } + case ScalarKind::Double: { + return createReaderImpl>( + decoders, descriptor); + } + case ScalarKind::UInt8: + case ScalarKind::UInt16: + case ScalarKind::UInt32: { + return createReaderImpl>( + decoders, descriptor); + } + case ScalarKind::UInt64: + case ScalarKind::String: + case ScalarKind::Binary: + case ScalarKind::Undefined: { + NIMBLE_NOT_SUPPORTED(fmt::format( + "Unsupported nimble scalar type: {}.", + toString(descriptor.scalarKind()))) + } + } + NIMBLE_UNREACHABLE(fmt::format( + "Should not have nimble scalar type: {}.", + toString(descriptor.scalarKind()))) + } +}; + +class StringFieldReader final : public FieldReader { + public: + StringFieldReader( + velox::memory::MemoryPool& pool, + velox::TypePtr type, + Decoder* decoder, + std::vector& buffer) + : FieldReader{pool, std::move(type), decoder}, buffer_{buffer} {} + + void next( + uint32_t count, + velox::VectorPtr& output, + const bits::Bitmap* scatterBitmap) final { + auto rowCount = scatterCount(count, scatterBitmap); + auto vector = + VectorInitializer>::initialize( + &pool_, output, type_, rowCount); + vector->clearStringBuffers(); + vector->resize(rowCount); + buffer_.resize(rowCount); + + auto nonNullCount = decoder_->next( + count, + buffer_.data(), + [&]() { return paddedNulls(vector, rowCount); }, + scatterBitmap); + + if (nonNullCount == rowCount) { + vector->resetNulls(); + for (uint32_t i = 0; i < rowCount; ++i) { + vector->set( + i, + // @lint-ignore CLANGTIDY facebook-hte-MemberUncheckedArrayBounds + {buffer_[i].data(), static_cast(buffer_[i].length())}); + } + } else { + vector->setNullCount(rowCount - nonNullCount); + for (uint32_t i = 0; i < rowCount; ++i) { + if (!vector->isNullAt(i)) { + vector->set( + i, + // @lint-ignore CLANGTIDY facebook-hte-MemberUncheckedArrayBounds + {buffer_[i].data(), static_cast(buffer_[i].length())}); + } + } + } + } + + void skip(uint32_t count) final { + decoder_->skip(count); + } + + private: + std::vector& buffer_; +}; + +class StringFieldReaderFactory final : public FieldReaderFactory { + public: + StringFieldReaderFactory( + velox::memory::MemoryPool& pool, + velox::TypePtr veloxType, + const Type* type) + : FieldReaderFactory{pool, std::move(veloxType), type} {} + + std::unique_ptr createReader( + const folly::F14FastMap>& decoders) + final { + return createReaderImpl( + decoders, nimbleType_->asScalar().scalarDescriptor(), wrap(buffer_)); + } + + private: + std::vector buffer_; +}; + +class MultiValueFieldReader : public FieldReader { + public: + using FieldReader::FieldReader; + + protected: + template + velox::vector_size_t loadOffsets( + uint32_t count, + velox::VectorPtr& output, + const bits::Bitmap* scatterBitmap, + velox::vector_size_t allocationSize = 0, + Args&&... args) { + auto rowCount = scatterCount(count, scatterBitmap); + if (allocationSize == 0) { + allocationSize = rowCount; + } + NIMBLE_ASSERT( + allocationSize >= rowCount, + fmt::format( + "readCount should be less than allocationSize. {} vs {}", + allocationSize, + rowCount)); + + auto vector = VectorInitializer::initialize( + &pool_, output, type_, allocationSize, std::forward(args)...); + vector->resize(allocationSize); + + velox::vector_size_t* sizes = + vector->mutableSizes(allocationSize) + ->template asMutable(); + velox::vector_size_t* offsets = + vector->mutableOffsets(allocationSize) + ->template asMutable(); + + auto nonNullCount = decoder_->next( + count, + sizes, + [&]() { return paddedNulls(vector, allocationSize); }, + scatterBitmap); + + size_t childrenRows = 0; + if (nonNullCount == rowCount) { + vector->resetNulls(); + for (uint32_t i = 0; i < rowCount; ++i) { + offsets[i] = static_cast(childrenRows); + childrenRows += sizes[i]; + } + } else { + vector->setNullCount(rowCount - nonNullCount); + + auto notNulls = reinterpret_cast(vector->rawNulls()); + for (uint32_t i = 0; i < rowCount; ++i) { + offsets[i] = static_cast(childrenRows); + if (bits::getBit(i, notNulls)) { + childrenRows += sizes[i]; + } else { + sizes[i] = 0; + } + } + } + + NIMBLE_CHECK( + childrenRows <= std::numeric_limits::max(), + fmt::format("Unsupported children count: {}", childrenRows)); + return static_cast(childrenRows); + } + + uint32_t skipLengths(uint32_t count) { + size_t childrenCount = 0; + std::array sizes; + + constexpr auto byteSize = nullBytes(kSkipBatchSize); + std::array nulls; + + while (count > 0) { + auto readSize = std::min(count, kSkipBatchSize); + auto nonNullCount = decoder_->next( + readSize, + sizes.data(), + [&]() { return nulls.data(); }, + /* scatterBitmap */ nullptr); + + if (nonNullCount == readSize) { + for (uint32_t i = 0; i < readSize; ++i) { + childrenCount += sizes[i]; + } + } else { + for (uint32_t i = 0; i < readSize; ++i) { + if (bits::getBit(i, nulls.data())) { + childrenCount += sizes[i]; + } + } + } + count -= readSize; + } + + NIMBLE_CHECK( + childrenCount <= std::numeric_limits::max(), + fmt::format("Unsupported children count: {}", childrenCount)); + return static_cast(childrenCount); + } +}; + +class ArrayFieldReader final : public MultiValueFieldReader { + public: + ArrayFieldReader( + velox::memory::MemoryPool& pool, + velox::TypePtr type, + Decoder* decoder, + std::unique_ptr elementsReader) + : MultiValueFieldReader{pool, std::move(type), decoder}, + elementsReader_{std::move(elementsReader)} {} + + void next( + uint32_t count, + velox::VectorPtr& output, + const bits::Bitmap* scatterBitmap) final { + auto childrenRows = this->template loadOffsets( + count, output, scatterBitmap); + + // As the fields are aligned by lengths decoder so no need to pass scatter + // to elements + elementsReader_->next( + childrenRows, + static_cast(*output).elements(), + /* scatterBitmap */ nullptr); + } + + void skip(uint32_t count) final { + auto childrenCount = this->skipLengths(count); + if (childrenCount > 0) { + elementsReader_->skip(childrenCount); + } + } + + void reset() final { + FieldReader::reset(); + elementsReader_->reset(); + } + + private: + std::unique_ptr elementsReader_; +}; + +class ArrayFieldReaderFactory final : public FieldReaderFactory { + public: + // Here the index is the index of the array lengths. + ArrayFieldReaderFactory( + velox::memory::MemoryPool& pool, + velox::TypePtr veloxType, + const Type* type, + std::unique_ptr elements) + : FieldReaderFactory{pool, std::move(veloxType), type}, + elements_{std::move(elements)} {} + + std::unique_ptr createReader( + const folly::F14FastMap>& decoders) + final { + return createReaderImpl( + decoders, nimbleType_->asArray().lengthsDescriptor(), [&]() { + return elements_->createReader(decoders); + }); + } + + private: + std::unique_ptr elements_; +}; + +class ArrayWithOffsetsFieldReader final : public MultiValueFieldReader { + public: + using OffsetType = uint32_t; + + ArrayWithOffsetsFieldReader( + velox::memory::MemoryPool& pool, + velox::TypePtr type, + Decoder* decoder, + Decoder* offsetDecoder, + std::unique_ptr elementsReader) + : MultiValueFieldReader{pool, std::move(type), decoder}, + offsetDecoder_{offsetDecoder}, + elementsReader_{std::move(elementsReader)}, + cached_{false}, + cachedValue_{nullptr}, + cachedIndex_{0}, + cachedSize_{0}, + cachedLazyLoad_{false}, + cachedLazyChildrenRows_{0} { + VectorInitializer::initialize( + &pool_, cachedValue_, type_, 1); + } + + void next( + uint32_t count, + velox::VectorPtr& output, + const bits::Bitmap* scatterBitmap) final { + auto rowCount = scatterCount(count, scatterBitmap); + // read the offsets/indices which is one value per rowCount + // and filter out deduped arrays to be read + uint32_t nonNullCount; + + auto dictionaryVector = + verifyVectorState>(output); + + if (dictionaryVector) { + dictionaryVector->resize(rowCount); + resetIfNotWritable(output, dictionaryVector->indices()); + } else { + velox::VectorPtr child; + VectorInitializer::initialize( + &pool_, child, type_, rowCount); + auto indices = + velox::AlignedBuffer::allocate(rowCount, &pool_); + + // Note: when creating a dictionary vector, it validates the vector (in + // debug builds) for correctness. Therefore, we allocate all the buffers + // above with the right size, but we "resize" them to zero, before + // creating the dictionary vector, to avoid failing this validation. + // We will later resize the vector to the correct size. + // These resize operations are "no cost" operations, as shrinking a + // vector/buffer doesn't free its memory, and resizing to the orginal size + // doesn't allocate, as capacity is guaranteed to be enough. + child->resize(0); + indices->setSize(0); + + output = velox::BaseVector::wrapInDictionary( + /* nulls */ nullptr, + /* indices */ std::move(indices), + /* size */ 0, + /* values */ std::move(child)); + dictionaryVector = + output->as>(); + dictionaryVector->resize(rowCount); + } + + void* nullsPtr = nullptr; + uint32_t dedupCount = getIndicesDeduplicated( + dictionaryVector->indices()->asMutable(), + [&]() { + // The pointer will be initialized ONLY if the data is nullable. + // Otherwise, it will remain nullptr, and this is handled below. + nullsPtr = paddedNulls(dictionaryVector, rowCount); + return nullsPtr; + }, + nonNullCount, + count, + scatterBitmap); + + bool hasNulls = nonNullCount != rowCount; + auto indices = dictionaryVector->indices()->asMutable(); + NIMBLE_DASSERT(indices, "Indices missing."); + + // Returns the first non-null index or -1 (if all are null). + auto baseIndex = findFirstBit(rowCount, hasNulls, nullsPtr, indices); + + bool cachedLocally = rowCount > 0 && cached_ && (baseIndex == cachedIndex_); + + // Initializes sizes and offsets in the vector. + auto childrenRows = loadOffsets( + dedupCount - cachedLocally, + dictionaryVector->valueVector(), + /* scatterBitmap */ nullptr, + dedupCount); + + if (cached_ && cachedLazyLoad_) { + if (cachedLocally) { + elementsReader_->next( + cachedLazyChildrenRows_, + static_cast(*cachedValue_).elements(), + /* scatterBitmap */ nullptr); + } else { + elementsReader_->skip(cachedLazyChildrenRows_); + } + cachedLazyLoad_ = false; + } + + elementsReader_->next( + childrenRows, + static_cast(*dictionaryVector->valueVector()) + .elements(), + /* scatterBitmap */ nullptr); + + if (cachedLocally) { + auto vector = static_cast( + dictionaryVector->valueVector().get()); + + // Copy elements from cache + const auto cacheIdx = static_cast(dedupCount) - 1; + velox::BaseVector::CopyRange cacheRange{ + 0, static_cast(cacheIdx), 1}; + vector->copyRanges(cachedValue_.get(), folly::Range(&cacheRange, 1)); + + // copyRanges overwrites offsets from the source array and must be reset + OffsetType* sizes = + vector->mutableSizes(dedupCount)->template asMutable(); + OffsetType* offsets = + vector->mutableOffsets(dedupCount)->template asMutable(); + + size_t rows = 0; + if (cacheIdx > 0) { + rows = offsets[cacheIdx - 1] + sizes[cacheIdx - 1]; + } + + sizes[cacheIdx] = cachedSize_; + offsets[cacheIdx] = static_cast(rows); + + if (hasNulls) { + vector->setNull(cacheIdx, false); + } + } + + // Cache last item + if (dedupCount > 0) { + const auto& values = dictionaryVector->valueVector(); + auto idxToCache = std::max( + 0, static_cast(dedupCount - 1 - cachedLocally)); + velox::BaseVector::CopyRange cacheRange{ + static_cast(idxToCache), 0, 1}; + + cachedValue_->prepareForReuse(); + cachedValue_->copyRanges(values.get(), folly::Range(&cacheRange, 1)); + + // Get the index for this last element which must be non-null + cachedIndex_ = indices[findLastBit(rowCount, hasNulls, nullsPtr)]; + + cachedSize_ = + static_cast(*values).sizeAt(idxToCache); + cached_ = true; + cachedLazyLoad_ = false; + } + + // normalize the indices if not all null + if (nonNullCount > 0) { + if (hasNulls) { + NIMBLE_DASSERT(nullsPtr, "Nulls buffer missing."); + for (OffsetType idx = 0; idx < rowCount; idx++) { + if (velox::bits::isBitNull( + static_cast(nullsPtr), idx)) { + continue; + } + + indices[idx] = indices[idx] - baseIndex; + } + } else { + for (OffsetType idx = 0; idx < rowCount; idx++) { + indices[idx] = indices[idx] - baseIndex; + } + } + } + + // update the indices as per cached and null locations + if (hasNulls) { + dictionaryVector->setNullCount(nonNullCount != rowCount); + NIMBLE_DASSERT(nullsPtr, "Nulls buffer missing."); + for (OffsetType idx = 0; idx < rowCount; idx++) { + if (velox::bits::isBitNull( + static_cast(nullsPtr), idx)) { + indices[idx] = dedupCount - 1; + } else { + if (indices[idx] == 0 && cachedLocally) { // cached index + indices[idx] = dedupCount - 1; + } else { + indices[idx] -= cachedLocally; + } + } + } + } else { + dictionaryVector->resetNulls(); + for (OffsetType idx = 0; idx < rowCount; idx++) { + if (indices[idx] == 0 && cachedLocally) { // cached index + indices[idx] = dedupCount - 1; + } else { + indices[idx] -= cachedLocally; + } + } + } + } + + void skip(uint32_t count) final { + // read the offsets/indices which is one value per rowCount + // and filter out deduped arrays to be read + std::array indices; + std::array nulls; + void* nullsPtr = nulls.data(); + uint32_t nonNullCount; + + while (count > 0) { + auto batchedRowCount = std::min(count, kSkipBatchSize); + uint32_t dedupCount = getIndicesDeduplicated( + indices.data(), + [&]() { return nullsPtr; }, + nonNullCount, + batchedRowCount); + + bool hasNulls = nonNullCount != batchedRowCount; + + // baseIndex is the first non-null index + auto baseIndex = + findFirstBit(batchedRowCount, hasNulls, nullsPtr, indices.data()); + + bool cachedLocally = cached_ && (baseIndex == cachedIndex_); + if (cachedLocally) { + dedupCount--; + } + + // skip all the children except the last one + if (dedupCount > 0) { + auto childrenRows = + cached_ && cachedLazyLoad_ ? cachedLazyChildrenRows_ : 0; + childrenRows += this->skipLengths(dedupCount - 1); + if (childrenRows > 0) { + elementsReader_->skip(childrenRows); + } + + /// cache the last child + + // get the index for this last element which must be non-null + cachedIndex_ = + indices[findLastBit(batchedRowCount, hasNulls, nullsPtr)]; + cached_ = true; + cachedLazyLoad_ = true; + cachedLazyChildrenRows_ = + loadOffsets(1, cachedValue_, nullptr); + + cachedSize_ = static_cast(*cachedValue_).sizeAt(0); + } + + count -= batchedRowCount; + } + } + + void reset() final { + FieldReader::reset(); + offsetDecoder_->reset(); + cached_ = false; + elementsReader_->reset(); + } + + private: + Decoder* offsetDecoder_; + std::unique_ptr elementsReader_; + bool cached_; + velox::VectorPtr cachedValue_; + OffsetType cachedIndex_; + uint32_t cachedSize_; + bool cachedLazyLoad_; + uint32_t cachedLazyChildrenRows_; + + static inline OffsetType findLastBit( + uint32_t rowCount, + bool hasNulls, + const void* FOLLY_NULLABLE nulls) { + if (!hasNulls) { + return rowCount - 1; + } + + NIMBLE_DASSERT(nulls, "Nulls buffer missing."); + auto index = velox::bits::findLastBit( + static_cast(nulls), 0, rowCount); + if (index == -1) { + return rowCount - 1; + } + + return index; + } + + static inline int32_t findFirstBit( + uint32_t rowCount, + bool hasNulls, + const void* FOLLY_NULLABLE nulls, + const OffsetType* indices) { + if (!hasNulls) { + return indices[0]; + } + + NIMBLE_DASSERT(nulls, "Nulls buffer missing."); + auto index = velox::bits::findFirstBit( + static_cast(nulls), 0, rowCount); + + if (index == -1) { + return -1; + } + + return indices[index]; + } + + uint32_t getIndicesDeduplicated( + OffsetType* indices, + std::function nulls, + uint32_t& nonNullCount, + uint32_t count, + const bits::Bitmap* scatterBitmap = nullptr) { + auto rowCount = scatterCount(count, scatterBitmap); + // OffsetType* indices = dictIndices->asMutable(); + void* nullsPtr; + + nonNullCount = offsetDecoder_->next( + count, + indices, + [&]() { + nullsPtr = nulls(); + return nullsPtr; + }, + scatterBitmap); + + // remove duplicated indices and calculate unique count + uint32_t uniqueCount = 0; + uint32_t prevIdx = 0; + bool hasNulls = nonNullCount != rowCount; + + if (hasNulls) { + NIMBLE_DASSERT( + nullsPtr != nullptr, + "Data contain nulls but nulls buffer is not initialized."); + + for (uint32_t idx = 0; idx < rowCount; idx++) { + if (velox::bits::isBitNull( + static_cast(nullsPtr), idx)) { + indices[idx] = 0; + continue; + } + + if (uniqueCount == 0 || indices[idx] != indices[prevIdx]) { + uniqueCount++; + } + prevIdx = idx; + } + } else { + for (uint32_t idx = 0; idx < rowCount; idx++) { + if (uniqueCount == 0 || indices[idx] != indices[prevIdx]) { + uniqueCount++; + } + prevIdx = idx; + } + } + + return uniqueCount; + } +}; + +class ArrayWithOffsetsFieldReaderFactory final : public FieldReaderFactory { + public: + // Here the index is the index of the array lengths. + ArrayWithOffsetsFieldReaderFactory( + velox::memory::MemoryPool& pool, + velox::TypePtr veloxType, + const Type* type, + std::unique_ptr elements) + : FieldReaderFactory{pool, std::move(veloxType), type}, + elements_{std::move(elements)} {} + + std::unique_ptr createReader( + const folly::F14FastMap>& decoders) + final { + return createReaderImpl( + decoders, + nimbleType_->asArrayWithOffsets().lengthsDescriptor(), + [&]() { + return getDecoder( + decoders, nimbleType_->asArrayWithOffsets().offsetsDescriptor()); + }, + [&]() { return elements_->createReader(decoders); }); + } + + private: + std::unique_ptr elements_; +}; + +class MapFieldReader final : public MultiValueFieldReader { + public: + MapFieldReader( + velox::memory::MemoryPool& pool, + velox::TypePtr type, + Decoder* decoder, + std::unique_ptr keysReader, + std::unique_ptr valuesReader) + : MultiValueFieldReader{pool, std::move(type), decoder}, + keysReader_{std::move(keysReader)}, + valuesReader_{std::move(valuesReader)} {} + + void next( + uint32_t count, + velox::VectorPtr& output, + const bits::Bitmap* scatterBitmap) final { + auto childrenRows = this->template loadOffsets( + count, output, scatterBitmap); + + // As the field is aligned by lengths decoder then no need to pass + // scatterBitmap to keys and values + auto& mapVector = static_cast(*output); + keysReader_->next( + childrenRows, mapVector.mapKeys(), /* scatterBitmap */ nullptr); + valuesReader_->next( + childrenRows, mapVector.mapValues(), /* scatterBitmap */ nullptr); + } + + void skip(uint32_t count) final { + auto childrenCount = this->skipLengths(count); + if (childrenCount > 0) { + keysReader_->skip(childrenCount); + valuesReader_->skip(childrenCount); + } + } + + void reset() final { + FieldReader::reset(); + keysReader_->reset(); + valuesReader_->reset(); + } + + private: + std::unique_ptr keysReader_; + std::unique_ptr valuesReader_; +}; + +class MapFieldReaderFactory final : public FieldReaderFactory { + public: + // Here the index is the index of the array lengths. + MapFieldReaderFactory( + velox::memory::MemoryPool& pool, + velox::TypePtr veloxType, + const Type* type, + std::unique_ptr keys, + std::unique_ptr values) + : FieldReaderFactory{pool, std::move(veloxType), type}, + keys_{std::move(keys)}, + values_{std::move(values)} {} + + std::unique_ptr createReader( + const folly::F14FastMap>& decoders) + final { + return createReaderImpl( + decoders, + nimbleType_->asMap().lengthsDescriptor(), + [&]() { return keys_->createReader(decoders); }, + [&]() { return values_->createReader(decoders); }); + } + + private: + std::unique_ptr keys_; + std::unique_ptr values_; +}; + +// Read values from boolean decoder and return number of true values. +template +uint32_t readBooleanValues( + Decoder* decoder, + bool* buffer, + uint32_t count, + TrueHandler handler) { + decoder->next(count, buffer); + + uint32_t trueCount = 0; + for (uint32_t i = 0; i < count; ++i) { + if (buffer[i]) { + handler(i); + ++trueCount; + } + } + return trueCount; +} + +uint32_t readBooleanValues(Decoder* decoder, bool* buffer, uint32_t count) { + return readBooleanValues(decoder, buffer, count, [](auto /* ignored */) {}); +} + +template +class RowFieldReader final : public FieldReader { + public: + RowFieldReader( + velox::memory::MemoryPool& pool, + velox::TypePtr type, + Decoder* decoder, + std::vector> childrenReaders, + Vector& boolBuffer, + folly::Executor* executor) + : FieldReader{pool, std::move(type), decoder}, + childrenReaders_{std::move(childrenReaders)}, + boolBuffer_{boolBuffer}, + executor_{executor} {} + + void next( + uint32_t count, + velox::VectorPtr& output, + const bits::Bitmap* scatterBitmap) final { + auto rowCount = scatterCount(count, scatterBitmap); + auto vector = VectorInitializer::initialize( + &pool_, output, type_, rowCount); + vector->children().resize(childrenReaders_.size()); + vector->unsafeResize(rowCount); + const void* childrenBits = nullptr; + uint32_t childrenCount = 0; + + if constexpr (hasNull) { + zeroNulls(vector, rowCount); + // if it is a scattered read case then we can't read the rowCount + // values from the nulls, we count the set value in scatterBitmap and + // read only those values, if there is no scatter then we can read + // rowCount values + boolBuffer_.resize(count); + decoder_->next(count, boolBuffer_.data()); + + auto* nullBuffer = paddedNulls(vector, rowCount); + bits::BitmapBuilder nullBits{nullBuffer, rowCount}; + if (scatterBitmap) { + uint32_t boolBufferOffset = 0; + for (uint32_t i = 0; i < rowCount; ++i) { + if (scatterBitmap->test(i) && boolBuffer_[boolBufferOffset++]) { + nullBits.set(i); + ++childrenCount; + } + } + } else { + for (uint32_t i = 0; i < rowCount; ++i) { + if (boolBuffer_[i]) { + nullBits.set(i); + ++childrenCount; + } + } + } + if (UNLIKELY(childrenCount == rowCount)) { + vector->resetNulls(); + } else { + vector->setNullCount(rowCount - childrenCount); + childrenBits = nullBuffer; + } + } else { + childrenCount = count; + if (scatterBitmap) { + auto requiredBytes = bits::bytesRequired(rowCount); + auto* nullBuffer = paddedNulls(vector, rowCount); + // @lint-ignore CLANGSECURITY facebook-security-vulnerable-memcpy + std::memcpy( + nullBuffer, + static_cast(scatterBitmap->bits()), + requiredBytes); + vector->setNullCount(rowCount - count); + childrenBits = scatterBitmap->bits(); + } else { + vector->resetNulls(); + } + } + + bits::Bitmap childrenBitmap{childrenBits, rowCount}; + auto bitmapPtr = childrenBits ? &childrenBitmap : nullptr; + + if (executor_) { + for (uint32_t i = 0; i < childrenReaders_.size(); ++i) { + auto& reader = childrenReaders_[i]; + if (reader) { + executor_->add([childrenCount, + bitmapPtr, + &reader, + &child = vector->childAt(i)]() { + reader->next(childrenCount, child, bitmapPtr); + }); + } + } + } else { + for (uint32_t i = 0; i < childrenReaders_.size(); ++i) { + auto& reader = childrenReaders_[i]; + if (reader) { + reader->next(childrenCount, vector->childAt(i), bitmapPtr); + } + } + } + } + + void skip(uint32_t count) final { + uint32_t childRowCount = count; + if constexpr (hasNull) { + std::array buffer; + childRowCount = 0; + while (count > 0) { + auto readSize = std::min(count, kSkipBatchSize); + childRowCount += readBooleanValues(decoder_, buffer.data(), readSize); + count -= readSize; + } + } + + if (childRowCount > 0) { + for (auto& reader : childrenReaders_) { + if (reader) { + reader->skip(childRowCount); + } + } + } + } + + void reset() final { + FieldReader::reset(); + for (auto& reader : childrenReaders_) { + if (reader) { + reader->reset(); + } + } + } + + private: + std::vector> childrenReaders_; + Vector& boolBuffer_; + folly::Executor* executor_; +}; + +class RowFieldReaderFactory final : public FieldReaderFactory { + public: + // Here the index is the index of the null decoder. + RowFieldReaderFactory( + velox::memory::MemoryPool& pool, + velox::TypePtr veloxType, + const Type* type, + std::vector> children, + folly::Executor* executor) + : FieldReaderFactory{pool, std::move(veloxType), type}, + children_{std::move(children)}, + boolBuffer_{&pool_}, + executor_{executor} {} + + std::unique_ptr createReader( + const folly::F14FastMap>& decoders) + final { + auto nulls = getDecoder(decoders, nimbleType_->asRow().nullsDescriptor()); + + std::vector> childrenReaders(children_.size()); + for (uint32_t i = 0; i < children_.size(); ++i) { + auto& child = children_[i]; + if (child) { + // @lint-ignore CLANGTIDY facebook-hte-LocalUncheckedArrayBounds + childrenReaders[i] = child->createReader(decoders); + } + } + + if (!nulls) { + return std::make_unique>( + pool_, + veloxType_, + nulls, + std::move(childrenReaders), + boolBuffer_, + executor_); + } + + return std::make_unique>( + pool_, + veloxType_, + nulls, + std::move(childrenReaders), + boolBuffer_, + executor_); + } + + private: + std::vector> children_; + Vector boolBuffer_; + folly::Executor* executor_; +}; + +// Represent a keyed value node for flat map +// Before reading the value, InMap vectors we need to call load() +template +class FlatMapKeyNode { + public: + FlatMapKeyNode( + velox::memory::MemoryPool& memoryPool, + std::unique_ptr valueReader, + Decoder* inMapDecoder, + const velox::dwio::common::flatmap::KeyValue& key) + : valueReader_{std::move(valueReader)}, + inMapDecoder_{inMapDecoder}, + inMapData_{&memoryPool}, + key_{key}, + mergedNulls_{&memoryPool} {} + + ~FlatMapKeyNode() = default; + + void readAsChild( + velox::VectorPtr& vector, + uint32_t numValues, + uint32_t nonNullValues, + const Vector& mapNulls, + Vector* mergedNulls = nullptr) { + if (!mergedNulls) { + mergedNulls = &mergedNulls_; + } + auto nonNullCount = + mergeNulls(numValues, nonNullValues, mapNulls, *mergedNulls); + bits::Bitmap bitmap{mergedNulls->data(), numValues}; + valueReader_->next(nonNullCount, vector, &bitmap); + NIMBLE_DCHECK(numValues == vector->size(), "Items not loaded"); + } + + uint32_t readInMapData(uint32_t numValues) { + inMapData_.resize(numValues); + numValues_ = readBooleanValues(inMapDecoder_, inMapData_.data(), numValues); + return numValues_; + } + + void loadValues(velox::VectorPtr& values) { + valueReader_->next(numValues_, values, /* scatterBitmap */ nullptr); + NIMBLE_DCHECK(numValues_ == values->size(), "Items not loaded"); + } + + void skip(uint32_t numValues) { + auto numItems = readInMapData(numValues); + if (numItems > 0) { + valueReader_->skip(numItems); + } + } + + const velox::dwio::common::flatmap::KeyValue& key() const { + return key_; + } + + bool inMap(uint32_t index) const { + return inMapData_[index]; + } + + void reset() { + inMapDecoder_->reset(); + valueReader_->reset(); + } + + private: + // Merge the mapNulls and inMapData into mergedNulls + uint32_t mergeNulls( + uint32_t numValues, + uint32_t nonNullMaps, + const Vector& mapNulls, + Vector& mergedNulls) { + const auto numItems = readInMapData(nonNullMaps); + auto requiredBytes = bits::bytesRequired(numValues); + mergedNulls.resize(requiredBytes); + memset(mergedNulls.data(), 0, requiredBytes); + if (numItems == 0) { + return 0; + } + + if (nonNullMaps == numValues) { + // All values are nonNull + bits::packBitmap(inMapData_, mergedNulls.data()); + return numItems; + } + uint32_t inMapOffset = 0; + for (uint32_t i = 0; i < numValues; ++i) { + if (mapNulls[i] && inMapData_[inMapOffset++]) { + bits::setBit(i, mergedNulls.data()); + } + } + return numItems; + } + + std::unique_ptr valueReader_; + Decoder* inMapDecoder_; + Vector inMapData_; + const velox::dwio::common::flatmap::KeyValue& key_; + uint32_t numValues_; + // nulls buffer used in parallel read cases. + Vector mergedNulls_; +}; + +template +class FlatMapFieldReaderBase : public FieldReader { + public: + FlatMapFieldReaderBase( + velox::memory::MemoryPool& pool, + velox::TypePtr type, + Decoder* decoder, + std::vector>> keyNodes, + Vector& boolBuffer) + : FieldReader{pool, std::move(type), decoder}, + keyNodes_{std::move(keyNodes)}, + boolBuffer_{boolBuffer} {} + + uint32_t loadNulls(uint32_t rowCount, velox::BaseVector* vector) { + if constexpr (hasNull) { + zeroNulls(vector, rowCount); + auto* nullBuffer = paddedNulls(vector, rowCount); + bits::BitmapBuilder bitmap{nullBuffer, rowCount}; + + boolBuffer_.resize(rowCount); + auto nonNullCount = readBooleanValues( + decoder_, boolBuffer_.data(), rowCount, [&](auto i) { + bitmap.set(i); + }); + + if (UNLIKELY(nonNullCount == rowCount)) { + vector->resetNulls(); + } else { + vector->setNullCount(rowCount - nonNullCount); + } + return nonNullCount; + } else { + vector->resetNulls(); + return rowCount; + } + } + + void skip(uint32_t count) final { + uint32_t nonNullCount = count; + + if constexpr (hasNull) { + std::array buffer; + nonNullCount = 0; + while (count > 0) { + auto readSize = std::min(count, kSkipBatchSize); + nonNullCount += readBooleanValues(decoder_, buffer.data(), readSize); + count -= readSize; + } + } + + if (nonNullCount > 0) { + for (auto& node : keyNodes_) { + if (node) { + node->skip(nonNullCount); + } + } + } + } + + void reset() final { + FieldReader::reset(); + for (auto& node : keyNodes_) { + if (node) { + node->reset(); + } + } + } + + protected: + std::vector>> keyNodes_; + Vector& boolBuffer_; +}; + +template +class FlatMapFieldReaderFactoryBase : public FieldReaderFactory { + public: + FlatMapFieldReaderFactoryBase( + velox::memory::MemoryPool& pool, + velox::TypePtr veloxType, + const Type* type, + std::vector inMapDescriptors, + std::vector> valueReaders, + const std::vector& selectedChildren) + : FieldReaderFactory{pool, std::move(veloxType), type}, + inMapDescriptors_{std::move(inMapDescriptors)}, + valueReaders_{std::move(valueReaders)}, + boolBuffer_{&pool_} { + // inMapTypes contains all projected children, including those that don't + // exist in the schema. selectedChildren and valuesReaders only contain + // those that also exist in the schema. + NIMBLE_ASSERT( + inMapDescriptors_.size() >= valueReaders_.size(), + "Value and inMaps size mismatch!"); + NIMBLE_ASSERT( + selectedChildren.size() == valueReaders_.size(), + "Selected children and value readers size mismatch!"); + + auto& flatMap = type->asFlatMap(); + keyValues_.reserve(selectedChildren.size()); + for (auto childIdx : selectedChildren) { + keyValues_.push_back(velox::dwio::common::flatmap::parseKeyValue( + flatMap.nameAt(childIdx))); + } + } + + template < + template + typename ReaderT, + bool includeMissing, + typename... Args> + std::unique_ptr createFlatMapReader( + const folly::F14FastMap>& decoders, + Args&&... args) { + auto nulls = + getDecoder(decoders, nimbleType_->asFlatMap().nullsDescriptor()); + + std::vector>> keyNodes; + keyNodes.reserve(valueReaders_.size()); + uint32_t childIdx = 0; + for (auto inMapDescriptor : inMapDescriptors_) { + if (inMapDescriptor) { + auto currentIdx = childIdx++; + if (auto decoder = getDecoder(decoders, *inMapDescriptor)) { + keyNodes.push_back(std::make_unique>( + pool_, + // @lint-ignore CLANGTIDY facebook-hte-MemberUncheckedArrayBounds + valueReaders_[currentIdx]->createReader(decoders), + decoder, + // @lint-ignore CLANGTIDY facebook-hte-MemberUncheckedArrayBounds + keyValues_[currentIdx])); + continue; + } + } + + if constexpr (includeMissing) { + keyNodes.push_back(nullptr); + } + } + + if (!nulls) { + return std::make_unique>( + pool_, + this->veloxType_, + nulls, + std::move(keyNodes), + boolBuffer_, + std::forward(args)...); + } + + return std::make_unique>( + pool_, + this->veloxType_, + nulls, + std::move(keyNodes), + boolBuffer_, + std::forward(args)...); + } + + protected: + std::vector inMapDescriptors_; + std::vector> valueReaders_; + std::vector> keyValues_; + Vector boolBuffer_; +}; + +template +class StructFlatMapFieldReader : public FlatMapFieldReaderBase { + public: + StructFlatMapFieldReader( + velox::memory::MemoryPool& pool, + velox::TypePtr type, + Decoder* decoder, + std::vector>> keyNodes, + Vector& boolBuffer, + Vector& mergedNulls, + folly::Executor* executor) + : FlatMapFieldReaderBase( + pool, + std::move(type), + decoder, + std::move(keyNodes), + boolBuffer), + mergedNulls_{mergedNulls}, + executor_{executor} {} + + void next( + uint32_t rowCount, + velox::VectorPtr& output, + const bits::Bitmap* scatterBitmap) final { + NIMBLE_ASSERT(scatterBitmap == nullptr, "unexpected scatterBitmap"); + auto vector = VectorInitializer::initialize( + &this->pool_, output, this->type_, rowCount); + vector->unsafeResize(rowCount); + uint32_t nonNullCount = this->loadNulls(rowCount, vector); + + if (executor_) { + for (uint32_t i = 0; i < this->keyNodes_.size(); ++i) { + if (this->keyNodes_[i] == nullptr) { + this->ensureNullConstant( + rowCount, vector->childAt(i), this->type_->childAt(i)); + } else { + executor_->add([this, + rowCount, + nonNullCount, + &node = this->keyNodes_[i], + &child = vector->childAt(i)] { + node->readAsChild(child, rowCount, nonNullCount, this->boolBuffer_); + }); + } + } + } else { + for (uint32_t i = 0; i < this->keyNodes_.size(); ++i) { + if (this->keyNodes_[i] == nullptr) { + this->ensureNullConstant( + rowCount, vector->childAt(i), this->type_->childAt(i)); + } else { + this->keyNodes_[i]->readAsChild( + vector->childAt(i), + rowCount, + nonNullCount, + this->boolBuffer_, + &mergedNulls_); + } + } + } + } + + private: + Vector& mergedNulls_; + folly::Executor* executor_; +}; + +template +class StructFlatMapFieldReaderFactory final + : public FlatMapFieldReaderFactoryBase { + template + using ReaderType = StructFlatMapFieldReader; + + public: + StructFlatMapFieldReaderFactory( + velox::memory::MemoryPool& pool, + velox::TypePtr veloxType, + const Type* type, + std::vector inMapDescriptors, + std::vector> valueReaders, + const std::vector& selectedChildren, + folly::Executor* executor) + : FlatMapFieldReaderFactoryBase( + pool, + std::move(veloxType), + type, + std::move(inMapDescriptors), + std::move(valueReaders), + selectedChildren), + mergedNulls_{&this->pool_}, + executor_{executor} { + NIMBLE_ASSERT(this->nimbleType_->isFlatMap(), "Type should be a flat map."); + } + + std::unique_ptr createReader( + const folly::F14FastMap>& decoders) + final { + return this->template createFlatMapReader( + decoders, mergedNulls_, executor_); + } + + private: + Vector mergedNulls_; + folly::Executor* executor_; +}; + +template +class MergedFlatMapFieldReader final + : public FlatMapFieldReaderBase { + public: + MergedFlatMapFieldReader( + velox::memory::MemoryPool& pool, + velox::TypePtr type, + Decoder* decoder, + std::vector>> keyNodes, + Vector& boolBuffer) + : FlatMapFieldReaderBase( + pool, + std::move(type), + decoder, + std::move(keyNodes), + boolBuffer) {} + + void next( + uint32_t rowCount, + velox::VectorPtr& output, + const bits::Bitmap* scatterBitmap) final { + NIMBLE_ASSERT(scatterBitmap == nullptr, "unexpected scatterBitmap"); + auto vector = VectorInitializer::initialize( + &this->pool_, output, this->type_, rowCount); + vector->resize(rowCount); + velox::VectorPtr& keysVector = vector->mapKeys(); + // Check the refCount for key vector + auto flatKeysVector = VectorInitializer>::initialize( + &this->pool_, + keysVector, + std::static_pointer_cast(this->type_)->keyType(), + rowCount); + + velox::BufferPtr offsets = vector->mutableOffsets(rowCount); + velox::BufferPtr lengths = vector->mutableSizes(rowCount); + uint32_t nonNullCount = this->loadNulls(rowCount, vector); + + nodes_.clear(); + size_t totalChildren = 0; + for (auto& node : this->keyNodes_) { + auto numValues = node->readInMapData(nonNullCount); + if (numValues > 0) { + nodes_.push_back(node.get()); + totalChildren += numValues; + } + } + + velox::VectorPtr nodeValues; + velox::VectorPtr& valuesVector = vector->mapValues(); + if (totalChildren > 0) { + keysVector->resize(totalChildren, false); + velox::BaseVector::prepareForReuse(valuesVector, totalChildren); + } + + auto* offsetsPtr = offsets->asMutable(); + auto* lengthsPtr = lengths->asMutable(); + initRowWiseInMap(rowCount); + initOffsets(rowCount, offsetsPtr, lengthsPtr); + + // Always access inMap and value streams node-wise to avoid large striding + // through the memory and destroying CPU cache performance. + // + // Index symbology used in this class: + // i : Row index + // j : Node index + for (size_t j = 0; j < nodes_.size(); ++j) { + copyRanges_.clear(); + for (velox::vector_size_t i = 0; i < rowCount; ++i) { + if (!velox::bits::isBitSet( + rowWiseInMap_.data(), j + i * nodes_.size())) { + continue; + } + const velox::vector_size_t sourceIndex = copyRanges_.size(); + copyRanges_.push_back({sourceIndex, offsetsPtr[i], 1}); + flatKeysVector->set(offsetsPtr[i], nodes_[j]->key().get()); + ++offsetsPtr[i]; + } + nodes_[j]->loadValues(nodeValues); + valuesVector->copyRanges(nodeValues.get(), copyRanges_); + } + if (rowCount > 0) { + NIMBLE_ASSERT( + offsetsPtr[rowCount - 1] == totalChildren, + "Total map entry size mismatch"); + // We updated `offsetsPtr' during the copy process, so that now it was + // shifted to the left by 1 element (i.e. offsetsPtr[i] is really + // offsetsPtr[i+1]). Need to restore the values back to their correct + // positions. + std::copy_backward( + offsetsPtr, offsetsPtr + rowCount - 1, offsetsPtr + rowCount); + offsetsPtr[0] = 0; + } + + // Reset the updated value vector to result + vector->setKeysAndValues(std::move(keysVector), std::move(valuesVector)); + } + + private: + void initRowWiseInMap(velox::vector_size_t rowCount) { + rowWiseInMap_.resize(velox::bits::nwords(nodes_.size() * rowCount)); + std::fill(rowWiseInMap_.begin(), rowWiseInMap_.end(), 0); + for (size_t j = 0; j < nodes_.size(); ++j) { + uint32_t inMapIndex = 0; + for (velox::vector_size_t i = 0; i < rowCount; ++i) { + const bool isNull = hasNull && !this->boolBuffer_[i]; + if (!isNull && nodes_[j]->inMap(inMapIndex)) { + velox::bits::setBit(rowWiseInMap_.data(), j + i * nodes_.size()); + } + inMapIndex += !isNull; + } + } + } + + void initOffsets( + velox::vector_size_t rowCount, + velox::vector_size_t* offsets, + velox::vector_size_t* lengths) { + velox::vector_size_t offset = 0; + for (velox::vector_size_t i = 0; i < rowCount; ++i) { + offsets[i] = offset; + lengths[i] = velox::bits::countBits( + rowWiseInMap_.data(), i * nodes_.size(), (i + 1) * nodes_.size()); + offset += lengths[i]; + } + } + + // All the nodes that is selected to be read. + std::vector*> nodes_; + + // In-map mask (1 bit per value), organized in row first layout. + std::vector rowWiseInMap_; + + // Copy ranges from one node values into the merged values. + std::vector copyRanges_; +}; + +template +class MergedFlatMapFieldReaderFactory final + : public FlatMapFieldReaderFactoryBase { + template + using ReaderType = MergedFlatMapFieldReader; + + public: + using FlatMapFieldReaderFactoryBase::FlatMapFieldReaderFactoryBase; + + std::unique_ptr createReader( + const folly::F14FastMap>& decoders) + final { + return this->template createFlatMapReader(decoders); + } +}; + +std::unique_ptr createFlatMapReaderFactory( + velox::memory::MemoryPool& pool, + velox::TypeKind keyKind, + velox::TypePtr veloxType, + const Type* type, + std::vector inMapDescriptors, + std::vector> valueReaders, + const std::vector& selectedChildren, + bool flatMapAsStruct, + folly::Executor* executor) { + switch (keyKind) { +#define SCALAR_CASE(veloxKind, fieldType) \ + case velox::TypeKind::veloxKind: { \ + if (flatMapAsStruct) { \ + return std::make_unique>( \ + pool, \ + std::move(veloxType), \ + type, \ + std::move(inMapDescriptors), \ + std::move(valueReaders), \ + selectedChildren, \ + executor); \ + } else { \ + return std::make_unique>( \ + pool, \ + std::move(veloxType), \ + type, \ + std::move(inMapDescriptors), \ + std::move(valueReaders), \ + selectedChildren); \ + } \ + } + + SCALAR_CASE(TINYINT, int8_t); + SCALAR_CASE(SMALLINT, int16_t); + SCALAR_CASE(INTEGER, int32_t); + SCALAR_CASE(BIGINT, int64_t); + SCALAR_CASE(VARCHAR, velox::StringView); + SCALAR_CASE(VARBINARY, velox::StringView); +#undef SCALAR_CASE + + default: + NIMBLE_NOT_SUPPORTED( + fmt::format("Not supported flatmap key type: {} ", keyKind)); + } +} + +std::shared_ptr createFlatType( + const std::vector& selectedFeatures, + const velox::TypePtr& veloxType) { + NIMBLE_ASSERT( + !selectedFeatures.empty(), + "Empty feature selection not allowed for struct encoding."); + + auto& valueType = veloxType->asMap().valueType(); + return velox::ROW( + std::vector(selectedFeatures), + std::vector>( + selectedFeatures.size(), valueType)); +} + +velox::TypePtr inferType( + const FieldReaderParams& params, + const std::string& name, + const velox::TypePtr& type, + size_t level) { + // Special case for flatmaps. If the flatmap field is missing, still need to + // honor the "as struct" intent by returning row instead of map. + if (level == 1 && params.readFlatMapFieldAsStruct.count(name) > 0) { + NIMBLE_CHECK( + type->kind() == velox::TypeKind::MAP, + "Unexpected type kind of flat maps."); + auto it = params.flatMapFeatureSelector.find(name); + NIMBLE_CHECK( + it != params.flatMapFeatureSelector.end() && + !it->second.features.empty(), + fmt::format( + "Flat map feature selection for map '{}' has empty feature set.", + name)); + NIMBLE_CHECK( + it->second.mode == SelectionMode::Include, + "Flta map exclusion list is not supported when flat map field is missing."); + + return createFlatType(it->second.features, type); + } + return type; +} + +std::unique_ptr createFieldReaderFactory( + const FieldReaderParams& parameters, + velox::memory::MemoryPool& pool, + const std::shared_ptr& nimbleType, + const std::shared_ptr& veloxType, + std::vector& offsets, + const std::function& isSelected, + folly::Executor* executor, + size_t level = 0, + const std::string* name = nullptr) { + auto veloxKind = veloxType->type()->kind(); + // compatibleKinds are the types that can be upcasted to nimbleType + auto checkType = [&nimbleType]( + const std::vector& compatibleKinds) { + return std::any_of( + compatibleKinds.begin(), + compatibleKinds.end(), + [&nimbleType](ScalarKind k) { + return nimbleType->asScalar().scalarDescriptor().scalarKind() == k; + }); + }; + +// Assuming no-upcastingg is the most common case, putting the largest type size +// at begging so the compatibility check can finish quicker. +#define BOOLEAN_COMPATIBLE \ + { ScalarKind::Bool } +#define TINYINT_COMPATIBLE \ + { ScalarKind::Int8, ScalarKind::Bool } +#define SMALLINT_COMPATIBLE \ + { ScalarKind::Int16, ScalarKind::Int8, ScalarKind::Bool } +#define INTEGER_COMPATIBLE \ + { ScalarKind::Int32, ScalarKind::Int16, ScalarKind::Int8, ScalarKind::Bool } +#define BIGINT_COMPATIBLE \ + { \ + ScalarKind::Int64, ScalarKind::Int32, ScalarKind::Int16, ScalarKind::Int8, \ + ScalarKind::Bool \ + } +#define FLOAT_COMPATIBLE \ + { ScalarKind::Float } +#define DOUBLE_COMPATIBLE \ + { ScalarKind::Double, ScalarKind::Float } + + switch (veloxKind) { +#define SCALAR_CASE(veloxKind, cppType, compitableKinds) \ + case velox::TypeKind::veloxKind: { \ + NIMBLE_CHECK( \ + nimbleType->isScalar() && checkType(compitableKinds), \ + "Provided schema doesn't match file schema."); \ + offsets.push_back(nimbleType->asScalar().scalarDescriptor().offset()); \ + return std::make_unique>( \ + pool, veloxType->type(), nimbleType.get()); \ + } + + SCALAR_CASE(BOOLEAN, bool, BOOLEAN_COMPATIBLE); + SCALAR_CASE(TINYINT, int8_t, TINYINT_COMPATIBLE); + SCALAR_CASE(SMALLINT, int16_t, SMALLINT_COMPATIBLE); + SCALAR_CASE(INTEGER, int32_t, INTEGER_COMPATIBLE); + SCALAR_CASE(BIGINT, int64_t, BIGINT_COMPATIBLE); + SCALAR_CASE(REAL, float, FLOAT_COMPATIBLE); + SCALAR_CASE(DOUBLE, double, DOUBLE_COMPATIBLE); +#undef SCALAR_CASE + + case velox::TypeKind::VARCHAR: + case velox::TypeKind::VARBINARY: { + NIMBLE_CHECK( + nimbleType->isScalar() && + (veloxKind == velox::TypeKind::VARCHAR && + nimbleType->asScalar().scalarDescriptor().scalarKind() == + ScalarKind::String || + veloxKind == velox::TypeKind::VARBINARY && + nimbleType->asScalar().scalarDescriptor().scalarKind() == + ScalarKind::Binary), + "Provided schema doesn't match file schema."); + offsets.push_back(nimbleType->asScalar().scalarDescriptor().offset()); + return std::make_unique( + pool, veloxType->type(), nimbleType.get()); + } + case velox::TypeKind::ARRAY: { + NIMBLE_CHECK( + nimbleType->isArray() || nimbleType->isArrayWithOffsets(), + "Provided schema doesn't match file schema."); + NIMBLE_ASSERT( + veloxType->size() == 1, + "Velox array type should have exactly one child."); + if (nimbleType->isArray()) { + auto& nimbleArray = nimbleType->asArray(); + auto& elementType = veloxType->childAt(0); + offsets.push_back(nimbleArray.lengthsDescriptor().offset()); + auto elements = isSelected(elementType->id()) + ? createFieldReaderFactory( + parameters, + pool, + nimbleArray.elements(), + elementType, + offsets, + isSelected, + executor, + level + 1) + : std::make_unique( + pool, elementType->type()); + return std::make_unique( + pool, veloxType->type(), nimbleType.get(), std::move(elements)); + } else { + auto& nimbleArrayWithOffsets = nimbleType->asArrayWithOffsets(); + offsets.push_back(nimbleArrayWithOffsets.lengthsDescriptor().offset()); + offsets.push_back(nimbleArrayWithOffsets.offsetsDescriptor().offset()); + + auto& elementType = veloxType->childAt(0); + auto elements = isSelected(elementType->id()) + ? createFieldReaderFactory( + parameters, + pool, + nimbleArrayWithOffsets.elements(), + elementType, + offsets, + isSelected, + executor, + level + 1) + : std::make_unique( + pool, elementType->type()); + return std::make_unique( + pool, veloxType->type(), nimbleType.get(), std::move(elements)); + } + } + case velox::TypeKind::ROW: { + NIMBLE_CHECK( + nimbleType->isRow(), "Provided schema doesn't match file schema."); + + auto& nimbleRow = nimbleType->asRow(); + auto& veloxRow = veloxType->type()->as(); + std::vector> children; + std::vector childTypes; + children.reserve(veloxType->size()); + childTypes.reserve(veloxType->size()); + offsets.push_back(nimbleRow.nullsDescriptor().offset()); + + for (auto i = 0; i < veloxType->size(); ++i) { + auto& child = veloxType->childAt(i); + std::unique_ptr factory; + if (isSelected(child->id())) { + if (i < nimbleRow.childrenCount()) { + factory = createFieldReaderFactory( + parameters, + pool, + nimbleRow.childAt(i), + child, + offsets, + isSelected, + executor, + level + 1, + &veloxRow.nameOf(i)); + } else { + factory = std::make_unique( + pool, + inferType( + parameters, + veloxRow.nameOf(i), + veloxRow.childAt(i), + level + 1)); + } + } + childTypes.push_back(factory ? factory->veloxType() : child->type()); + children.push_back(std::move(factory)); + } + + // Underlying reader may return a different vector type than what + // specified, (eg. flat map read as struct). So create new ROW type based + // on children types. Note this special logic is only for Row type based + // on the constraint that flatmap can only be top level fields. + return std::make_unique( + pool, + velox::ROW( + std::vector(veloxRow.names()), + std::move(childTypes)), + nimbleType.get(), + std::move(children), + executor); + } + case velox::TypeKind::MAP: { + NIMBLE_CHECK( + nimbleType->isMap() || nimbleType->isFlatMap(), + "Provided schema doesn't match file schema."); + NIMBLE_ASSERT( + veloxType->size() == 2, + "Velox map type should have exactly two children."); + + if (nimbleType->isMap()) { + const auto& nimbleMap = nimbleType->asMap(); + auto& keyType = veloxType->childAt(0); + offsets.push_back(nimbleMap.lengthsDescriptor().offset()); + auto keys = isSelected(keyType->id()) + ? createFieldReaderFactory( + parameters, + pool, + nimbleMap.keys(), + keyType, + offsets, + isSelected, + executor, + level + 1) + : std::make_unique(pool, keyType->type()); + auto& valueType = veloxType->childAt(1); + auto values = isSelected(valueType->id()) + ? createFieldReaderFactory( + parameters, + pool, + nimbleMap.values(), + valueType, + offsets, + isSelected, + executor, + level + 1) + : std::make_unique(pool, valueType->type()); + return std::make_unique( + pool, + veloxType->type(), + nimbleType.get(), + std::move(keys), + std::move(values)); + } else { + auto& nimbleFlatMap = nimbleType->asFlatMap(); + offsets.push_back(nimbleFlatMap.nullsDescriptor().offset()); + NIMBLE_CHECK( + level == 1 && name != nullptr, + "Flat map is only supported as top level fields"); + auto flatMapAsStruct = + parameters.readFlatMapFieldAsStruct.count(*name) > 0; + + // Extract features only when flat map is not empty. When flatmap is + // empty, writer creates dummy child with empty name to carry schema + // information. We need to capture actual children count here. + auto childrenCount = nimbleFlatMap.childrenCount(); + if (childrenCount == 1 && nimbleFlatMap.nameAt(0).empty()) { + childrenCount = 0; + } + + folly::F14FastMap namesToIndices; + + auto featuresIt = parameters.flatMapFeatureSelector.find(*name); + auto hasFeatureSelection = + featuresIt != parameters.flatMapFeatureSelector.end(); + if (hasFeatureSelection) { + NIMBLE_CHECK( + !featuresIt->second.features.empty(), + fmt::format( + "Flat map feature selection for map '{}' has empty feature set.", + *name)); + + if (featuresIt->second.mode == SelectionMode::Include) { + // We have valid feature projection. Build name -> index lookup + // table. + namesToIndices.reserve(childrenCount); + for (auto i = 0; i < childrenCount; ++i) { + namesToIndices.emplace(nimbleFlatMap.nameAt(i), i); + } + } else { + NIMBLE_CHECK( + !flatMapAsStruct, + fmt::format( + "Exclusion can only be applied when flat map is returned as a regular map.")); + } + } else { + // Not specifying features for a flat map is only allowed when + // reconstructing a map column. For struct encoding, we require the + // caller to provide feature selection, as it dictates the order of + // the returned features. + NIMBLE_CHECK( + !flatMapAsStruct, + fmt::format( + "Flat map '{}' is configured to be returned as a struct, but feature selection is missing. " + "Feature selection is used to define the order of the features in the returned struct.", + *name)); + } + + auto actualType = veloxType->type(); + auto& valueType = veloxType->childAt(1); + std::vector selectedChildren; + std::vector inMapDescriptors; + + if (flatMapAsStruct) { + // When reading as struct, all children appear in the feature + // selection will need to be in the result even if they don't exist in + // the schema. + auto& features = featuresIt->second.features; + selectedChildren.reserve(features.size()); + inMapDescriptors.reserve(features.size()); + actualType = createFlatType(features, veloxType->type()); + + for (const auto& feature : features) { + auto it = namesToIndices.find(feature); + if (it != namesToIndices.end()) { + auto childIdx = it->second; + selectedChildren.push_back(childIdx); + auto* inMapDescriptor = + &nimbleFlatMap.inMapDescriptorAt(childIdx); + inMapDescriptors.push_back(inMapDescriptor); + offsets.push_back(inMapDescriptor->offset()); + } else { + inMapDescriptors.push_back(nullptr); + } + } + } else if (childrenCount > 0) { + // When reading as regular map, projection only matters if the map is + // not empty. + if (!hasFeatureSelection) { + selectedChildren.reserve(childrenCount); + for (auto i = 0; i < childrenCount; ++i) { + selectedChildren.push_back(i); + } + } else { + auto& features = featuresIt->second.features; + if (featuresIt->second.mode == SelectionMode::Include) { + // Note this path is slightly different from "as struct" path as + // it doesn't need to add the missing children to the selection. + selectedChildren.reserve(features.size()); + for (auto& feature : features) { + auto it = namesToIndices.find(feature); + if (it != namesToIndices.end()) { + selectedChildren.push_back(it->second); + } + } + } else { + folly::F14FastSet exclusions( + features.begin(), features.end()); + selectedChildren.reserve(childrenCount); + for (auto i = 0; i < childrenCount; ++i) { + if (exclusions.count(nimbleFlatMap.nameAt(i)) == 0) { + selectedChildren.push_back(i); + } + } + } + } + + inMapDescriptors.reserve(selectedChildren.size()); + for (auto childIdx : selectedChildren) { + auto* inMapDescriptor = &nimbleFlatMap.inMapDescriptorAt(childIdx); + inMapDescriptors.push_back(inMapDescriptor); + offsets.push_back(inMapDescriptor->offset()); + } + } + + std::vector> valueReaders; + valueReaders.reserve(selectedChildren.size()); + for (auto childIdx : selectedChildren) { + valueReaders.push_back(createFieldReaderFactory( + parameters, + pool, + nimbleFlatMap.childAt(childIdx), + valueType, + offsets, + isSelected, + executor, + level + 1)); + } + + auto& keySelectionCallback = parameters.keySelectionCallback; + if (keySelectionCallback) { + keySelectionCallback( + {.totalKeys = childrenCount, + .selectedKeys = selectedChildren.size()}); + } + + return createFlatMapReaderFactory( + pool, + veloxType->childAt(0)->type()->kind(), + std::move(actualType), + nimbleType.get(), + std::move(inMapDescriptors), + std::move(valueReaders), + selectedChildren, + flatMapAsStruct, + executor); + } + } + default: + NIMBLE_NOT_SUPPORTED( + fmt::format("Unsupported type: {}", veloxType->type()->kindName())); + } +} + +} // namespace + +void FieldReader::ensureNullConstant( + uint32_t rowCount, + velox::VectorPtr& output, + const std::shared_ptr& type) const { + // If output is already single referenced null constant, resize. Otherwise, + // allocate new one. + if (output && output.use_count() == 1 && + output->encoding() == velox::VectorEncoding::Simple::CONSTANT && + output->isNullAt(0)) { + output->resize(rowCount); + } else { + output = velox::BaseVector::createNullConstant(type, rowCount, &pool_); + } +} + +void FieldReader::reset() { + if (decoder_) { + decoder_->reset(); + } +} + +std::unique_ptr FieldReaderFactory::createNullColumnReader() + const { + return std::make_unique(pool_, veloxType_); +} + +Decoder* FOLLY_NULLABLE FieldReaderFactory::getDecoder( + const folly::F14FastMap>& decoders, + const StreamDescriptor& streamDescriptor) const { + auto it = decoders.find(streamDescriptor.offset()); + if (it == decoders.end()) { + // It is possible that for a given offset, we don't have a matching + // decoder. Each stripe might see different amount of streams, so for all + // the unknown streams, there won't be a matching decoder. + return nullptr; + } + + return it->second.get(); +} + +template +std::unique_ptr FieldReaderFactory::createReaderImpl( + const folly::F14FastMap>& decoders, + const StreamDescriptor& nullsDescriptor, + Args&&... args) const { + auto decoder = getDecoder(decoders, nullsDescriptor); + if (!decoder) { + return createNullColumnReader(); + } + + return std::make_unique(pool_, veloxType_, decoder, args()...); +} + +std::unique_ptr FieldReaderFactory::create( + const FieldReaderParams& parameters, + velox::memory::MemoryPool& pool, + const std::shared_ptr& nimbleType, + const std::shared_ptr& veloxType, + std::vector& offsets, + const std::function& isSelected, + folly::Executor* executor) { + return createFieldReaderFactory( + parameters, pool, nimbleType, veloxType, offsets, isSelected, executor); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/FieldReader.h b/dwio/nimble/velox/FieldReader.h new file mode 100644 index 0000000..54947b8 --- /dev/null +++ b/dwio/nimble/velox/FieldReader.h @@ -0,0 +1,125 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include "dwio/nimble/velox/Decoder.h" +#include "dwio/nimble/velox/SchemaReader.h" +#include "velox/common/memory/MemoryPool.h" +#include "velox/dwio/common/FlatMapHelper.h" +#include "velox/dwio/common/TypeWithId.h" +#include "velox/vector/BaseVector.h" + +namespace facebook::nimble { + +enum class SelectionMode { + Include = 0, + Exclude = 1, +}; + +struct FeatureSelection { + std::vector features; + // When mode == Include, only features appearing in 'features' will be + // included in returned map, otherwise, + // all features from the file will be returned in the map, excluding + // the features appearing in 'features'. + SelectionMode mode{SelectionMode::Include}; +}; + +struct FieldReaderParams { + // Allow selecting subset of features to be included/excluded in flat maps. + // The key in the map is the flat map (top-level) column name. + folly::F14FastMap flatMapFeatureSelector; + + // Contains flatmap field name which we want to return as Struct + folly::F14FastSet readFlatMapFieldAsStruct; + + // Callback to populate feature projection stats when needed + std::function + keySelectionCallback{nullptr}; +}; + +class FieldReader { + public: + FieldReader( + velox::memory::MemoryPool& pool, + velox::TypePtr type, + Decoder* decoder) + : pool_{pool}, type_{std::move(type)}, decoder_{decoder} {} + + virtual ~FieldReader() = default; + + // Place the next X rows of data into the passed in output vector. + virtual void next( + uint32_t count, + velox::VectorPtr& output, + const bits::Bitmap* scatterBitmap) = 0; + + virtual void skip(uint32_t count) = 0; + + // Called at the end of stripe + virtual void reset(); + + protected: + void ensureNullConstant( + uint32_t count, + velox::VectorPtr& output, + const std::shared_ptr& type) const; + + velox::memory::MemoryPool& pool_; + const velox::TypePtr type_; + Decoder* decoder_; +}; + +class FieldReaderFactory { + public: + FieldReaderFactory( + velox::memory::MemoryPool& pool, + velox::TypePtr veloxType, + const Type* nimbleType) + : pool_{pool}, + veloxType_{std::move(veloxType)}, + nimbleType_{nimbleType} {} + + virtual ~FieldReaderFactory() = default; + + virtual std::unique_ptr createReader( + const folly::F14FastMap>& + decoders) = 0; + + const velox::TypePtr& veloxType() const { + return veloxType_; + } + + // Build a field reader factory tree. Will traverse the passed in types and + // create matching field readers. + static std::unique_ptr create( + const FieldReaderParams& parameters, + velox::memory::MemoryPool& pool, + const std::shared_ptr& nimbleType, + const std::shared_ptr& veloxType, + std::vector& offsets, + const std::function& isSelected = + [](auto) { return true; }, + folly::Executor* executor = nullptr); + + protected: + std::unique_ptr createNullColumnReader() const; + + Decoder* FOLLY_NULLABLE getDecoder( + const folly::F14FastMap>& decoders, + const StreamDescriptor& streamDescriptor) const; + + template + std::unique_ptr createReaderImpl( + const folly::F14FastMap>& decoders, + const StreamDescriptor& nullsDecriptor, + Args&&... args) const; + + velox::memory::MemoryPool& pool_; + const velox::TypePtr veloxType_; + const Type* nimbleType_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/FieldWriter.cpp b/dwio/nimble/velox/FieldWriter.cpp new file mode 100644 index 0000000..3660340 --- /dev/null +++ b/dwio/nimble/velox/FieldWriter.cpp @@ -0,0 +1,1280 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/FieldWriter.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/velox/SchemaBuilder.h" +#include "dwio/nimble/velox/SchemaTypes.h" +#include "velox/common/base/CompareFlags.h" +#include "velox/vector/ComplexVector.h" +#include "velox/vector/FlatVector.h" + +namespace facebook::nimble { + +class FieldWriterContext::LocalDecodedVector { + public: + explicit LocalDecodedVector(FieldWriterContext& context) + : context_(context), vector_(context_.getDecodedVector()) {} + + LocalDecodedVector(LocalDecodedVector&& other) noexcept + : context_{other.context_}, vector_{std::move(other.vector_)} {} + + LocalDecodedVector& operator=(LocalDecodedVector&& other) = delete; + + ~LocalDecodedVector() { + if (vector_) { + context_.releaseDecodedVector(std::move(vector_)); + } + } + + velox::DecodedVector& get() { + return *vector_; + } + + private: + FieldWriterContext& context_; + std::unique_ptr vector_; +}; + +namespace { + +template +struct NimbleTypeTraits {}; + +template <> +struct NimbleTypeTraits { + static constexpr ScalarKind scalarKind = ScalarKind::Bool; +}; + +template <> +struct NimbleTypeTraits { + static constexpr ScalarKind scalarKind = ScalarKind::Int8; +}; + +template <> +struct NimbleTypeTraits { + static constexpr ScalarKind scalarKind = ScalarKind::Int16; +}; + +template <> +struct NimbleTypeTraits { + static constexpr ScalarKind scalarKind = ScalarKind::Int32; +}; + +template <> +struct NimbleTypeTraits { + static constexpr ScalarKind scalarKind = ScalarKind::Int64; +}; + +template <> +struct NimbleTypeTraits { + static constexpr ScalarKind scalarKind = ScalarKind::Float; +}; + +template <> +struct NimbleTypeTraits { + static constexpr ScalarKind scalarKind = ScalarKind::Double; +}; + +template <> +struct NimbleTypeTraits { + static constexpr ScalarKind scalarKind = ScalarKind::String; +}; + +template <> +struct NimbleTypeTraits { + static constexpr ScalarKind scalarKind = ScalarKind::Binary; +}; + +template <> +struct NimbleTypeTraits { + static constexpr ScalarKind scalarKind = ScalarKind::Int64; +}; + +// Adapters to handle flat or decoded vector using same interfaces. +template +class Flat { + static constexpr auto kIsBool = std::is_same_v; + + public: + explicit Flat(const velox::VectorPtr& vector) + : vector_{vector}, nulls_{vector->rawNulls()} { + if constexpr (!kIsBool) { + if (auto casted = vector->asFlatVector()) { + values_ = casted->rawValues(); + } + } + } + + bool hasNulls() const { + return vector_->mayHaveNulls(); + } + + bool isNullAt(velox::vector_size_t index) const { + return velox::bits::isBitNull(nulls_, index); + } + + T valueAt(velox::vector_size_t index) const { + if constexpr (kIsBool) { + return static_cast*>(vector_.get()) + ->valueAtFast(index); + } else { + return values_[index]; + } + } + + velox::vector_size_t index(velox::vector_size_t index) const { + return index; + } + + private: + const velox::VectorPtr& vector_; + const uint64_t* nulls_; + const T* values_; +}; + +template +class Decoded { + public: + explicit Decoded(const velox::DecodedVector& decoded) : decoded_{decoded} {} + + bool hasNulls() const { + return decoded_.mayHaveNulls(); + } + + bool isNullAt(velox::vector_size_t index) const { + return decoded_.isNullAt(index); + } + + T valueAt(velox::vector_size_t index) const { + return decoded_.valueAt(index); + } + + velox::vector_size_t index(velox::vector_size_t index) const { + return decoded_.index(index); + } + + private: + const velox::DecodedVector& decoded_; +}; + +template +uint64_t iterateNonNulls( + const OrderedRanges& ranges, + nimble::Vector& nonNulls, + const Vector& vector, + const Consumer& consumer, + const IndexOp& indexOp) { + uint64_t nonNullCount = 0; + if (vector.hasNulls()) { + ranges.applyEach([&](auto offset) { + auto notNull = !vector.isNullAt(offset); + if constexpr (addNulls) { + nonNulls.push_back(notNull); + } + if (notNull) { + ++nonNullCount; + consumer(indexOp(offset)); + } + }); + } else { + ranges.applyEach([&](auto offset) { consumer(indexOp(offset)); }); + nonNullCount = ranges.size(); + } + return nonNullCount; +} + +template +uint64_t iterateNonNullIndices( + const OrderedRanges& ranges, + nimble::Vector& nonNulls, + const Vector& vector, + const Op& op) { + return iterateNonNulls( + ranges, nonNulls, vector, op, [&](auto offset) { + return vector.index(offset); + }); +} + +template +uint64_t iterateNonNullValues( + const OrderedRanges& ranges, + nimble::Vector& nonNulls, + const Vector& vector, + const Op& op) { + return iterateNonNulls(ranges, nonNulls, vector, op, [&](auto offset) { + return vector.valueAt(offset); + }); +} + +template +bool equalDecodedVectorIndices( + const velox::DecodedVector& vec, + velox::vector_size_t index, + velox::vector_size_t otherIndex) { + bool otherNull = vec.isNullAt(otherIndex); + bool thisNull = vec.isNullAt(index); + + if (thisNull && otherNull) { + return true; + } + + if (thisNull || otherNull) { + return false; + } + + return vec.valueAt(index) == vec.valueAt(otherIndex); +} + +template +bool compareDecodedVectorToCache( + const velox::DecodedVector& thisVec, + velox::vector_size_t index, + velox::FlatVector* cachedFlatVec, + velox::vector_size_t cacheIndex, + velox::CompareFlags flags) { + bool thisNull = thisVec.isNullAt(index); + bool otherNull = cachedFlatVec->isNullAt(cacheIndex); + + if (thisNull && otherNull) { + return true; + } + + if (thisNull || otherNull) { + return false; + } + + return thisVec.valueAt(index) == cachedFlatVec->valueAt(cacheIndex); +} + +template +std::string_view convert(const Vector& input) { + return { + reinterpret_cast(input.data()), input.size() * sizeof(T)}; +} + +template >> +struct IdentityConverter { + static T convert(T t, Buffer&, uint64_t&) { + return t; + } +}; + +struct StringConverter { + static std::string_view + convert(velox::StringView sv, Buffer& buffer, uint64_t& memoryUsed) { + memoryUsed += sv.size(); + return buffer.writeString({sv.data(), sv.size()}); + } +}; + +struct TimestampConverter { + static int64_t convert(velox::Timestamp ts, Buffer&, uint64_t&) { + return ts.toMillis(); + } +}; + +template < + velox::TypeKind K, + typename C = IdentityConverter::NativeType>> +class SimpleFieldWriter : public FieldWriter { + using SourceType = typename velox::TypeTraits::NativeType; + using TargetType = decltype(C::convert( + SourceType(), + std::declval(), + std::declval())); + + public: + explicit SimpleFieldWriter(FieldWriterContext& context) + : FieldWriter( + context, + context.schemaBuilder.createScalarTypeBuilder( + NimbleTypeTraits::scalarKind)), + valuesStream_{context.createNullableContentStreamData( + typeBuilder_->asScalar().scalarDescriptor())} {} + + void write(const velox::VectorPtr& vector, const OrderedRanges& ranges) + override { + auto size = ranges.size(); + auto& buffer = context_.stringBuffer(); + auto& data = valuesStream_.mutableData(); + + if (auto flat = vector->asFlatVector()) { + valuesStream_.ensureNullsCapacity(flat->mayHaveNulls(), size); + bool rangeCopied = false; + if (!flat->mayHaveNulls()) { + if constexpr ( + std::is_same_v> && + K != velox::TypeKind::BOOLEAN) { + // NOTE: this is currently expensive to grow during a long sequence of + // ingest + // operators. We currently achieve a balance via a buffer growth + // policy. Another factor that can help us reduce this cost is to also + // consider the stripe progress. However, naive progress based + // policies can't be combined with the size based policies, and are + // thus currently not included. + auto newSize = data.size() + size; + if (newSize > data.capacity()) { + auto newCapacity = + context_.inputBufferGrowthPolicy->getExtendedCapacity( + newSize, data.capacity()); + ++context_.inputBufferGrowthStats.count; + context_.inputBufferGrowthStats.itemCount += newCapacity; + data.reserve(newCapacity); + } + ranges.apply([&](auto offset, auto count) { + data.insert( + data.end(), + flat->rawValues() + offset, + flat->rawValues() + offset + count); + }); + rangeCopied = true; + } + } + + if (!rangeCopied) { + iterateNonNullValues( + ranges, + valuesStream_.mutableNonNulls(), + Flat{vector}, + [&](SourceType value) { + data.push_back( + C::convert(value, buffer, valuesStream_.extraMemory())); + }); + } + } else { + auto localDecoded = decode(vector, ranges); + auto& decoded = localDecoded.get(); + valuesStream_.ensureNullsCapacity(decoded.mayHaveNulls(), size); + iterateNonNullValues( + ranges, + valuesStream_.mutableNonNulls(), + Decoded{decoded}, + [&](SourceType value) { + data.push_back( + C::convert(value, buffer, valuesStream_.extraMemory())); + }); + } + } + + void reset() override { + valuesStream_.reset(); + } + + private: + NullableContentStreamData& valuesStream_; +}; + +class RowFieldWriter : public FieldWriter { + public: + RowFieldWriter( + FieldWriterContext& context, + const std::shared_ptr& type) + : FieldWriter{context, context.schemaBuilder.createRowTypeBuilder(type->size())}, + nullsStream_{context_.createNullsStreamData( + typeBuilder_->asRow().nullsDescriptor())} { + auto rowType = + std::dynamic_pointer_cast(type->type()); + + fields_.reserve(rowType->size()); + for (auto i = 0; i < rowType->size(); ++i) { + fields_.push_back(FieldWriter::create(context, type->childAt(i))); + typeBuilder_->asRow().addChild( + rowType->nameOf(i), fields_.back()->typeBuilder()); + } + } + + void write(const velox::VectorPtr& vector, const OrderedRanges& ranges) + override { + auto size = ranges.size(); + OrderedRanges childRanges; + const OrderedRanges* childRangesPtr; + const velox::RowVector* row = vector->as(); + + if (row) { + NIMBLE_CHECK(fields_.size() == row->childrenSize(), "schema mismatch"); + nullsStream_.ensureNullsCapacity(vector->mayHaveNulls(), size); + if (row->mayHaveNulls()) { + childRangesPtr = &childRanges; + iterateNonNullIndices( + ranges, + nullsStream_.mutableNonNulls(), + Flat{vector}, + [&](auto offset) { childRanges.add(offset, 1); }); + } else { + childRangesPtr = &ranges; + } + } else { + auto localDecoded = decode(vector, ranges); + auto& decoded = localDecoded.get(); + row = decoded.base()->as(); + NIMBLE_ASSERT(row, "Unexpected vector type"); + NIMBLE_CHECK(fields_.size() == row->childrenSize(), "schema mismatch"); + childRangesPtr = &childRanges; + nullsStream_.ensureNullsCapacity(decoded.mayHaveNulls(), size); + iterateNonNullIndices( + ranges, + nullsStream_.mutableNonNulls(), + Decoded{decoded}, + [&](auto offset) { childRanges.add(offset, 1); }); + } + for (auto i = 0; i < fields_.size(); ++i) { + fields_[i]->write(row->childAt(i), *childRangesPtr); + } + } + + void reset() override { + nullsStream_.reset(); + + for (auto& field : fields_) { + field->reset(); + } + } + + void close() override { + for (auto& field : fields_) { + field->close(); + } + } + + private: + std::vector> fields_; + NullsStreamData& nullsStream_; +}; + +class MultiValueFieldWriter : public FieldWriter { + public: + MultiValueFieldWriter( + FieldWriterContext& context, + std::shared_ptr typeBuilder) + : FieldWriter{context, std::move(typeBuilder)}, + lengthsStream_{context.createNullableContentStreamData( + static_cast(*typeBuilder_) + .lengthsDescriptor())} {} + + void reset() override { + lengthsStream_.reset(); + } + + protected: + template + const T* ingestLengths( + const velox::VectorPtr& vector, + const OrderedRanges& ranges, + OrderedRanges& childRanges) { + auto size = ranges.size(); + const T* casted = vector->as(); + const velox::vector_size_t* offsets; + const velox::vector_size_t* lengths; + auto& data = lengthsStream_.mutableData(); + + auto proc = [&](velox::vector_size_t index) { + auto length = lengths[index]; + if (length > 0) { + childRanges.add(offsets[index], length); + } + data.push_back(length); + }; + + if (casted) { + offsets = casted->rawOffsets(); + lengths = casted->rawSizes(); + + lengthsStream_.ensureNullsCapacity(casted->mayHaveNulls(), size); + iterateNonNullIndices( + ranges, lengthsStream_.mutableNonNulls(), Flat{vector}, proc); + } else { + auto localDecoded = decode(vector, ranges); + auto& decoded = localDecoded.get(); + casted = decoded.base()->as(); + NIMBLE_ASSERT(casted, "Unexpected vector type"); + offsets = casted->rawOffsets(); + lengths = casted->rawSizes(); + + lengthsStream_.ensureNullsCapacity(decoded.mayHaveNulls(), size); + iterateNonNullIndices( + ranges, lengthsStream_.mutableNonNulls(), Decoded{decoded}, proc); + } + + return casted; + } + + NullableContentStreamData& lengthsStream_; +}; + +class ArrayFieldWriter : public MultiValueFieldWriter { + public: + ArrayFieldWriter( + FieldWriterContext& context, + const std::shared_ptr& type) + : MultiValueFieldWriter{ + context, + context.schemaBuilder.createArrayTypeBuilder()} { + auto arrayType = + std::dynamic_pointer_cast(type->type()); + + NIMBLE_DASSERT(type->size() == 1, "Invalid array type."); + elements_ = FieldWriter::create(context, type->childAt(0)); + + typeBuilder_->asArray().setChildren(elements_->typeBuilder()); + } + + void write(const velox::VectorPtr& vector, const OrderedRanges& ranges) + override { + OrderedRanges childRanges; + auto array = ingestLengths(vector, ranges, childRanges); + if (childRanges.size() > 0) { + elements_->write(array->elements(), childRanges); + } + } + + void reset() override { + MultiValueFieldWriter::reset(); + elements_->reset(); + } + + void close() override { + elements_->close(); + } + + private: + std::unique_ptr elements_; +}; + +class MapFieldWriter : public MultiValueFieldWriter { + public: + MapFieldWriter( + FieldWriterContext& context, + const std::shared_ptr& type) + : MultiValueFieldWriter{ + context, + context.schemaBuilder.createMapTypeBuilder()} { + auto mapType = + std::dynamic_pointer_cast(type->type()); + + NIMBLE_DASSERT(type->size() == 2, "Invalid map type."); + keys_ = FieldWriter::create(context, type->childAt(0)); + values_ = FieldWriter::create(context, type->childAt(1)); + typeBuilder_->asMap().setChildren( + keys_->typeBuilder(), values_->typeBuilder()); + } + + void write(const velox::VectorPtr& vector, const OrderedRanges& ranges) + override { + OrderedRanges childRanges; + auto map = ingestLengths(vector, ranges, childRanges); + if (childRanges.size() > 0) { + keys_->write(map->mapKeys(), childRanges); + values_->write(map->mapValues(), childRanges); + } + } + + void reset() override { + MultiValueFieldWriter::reset(); + keys_->reset(); + values_->reset(); + } + + void close() override { + keys_->close(); + values_->close(); + } + + private: + std::unique_ptr keys_; + std::unique_ptr values_; +}; + +class FlatMapValueFieldWriter { + public: + FlatMapValueFieldWriter( + FieldWriterContext& context, + const StreamDescriptorBuilder& inMapDescriptor, + std::unique_ptr valueField) + : inMapDescriptor_{inMapDescriptor}, + valueField_{std::move(valueField)}, + inMapStream_{context.createContentStreamData(inMapDescriptor)} {} + + // Clear the ranges and extend the inMapBuffer + void prepare(uint32_t numValues) { + auto& data = inMapStream_.mutableData(); + data.reserve(data.size() + numValues); + std::fill(data.end(), data.end() + numValues, false); + } + + void add(velox::vector_size_t offset, uint32_t mapIndex) { + auto& data = inMapStream_.mutableData(); + auto index = mapIndex + data.size(); + NIMBLE_CHECK(data.empty() || !data[index], "Duplicated key"); + ranges_.add(offset, 1); + data[index] = true; + } + + void write(const velox::VectorPtr& vector, uint32_t mapCount) { + auto& data = inMapStream_.mutableData(); + data.update_size(data.size() + mapCount); + + if (ranges_.size() > 0) { + valueField_->write(vector, ranges_); + } + + ranges_.clear(); + } + + void backfill(uint32_t count, uint32_t reserve) { + inMapStream_.mutableData().resize(count, false); + prepare(reserve); + } + + void reset() { + inMapStream_.reset(); + valueField_->reset(); + } + + void close() { + valueField_->close(); + } + + private: + const StreamDescriptorBuilder& inMapDescriptor_; + std::unique_ptr valueField_; + ContentStreamData& inMapStream_; + OrderedRanges ranges_; +}; + +template +class FlatMapFieldWriter : public FieldWriter { + using KeyType = typename velox::TypeTraits::NativeType; + + public: + FlatMapFieldWriter( + FieldWriterContext& context, + const std::shared_ptr& type) + : FieldWriter( + context, + context.schemaBuilder.createFlatMapTypeBuilder( + NimbleTypeTraits::scalarKind)), + nullsStream_{context_.createNullsStreamData( + typeBuilder_->asFlatMap().nullsDescriptor())}, + valueType_{type->childAt(1)} {} + + void write(const velox::VectorPtr& vector, const OrderedRanges& ranges) + override { + auto size = ranges.size(); + const velox::vector_size_t* offsets; + const velox::vector_size_t* lengths; + uint32_t nonNullCount = 0; + OrderedRanges keyRanges; + + // Lambda that iterates keys of a map and records the offsets to write to + // particular value node. + auto processMap = [&](velox::vector_size_t index, auto& keysVector) { + for (auto begin = offsets[index], end = begin + lengths[index]; + begin < end; + ++begin) { + auto valueField = getValueFieldWriter(keysVector.valueAt(begin), size); + // Add the value to the buffer by recording its offset in the values + // vector. + valueField->add(begin, nonNullCount); + } + ++nonNullCount; + }; + + // Lambda that calculates child ranges + auto computeKeyRanges = [&](velox::vector_size_t index) { + keyRanges.add(offsets[index], lengths[index]); + }; + + // Lambda that iterates the vector + auto processVector = [&](const auto& map, const auto& vector) { + auto& mapKeys = map->mapKeys(); + if (auto flatKeys = mapKeys->template asFlatVector()) { + // Keys are flat. + Flat keysVector{mapKeys}; + iterateNonNullIndices( + ranges, nullsStream_.mutableNonNulls(), vector, [&](auto offset) { + processMap(offset, keysVector); + }); + } else { + // Keys are encoded. Decode. + iterateNonNullIndices( + ranges, nullsStream_.mutableNonNulls(), vector, computeKeyRanges); + auto localDecodedKeys = decode(mapKeys, keyRanges); + auto& decodedKeys = localDecodedKeys.get(); + Decoded keysVector{decodedKeys}; + iterateNonNullIndices( + ranges, nullsStream_.mutableNonNulls(), vector, [&](auto offset) { + processMap(offset, keysVector); + }); + } + }; + + // Reset existing value fields for next batch + for (auto& pair : currentValueFields_) { + pair.second->prepare(size); + } + + const velox::MapVector* map = vector->as(); + if (map) { + // Map is flat + offsets = map->rawOffsets(); + lengths = map->rawSizes(); + + nullsStream_.ensureNullsCapacity(map->mayHaveNulls(), size); + processVector(map, Flat{vector}); + } else { + // Map is encoded. Decode. + auto localDecodedMap = decode(vector, ranges); + auto& decodedMap = localDecodedMap.get(); + map = decodedMap.base()->template as(); + NIMBLE_ASSERT(map, "Unexpected vector type"); + offsets = map->rawOffsets(); + lengths = map->rawSizes(); + + nullsStream_.ensureNullsCapacity(decodedMap.mayHaveNulls(), size); + processVector(map, Decoded{decodedMap}); + } + + // Now actually ingest the map values + if (nonNullCount > 0) { + auto& values = map->mapValues(); + for (auto& pair : currentValueFields_) { + pair.second->write(values, nonNullCount); + } + } + nonNullCount_ += nonNullCount; + } + + void reset() override { + for (auto& field : currentValueFields_) { + field.second->reset(); + } + + nullsStream_.reset(); + nonNullCount_ = 0; + currentValueFields_.clear(); + } + + void close() override { + // Add dummy node so we can preserve schema of an empty flat map. + if (allValueFields_.empty()) { + auto valueField = FieldWriter::create(context_, valueType_); + typeBuilder_->asFlatMap().addChild("", valueField->typeBuilder()); + } else { + for (auto& pair : allValueFields_) { + pair.second->close(); + } + } + } + + private: + FlatMapValueFieldWriter* getValueFieldWriter(KeyType key, uint32_t size) { + auto it = currentValueFields_.find(key); + if (it != currentValueFields_.end()) { + return it->second; + } + + auto stringKey = folly::to(key); + NIMBLE_DASSERT( + !stringKey.empty(), "String key cannot be empty for flatmap"); + + // check whether the typebuilder for this key is already present + auto flatFieldIt = allValueFields_.find(key); + if (flatFieldIt == allValueFields_.end()) { + auto valueFieldWriter = FieldWriter::create(context_, valueType_); + const auto& inMapDescriptor = typeBuilder_->asFlatMap().addChild( + stringKey, valueFieldWriter->typeBuilder()); + if (context_.flatmapFieldAddedEventHandler) { + context_.flatmapFieldAddedEventHandler( + *typeBuilder_, stringKey, *valueFieldWriter->typeBuilder()); + } + auto flatMapValueField = std::make_unique( + context_, inMapDescriptor, std::move(valueFieldWriter)); + flatFieldIt = + allValueFields_.emplace(key, std::move(flatMapValueField)).first; + } + // TODO: assert on not having too many keys? + it = currentValueFields_.emplace(key, flatFieldIt->second.get()).first; + + // At this point we will have at max nonNullCount_ for the field which we + // backfill to false, later when ingest is completed (using scatter write) + // we update the inMapBuffer to nonNullCount which represnet the actual + // values written in file + it->second->backfill(nonNullCount_, size); + return it->second; + } + + NullsStreamData& nullsStream_; + // This map store the FlatMapValue fields used in current flush unit. + folly::F14FastMap currentValueFields_; + const std::shared_ptr& valueType_; + uint64_t nonNullCount_ = 0; + // This map store all FlatMapValue fields encountered by the VeloxWriter + // across the whole file. + folly::F14FastMap> + allValueFields_; +}; + +std::unique_ptr createFlatMapFieldWriter( + FieldWriterContext& context, + const std::shared_ptr& type) { + NIMBLE_DASSERT( + type->type()->kind() == velox::TypeKind::MAP, + "Unexpected flat-map field type."); + NIMBLE_DASSERT(type->size() == 2, "Invalid flat-map field type."); + auto kind = type->childAt(0)->type()->kind(); + switch (kind) { + case velox::TypeKind::TINYINT: + return std::make_unique>( + context, type); + case velox::TypeKind::SMALLINT: + return std::make_unique>( + context, type); + case velox::TypeKind::INTEGER: + return std::make_unique>( + context, type); + case velox::TypeKind::BIGINT: + return std::make_unique>( + context, type); + case velox::TypeKind::VARCHAR: + return std::make_unique>( + context, type); + case velox::TypeKind::VARBINARY: + return std::make_unique>( + context, type); + default: + NIMBLE_NOT_SUPPORTED(fmt::format( + "Unsupported flat map key type {}.", + type->childAt(0)->type()->toString())); + } +} + +template +class ArrayWithOffsetsFieldWriter : public FieldWriter { + using SourceType = typename velox::TypeTraits::NativeType; + using OffsetType = uint32_t; + + public: + ArrayWithOffsetsFieldWriter( + FieldWriterContext& context, + const std::shared_ptr& type) + : FieldWriter{context, context.schemaBuilder.createArrayWithOffsetsTypeBuilder()}, + offsetsStream_{context.createNullableContentStreamData( + typeBuilder_->asArrayWithOffsets().offsetsDescriptor())}, + lengthsStream_{context.createContentStreamData( + typeBuilder_->asArrayWithOffsets().lengthsDescriptor())}, + cached_(false), + cachedValue_(nullptr), + cachedSize_(0) { + elements_ = FieldWriter::create(context, type->childAt(0)); + + typeBuilder_->asArrayWithOffsets().setChildren(elements_->typeBuilder()); + + cachedValue_ = velox::ArrayVector::create( + type->type(), 1, context.bufferMemoryPool.get()); + } + + void write(const velox::VectorPtr& vector, const OrderedRanges& ranges) + override { + OrderedRanges childFilteredRanges; + auto array = ingestLengthsOffsets(vector, ranges, childFilteredRanges); + if (childFilteredRanges.size() > 0) { + elements_->write(array->elements(), childFilteredRanges); + } + } + + void reset() override { + offsetsStream_.reset(); + lengthsStream_.reset(); + elements_->reset(); + + cached_ = false; + nextOffset_ = 0; + } + + void close() override { + elements_->close(); + } + + private: + std::unique_ptr elements_; + NullableContentStreamData& + offsetsStream_; /** offsets for each data after dedup */ + ContentStreamData& + lengthsStream_; /** lengths of the each deduped data */ + OffsetType nextOffset_{0}; /** next available offset for dedup storing */ + + bool cached_; + velox::VectorPtr cachedValue_; + velox::vector_size_t cachedSize_; + + template + void ingestLengthsOffsetsByElements( + const velox::ArrayVector* array, + const Vector& iterableVector, + const OrderedRanges& ranges, + const OrderedRanges& childRanges, + OrderedRanges& filteredRanges) { + const velox::vector_size_t* rawOffsets = array->rawOffsets(); + const velox::vector_size_t* rawLengths = array->rawSizes(); + velox::vector_size_t prevIndex = -1; + velox::vector_size_t numLengths = 0; + auto& lengthsData = lengthsStream_.mutableData(); + auto& offsetsData = offsetsStream_.mutableData(); + auto& nonNulls = offsetsStream_.mutableNonNulls(); + + std::function + compareConsecutive; + + std::function compareToCache; + + /** dedup arrays by consecutive elements */ + auto dedupProc = [&](velox::vector_size_t index) { + auto const length = rawLengths[index]; + + bool match = false; + /// Don't compare on the first run + if (prevIndex >= 0) { + match = + (index == prevIndex || + (length == rawLengths[prevIndex] && + compareConsecutive(index, prevIndex))); + } else if (cached_) { // check cache here + match = (length == cachedSize_ && compareToCache(index)); + } + + if (!match) { + if (length > 0) { + filteredRanges.add(rawOffsets[index], length); + } + lengthsData.push_back(length); + ++numLengths; + ++nextOffset_; + } + + prevIndex = index; + offsetsData.push_back(nextOffset_ - 1); + }; + + auto& vectorElements = array->elements(); + if (auto flat = vectorElements->asFlatVector()) { + /** compare array at index and prevIndex to be equal */ + compareConsecutive = [&](velox::vector_size_t index, + velox::vector_size_t prevIndex) { + bool match = true; + velox::CompareFlags flags; + for (velox::vector_size_t idx = 0; idx < rawLengths[index]; ++idx) { + match = flat->compare( + flat, + rawOffsets[index] + idx, + rawOffsets[prevIndex] + idx, + flags) + .value_or(-1) == 0; + if (!match) { + break; + } + } + return match; + }; + + compareToCache = [&](velox::vector_size_t index) { + velox::CompareFlags flags; + return array->compare(cachedValue_.get(), index, 0, flags) + .value_or(-1) == 0; + }; + + iterateNonNullIndices(ranges, nonNulls, iterableVector, dedupProc); + } else { + auto localDecoded = decode(vectorElements, childRanges); + auto& decoded = localDecoded.get(); + /** compare array at index and prevIndex to be equal */ + compareConsecutive = [&](velox::vector_size_t index, + velox::vector_size_t prevIndex) { + bool match = true; + for (velox::vector_size_t idx = 0; idx < rawLengths[index]; ++idx) { + match = equalDecodedVectorIndices( + decoded, rawOffsets[index] + idx, rawOffsets[prevIndex] + idx); + if (!match) { + break; + } + } + return match; + }; + + auto cachedElements = + (cachedValue_->as())->elements(); + auto cachedFlat = cachedElements->asFlatVector(); + compareToCache = [&](velox::vector_size_t index) { + bool match = true; + velox::CompareFlags flags; + for (velox::vector_size_t idx = 0; idx < rawLengths[index]; ++idx) { + match = compareDecodedVectorToCache( + decoded, rawOffsets[index] + idx, cachedFlat, idx, flags); + if (!match) { + break; + } + } + return match; + }; + iterateNonNullIndices(ranges, nonNulls, iterableVector, dedupProc); + } + + // Copy the last valid element into the cache. + // Cache is saved across calls to write(), as long as the same FieldWriter + // object is used. + if (prevIndex != -1 && lengthsData.size() > 0) { + cached_ = true; + cachedSize_ = lengthsData[lengthsData.size() - 1]; + NIMBLE_ASSERT( + lengthsData[lengthsData.size() - 1] == rawLengths[prevIndex], + "Unexpected index: Prev index is not the last item in the list."); + cachedValue_->prepareForReuse(); + velox::BaseVector::CopyRange cacheRange{ + static_cast(prevIndex) /* source index*/, + 0 /* target index*/, + 1 /* count*/}; + cachedValue_->copyRanges(array, folly::Range(&cacheRange, 1)); + } + } + + const velox::ArrayVector* ingestLengthsOffsets( + const velox::VectorPtr& vector, + const OrderedRanges& ranges, + OrderedRanges& filteredRanges) { + auto size = ranges.size(); + const velox::ArrayVector* arrayVector = vector->as(); + const velox::vector_size_t* rawOffsets; + const velox::vector_size_t* rawLengths; + OrderedRanges childRanges; + + auto proc = [&](velox::vector_size_t index) { + auto length = rawLengths[index]; + if (length > 0) { + childRanges.add(rawOffsets[index], length); + } + }; + + if (arrayVector) { + rawOffsets = arrayVector->rawOffsets(); + rawLengths = arrayVector->rawSizes(); + + offsetsStream_.ensureNullsCapacity(arrayVector->mayHaveNulls(), size); + Flat iterableVector{vector}; + iterateNonNullIndices( + ranges, offsetsStream_.mutableNonNulls(), iterableVector, proc); + ingestLengthsOffsetsByElements( + arrayVector, iterableVector, ranges, childRanges, filteredRanges); + } else { + auto localDecoded = decode(vector, ranges); + auto& decoded = localDecoded.get(); + arrayVector = decoded.base()->template as(); + NIMBLE_ASSERT(arrayVector, "Unexpected vector type"); + rawOffsets = arrayVector->rawOffsets(); + rawLengths = arrayVector->rawSizes(); + + offsetsStream_.ensureNullsCapacity(decoded.mayHaveNulls(), size); + Decoded iterableVector{decoded}; + iterateNonNullIndices( + ranges, offsetsStream_.mutableNonNulls(), iterableVector, proc); + ingestLengthsOffsetsByElements( + arrayVector, iterableVector, ranges, childRanges, filteredRanges); + } + return arrayVector; + } +}; + +std::unique_ptr createArrayWithOffsetsFieldWriter( + FieldWriterContext& context, + const std::shared_ptr& type) { + NIMBLE_DASSERT( + type->type()->kind() == velox::TypeKind::ARRAY, + "Unexpected offset-array field type."); + NIMBLE_DASSERT(type->size() == 1, "Invalid offset-array field type."); + auto kind = type->childAt(0)->type()->kind(); + switch (kind) { + case velox::TypeKind::TINYINT: + return std::make_unique< + ArrayWithOffsetsFieldWriter>(context, type); + case velox::TypeKind::SMALLINT: + return std::make_unique< + ArrayWithOffsetsFieldWriter>( + context, type); + case velox::TypeKind::INTEGER: + return std::make_unique< + ArrayWithOffsetsFieldWriter>(context, type); + case velox::TypeKind::BIGINT: + return std::make_unique< + ArrayWithOffsetsFieldWriter>(context, type); + case velox::TypeKind::REAL: + return std::make_unique< + ArrayWithOffsetsFieldWriter>(context, type); + case velox::TypeKind::DOUBLE: + return std::make_unique< + ArrayWithOffsetsFieldWriter>(context, type); + default: + NIMBLE_NOT_SUPPORTED(fmt::format( + "Unsupported dedup array element type {}.", + type->childAt(0)->type()->toString())); + } +} + +} // namespace + +void NullsStreamData::ensureNullsCapacity( + bool mayHaveNulls, + velox::vector_size_t size) { + if (mayHaveNulls || hasNulls_) { + auto newSize = bufferedCount_ + size; + nonNulls_.reserve(newSize); + if (!hasNulls_) { + hasNulls_ = true; + std::fill(nonNulls_.data(), nonNulls_.data() + bufferedCount_, true); + nonNulls_.update_size(bufferedCount_); + } + if (!mayHaveNulls) { + std::fill( + nonNulls_.data() + bufferedCount_, nonNulls_.data() + newSize, true); + nonNulls_.update_size(newSize); + } + } + bufferedCount_ += size; +} + +FieldWriterContext::LocalDecodedVector +FieldWriterContext::getLocalDecodedVector() { + return LocalDecodedVector{*this}; +} + +velox::SelectivityVector& FieldWriterContext::getSelectivityVector( + velox::vector_size_t size) { + if (LIKELY(selectivity_.get() != nullptr)) { + selectivity_->resize(size); + } else { + selectivity_ = std::make_unique(size); + } + return *selectivity_; +} + +std::unique_ptr FieldWriterContext::getDecodedVector() { + if (decodedVectorPool_.empty()) { + return std::make_unique(); + } + auto vector = std::move(decodedVectorPool_.back()); + decodedVectorPool_.pop_back(); + return vector; +} + +void FieldWriterContext::releaseDecodedVector( + std::unique_ptr&& vector) { + decodedVectorPool_.push_back(std::move(vector)); +} + +FieldWriterContext::LocalDecodedVector FieldWriter::decode( + const velox::VectorPtr& vector, + const OrderedRanges& ranges) { + auto& selectivityVector = context_.getSelectivityVector(vector->size()); + // initialize selectivity vector + selectivityVector.clearAll(); + ranges.apply([&](auto offset, auto size) { + selectivityVector.setValidRange(offset, offset + size, true); + }); + selectivityVector.updateBounds(); + + auto localDecoded = context_.getLocalDecodedVector(); + localDecoded.get().decode(*vector, selectivityVector); + return localDecoded; +} + +std::unique_ptr FieldWriter::create( + FieldWriterContext& context, + const std::shared_ptr& type, + std::function typeAddedHandler) { + context.typeAddedHandler = std::move(typeAddedHandler); + return create(context, type); +} + +std::unique_ptr FieldWriter::create( + FieldWriterContext& context, + const std::shared_ptr& type) { + std::unique_ptr field; + switch (type->type()->kind()) { + case velox::TypeKind::BOOLEAN: { + field = std::make_unique>( + context); + break; + } + case velox::TypeKind::TINYINT: { + field = std::make_unique>( + context); + break; + } + case velox::TypeKind::SMALLINT: { + field = std::make_unique>( + context); + break; + } + case velox::TypeKind::INTEGER: { + field = std::make_unique>( + context); + break; + } + case velox::TypeKind::BIGINT: { + field = + std::make_unique>(context); + break; + } + case velox::TypeKind::REAL: { + field = + std::make_unique>(context); + break; + } + case velox::TypeKind::DOUBLE: { + field = + std::make_unique>(context); + break; + } + case velox::TypeKind::VARCHAR: { + field = std::make_unique< + SimpleFieldWriter>( + context); + break; + } + case velox::TypeKind::VARBINARY: { + field = std::make_unique< + SimpleFieldWriter>( + context); + break; + } + case velox::TypeKind::TIMESTAMP: { + field = std::make_unique< + SimpleFieldWriter>( + context); + break; + } + case velox::TypeKind::ROW: { + field = std::make_unique(context, type); + break; + } + case velox::TypeKind::ARRAY: { + field = context.dictionaryArrayNodeIds.contains(type->id()) + ? createArrayWithOffsetsFieldWriter(context, type) + : std::make_unique(context, type); + break; + } + case velox::TypeKind::MAP: { + field = context.flatMapNodeIds.contains(type->id()) + ? createFlatMapFieldWriter(context, type) + : std::make_unique(context, type); + break; + } + default: + NIMBLE_NOT_SUPPORTED( + fmt::format("Unsupported kind: {}.", type->type()->kind())); + } + + context.typeAddedHandler(*field->typeBuilder()); + + return field; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/FieldWriter.h b/dwio/nimble/velox/FieldWriter.h new file mode 100644 index 0000000..e1e21d1 --- /dev/null +++ b/dwio/nimble/velox/FieldWriter.h @@ -0,0 +1,334 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/velox/BufferGrowthPolicy.h" +#include "dwio/nimble/velox/OrderedRanges.h" +#include "dwio/nimble/velox/SchemaBuilder.h" +#include "velox/dwio/common/TypeWithId.h" +#include "velox/vector/DecodedVector.h" + +namespace facebook::nimble { + +// Stream data is a generic interface representing a stream of data, allowing +// generic access to the content to be used by writers +class StreamData { + public: + explicit StreamData(const StreamDescriptorBuilder& descriptor) + : descriptor_{descriptor} {} + + StreamData(const StreamData&) = delete; + StreamData(StreamData&&) = delete; + StreamData& operator=(const StreamData&) = delete; + StreamData& operator=(StreamData&&) = delete; + + virtual std::string_view data() const = 0; + virtual std::span nonNulls() const = 0; + virtual bool hasNulls() const = 0; + virtual bool empty() const = 0; + virtual uint64_t memoryUsed() const = 0; + + virtual void reset() = 0; + virtual void materialize() {} + + const StreamDescriptorBuilder& descriptor() const { + return descriptor_; + } + + virtual ~StreamData() = default; + + private: + const StreamDescriptorBuilder& descriptor_; +}; + +// Content only data stream. +// Used when a stream doesn't contain nulls. +template +class ContentStreamData final : public StreamData { + public: + ContentStreamData( + velox::memory::MemoryPool& memoryPool, + const StreamDescriptorBuilder& descriptor) + : StreamData(descriptor), data_{&memoryPool}, extraMemory_{0} {} + + inline virtual std::string_view data() const override { + return { + reinterpret_cast(data_.data()), data_.size() * sizeof(T)}; + } + + inline virtual std::span nonNulls() const override { + return {}; + } + + inline virtual bool hasNulls() const override { + return false; + } + + inline virtual bool empty() const override { + return data_.empty(); + } + + inline virtual uint64_t memoryUsed() const override { + return (data_.size() * sizeof(T)) + extraMemory_; + } + + inline Vector& mutableData() { + return data_; + } + + inline uint64_t& extraMemory() { + return extraMemory_; + } + + inline virtual void reset() override { + data_.clear(); + extraMemory_ = 0; + } + + private: + Vector data_; + uint64_t extraMemory_; +}; + +// Nulls only data stream. +// Used in cases where boolean data (representing nulls) is needed. +// NOTE: ContentStreamData can also be used to represent these data +// streams, however, for these "null streams", we have special optimizations, +// where if all data is non-null, we omit the stream. This class specialization +// helps with reusing enabling this optimization. +class NullsStreamData : public StreamData { + public: + NullsStreamData( + velox::memory::MemoryPool& memoryPool, + const StreamDescriptorBuilder& descriptor) + : StreamData(descriptor), + nonNulls_{&memoryPool}, + hasNulls_{false}, + bufferedCount_{0} {} + + inline virtual std::string_view data() const override { + return {}; + } + + inline virtual std::span nonNulls() const override { + return nonNulls_; + } + + inline virtual bool hasNulls() const override { + return hasNulls_; + } + + inline virtual bool empty() const override { + return nonNulls_.empty() && bufferedCount_ == 0; + } + + inline virtual uint64_t memoryUsed() const override { + return nonNulls_.size(); + } + + inline Vector& mutableNonNulls() { + return nonNulls_; + } + + inline virtual void reset() override { + nonNulls_.clear(); + hasNulls_ = false; + bufferedCount_ = 0; + } + + void materialize() override { + if (nonNulls_.size() < bufferedCount_) { + const auto offset = nonNulls_.size(); + nonNulls_.resize(bufferedCount_); + std::fill( + nonNulls_.data() + offset, nonNulls_.data() + bufferedCount_, true); + } + } + + void ensureNullsCapacity(bool mayHaveNulls, velox::vector_size_t size); + + protected: + Vector nonNulls_; + bool hasNulls_; + uint32_t bufferedCount_; +}; + +// Nullable content data stream. +// Used in all cases where data may contain nulls. +template +class NullableContentStreamData final : public NullsStreamData { + public: + NullableContentStreamData( + velox::memory::MemoryPool& memoryPool, + const StreamDescriptorBuilder& descriptor) + : NullsStreamData(memoryPool, descriptor), + data_{&memoryPool}, + extraMemory_{0} {} + + inline virtual std::string_view data() const override { + return { + reinterpret_cast(data_.data()), data_.size() * sizeof(T)}; + } + + inline virtual bool empty() const override { + return NullsStreamData::empty() && data_.empty(); + } + + inline virtual uint64_t memoryUsed() const override { + return (data_.size() * sizeof(T)) + extraMemory_ + + NullsStreamData::memoryUsed(); + } + + inline Vector& mutableData() { + return data_; + } + + inline uint64_t& extraMemory() { + return extraMemory_; + } + + inline virtual void reset() override { + NullsStreamData::reset(); + data_.clear(); + extraMemory_ = 0; + } + + private: + Vector data_; + uint64_t extraMemory_; +}; + +struct InputBufferGrowthStats { + std::atomic count{0}; + // realloc bytes would be interesting, but requires a bit more + // trouble to get. + std::atomic itemCount{0}; +}; + +struct FieldWriterContext { + class LocalDecodedVector; + + explicit FieldWriterContext( + velox::memory::MemoryPool& memoryPool, + std::unique_ptr reclaimer = nullptr) + : bufferMemoryPool{memoryPool.addLeafChild( + "field_writer_buffer", + true, + std::move(reclaimer))}, + inputBufferGrowthPolicy{ + DefaultInputBufferGrowthPolicy::withDefaultRanges()} { + resetStringBuffer(); + } + + std::shared_ptr bufferMemoryPool; + SchemaBuilder schemaBuilder; + + folly::F14FastSet flatMapNodeIds; + folly::F14FastSet dictionaryArrayNodeIds; + + std::unique_ptr inputBufferGrowthPolicy; + InputBufferGrowthStats inputBufferGrowthStats; + + std::function + flatmapFieldAddedEventHandler; + + std::function typeAddedHandler = + [](const TypeBuilder&) {}; + + LocalDecodedVector getLocalDecodedVector(); + velox::SelectivityVector& getSelectivityVector(velox::vector_size_t size); + + Buffer& stringBuffer() { + return *buffer_; + } + + // Reset writer context for use by next stripe. + void resetStringBuffer() { + buffer_ = std::make_unique(*bufferMemoryPool); + } + + const std::vector>& streams() { + return streams_; + } + + template + NullsStreamData& createNullsStreamData( + const StreamDescriptorBuilder& descriptor) { + return static_cast(*streams_.emplace_back( + std::make_unique(*bufferMemoryPool, descriptor))); + } + + template + ContentStreamData& createContentStreamData( + const StreamDescriptorBuilder& descriptor) { + return static_cast&>(*streams_.emplace_back( + std::make_unique>(*bufferMemoryPool, descriptor))); + } + + template + NullableContentStreamData& createNullableContentStreamData( + const StreamDescriptorBuilder& descriptor) { + return static_cast&>( + *streams_.emplace_back(std::make_unique>( + *bufferMemoryPool, descriptor))); + } + + private: + std::unique_ptr getDecodedVector(); + void releaseDecodedVector(std::unique_ptr&& vector); + + std::unique_ptr buffer_; + std::vector> decodedVectorPool_; + std::unique_ptr selectivity_; + std::vector> streams_; +}; + +using OrderedRanges = range_helper::OrderedRanges; + +class FieldWriter { + public: + FieldWriter( + FieldWriterContext& context, + std::shared_ptr typeBuilder) + : context_{context}, typeBuilder_{std::move(typeBuilder)} {} + + virtual ~FieldWriter() = default; + + // Writes the vector to internal buffers. + virtual void write( + const velox::VectorPtr& vector, + const OrderedRanges& ranges) = 0; + + // Clears interanl state and any accumulated data in internal buffers. + virtual void reset() = 0; + + // Called when all writes are done, allowing field writers to finalize + // internal state. + virtual void close() {} + + const std::shared_ptr& typeBuilder() { + return typeBuilder_; + } + + static std::unique_ptr create( + FieldWriterContext& context, + const std::shared_ptr& type); + + static std::unique_ptr create( + FieldWriterContext& context, + const std::shared_ptr& type, + std::function typeAddedHandler); + + protected: + FieldWriterContext& context_; + std::shared_ptr typeBuilder_; + + FieldWriterContext::LocalDecodedVector decode( + const velox::VectorPtr& vector, + const OrderedRanges& ranges); +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/FlatMapLayoutPlanner.cpp b/dwio/nimble/velox/FlatMapLayoutPlanner.cpp new file mode 100644 index 0000000..d1ce258 --- /dev/null +++ b/dwio/nimble/velox/FlatMapLayoutPlanner.cpp @@ -0,0 +1,191 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/FlatMapLayoutPlanner.h" +#include + +namespace facebook::nimble { + +namespace { + +void appendAllNestedStreams( + const TypeBuilder& type, + std::vector& childrenOffsets) { + switch (type.kind()) { + case Kind::Scalar: { + childrenOffsets.push_back(type.asScalar().scalarDescriptor().offset()); + break; + } + case Kind::Row: { + auto& row = type.asRow(); + childrenOffsets.push_back(row.nullsDescriptor().offset()); + for (auto i = 0; i < row.childrenCount(); ++i) { + appendAllNestedStreams(row.childAt(i), childrenOffsets); + } + break; + } + case Kind::Array: { + auto& array = type.asArray(); + childrenOffsets.push_back(array.lengthsDescriptor().offset()); + appendAllNestedStreams(array.elements(), childrenOffsets); + break; + } + case Kind::ArrayWithOffsets: { + auto& arrayWithOffsets = type.asArrayWithOffsets(); + childrenOffsets.push_back(arrayWithOffsets.offsetsDescriptor().offset()); + childrenOffsets.push_back(arrayWithOffsets.lengthsDescriptor().offset()); + appendAllNestedStreams(arrayWithOffsets.elements(), childrenOffsets); + break; + } + case Kind::Map: { + auto& map = type.asMap(); + childrenOffsets.push_back(map.lengthsDescriptor().offset()); + appendAllNestedStreams(map.keys(), childrenOffsets); + appendAllNestedStreams(map.values(), childrenOffsets); + break; + } + case Kind::FlatMap: { + auto& flatMap = type.asFlatMap(); + childrenOffsets.push_back(flatMap.nullsDescriptor().offset()); + for (auto i = 0; i < flatMap.childrenCount(); ++i) { + childrenOffsets.push_back(flatMap.inMapDescriptorAt(i).offset()); + appendAllNestedStreams(flatMap.childAt(i), childrenOffsets); + } + break; + } + } +} + +} // namespace + +FlatMapLayoutPlanner::FlatMapLayoutPlanner( + std::function()> typeResolver, + std::vector>> flatMapFeatureOrder) + : typeResolver_{std::move(typeResolver)}, + flatMapFeatureOrder_{std::move(flatMapFeatureOrder)} { + NIMBLE_ASSERT(typeResolver_ != nullptr, "typeResolver is not supplied"); +} + +std::vector FlatMapLayoutPlanner::getLayout( + std::vector&& streams) { + auto type = typeResolver_(); + NIMBLE_ASSERT( + type->kind() == Kind::Row, + "Flat map layout planner requires row as the schema root."); + auto& root = type->asRow(); + + // Layout logic: + // 1. Root stream (Row nulls) is always first + // 2. Later, all flat maps included in config are layed out. + // For each map, we layout all the features included in the config for + // that map, in the order they appeared in the config. For each feature, we + // first add it's in-map stream and then all the value streams for that + // feature (if the value is a complex type, we add all the nested streams + // for this complex type together). + // 3. We then layout all the other "leftover" streams, in "schema order". This + // guarantees that all "related" streams are next to each other. + // Leftover streams include all streams belonging to other columns, and all + // flat map features not included in the config. + + // This vector is going to hold all the ordered flat-map streams contained in + // the config + std::vector orderedFlatMapOffsets; + orderedFlatMapOffsets.reserve(flatMapFeatureOrder_.size() * 3); + + for (const auto& flatMapFeatures : flatMapFeatureOrder_) { + NIMBLE_CHECK( + std::get<0>(flatMapFeatures) < root.childrenCount(), + fmt::format( + "Column ordinal {} for feature ordering is out of range. " + "Top-level row has {} columns.", + std::get<0>(flatMapFeatures), + root.childrenCount())); + auto& column = root.childAt(std::get<0>(flatMapFeatures)); + NIMBLE_CHECK( + column.kind() == Kind::FlatMap, + fmt::format( + "Column '{}' for feature ordering is not a flat map.", + root.nameAt(std::get<0>(flatMapFeatures)))); + + auto& flatMap = column.asFlatMap(); + + // For each flat map, first we push the flat map nulls stream. + orderedFlatMapOffsets.push_back(flatMap.nullsDescriptor().offset()); + + // Build a lookup table from feature name to its schema offset. + std::unordered_map flatMapNamedOrdinals; + flatMapNamedOrdinals.reserve(flatMap.childrenCount()); + for (auto i = 0; i < flatMap.childrenCount(); ++i) { + flatMapNamedOrdinals.insert({flatMap.nameAt(i), i}); + } + + // For every ordered feature, check if we have streams for it and it, along + // with all its nested streams to the ordered stream list. + for (const auto& feature : std::get<1>(flatMapFeatures)) { + auto it = flatMapNamedOrdinals.find(folly::to(feature)); + if (it == flatMapNamedOrdinals.end()) { + continue; + } + + auto ordinal = it->second; + auto& inMapDescriptor = flatMap.inMapDescriptorAt(ordinal); + orderedFlatMapOffsets.push_back(inMapDescriptor.offset()); + appendAllNestedStreams(flatMap.childAt(ordinal), orderedFlatMapOffsets); + } + } + + // This vector is going to hold all the streams, ordered based on schema + // order. This will include streams that already appear in + // 'orderedFlatMapOffsets'. Later, while laying out the final stream oder, + // we'll de-dup these streams. + std::vector orderedAllOffsets; + appendAllNestedStreams(root, orderedAllOffsets); + + // Build a lookup table from type builders to their streams + std::unordered_map offsetsToStreams; + offsetsToStreams.reserve(streams.size()); + std::transform( + streams.begin(), + streams.end(), + std::inserter(offsetsToStreams, offsetsToStreams.begin()), + [](auto& stream) { return std::make_pair(stream.offset, &stream); }); + + std::vector layout; + layout.reserve(streams.size()); + + auto tryAppendStream = [&offsetsToStreams, &layout](uint32_t offset) { + auto it = offsetsToStreams.find(offset); + if (it != offsetsToStreams.end()) { + layout.emplace_back(std::move(*it->second)); + offsetsToStreams.erase(it); + } + }; + + // At this point we have ordered all the type builders, so now we are going to + // try and find matching streams for each type builder and append them to the + // final ordered stream list. + + // First add the root row's null stream + tryAppendStream(root.nullsDescriptor().offset()); + + // Then, add all ordered flat maps + for (auto offset : orderedFlatMapOffsets) { + tryAppendStream(offset); + } + + // Then add all "leftover" streams, in schema order. + // 'tryAppendStream' will de-dup streams that were already added in previous + // steps. + for (auto offset : orderedAllOffsets) { + tryAppendStream(offset); + } + + NIMBLE_ASSERT( + streams.size() == layout.size(), + fmt::format( + "Stream count mismatch. Input size: {}, output size: {}.", + streams.size(), + layout.size())); + + return layout; +} +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/FlatMapLayoutPlanner.h b/dwio/nimble/velox/FlatMapLayoutPlanner.h new file mode 100644 index 0000000..cef47ed --- /dev/null +++ b/dwio/nimble/velox/FlatMapLayoutPlanner.h @@ -0,0 +1,23 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/tablet/Tablet.h" +#include "dwio/nimble/velox/SchemaBuilder.h" + +namespace facebook::nimble { + +class FlatMapLayoutPlanner : public LayoutPlanner { + public: + FlatMapLayoutPlanner( + std::function()> typeResolver, + std::vector>> + flatMapFeatureOrder); + + virtual std::vector getLayout(std::vector&& streams) override; + + private: + std::function()> typeResolver_; + std::vector>> flatMapFeatureOrder_; +}; +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/FlushPolicy.cpp b/dwio/nimble/velox/FlushPolicy.cpp new file mode 100644 index 0000000..4831d45 --- /dev/null +++ b/dwio/nimble/velox/FlushPolicy.cpp @@ -0,0 +1,17 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/FlushPolicy.h" + +namespace facebook::nimble { + +FlushDecision RawStripeSizeFlushPolicy::shouldFlush( + const StripeProgress& stripeProgress) { + return stripeProgress.rawStripeSize >= rawStripeSize_ ? FlushDecision::Stripe + : FlushDecision::None; +} + +void RawStripeSizeFlushPolicy::onClose() { + // No-op +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/FlushPolicy.h b/dwio/nimble/velox/FlushPolicy.h new file mode 100644 index 0000000..15ec3ce --- /dev/null +++ b/dwio/nimble/velox/FlushPolicy.h @@ -0,0 +1,65 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include + +namespace facebook::nimble { + +struct StripeProgress { + // Size of the stripe data when it's fully decompressed and decoded + const uint64_t rawStripeSize; + // Size of the stripe after buffered data is encoded and optionally compressed + const uint64_t stripeSize; + // Size of the allocated buffer in the writer + const uint64_t bufferSize; +}; + +enum class FlushDecision : uint8_t { + None = 0, + Stripe = 1, + Chunk = 2, +}; + +class FlushPolicy { + public: + virtual ~FlushPolicy() = default; + virtual FlushDecision shouldFlush(const StripeProgress& stripeProgress) = 0; + // Required for memory pressure coordination for now. Will remove in the + // future. + virtual void onClose() = 0; +}; + +class RawStripeSizeFlushPolicy final : public FlushPolicy { + public: + explicit RawStripeSizeFlushPolicy(uint64_t rawStripeSize) + : rawStripeSize_{rawStripeSize} {} + + FlushDecision shouldFlush(const StripeProgress& stripeProgress) override; + + void onClose() override; + + private: + const uint64_t rawStripeSize_; +}; + +class LambdaFlushPolicy : public FlushPolicy { + public: + explicit LambdaFlushPolicy( + std::function lambda) + : lambda_{lambda} {} + + FlushDecision shouldFlush(const StripeProgress& stripeProgress) override { + return lambda_(stripeProgress); + } + + void onClose() override { + // No-op + } + + private: + std::function lambda_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/Metadata.fbs b/dwio/nimble/velox/Metadata.fbs new file mode 100644 index 0000000..96e776a --- /dev/null +++ b/dwio/nimble/velox/Metadata.fbs @@ -0,0 +1,14 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +namespace facebook.nimble.serialization; + +table MetadataEntry { + key:string; + value:string; +} + +table Metadata { + entries:[MetadataEntry]; +} + +root_type Metadata; diff --git a/dwio/nimble/velox/OrderedRanges.h b/dwio/nimble/velox/OrderedRanges.h new file mode 100644 index 0000000..2f30b98 --- /dev/null +++ b/dwio/nimble/velox/OrderedRanges.h @@ -0,0 +1,76 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include +#include + +namespace facebook::nimble::range_helper { + +// Range helper used in cases that order need to be maintained. It doesn't +// handle overlaps. +// Suppose we have: +// offsets: [0, 1, 100, 50] +// sizes: [1, 2, 2, 50] +// The result will be: +// [0, 3), [100, 102), [50, 100) +// Time complexity of `add()` is O(1). +template +class OrderedRanges { + public: + template + inline void apply(F f) const { + for (auto& range : ranges_) { + f(std::get<0>(range), std::get<1>(range)); + } + } + + template + inline void applyEach(F f) const { + for (auto& range : ranges_) { + for (auto offset = std::get<0>(range), end = offset + std::get<1>(range); + offset < end; + ++offset) { + f(offset); + } + } + } + + inline void add(T offset, T size) { + size_ += size; + if (ranges_.size() > 0) { + auto& last = ranges_.back(); + auto& end = std::get<1>(last); + if (std::get<0>(last) + end == offset) { + end += size; + return; + } + } + ranges_.emplace_back(offset, size); + } + + inline T size() const { + return size_; + } + + inline void clear() { + ranges_.clear(); + size_ = 0; + } + + static OrderedRanges of(T offset, T size) { + OrderedRanges r; + r.add(offset, size); + return r; + } + + const std::vector>& ranges() const { + return ranges_; + } + + private: + std::vector> ranges_; + T size_{0}; +}; + +} // namespace facebook::nimble::range_helper diff --git a/dwio/nimble/velox/Schema.fbs b/dwio/nimble/velox/Schema.fbs new file mode 100644 index 0000000..1d6f831 --- /dev/null +++ b/dwio/nimble/velox/Schema.fbs @@ -0,0 +1,49 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +namespace facebook.nimble.serialization; + +enum Kind:uint8 { + Int8 = 0, + UInt8, + Int16, + UInt16, + Int32, + UInt32, + Int64, + UInt64, + Float, + Double, + Bool, + String, + Binary, + Row = 50, + Array, + Map, + ArrayWithOffsets, + FlatMapInt8 = 100, + FlatMapUInt8, + FlatMapInt16, + FlatMapUInt16, + FlatMapInt32, + FlatMapUInt32, + FlatMapInt64, + FlatMapUInt64, + FlatMapFloat, + FlatMapDouble, + FlatMapBool, + FlatMapString, + FlatMapBinary, +} + +table SchemaNode { + kind:Kind; + children:uint32; + name:string; + offset:uint32; +} + +table Schema { + nodes:[SchemaNode]; +} + +root_type Schema; diff --git a/dwio/nimble/velox/SchemaBuilder.cpp b/dwio/nimble/velox/SchemaBuilder.cpp new file mode 100644 index 0000000..23f538d --- /dev/null +++ b/dwio/nimble/velox/SchemaBuilder.cpp @@ -0,0 +1,559 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/SchemaBuilder.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/velox/SchemaTypes.h" + +namespace facebook::nimble { + +TypeBuilder::TypeBuilder(SchemaBuilder& schemaBuilder, Kind kind) + : schemaBuilder_{schemaBuilder}, kind_{kind} {} + +ScalarTypeBuilder& TypeBuilder::asScalar() { + return dynamic_cast(*this); +} + +ArrayTypeBuilder& TypeBuilder::asArray() { + return dynamic_cast(*this); +} + +MapTypeBuilder& TypeBuilder::asMap() { + return dynamic_cast(*this); +} + +RowTypeBuilder& TypeBuilder::asRow() { + return dynamic_cast(*this); +} + +FlatMapTypeBuilder& TypeBuilder::asFlatMap() { + return dynamic_cast(*this); +} + +ArrayWithOffsetsTypeBuilder& TypeBuilder::asArrayWithOffsets() { + return dynamic_cast(*this); +} + +const ScalarTypeBuilder& TypeBuilder::asScalar() const { + return dynamic_cast(*this); +} + +const ArrayTypeBuilder& TypeBuilder::asArray() const { + return dynamic_cast(*this); +} + +const MapTypeBuilder& TypeBuilder::asMap() const { + return dynamic_cast(*this); +} + +const RowTypeBuilder& TypeBuilder::asRow() const { + return dynamic_cast(*this); +} + +const FlatMapTypeBuilder& TypeBuilder::asFlatMap() const { + return dynamic_cast(*this); +} + +const ArrayWithOffsetsTypeBuilder& TypeBuilder::asArrayWithOffsets() const { + return dynamic_cast(*this); +} + +Kind TypeBuilder::kind() const { + return kind_; +} + +ScalarTypeBuilder::ScalarTypeBuilder( + SchemaBuilder& schemaBuilder, + ScalarKind scalarKind) + : TypeBuilder{schemaBuilder, Kind::Scalar}, + scalarDescriptor_{schemaBuilder_.allocateStreamOffset(), scalarKind} {} + +const StreamDescriptorBuilder& ScalarTypeBuilder::scalarDescriptor() const { + return scalarDescriptor_; +} + +LengthsTypeBuilder::LengthsTypeBuilder(SchemaBuilder& schemaBuilder, Kind kind) + : TypeBuilder(schemaBuilder, kind), + lengthsDescriptor_{ + schemaBuilder_.allocateStreamOffset(), + ScalarKind::UInt32} {} + +const StreamDescriptorBuilder& LengthsTypeBuilder::lengthsDescriptor() const { + return lengthsDescriptor_; +} + +ArrayTypeBuilder::ArrayTypeBuilder(SchemaBuilder& schemaBuilder) + : LengthsTypeBuilder(schemaBuilder, Kind::Array) {} + +const TypeBuilder& ArrayTypeBuilder::elements() const { + return *elements_; +} + +void ArrayTypeBuilder::setChildren(std::shared_ptr elements) { + NIMBLE_ASSERT(!elements_, "ArrayTypeBuilder elements already initialized."); + schemaBuilder_.registerChild(elements); + elements_ = std::move(elements); +} + +MapTypeBuilder::MapTypeBuilder(SchemaBuilder& schemaBuilder) + : LengthsTypeBuilder{schemaBuilder, Kind::Map} {} + +const TypeBuilder& MapTypeBuilder::keys() const { + return *keys_; +} + +const TypeBuilder& MapTypeBuilder::values() const { + return *values_; +} + +void MapTypeBuilder::setChildren( + std::shared_ptr keys, + std::shared_ptr values) { + NIMBLE_ASSERT(!keys_, "MapTypeBuilder keys already initialized."); + NIMBLE_ASSERT(!values_, "MapTypeBuilder values already initialized."); + schemaBuilder_.registerChild(keys); + schemaBuilder_.registerChild(values); + keys_ = std::move(keys); + values_ = std::move(values); +} + +ArrayWithOffsetsTypeBuilder::ArrayWithOffsetsTypeBuilder( + SchemaBuilder& schemaBuilder) + : TypeBuilder(schemaBuilder, Kind::ArrayWithOffsets), + offsetsDescriptor_{ + schemaBuilder_.allocateStreamOffset(), + ScalarKind::UInt32}, + lengthsDescriptor_{ + schemaBuilder_.allocateStreamOffset(), + ScalarKind::UInt32} {} + +const StreamDescriptorBuilder& ArrayWithOffsetsTypeBuilder::offsetsDescriptor() + const { + return offsetsDescriptor_; +} + +const StreamDescriptorBuilder& ArrayWithOffsetsTypeBuilder::lengthsDescriptor() + const { + return lengthsDescriptor_; +} + +const TypeBuilder& ArrayWithOffsetsTypeBuilder::elements() const { + return *elements_; +} + +void ArrayWithOffsetsTypeBuilder::setChildren( + std::shared_ptr elements) { + NIMBLE_ASSERT( + !elements_, "ArrayWithOffsetsTypeBuilder elements already initialized."); + schemaBuilder_.registerChild(elements); + elements_ = std::move(elements); +} + +RowTypeBuilder::RowTypeBuilder( + SchemaBuilder& schemaBuilder, + size_t childrenCount) + : TypeBuilder{schemaBuilder, Kind::Row}, + nullsDescriptor_{ + schemaBuilder_.allocateStreamOffset(), + ScalarKind::Bool} { + names_.reserve(childrenCount); + children_.reserve(childrenCount); +} + +const StreamDescriptorBuilder& RowTypeBuilder::nullsDescriptor() const { + return nullsDescriptor_; +} + +size_t RowTypeBuilder::childrenCount() const { + return children_.size(); +} + +const TypeBuilder& RowTypeBuilder::childAt(size_t index) const { + NIMBLE_ASSERT( + index < children_.size(), + fmt::format( + "Index out of range. index: {}, size: {}.", index, children_.size())); + return *children_[index]; +} + +const std::string& RowTypeBuilder::nameAt(size_t index) const { + NIMBLE_ASSERT( + index < children_.size(), + fmt::format( + "Index out of range. index: {}, size: {}.", index, children_.size())); + return names_[index]; +} + +void RowTypeBuilder::addChild( + std::string name, + std::shared_ptr child) { + NIMBLE_DASSERT( + children_.size() < children_.capacity(), + fmt::format( + "Registering more row children than expected. Capacity: {}", + children_.capacity())); + schemaBuilder_.registerChild(child); + names_.push_back(std::move(name)); + children_.push_back(std::move(child)); +} + +FlatMapTypeBuilder::FlatMapTypeBuilder( + SchemaBuilder& schemaBuilder, + ScalarKind keyScalarKind) + : TypeBuilder{schemaBuilder, Kind::FlatMap}, + keyScalarKind_{keyScalarKind}, + nullsDescriptor_{ + schemaBuilder_.allocateStreamOffset(), + ScalarKind::Bool} {} + +const StreamDescriptorBuilder& FlatMapTypeBuilder::nullsDescriptor() const { + return nullsDescriptor_; +} + +const StreamDescriptorBuilder& FlatMapTypeBuilder::inMapDescriptorAt( + size_t index) const { + NIMBLE_ASSERT( + index < inMapDescriptors_.size(), + fmt::format( + "Index out of range. index: {}, size: {}.", + index, + inMapDescriptors_.size())); + return *inMapDescriptors_[index]; +} + +const TypeBuilder& FlatMapTypeBuilder::childAt(size_t index) const { + NIMBLE_ASSERT( + index < children_.size(), + fmt::format( + "Index out of range. index: {}, size: {}.", index, children_.size())); + return *children_[index]; +} + +ScalarKind FlatMapTypeBuilder::keyScalarKind() const { + return keyScalarKind_; +} + +size_t FlatMapTypeBuilder::childrenCount() const { + return children_.size(); +} + +const std::string& FlatMapTypeBuilder::nameAt(size_t index) const { + NIMBLE_ASSERT( + index < names_.size(), + fmt::format( + "Index out of range. index: {}, size: {}.", index, names_.size())); + return names_[index]; +} + +const StreamDescriptorBuilder& FlatMapTypeBuilder::addChild( + std::string name, + std::shared_ptr child) { + auto& inMapDescriptor = + inMapDescriptors_.emplace_back(std::make_unique( + schemaBuilder_.allocateStreamOffset(), ScalarKind::Bool)); + + schemaBuilder_.registerChild(child); + + names_.push_back(std::move(name)); + children_.push_back(std::move(child)); + + return *inMapDescriptor; +} + +std::shared_ptr SchemaBuilder::createScalarTypeBuilder( + ScalarKind scalarKind) { + struct MakeSharedEnabler : public ScalarTypeBuilder { + MakeSharedEnabler(SchemaBuilder& schemaBuilder, ScalarKind scalarKind) + : ScalarTypeBuilder{schemaBuilder, scalarKind} {} + }; + + auto type = std::make_shared(*this, scalarKind); + + // This new type builder is not attached to a parent, therefore it is a new + // tree "root" (as of now), so we add it to the roots list. + roots_.insert(type); + return type; +} + +std::shared_ptr SchemaBuilder::createArrayTypeBuilder() { + struct MakeSharedEnabler : public ArrayTypeBuilder { + explicit MakeSharedEnabler(SchemaBuilder& schemaBuilder) + : ArrayTypeBuilder{schemaBuilder} {} + }; + auto type = std::make_shared(*this); + + // This new type builder is not attached to a parent, therefore it is a new + // tree "root" (as of now), so we add it to the roots list. + roots_.insert(type); + return type; +} + +std::shared_ptr +SchemaBuilder::createArrayWithOffsetsTypeBuilder() { + struct MakeSharedEnabler : public ArrayWithOffsetsTypeBuilder { + explicit MakeSharedEnabler(SchemaBuilder& schemaBuilder) + : ArrayWithOffsetsTypeBuilder{schemaBuilder} {} + }; + + auto type = std::make_shared(*this); + + // This new type builder is not attached to a parent, therefore it is a new + // tree "root" (as of now), so we add it to the roots list. + roots_.insert(type); + return type; +} + +std::shared_ptr SchemaBuilder::createRowTypeBuilder( + size_t childrenCount) { + struct MakeSharedEnabler : public RowTypeBuilder { + MakeSharedEnabler(SchemaBuilder& schemaBuilder, size_t childrenCount) + : RowTypeBuilder{schemaBuilder, childrenCount} {} + }; + auto type = std::make_shared(*this, childrenCount); + + // This new type builder is not attached to a parent, therefore it is a new + // tree "root" (as of now), so we add it to the roots list. + roots_.insert(type); + return type; +} + +std::shared_ptr SchemaBuilder::createMapTypeBuilder() { + struct MakeSharedEnabler : public MapTypeBuilder { + explicit MakeSharedEnabler(SchemaBuilder& schemaBuilder) + : MapTypeBuilder{schemaBuilder} {} + }; + auto type = std::make_shared(*this); + + // This new type builder is not attached to a parent, therefore it is a new + // tree "root" (as of now), so we add it to the roots list. + roots_.insert(type); + return type; +} + +std::shared_ptr SchemaBuilder::createFlatMapTypeBuilder( + ScalarKind keyScalarKind) { + struct MakeSharedEnabler : public FlatMapTypeBuilder { + MakeSharedEnabler(SchemaBuilder& schemaBuilder, ScalarKind keyScalarKind) + : FlatMapTypeBuilder(schemaBuilder, keyScalarKind) {} + }; + auto type = std::make_shared(*this, keyScalarKind); + + // This new type builder is not attached to a parent, therefore it is a new + // tree "root" (as of now), so we add it to the roots list. + roots_.insert(type); + return type; +} + +offset_size SchemaBuilder::nodeCount() const { + return currentOffset_; +} + +const std::shared_ptr& SchemaBuilder::getRoot() const { + // When retreiving schema nodes, we return a vector ordered based on the + // schema tree DFS order. To be able to flatten the schema tree to a flat + // ordered vector, we need to guarantee that the schema tree has a single root + // node, where we start traversing from. + NIMBLE_ASSERT( + roots_.size() == 1, + fmt::format( + "Unable to determine schema root. List of roots contain {} entries.", + roots_.size())); + return *roots_.cbegin(); +} + +void SchemaBuilder::registerChild(const std::shared_ptr& type) { + // If we try to attach a node to a parent, but this node doesn't exist in the + // roots list, it means that either this node was already attached to a parent + // before (and therefore was removed from the roots list), or the node was + // created using a different schema builder instance (and therefore belongs to + // a roots list in the other schema builder instance). + NIMBLE_ASSERT( + roots_.find(type) != roots_.end(), + "Child type not found. This can happen if child is registered more than once, " + "or if a different Schema Builder was used to create the child."); + + // Now that the node is attached to a parent, it is no longer a "root" of a + // tree, and should be removed from the list of roots. + roots_.erase(type); +} + +offset_size SchemaBuilder::allocateStreamOffset() { + return currentOffset_++; +} + +void SchemaBuilder::addNode( + std::vector>& nodes, + const TypeBuilder& type, + std::optional name) const { + switch (type.kind()) { + case Kind::Scalar: { + const auto& scalar = type.asScalar(); + nodes.push_back(std::make_unique( + type.kind(), + scalar.scalarDescriptor().offset(), + scalar.scalarDescriptor().scalarKind(), + std::move(name))); + break; + } + case Kind::Array: { + const auto& array = type.asArray(); + nodes.push_back(std::make_unique( + type.kind(), + array.lengthsDescriptor().offset(), + ScalarKind::UInt32, + std::move(name))); + addNode(nodes, array.elements()); + break; + } + case Kind::ArrayWithOffsets: { + const auto& array = type.asArrayWithOffsets(); + nodes.push_back(std::make_unique( + type.kind(), + array.lengthsDescriptor().offset(), + ScalarKind::UInt32, + std::move(name))); + nodes.push_back(std::make_unique( + Kind::Scalar, + array.offsetsDescriptor().offset(), + ScalarKind::UInt32, + std::nullopt)); + addNode(nodes, array.elements()); + break; + } + case Kind::Row: { + auto& row = type.asRow(); + nodes.push_back(std::make_unique( + type.kind(), + row.nullsDescriptor().offset(), + ScalarKind::Bool, + std::move(name), + row.childrenCount())); + for (auto i = 0; i < row.childrenCount(); ++i) { + addNode(nodes, row.childAt(i), row.nameAt(i)); + } + break; + } + case Kind::Map: { + const auto& map = type.asMap(); + nodes.push_back(std::make_unique( + type.kind(), + map.lengthsDescriptor().offset(), + ScalarKind::UInt32, + std::move(name))); + addNode(nodes, map.keys()); + addNode(nodes, map.values()); + break; + } + case Kind::FlatMap: { + auto& map = type.asFlatMap(); + + size_t childrenSize = map.childrenCount(); + nodes.push_back(std::make_unique( + type.kind(), + map.nullsDescriptor().offset(), + map.keyScalarKind(), + std::move(name), + childrenSize)); + NIMBLE_ASSERT( + map.inMapDescriptors_.size() == childrenSize, + "Flat map in-maps collection size and children collection size should be the same."); + for (size_t i = 0; i < childrenSize; ++i) { + nodes.push_back(std::make_unique( + Kind::Scalar, + map.inMapDescriptorAt(i).offset(), + ScalarKind::Bool, + map.nameAt(i))); + addNode(nodes, map.childAt(i)); + } + + break; + } + + default: + NIMBLE_UNREACHABLE( + fmt::format("Unknown type kind: {}.", toString(type.kind()))); + } +} + +std::vector> SchemaBuilder::getSchemaNodes() + const { + auto& root = getRoot(); + std::vector> nodes; + nodes.reserve(currentOffset_); + addNode(nodes, *root); + return nodes; +} + +void printType( + std::ostream& out, + const TypeBuilder& builder, + uint32_t indentation, + const std::optional& name = std::nullopt) { + out << std::string(indentation, ' ') + << (name.has_value() ? name.value() + ":" : ""); + + switch (builder.kind()) { + case Kind::Scalar: { + const auto& scalar = builder.asScalar(); + out << "[" << scalar.scalarDescriptor().offset() << "]" + << toString(scalar.scalarDescriptor().scalarKind()) << "\n"; + break; + } + case Kind::Array: { + const auto& array = builder.asArray(); + out << "[" << array.lengthsDescriptor().offset() << "]ARRAY\n"; + printType(out, array.elements(), indentation + 2, "elements"); + out << "\n"; + break; + } + case Kind::Map: { + const auto& map = builder.asMap(); + out << "[" << map.lengthsDescriptor().offset() << "]MAP\n"; + printType(out, map.keys(), indentation + 2, "keys"); + printType(out, map.values(), indentation + 2, "values"); + out << "\n"; + break; + } + case Kind::Row: { + const auto& row = builder.asRow(); + out << "[" << row.nullsDescriptor().offset() << "]ROW\n"; + for (auto i = 0; i < row.childrenCount(); ++i) { + printType(out, row.childAt(i), indentation + 2, row.nameAt(i)); + } + out << "\n"; + break; + } + case Kind::FlatMap: { + const auto& map = builder.asFlatMap(); + out << "[" << map.nullsDescriptor().offset() << "]FLATMAP\n"; + for (auto i = 0; i < map.childrenCount(); ++i) { + printType( + out, + map.childAt(i), + indentation + 2, + // @lint-ignore CLANGTIDY facebook-hte-LocalUncheckedArrayBounds + map.nameAt(i)); + } + out << "\n"; + break; + } + case Kind::ArrayWithOffsets: { + const auto& array = builder.asArrayWithOffsets(); + out << "[" << array.offsetsDescriptor().offset() << "," + << array.lengthsDescriptor().offset() << "]OFFSETARRAY\n"; + printType(out, array.elements(), indentation + 2, "elements"); + out << "\n"; + break; + } + } +} + +std::ostream& operator<<(std::ostream& out, const SchemaBuilder& schema) { + size_t index = 0; + for (const auto& root : schema.roots_) { + out << "Root " << index++ << ":\n"; + printType(out, *root, 2); + } + + return out; +} +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/SchemaBuilder.h b/dwio/nimble/velox/SchemaBuilder.h new file mode 100644 index 0000000..cfa6584 --- /dev/null +++ b/dwio/nimble/velox/SchemaBuilder.h @@ -0,0 +1,303 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/velox/SchemaTypes.h" +#include "folly/container/F14Set.h" + +// SchemaBuilder is used to construct a tablet compatible schema. +// The table schema translates logical type tree to its underlying Nimble +// streams stored in the file. Each stream has a unique offset assigned to it, +// which is later used when storing stream related data in the tablet footer. +// The class supports primitive types (scalar types) and complex types: rows, +// arrays, maps and flat maps. +// +// This class serves two purposes: +// 1. Allow constructing parent schema nodes, before creating their children, +// and later allow attaching children to parent nodes. This is needed because +// writers usually construct their inner state top to bottom (recursively), so +// this allows for a much easier integration with writers. +// 2. When writing an Nimble file, not all inner streams are known ahead of +// time. This is because FlatMaps may append new streams when new data arrives, +// and new keys are encountered. This class abstracts the allocation of tablet +// offsets (indices to streams) and makes sure they stay consistent across +// stripes, even while the schema keeps changing. +// NOTE: The offsets allocated by this class are used only during write time, +// and are different than the final serialized offsets. During footer +// serialization, "build time offsets" (offsets managed by this class) are +// translated to "schema ordered" offsets (rely on the final flattened schema +// tree layout). +namespace facebook::nimble { + +class SchemaBuilder; +class ScalarTypeBuilder; +class ArrayTypeBuilder; +class MapTypeBuilder; +class RowTypeBuilder; +class FlatMapTypeBuilder; +class ArrayWithOffsetsTypeBuilder; + +class StreamContext { + public: + StreamContext() = default; + StreamContext(const StreamContext&) = delete; + StreamContext(StreamContext&&) = delete; + StreamContext& operator=(const StreamContext&) = delete; + StreamContext& operator=(StreamContext&&) = delete; + virtual ~StreamContext() = default; +}; + +class StreamDescriptorBuilder : public StreamDescriptor { + public: + StreamDescriptorBuilder(offset_size offset, ScalarKind scalarKind) + : StreamDescriptor(offset, scalarKind) {} + + void setContext(std::unique_ptr&& context) const { + context_ = std::move(context); + } + + template + T* context() const { + return dynamic_cast(context_.get()); + } + + private: + mutable std::unique_ptr context_; +}; + +class TypeBuilderContext { + public: + virtual ~TypeBuilderContext() = default; +}; + +class TypeBuilder { + public: + void setContext(std::unique_ptr&& context) const { + context_ = std::move(context); + } + + template + const T* context() const { + return dynamic_cast(context_.get()); + } + + Kind kind() const; + + ScalarTypeBuilder& asScalar(); + ArrayTypeBuilder& asArray(); + MapTypeBuilder& asMap(); + RowTypeBuilder& asRow(); + FlatMapTypeBuilder& asFlatMap(); + ArrayWithOffsetsTypeBuilder& asArrayWithOffsets(); + const ScalarTypeBuilder& asScalar() const; + const ArrayTypeBuilder& asArray() const; + const MapTypeBuilder& asMap() const; + const RowTypeBuilder& asRow() const; + const FlatMapTypeBuilder& asFlatMap() const; + const ArrayWithOffsetsTypeBuilder& asArrayWithOffsets() const; + + protected: + TypeBuilder(SchemaBuilder& schemaBuilder, Kind kind); + virtual ~TypeBuilder() = default; + + SchemaBuilder& schemaBuilder_; + + private: + Kind kind_; + mutable std::unique_ptr context_; + + friend class SchemaBuilder; +}; + +class ScalarTypeBuilder : public TypeBuilder { + public: + const StreamDescriptorBuilder& scalarDescriptor() const; + + private: + ScalarTypeBuilder(SchemaBuilder& schemaBuilder, ScalarKind scalarKind); + + StreamDescriptorBuilder scalarDescriptor_; + + friend class SchemaBuilder; +}; + +class LengthsTypeBuilder : public TypeBuilder { + public: + const StreamDescriptorBuilder& lengthsDescriptor() const; + + protected: + LengthsTypeBuilder(SchemaBuilder& schemaBuilder, Kind kind); + + private: + StreamDescriptorBuilder lengthsDescriptor_; +}; + +class ArrayTypeBuilder : public LengthsTypeBuilder { + public: + const TypeBuilder& elements() const; + void setChildren(std::shared_ptr elements); + + private: + explicit ArrayTypeBuilder(SchemaBuilder& schemaBuilder); + + std::shared_ptr elements_; + + friend class SchemaBuilder; +}; + +class MapTypeBuilder : public LengthsTypeBuilder { + public: + const TypeBuilder& keys() const; + const TypeBuilder& values() const; + + void setChildren( + std::shared_ptr keys, + std::shared_ptr values); + + private: + explicit MapTypeBuilder(SchemaBuilder& schemaBuilder); + + std::shared_ptr keys_; + std::shared_ptr values_; + + friend class SchemaBuilder; +}; + +class RowTypeBuilder : public TypeBuilder { + public: + const StreamDescriptorBuilder& nullsDescriptor() const; + size_t childrenCount() const; + const TypeBuilder& childAt(size_t index) const; + const std::string& nameAt(size_t index) const; + void addChild(std::string name, std::shared_ptr child); + + private: + RowTypeBuilder(SchemaBuilder& schemaBuilder, size_t childrenCount); + + StreamDescriptorBuilder nullsDescriptor_; + std::vector names_; + std::vector> children_; + + friend class SchemaBuilder; +}; + +class FlatMapTypeBuilder : public TypeBuilder { + public: + const StreamDescriptorBuilder& nullsDescriptor() const; + const StreamDescriptorBuilder& inMapDescriptorAt(size_t index) const; + size_t childrenCount() const; + const TypeBuilder& childAt(size_t index) const; + const std::string& nameAt(size_t index) const; + ScalarKind keyScalarKind() const; + + const StreamDescriptorBuilder& addChild( + std::string name, + std::shared_ptr child); + + private: + FlatMapTypeBuilder(SchemaBuilder& schemaBuilder, ScalarKind keyScalarKind); + + ScalarKind keyScalarKind_; + StreamDescriptorBuilder nullsDescriptor_; + std::vector names_; + std::vector> inMapDescriptors_; + std::vector> children_; + + friend class SchemaBuilder; +}; + +class ArrayWithOffsetsTypeBuilder : public TypeBuilder { + public: + const StreamDescriptorBuilder& offsetsDescriptor() const; + const StreamDescriptorBuilder& lengthsDescriptor() const; + const TypeBuilder& elements() const; + void setChildren(std::shared_ptr elements); + + private: + explicit ArrayWithOffsetsTypeBuilder(SchemaBuilder& schemaBuilder); + + StreamDescriptorBuilder offsetsDescriptor_; + StreamDescriptorBuilder lengthsDescriptor_; + std::shared_ptr elements_; + + friend class SchemaBuilder; +}; + +class SchemaBuilder { + public: + // Create a builder representing a scalar type with kind |scalarKind|. + std::shared_ptr createScalarTypeBuilder( + ScalarKind scalarKind); + + // Create an array builder + std::shared_ptr createArrayTypeBuilder(); + + // Create an array with offsets builder + std::shared_ptr + createArrayWithOffsetsTypeBuilder(); + + // Create a row (struct) builder, with fixed number of children + // |childrenCount|. + std::shared_ptr createRowTypeBuilder(size_t childrenCount); + + // Create a map builder. + std::shared_ptr createMapTypeBuilder(); + + // Create a flat map builder. |keyScalarKind| captures the type of the map + // key. + std::shared_ptr createFlatMapTypeBuilder( + ScalarKind keyScalarKind); + + // Retrieves all the nodes CURRENTLY known to the schema builder. + // If more nodes are added to the schema builder later on, following calls to + // this method will return existing and new nodes. It is guaranteed that the + // allocated offsets for each node will stay the same across invocations. + // NOTE: The schema must be in a valid state when calling this method. Valid + // state means that the schema is a valid tree, with a single root node. This + // means that all created builders were attached to their parents by calling + // addChild/setChildren. + std::vector> getSchemaNodes() const; + + offset_size nodeCount() const; + + const std::shared_ptr& getRoot() const; + + private: + void registerChild(const std::shared_ptr& type); + offset_size allocateStreamOffset(); + + void addNode( + std::vector>& nodes, + const TypeBuilder& type, + std::optional name = std::nullopt) const; + + // Schema builder is building a tree of types. As a tree, it should have a + // single root. We use |roots_| to track that the callers don't forget to + // attach every created type to a parent (and thus constructing a valid tree). + // Every time a new type is constructed (and still wasn't attached to a + // parent), we add it to |roots_|, to indicate that a new root (unattached + // node) exists. Every time a node is attached to a parent, we remove it from + // |roots_|, as it is no longer a free standing node. + // We also use |roots_| to identify two other misuses: + // 1. If a node was created using a different schema builder instance and then + // was attached to a parent of a different schema builder. + // 2. Attaching a node more than once to a parent. + folly::F14FastSet> roots_; + offset_size currentOffset_ = 0; + + friend class ScalarTypeBuilder; + friend class LengthsTypeBuilder; + friend class ArrayTypeBuilder; + friend class ArrayWithOffsetsTypeBuilder; + friend class RowTypeBuilder; + friend class MapTypeBuilder; + friend class FlatMapTypeBuilder; + friend std::ostream& operator<<( + std::ostream& out, + const SchemaBuilder& schema); +}; + +std::ostream& operator<<(std::ostream& out, const SchemaBuilder& schema); + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/SchemaReader.cpp b/dwio/nimble/velox/SchemaReader.cpp new file mode 100644 index 0000000..5afd97e --- /dev/null +++ b/dwio/nimble/velox/SchemaReader.cpp @@ -0,0 +1,521 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/SchemaReader.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/velox/SchemaTypes.h" +#include "folly/container/F14Map.h" + +namespace facebook::nimble { + +namespace { + +inline std::string getKindName(Kind kind) { + static folly::F14FastMap names{ + {Kind::Scalar, "Scalar"}, + {Kind::Row, "Row"}, + {Kind::Array, "Array"}, + {Kind::Map, "Map"}, + {Kind::FlatMap, "FlatMap"}, + }; + + auto it = names.find(kind); + if (UNLIKELY(it == names.end())) { + return folly::to(""); + } + + return it->second; +} + +struct NamedType { + std::shared_ptr type; + std::optional name; +}; + +} // namespace + +Type::Type(Kind kind) : kind_{kind} {} + +Kind Type::kind() const { + return kind_; +} + +bool Type::isScalar() const { + return kind_ == Kind::Scalar; +} + +bool Type::isRow() const { + return kind_ == Kind::Row; +} + +bool Type::isArray() const { + return kind_ == Kind::Array; +} + +bool Type::isArrayWithOffsets() const { + return kind_ == Kind::ArrayWithOffsets; +} + +bool Type::isMap() const { + return kind_ == Kind::Map; +} + +bool Type::isFlatMap() const { + return kind_ == Kind::FlatMap; +} + +const ScalarType& Type::asScalar() const { + NIMBLE_ASSERT( + isScalar(), + fmt::format( + "Cannot cast to Scalar. Current type is {}.", getKindName(kind_))); + return dynamic_cast(*this); +} + +const RowType& Type::asRow() const { + NIMBLE_ASSERT( + isRow(), + fmt::format( + "Cannot cast to Row. Current type is {}.", getKindName(kind_))); + return dynamic_cast(*this); +} + +const ArrayType& Type::asArray() const { + NIMBLE_ASSERT( + isArray(), + fmt::format( + "Cannot cast to Array. Current type is {}.", getKindName(kind_))); + return dynamic_cast(*this); +} + +const ArrayWithOffsetsType& Type::asArrayWithOffsets() const { + NIMBLE_ASSERT( + isArrayWithOffsets(), + fmt::format( + "Cannot cast to ArrayWithOffsets. Current type is {}.", + getKindName(kind_))); + return dynamic_cast(*this); +} + +const MapType& Type::asMap() const { + NIMBLE_ASSERT( + isMap(), + fmt::format( + "Cannot cast to Map. Current type is {}.", getKindName(kind_))); + return dynamic_cast(*this); +} + +const FlatMapType& Type::asFlatMap() const { + NIMBLE_ASSERT( + isFlatMap(), + fmt::format( + "Cannot cast to FlatMap. Current type is {}.", getKindName(kind_))); + return dynamic_cast(*this); +} + +ScalarType::ScalarType(StreamDescriptor scalarDescriptor) + : Type(Kind::Scalar), scalarDescriptor_{std::move(scalarDescriptor)} {} + +const StreamDescriptor& ScalarType::scalarDescriptor() const { + return scalarDescriptor_; +} + +ArrayType::ArrayType( + StreamDescriptor lengthsDescriptor, + std::shared_ptr elements) + : Type(Kind::Array), + lengthsDescriptor_{std::move(lengthsDescriptor)}, + elements_{std::move(elements)} {} + +const StreamDescriptor& ArrayType::lengthsDescriptor() const { + return lengthsDescriptor_; +} + +const std::shared_ptr& ArrayType::elements() const { + return elements_; +} + +MapType::MapType( + StreamDescriptor lengthsDescriptor, + std::shared_ptr keys, + std::shared_ptr values) + : Type(Kind::Map), + lengthsDescriptor_{std::move(lengthsDescriptor)}, + keys_{std::move(keys)}, + values_{std::move(values)} {} + +const StreamDescriptor& MapType::lengthsDescriptor() const { + return lengthsDescriptor_; +} + +const std::shared_ptr& MapType::keys() const { + return keys_; +} + +const std::shared_ptr& MapType::values() const { + return values_; +} + +RowType::RowType( + StreamDescriptor nullsDescriptor, + std::vector names, + std::vector> children) + : Type(Kind::Row), + nullsDescriptor_{std::move(nullsDescriptor)}, + names_{std::move(names)}, + children_{std::move(children)} { + NIMBLE_ASSERT( + names_.size() == children_.size(), + fmt::format( + "Size mismatch. names: {} vs children: {}.", + names_.size(), + children_.size())); +} + +const StreamDescriptor& RowType::nullsDescriptor() const { + return nullsDescriptor_; +} + +size_t RowType::childrenCount() const { + return children_.size(); +} + +const std::shared_ptr& RowType::childAt(size_t index) const { + NIMBLE_ASSERT( + index < children_.size(), + fmt::format( + "Index out of range. index: {}, size: {}.", index, children_.size())); + return children_[index]; +} + +const std::string& RowType::nameAt(size_t index) const { + NIMBLE_ASSERT( + index < children_.size(), + fmt::format( + "Index out of range. index: {}, size: {}.", index, children_.size())); + return names_[index]; +} + +FlatMapType::FlatMapType( + StreamDescriptor nullsDescriptor, + ScalarKind keyScalarKind, + std::vector names, + std::vector> inMapDescriptors, + std::vector> children) + : Type(Kind::FlatMap), + nullsDescriptor_{nullsDescriptor}, + keyScalarKind_{keyScalarKind}, + names_{std::move(names)}, + inMapDescriptors_{std::move(inMapDescriptors)}, + children_{std::move(children)} { + NIMBLE_ASSERT( + names_.size() == children_.size() && + inMapDescriptors_.size() == children_.size(), + fmt::format( + "Size mismatch. names: {} vs inMaps: {} vs children: {}.", + names_.size(), + inMapDescriptors_.size(), + children_.size())); +} + +const StreamDescriptor& FlatMapType::nullsDescriptor() const { + return nullsDescriptor_; +} + +const StreamDescriptor& FlatMapType::inMapDescriptorAt(size_t index) const { + NIMBLE_ASSERT( + index < inMapDescriptors_.size(), + fmt::format( + "Index out of range. index: {}, size: {}.", + index, + inMapDescriptors_.size())); + return *inMapDescriptors_[index]; +} + +ScalarKind FlatMapType::keyScalarKind() const { + return keyScalarKind_; +} + +size_t FlatMapType::childrenCount() const { + return children_.size(); +} + +const std::shared_ptr& FlatMapType::childAt(size_t index) const { + NIMBLE_ASSERT( + index < children_.size(), + fmt::format( + "Index out of range. index: {}, size: {}.", index, children_.size())); + return children_[index]; +} + +const std::string& FlatMapType::nameAt(size_t index) const { + NIMBLE_ASSERT( + index < names_.size(), + fmt::format( + "Index out of range. index: {}, size: {}.", index, names_.size())); + return names_[index]; +} + +ArrayWithOffsetsType::ArrayWithOffsetsType( + StreamDescriptor offsetsDescriptor, + StreamDescriptor lengthsDescriptor, + std::shared_ptr elements) + : Type(Kind::ArrayWithOffsets), + offsetsDescriptor_{std::move(offsetsDescriptor)}, + lengthsDescriptor_{std::move(lengthsDescriptor)}, + elements_{std::move(elements)} {} + +const StreamDescriptor& ArrayWithOffsetsType::offsetsDescriptor() const { + return offsetsDescriptor_; +} + +const StreamDescriptor& ArrayWithOffsetsType::lengthsDescriptor() const { + return lengthsDescriptor_; +} + +const std::shared_ptr& ArrayWithOffsetsType::elements() const { + return elements_; +} + +NamedType getType( + offset_size& index, + const std::vector>& nodes) { + NIMBLE_DASSERT(index < nodes.size(), "Index out of range."); + auto& node = nodes[index++]; + auto offset = node->offset(); + auto kind = node->kind(); + switch (kind) { + case Kind::Scalar: { + return { + .type = std::make_shared( + StreamDescriptor{offset, node->scalarKind()}), + .name = node->name()}; + } + case Kind::Array: { + auto elements = getType(index, nodes).type; + return { + .type = std::make_shared( + StreamDescriptor{offset, ScalarKind::UInt32}, + std::move(elements)), + .name = node->name()}; + } + case Kind::Map: { + auto keys = getType(index, nodes).type; + auto values = getType(index, nodes).type; + return { + .type = std::make_shared( + StreamDescriptor{offset, ScalarKind::UInt32}, + std::move(keys), + std::move(values)), + .name = node->name()}; + } + case Kind::Row: { + auto childrenCount = node->childrenCount(); + std::vector names{childrenCount}; + std::vector> children{childrenCount}; + for (auto i = 0; i < childrenCount; ++i) { + auto namedType = getType(index, nodes); + NIMBLE_ASSERT( + namedType.name.has_value(), "Row fields must have names."); + children[i] = namedType.type; + names[i] = namedType.name.value(); + } + return { + .type = std::make_shared( + StreamDescriptor{offset, ScalarKind::Bool}, + std::move(names), + std::move(children)), + .name = node->name()}; + } + case Kind::FlatMap: { + auto childrenCount = node->childrenCount(); + + std::vector names{childrenCount}; + std::vector> inMapDescriptors{ + childrenCount}; + std::vector> children{childrenCount}; + + for (auto i = 0; i < childrenCount; ++i) { + NIMBLE_DASSERT(index < nodes.size(), "Unexpected node index."); + auto& inMapNode = nodes[index++]; + NIMBLE_ASSERT( + inMapNode->kind() == Kind::Scalar && + inMapNode->scalarKind() == ScalarKind::Bool, + "Flat map in-map field must have a boolean scalar type."); + NIMBLE_ASSERT( + inMapNode->name().has_value(), "Flat map fields must have names."); + auto field = getType(index, nodes); + names[i] = inMapNode->name().value(); + inMapDescriptors[i] = std::make_unique( + inMapNode->offset(), inMapNode->scalarKind()); + children[i] = field.type; + } + return { + .type = std::make_shared( + StreamDescriptor{offset, ScalarKind::Bool}, + node->scalarKind(), + std::move(names), + std::move(inMapDescriptors), + std::move(children)), + .name = node->name()}; + } + case Kind::ArrayWithOffsets: { + auto& offsetsNode = nodes[index++]; + NIMBLE_ASSERT( + offsetsNode->kind() == Kind::Scalar && + offsetsNode->scalarKind() == ScalarKind::UInt32, + "Array with offsets field must have a uint32 scalar type."); + auto elements = getType(index, nodes).type; + return { + .type = std::make_shared( + StreamDescriptor{ + offsetsNode->offset(), offsetsNode->scalarKind()}, + StreamDescriptor{offset, ScalarKind::UInt32}, + std::move(elements)), + .name = node->name()}; + } + + default: { + NIMBLE_UNREACHABLE(fmt::format("Unknown node kind: ", toString(kind))); + } + } +} + +std::shared_ptr SchemaReader::getSchema( + const std::vector>& nodes) { + offset_size index = 0; + auto namedType = getType(index, nodes); + return namedType.type; +} + +void traverseSchema( + size_t& index, + uint32_t level, + const std::shared_ptr& type, + const std::function< + void(uint32_t, const Type&, const SchemaReader::NodeInfo&)>& visitor, + const SchemaReader::NodeInfo& info) { + visitor(level, *type, info); + ++index; + switch (type->kind()) { + case Kind::Scalar: + break; + case Kind::Row: { + auto& row = type->asRow(); + auto childrenCount = row.childrenCount(); + for (size_t i = 0; i < childrenCount; ++i) { + traverseSchema( + index, + level + 1, + row.childAt(i), + visitor, + {.name = row.nameAt(i), + .parentType = type.get(), + .placeInSibling = i}); + } + break; + } + case Kind::Array: { + auto& array = type->asArray(); + traverseSchema( + index, + level + 1, + array.elements(), + visitor, + {.name = "elements", .parentType = type.get()}); + break; + } + case Kind::ArrayWithOffsets: { + auto& arrayWithOffsets = type->asArrayWithOffsets(); + traverseSchema( + index, + level + 1, + arrayWithOffsets.elements(), + visitor, + {.name = "elements", .parentType = type.get()}); + + break; + } + case Kind::Map: { + auto& map = type->asMap(); + traverseSchema( + index, + level + 1, + map.keys(), + visitor, + {.name = "keys", .parentType = type.get(), .placeInSibling = 0}); + traverseSchema( + index, + level + 1, + map.values(), + visitor, + {.name = "values", .parentType = type.get(), .placeInSibling = 1}); + break; + } + case Kind::FlatMap: { + auto& map = type->asFlatMap(); + for (size_t i = 0; i < map.childrenCount(); ++i) { + traverseSchema( + index, + level + 1, + map.childAt(i), + visitor, + {.name = map.nameAt(i), + .parentType = type.get(), + .placeInSibling = i}); + } + break; + } + } +} + +void SchemaReader::traverseSchema( + const std::shared_ptr& root, + std::function + visitor) { + size_t index = 0; + nimble::traverseSchema( + index, 0, root, visitor, {.name = "root", .parentType = nullptr}); +} + +std::ostream& operator<<( + std::ostream& out, + const std::shared_ptr& root) { + SchemaReader::traverseSchema( + root, + [&out](uint32_t level, const Type& type, const SchemaReader::NodeInfo&) { + out << std::string((std::basic_string::size_type)level * 2, ' '); + if (type.isScalar()) { + auto& scalar = type.asScalar(); + out << "[" << scalar.scalarDescriptor().offset() << "]" + << toString(scalar.scalarDescriptor().scalarKind()) << "\n"; + } else if (type.isArray()) { + out << "[" << type.asArray().lengthsDescriptor().offset() << "]" + << "ARRAY\n"; + } else if (type.isArrayWithOffsets()) { + const auto& array = type.asArrayWithOffsets(); + out << "[o:" << array.offsetsDescriptor().offset() + << ",l:" << array.lengthsDescriptor().offset() << "]" + << "OFFSETARRAY\n"; + } else if (type.isRow()) { + auto& row = type.asRow(); + out << "[" << row.nullsDescriptor().offset() << "]" << "ROW["; + for (auto i = 0; i < row.childrenCount(); ++i) { + out << row.nameAt(i) << (i < row.childrenCount() - 1 ? "," : ""); + } + out << "]\n"; + } else if (type.isMap()) { + out << "[" << type.asMap().lengthsDescriptor().offset() << "]" + << "MAP\n"; + } else if (type.isFlatMap()) { + auto& map = type.asFlatMap(); + out << "[" << map.nullsDescriptor().offset() << "]" << "FLATMAP["; + for (auto i = 0; i < map.childrenCount(); ++i) { + out << map.nameAt(i) << (i < map.childrenCount() - 1 ? "," : ""); + } + out << "]\n"; + } + }); + return out; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/SchemaReader.h b/dwio/nimble/velox/SchemaReader.h new file mode 100644 index 0000000..358684d --- /dev/null +++ b/dwio/nimble/velox/SchemaReader.h @@ -0,0 +1,187 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include +#include + +#include "dwio/nimble/velox/SchemaTypes.h" + +// Schema reader provides a strongly typed, tree like, reader friendly facade on +// top of a flat tablet schema. +// Schema stored in a tablet footer is a DFS representation of a type tree. +// Reconstructing a tree out of this flat representation require deep knowledge +// of how each complex type (rows, arrays, maps and flat maps) are laid out. +// Using this class, it is possible to reconstruct an easy to use tree +// representation on the logical type tree. +// The main usage of this class is to allow efficient retrieval of subsets of +// the schema tree (used in column and feature projection), as it can +// efficiently skip full sections of the schema tree (which cannot be done +// efficiently using a flat schema representation, that relies on schema node +// order for decoding). +namespace facebook::nimble { + +class ScalarType; +class RowType; +class ArrayType; +class ArrayWithOffsetsType; +class MapType; +class FlatMapType; + +class Type { + public: + Kind kind() const; + + bool isScalar() const; + bool isRow() const; + bool isArray() const; + bool isArrayWithOffsets() const; + bool isMap() const; + bool isFlatMap() const; + + const ScalarType& asScalar() const; + const RowType& asRow() const; + const ArrayType& asArray() const; + const ArrayWithOffsetsType& asArrayWithOffsets() const; + const MapType& asMap() const; + const FlatMapType& asFlatMap() const; + + protected: + explicit Type(Kind kind); + + virtual ~Type() = default; + + private: + Type(const Type&) = delete; + Type(Type&&) = delete; + + Kind kind_; +}; + +class ScalarType : public Type { + public: + explicit ScalarType(StreamDescriptor scalarDescriptor); + + const StreamDescriptor& scalarDescriptor() const; + + private: + StreamDescriptor scalarDescriptor_; +}; + +class ArrayType : public Type { + public: + ArrayType( + StreamDescriptor lengthsDescriptor, + std::shared_ptr elements); + + const StreamDescriptor& lengthsDescriptor() const; + const std::shared_ptr& elements() const; + + protected: + StreamDescriptor lengthsDescriptor_; + std::shared_ptr elements_; +}; + +class MapType : public Type { + public: + MapType( + StreamDescriptor lengthsDescriptor, + std::shared_ptr keys, + std::shared_ptr values); + + const StreamDescriptor& lengthsDescriptor() const; + const std::shared_ptr& keys() const; + const std::shared_ptr& values() const; + + protected: + StreamDescriptor lengthsDescriptor_; + std::shared_ptr keys_; + std::shared_ptr values_; +}; + +class RowType : public Type { + public: + RowType( + StreamDescriptor nullsDescriptor, + std::vector names, + std::vector> children); + + const StreamDescriptor& nullsDescriptor() const; + size_t childrenCount() const; + const std::shared_ptr& childAt(size_t index) const; + const std::string& nameAt(size_t index) const; + + protected: + StreamDescriptor nullsDescriptor_; + std::vector names_; + std::vector> children_; +}; + +class FlatMapType : public Type { + public: + FlatMapType( + StreamDescriptor nullsDescriptor, + ScalarKind keyScalarKind, + std::vector names, + std::vector> inMapDescriptors, + std::vector> children); + + const StreamDescriptor& nullsDescriptor() const; + const StreamDescriptor& inMapDescriptorAt(size_t index) const; + ScalarKind keyScalarKind() const; + size_t childrenCount() const; + const std::shared_ptr& childAt(size_t index) const; + const std::string& nameAt(size_t index) const; + + private: + StreamDescriptor nullsDescriptor_; + ScalarKind keyScalarKind_; + std::vector names_; + std::vector> inMapDescriptors_; + std::vector> children_; +}; + +class ArrayWithOffsetsType : public Type { + public: + ArrayWithOffsetsType( + StreamDescriptor offsetsDescriptor, + StreamDescriptor lengthsDescriptor, + std::shared_ptr elements); + + const StreamDescriptor& offsetsDescriptor() const; + const StreamDescriptor& lengthsDescriptor() const; + const std::shared_ptr& elements() const; + + protected: + StreamDescriptor offsetsDescriptor_; + StreamDescriptor lengthsDescriptor_; + std::shared_ptr elements_; +}; + +class SchemaReader { + public: + // Construct type tree from an ordered list of schema nodes. + // If the schema nodes are ordered incorrectly the behavior is undefined. + static std::shared_ptr getSchema( + const std::vector>& nodes); + + struct NodeInfo { + std::string_view name; + const Type* parentType; + size_t placeInSibling = 0; + }; + + static void traverseSchema( + const std::shared_ptr& root, + std::function visitor); +}; + +std::ostream& operator<<( + std::ostream& out, + const std::shared_ptr& root); + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/SchemaSerialization.cpp b/dwio/nimble/velox/SchemaSerialization.cpp new file mode 100644 index 0000000..acaad08 --- /dev/null +++ b/dwio/nimble/velox/SchemaSerialization.cpp @@ -0,0 +1,213 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/SchemaSerialization.h" +#include "dwio/nimble/velox/SchemaReader.h" + +namespace facebook::nimble { + +namespace { + +constexpr uint32_t kInitialSchemaSectionSize = 1 << 20; // 1MB + +serialization::Kind nodeToSerializationKind(const SchemaNode* node) { + switch (node->kind()) { + case Kind::Scalar: { + switch (node->scalarKind()) { + case ScalarKind::Int8: + return serialization::Kind_Int8; + case ScalarKind::UInt8: + return serialization::Kind_UInt8; + case ScalarKind::Int16: + return serialization::Kind_Int16; + case ScalarKind::UInt16: + return serialization::Kind_UInt16; + case ScalarKind::Int32: + return serialization::Kind_Int32; + case ScalarKind::UInt32: + return serialization::Kind_UInt32; + case ScalarKind::Int64: + return serialization::Kind_Int64; + case ScalarKind::UInt64: + return serialization::Kind_UInt64; + case ScalarKind::Float: + return serialization::Kind_Float; + case ScalarKind::Double: + return serialization::Kind_Double; + case ScalarKind::Bool: + return serialization::Kind_Bool; + case ScalarKind::String: + return serialization::Kind_String; + case ScalarKind::Binary: + return serialization::Kind_Binary; + default: + NIMBLE_UNREACHABLE(fmt::format( + "Unknown scalar kind {}.", toString(node->scalarKind()))); + } + } + case Kind::Array: + return serialization::Kind_Array; + case Kind::ArrayWithOffsets: + return serialization::Kind_ArrayWithOffsets; + case Kind::Row: + return serialization::Kind_Row; + case Kind::Map: + return serialization::Kind_Map; + case Kind::FlatMap: { + switch (node->scalarKind()) { + case ScalarKind::Int8: + return serialization::Kind_FlatMapInt8; + case ScalarKind::UInt8: + return serialization::Kind_FlatMapUInt8; + case ScalarKind::Int16: + return serialization::Kind_FlatMapInt16; + case ScalarKind::UInt16: + return serialization::Kind_FlatMapUInt16; + case ScalarKind::Int32: + return serialization::Kind_FlatMapInt32; + case ScalarKind::UInt32: + return serialization::Kind_FlatMapUInt32; + case ScalarKind::Int64: + return serialization::Kind_FlatMapInt64; + case ScalarKind::UInt64: + return serialization::Kind_FlatMapUInt64; + case ScalarKind::Float: + return serialization::Kind_FlatMapFloat; + case ScalarKind::Double: + return serialization::Kind_FlatMapDouble; + case ScalarKind::Bool: + return serialization::Kind_FlatMapBool; + case ScalarKind::String: + return serialization::Kind_FlatMapString; + case ScalarKind::Binary: + return serialization::Kind_FlatMapBinary; + default: + NIMBLE_UNREACHABLE(fmt::format( + "Unknown flat map key kind {}.", toString(node->scalarKind()))); + } + } + default: + NIMBLE_UNREACHABLE( + fmt::format("Unknown node kind {}.", toString(node->kind()))); + } +} + +std::pair serializationNodeToKind( + const serialization::SchemaNode* node) { + switch (node->kind()) { + case nimble::serialization::Kind_Int8: + return {Kind::Scalar, ScalarKind::Int8}; + case nimble::serialization::Kind_UInt8: + return {Kind::Scalar, ScalarKind::UInt8}; + case nimble::serialization::Kind_Int16: + return {Kind::Scalar, ScalarKind::Int16}; + case nimble::serialization::Kind_UInt16: + return {Kind::Scalar, ScalarKind::UInt16}; + case nimble::serialization::Kind_Int32: + return {Kind::Scalar, ScalarKind::Int32}; + case nimble::serialization::Kind_UInt32: + return {Kind::Scalar, ScalarKind::UInt32}; + case nimble::serialization::Kind_Int64: + return {Kind::Scalar, ScalarKind::Int64}; + case nimble::serialization::Kind_UInt64: + return {Kind::Scalar, ScalarKind::UInt64}; + case nimble::serialization::Kind_Float: + return {Kind::Scalar, ScalarKind::Float}; + case nimble::serialization::Kind_Double: + return {Kind::Scalar, ScalarKind::Double}; + case nimble::serialization::Kind_Bool: + return {Kind::Scalar, ScalarKind::Bool}; + case nimble::serialization::Kind_String: + return {Kind::Scalar, ScalarKind::String}; + case nimble::serialization::Kind_Binary: + return {Kind::Scalar, ScalarKind::Binary}; + case nimble::serialization::Kind_Row: + return {Kind::Row, ScalarKind::Undefined}; + case nimble::serialization::Kind_Array: + return {Kind::Array, ScalarKind::Undefined}; + case nimble::serialization::Kind_ArrayWithOffsets: + return {Kind::ArrayWithOffsets, ScalarKind::Undefined}; + case nimble::serialization::Kind_Map: + return {Kind::Map, ScalarKind::Undefined}; + case nimble::serialization::Kind_FlatMapInt8: + return {Kind::FlatMap, ScalarKind::Int8}; + case nimble::serialization::Kind_FlatMapUInt8: + return {Kind::FlatMap, ScalarKind::UInt8}; + case nimble::serialization::Kind_FlatMapInt16: + return {Kind::FlatMap, ScalarKind::Int16}; + case nimble::serialization::Kind_FlatMapUInt16: + return {Kind::FlatMap, ScalarKind::UInt16}; + case nimble::serialization::Kind_FlatMapInt32: + return {Kind::FlatMap, ScalarKind::Int32}; + case nimble::serialization::Kind_FlatMapUInt32: + return {Kind::FlatMap, ScalarKind::UInt32}; + case nimble::serialization::Kind_FlatMapInt64: + return {Kind::FlatMap, ScalarKind::Int64}; + case nimble::serialization::Kind_FlatMapUInt64: + return {Kind::FlatMap, ScalarKind::UInt64}; + case nimble::serialization::Kind_FlatMapFloat: + return {Kind::FlatMap, ScalarKind::Float}; + case nimble::serialization::Kind_FlatMapDouble: + return {Kind::FlatMap, ScalarKind::Double}; + case nimble::serialization::Kind_FlatMapBool: + return {Kind::FlatMap, ScalarKind::Bool}; + case nimble::serialization::Kind_FlatMapString: + return {Kind::FlatMap, ScalarKind::String}; + case nimble::serialization::Kind_FlatMapBinary: + return {Kind::FlatMap, ScalarKind::Binary}; + default: + NIMBLE_NOT_SUPPORTED(fmt::format( + "Unknown schema node kind {}.", + nimble::serialization::EnumNameKind(node->kind()))); + } +} + +} // namespace + +SchemaSerializer::SchemaSerializer() : builder_{kInitialSchemaSectionSize} {} + +std::string_view SchemaSerializer::serialize( + const SchemaBuilder& schemaBuilder) { + auto nodes = schemaBuilder.getSchemaNodes(); + builder_.Clear(); + auto schema = + builder_.CreateVector>( + nodes.size(), [this, &nodes](size_t i) { + auto& node = nodes[i]; + return serialization::CreateSchemaNode( + builder_, + nodeToSerializationKind(node.get()), + node->childrenCount(), + node->name().has_value() + ? builder_.CreateString(node->name().value()) + : 0, + node->offset()); + }); + + builder_.Finish(serialization::CreateSchema(builder_, schema)); + return { + reinterpret_cast(builder_.GetBufferPointer()), + builder_.GetSize()}; +} + +std::shared_ptr SchemaDeserializer::deserialize( + std::string_view input) { + auto schema = flatbuffers::GetRoot(input.data()); + auto nodeCount = schema->nodes()->size(); + std::vector> nodes(nodeCount); + + for (auto i = 0; i < nodeCount; ++i) { + auto* node = schema->nodes()->Get(i); + auto kind = serializationNodeToKind(node); + nodes[i] = std::make_unique( + kind.first, + node->offset(), + kind.second, + node->name() ? std::optional(node->name()->str()) + : std::nullopt, + node->children()); + } + + return SchemaReader::getSchema(nodes); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/SchemaSerialization.h b/dwio/nimble/velox/SchemaSerialization.h new file mode 100644 index 0000000..ab61e30 --- /dev/null +++ b/dwio/nimble/velox/SchemaSerialization.h @@ -0,0 +1,26 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/velox/SchemaBuilder.h" +#include "dwio/nimble/velox/SchemaGenerated.h" +#include "dwio/nimble/velox/SchemaReader.h" + +namespace facebook::nimble { + +class SchemaSerializer { + public: + SchemaSerializer(); + + std::string_view serialize(const SchemaBuilder& builder); + + private: + flatbuffers::FlatBufferBuilder builder_; +}; + +class SchemaDeserializer { + public: + static std::shared_ptr deserialize(std::string_view schema); +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/SchemaTypes.cpp b/dwio/nimble/velox/SchemaTypes.cpp new file mode 100644 index 0000000..fb48042 --- /dev/null +++ b/dwio/nimble/velox/SchemaTypes.cpp @@ -0,0 +1,52 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/SchemaTypes.h" + +#include "dwio/nimble/common/Exceptions.h" + +#include + +namespace facebook::nimble { +std::string toString(ScalarKind kind) { + switch (kind) { +#define CASE(KIND) \ + case ScalarKind::KIND: { \ + return #KIND; \ + } + CASE(Int8); + CASE(UInt8); + CASE(Int16); + CASE(UInt16); + CASE(Int32); + CASE(UInt32); + CASE(Int64); + CASE(UInt64); + CASE(Float); + CASE(Double); + CASE(Bool); + CASE(String); + CASE(Binary); + CASE(Undefined); +#undef CASE + } + NIMBLE_UNREACHABLE(fmt::format("Unknown: {}.", static_cast(kind))); +} + +std::string toString(Kind kind) { + switch (kind) { +#define CASE(KIND) \ + case Kind::KIND: { \ + return #KIND; \ + } + CASE(Scalar); + CASE(Row); + CASE(Array); + CASE(ArrayWithOffsets); + CASE(Map); + CASE(FlatMap); +#undef CASE + } + NIMBLE_UNREACHABLE(fmt::format("Unknown: {}.", static_cast(kind))); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/SchemaTypes.h b/dwio/nimble/velox/SchemaTypes.h new file mode 100644 index 0000000..b6f9010 --- /dev/null +++ b/dwio/nimble/velox/SchemaTypes.h @@ -0,0 +1,102 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include + +namespace facebook::nimble { + +using offset_size = uint32_t; + +enum class ScalarKind : uint8_t { + Int8, + UInt8, + Int16, + UInt16, + Int32, + UInt32, + Int64, + UInt64, + Float, + Double, + Bool, + String, + Binary, + + Undefined = 255, +}; + +enum class Kind : uint8_t { + Scalar, + Row, + Array, + ArrayWithOffsets, + Map, + FlatMap, +}; + +std::string toString(ScalarKind kind); +std::string toString(Kind kind); + +class StreamDescriptor { + public: + StreamDescriptor(offset_size offset, ScalarKind scalarKind) + : offset_{offset}, scalarKind_{scalarKind} {} + + offset_size offset() const { + return offset_; + } + + ScalarKind scalarKind() const { + return scalarKind_; + } + + private: + offset_size offset_; + ScalarKind scalarKind_; +}; + +class SchemaNode { + public: + SchemaNode( + Kind kind, + offset_size offset, + ScalarKind scalarKind, + std::optional name = std::nullopt, + size_t childrenCount = 0) + : kind_{kind}, + offset_{offset}, + name_{std::move(name)}, + scalarKind_{scalarKind}, + childrenCount_{childrenCount} {} + + Kind kind() const { + return kind_; + } + + size_t childrenCount() const { + return childrenCount_; + } + + offset_size offset() const { + return offset_; + } + + std::optional name() const { + return name_; + } + + ScalarKind scalarKind() const { + return scalarKind_; + } + + private: + Kind kind_; + offset_size offset_; + std::optional name_; + ScalarKind scalarKind_; + size_t childrenCount_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/SchemaUtils.cpp b/dwio/nimble/velox/SchemaUtils.cpp new file mode 100644 index 0000000..a83877b --- /dev/null +++ b/dwio/nimble/velox/SchemaUtils.cpp @@ -0,0 +1,154 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/SchemaUtils.h" +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/velox/SchemaBuilder.h" + +namespace facebook::nimble { + +namespace { + +velox::TypePtr convertToVeloxScalarType(ScalarKind scalarKind) { + switch (scalarKind) { + case ScalarKind::Int8: + return velox::ScalarType::create(); + case ScalarKind::Int16: + return velox::ScalarType::create(); + case ScalarKind::Int32: + return velox::ScalarType::create(); + case ScalarKind::Int64: + return velox::ScalarType::create(); + case ScalarKind::Float: + return velox::ScalarType::create(); + case ScalarKind::Double: + return velox::ScalarType::create(); + case ScalarKind::Bool: + return velox::ScalarType::create(); + case ScalarKind::String: + return velox::ScalarType::create(); + case ScalarKind::Binary: + return velox::ScalarType::create(); + case ScalarKind::UInt8: + case ScalarKind::UInt16: + case ScalarKind::UInt32: + case ScalarKind::UInt64: + case ScalarKind::Undefined: + NIMBLE_NOT_SUPPORTED(fmt::format( + "Scalar kind {} is not supported by Velox.", toString(scalarKind))); + } + NIMBLE_UNREACHABLE( + fmt::format("Unknown scalarKind: {}.", toString(scalarKind))); +} + +std::shared_ptr convertToNimbleType( + SchemaBuilder& builder, + const velox::Type& type) { + switch (type.kind()) { + case velox::TypeKind::BOOLEAN: + return builder.createScalarTypeBuilder(ScalarKind::Bool); + case velox::TypeKind::TINYINT: + return builder.createScalarTypeBuilder(ScalarKind::Int8); + case velox::TypeKind::SMALLINT: + return builder.createScalarTypeBuilder(ScalarKind::Int16); + case velox::TypeKind::INTEGER: + return builder.createScalarTypeBuilder(ScalarKind::Int32); + case velox::TypeKind::BIGINT: + return builder.createScalarTypeBuilder(ScalarKind::Int64); + case velox::TypeKind::REAL: + return builder.createScalarTypeBuilder(ScalarKind::Float); + case velox::TypeKind::DOUBLE: + return builder.createScalarTypeBuilder(ScalarKind::Double); + case velox::TypeKind::VARCHAR: + return builder.createScalarTypeBuilder(ScalarKind::String); + case velox::TypeKind::VARBINARY: + return builder.createScalarTypeBuilder(ScalarKind::Binary); + case velox::TypeKind::ARRAY: { + auto& arrayType = type.asArray(); + auto nimbleType = builder.createArrayTypeBuilder(); + nimbleType->setChildren( + convertToNimbleType(builder, *arrayType.elementType())); + return nimbleType; + } + case velox::TypeKind::MAP: { + auto& mapType = type.asMap(); + auto nimbleType = builder.createMapTypeBuilder(); + auto key = convertToNimbleType(builder, *mapType.keyType()); + auto value = convertToNimbleType(builder, *mapType.valueType()); + nimbleType->setChildren(key, value); + return nimbleType; + } + case velox::TypeKind::ROW: { + auto& rowType = type.asRow(); + auto nimbleType = builder.createRowTypeBuilder(rowType.size()); + for (size_t i = 0; i < rowType.size(); ++i) { + nimbleType->addChild( + rowType.nameOf(i), + convertToNimbleType(builder, *rowType.childAt(i))); + } + return nimbleType; + } + default: + NIMBLE_NOT_SUPPORTED( + fmt::format("Unsupported type kind {}.", type.kind())); + } +} + +} // namespace + +velox::TypePtr convertToVeloxType(const Type& type) { + switch (type.kind()) { + case Kind::Scalar: { + return convertToVeloxScalarType( + type.asScalar().scalarDescriptor().scalarKind()); + } + case Kind::Row: { + const auto& rowType = type.asRow(); + std::vector names; + std::vector children; + names.reserve(rowType.childrenCount()); + children.reserve(rowType.childrenCount()); + for (size_t i = 0; i < rowType.childrenCount(); ++i) { + names.push_back(rowType.nameAt(i)); + children.push_back(convertToVeloxType(*rowType.childAt(i))); + } + return std::make_shared( + std::move(names), std::move(children)); + } + case Kind::Array: { + const auto& arrayType = type.asArray(); + return std::make_shared( + convertToVeloxType(*arrayType.elements())); + } + case Kind::ArrayWithOffsets: { + const auto& arrayWithOffsetsType = type.asArrayWithOffsets(); + return std::make_shared( + convertToVeloxType(*arrayWithOffsetsType.elements())); + } + case Kind::Map: { + const auto& mapType = type.asMap(); + return std::make_shared( + convertToVeloxType(*mapType.keys()), + convertToVeloxType(*mapType.values())); + } + case Kind::FlatMap: { + const auto& flatMapType = type.asFlatMap(); + return std::make_shared( + convertToVeloxScalarType(flatMapType.keyScalarKind()), + // When Flatmap is empty, we always insert dummy key/value + // to it, so it is guaranteed that flatMapType.childAt(0) + // is always valid. + convertToVeloxType(*flatMapType.childAt(0))); + } + default: + NIMBLE_UNREACHABLE( + fmt::format("Unknown type kind {}.", toString(type.kind()))); + } +} + +std::shared_ptr convertToNimbleType(const velox::Type& type) { + SchemaBuilder builder; + convertToNimbleType(builder, type); + return SchemaReader::getSchema(builder.getSchemaNodes()); +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/SchemaUtils.h b/dwio/nimble/velox/SchemaUtils.h new file mode 100644 index 0000000..2513e00 --- /dev/null +++ b/dwio/nimble/velox/SchemaUtils.h @@ -0,0 +1,14 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/velox/SchemaReader.h" +#include "velox/type/Type.h" + +namespace facebook::nimble { + +velox::TypePtr convertToVeloxType(const Type& type); + +std::shared_ptr convertToNimbleType(const velox::Type& type); + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/Serializer.cpp b/dwio/nimble/velox/Serializer.cpp new file mode 100644 index 0000000..1ee5c0d --- /dev/null +++ b/dwio/nimble/velox/Serializer.cpp @@ -0,0 +1,150 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/Serializer.h" +#include +#include +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/velox/SchemaBuilder.h" +#include "dwio/nimble/velox/SchemaUtils.h" + +namespace facebook::nimble { + +namespace { + +ScalarKind getScalarKind(const Type& type) { + switch (type.kind()) { + case Kind::Scalar: + return type.asScalar().scalarDescriptor().scalarKind(); + case Kind::Row: + case Kind::Array: + case Kind::ArrayWithOffsets: + case Kind::Map: + case Kind::FlatMap: + return ScalarKind::Undefined; + } +} + +void writeMissingStreams( + Vector& buffer, + uint32_t prevIndex, + uint32_t endIndex) { + if (endIndex != prevIndex + 1) { + buffer.extend((endIndex - prevIndex - 1) * sizeof(uint32_t), 0); + } +} + +} // namespace + +std::string_view Serializer::serialize( + const velox::VectorPtr& vector, + const OrderedRanges& ranges) { + buffer_.resize(sizeof(uint32_t)); + auto pos = buffer_.data(); + encoding::writeUint32(ranges.size(), pos); + writer_->write(vector, ranges); + uint32_t lastStream = 0xffffffff; + + for (auto& streamData : context_.streams()) { + auto stream = streamData->descriptor().offset(); + auto nonNulls = streamData->nonNulls(); + auto data = streamData->data(); + + if (data.empty() && nonNulls.empty()) { + continue; + } + + // Current implementation has strong assumption that schema traversal is + // pre-order, hence handles types with offsets in increasing order. This + // assumption may not be true if schema changes based on data shape (ie. + // flatmap). We need to implement different ways of maintaining the order if + // need to support that. + NIMBLE_CHECK( + lastStream + 1 <= stream, + fmt::format("unexpected stream offset {}", stream)); + // We expect streams to arrive in ascending offset order. If there is an + // offset gap, it means that those streams are missing (and will not show up + // later on), so we fill zeros for all of the missing streams. + writeMissingStreams(buffer_, lastStream, stream); + lastStream = stream; + + NIMBLE_CHECK( + nonNulls.empty() || + std::all_of( + nonNulls.begin(), + nonNulls.end(), + [](bool notNull) { return notNull; }), + "nulls not supported"); + auto oldSize = buffer_.size(); + auto scalarKind = streamData->descriptor().scalarKind(); + if (scalarKind == ScalarKind::String || scalarKind == ScalarKind::Binary) { + // TODO: handle string compression + const auto strData = + reinterpret_cast(data.data()); + const auto strDataEnd = + reinterpret_cast(data.end()); + uint32_t size = 0; + for (auto sv = strData; sv < strDataEnd; ++sv) { + size += (sv->size() + sizeof(uint32_t)); + } + buffer_.resize(oldSize + size + sizeof(uint32_t)); + auto pos = buffer_.data() + oldSize; + encoding::writeUint32(size, pos); + for (auto sv = strData; sv < strDataEnd; ++sv) { + encoding::writeString(*sv, pos); + } + } else { + // Size prefix + compression type + actual content + const auto size = data.size(); + buffer_.resize(oldSize + size + sizeof(uint32_t) + 1); + + auto compression = options_.compressionType; + bool writeUncompressed = true; + if (compression != CompressionType::Uncompressed && + size >= options_.compressionThreshold) { + auto pos = buffer_.data() + oldSize + sizeof(uint32_t); + encoding::writeChar(static_cast(compression), pos); + // TODO: share compression implementation + switch (compression) { + case CompressionType::Zstd: { + auto ret = ZSTD_compress( + pos, size, data.data(), size, options_.compressionLevel); + if (ZSTD_isError(ret)) { + NIMBLE_ASSERT( + ZSTD_getErrorCode(ret) == + ZSTD_ErrorCode::ZSTD_error_dstSize_tooSmall, + "zstd error"); + // fall back to uncompressed + } else { + pos = buffer_.data() + oldSize; + encoding::writeUint32(ret + 1, pos); + // reflect the compressed size + buffer_.resize(oldSize + ret + sizeof(uint32_t) + 1); + writeUncompressed = false; + } + break; + } + default: + NIMBLE_NOT_SUPPORTED(fmt::format( + "Unsupported compression {}", toString(compression))); + } + } + + if (writeUncompressed) { + auto pos = buffer_.data() + oldSize; + encoding::writeUint32(size + 1, pos); + encoding::writeChar( + static_cast(CompressionType::Uncompressed), pos); + std::copy(data.data(), data.end(), pos); + } + } + } + + writer_->reset(); + + // Write missing streams similar to above + writeMissingStreams(buffer_, lastStream, context_.schemaBuilder.nodeCount()); + return {buffer_.data(), buffer_.size()}; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/Serializer.h b/dwio/nimble/velox/Serializer.h new file mode 100644 index 0000000..b491315 --- /dev/null +++ b/dwio/nimble/velox/Serializer.h @@ -0,0 +1,47 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/velox/FieldWriter.h" +#include "velox/dwio/common/TypeWithId.h" +#include "velox/vector/BaseVector.h" + +namespace facebook::nimble { + +struct SerializerOptions { + CompressionType compressionType{CompressionType::Uncompressed}; + uint32_t compressionThreshold{0}; + int32_t compressionLevel{0}; +}; + +class Serializer { + public: + Serializer( + SerializerOptions options, + velox::memory::MemoryPool& pool, + const std::shared_ptr& type) + : options_{std::move(options)}, + context_{pool}, + writer_{FieldWriter::create( + context_, + velox::dwio::common::TypeWithId::create(type))}, + buffer_{context_.bufferMemoryPool.get()} {} + + std::string_view serialize( + const velox::VectorPtr& vector, + const OrderedRanges& ranges); + + const SchemaBuilder& schemaBuilder() const { + return context_.schemaBuilder; + } + + private: + SerializerOptions options_; + FieldWriterContext context_; + std::unique_ptr writer_; + Vector buffer_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/StreamLabels.cpp b/dwio/nimble/velox/StreamLabels.cpp new file mode 100644 index 0000000..80f66c4 --- /dev/null +++ b/dwio/nimble/velox/StreamLabels.cpp @@ -0,0 +1,177 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/StreamLabels.h" + +#include "dwio/nimble/common/Exceptions.h" + +namespace facebook::nimble { + +namespace { +void addLabels( + const std::shared_ptr& node, + std::vector& labels, + std::vector& offsetToLabel, + size_t labelIndex, + const std::string& name) { + switch (node->kind()) { + case Kind::Scalar: { + const auto& scalar = node->asScalar(); + const auto offset = scalar.scalarDescriptor().offset(); + NIMBLE_DASSERT(labelIndex < labels.size(), "Unexpected label index."); + NIMBLE_DASSERT(offsetToLabel.size() > offset, "Unexpected offset."); + labels.push_back(labels[labelIndex] + name); + offsetToLabel[offset] = labels.size() - 1; + break; + } + case Kind::Array: { + const auto& array = node->asArray(); + const auto offset = array.lengthsDescriptor().offset(); + NIMBLE_DASSERT(labelIndex < labels.size(), "Unexpected label index."); + NIMBLE_DASSERT(offsetToLabel.size() > offset, "Unexpected offset."); + labels.push_back(labels[labelIndex] + name); + labelIndex = labels.size() - 1; + offsetToLabel[offset] = labelIndex; + addLabels(array.elements(), labels, offsetToLabel, labelIndex, ""); + break; + } + case Kind::Map: { + const auto& map = node->asMap(); + const auto offset = map.lengthsDescriptor().offset(); + NIMBLE_DASSERT(labelIndex < labels.size(), "Unexpected label index."); + NIMBLE_DASSERT(offsetToLabel.size() > offset, "Unexpected offset."); + labels.push_back(labels[labelIndex] + name); + labelIndex = labels.size() - 1; + offsetToLabel[offset] = labelIndex; + addLabels(map.keys(), labels, offsetToLabel, labelIndex, ""); + addLabels(map.values(), labels, offsetToLabel, labelIndex, ""); + break; + } + case Kind::Row: { + const auto& row = node->asRow(); + const auto offset = row.nullsDescriptor().offset(); + NIMBLE_DASSERT(labelIndex < labels.size(), "Unexpected label index."); + NIMBLE_DASSERT(offsetToLabel.size() > offset, "Unexpected offset."); + labels.push_back(labels[labelIndex] + "/"); + labelIndex = labels.size() - 1; + offsetToLabel[offset] = labelIndex; + for (auto i = 0; i < row.childrenCount(); ++i) { + addLabels( + row.childAt(i), + labels, + offsetToLabel, + labelIndex, + folly::to(i)); + } + break; + } + case Kind::FlatMap: { + const auto& map = node->asFlatMap(); + const auto offset = map.nullsDescriptor().offset(); + NIMBLE_DASSERT(labelIndex < labels.size(), "Unexpected label index."); + NIMBLE_DASSERT(offsetToLabel.size() > offset, "Unexpected offset."); + labels.push_back(labels[labelIndex] + name); + labelIndex = labels.size() - 1; + offsetToLabel[offset] = labelIndex; + for (auto i = 0; i < map.childrenCount(); ++i) { + const auto inMapOffset = map.inMapDescriptorAt(i).offset(); + NIMBLE_DASSERT( + offsetToLabel.size() > inMapOffset, "Unexpected offset."); + labels.push_back(labels[labelIndex] + "/" + map.nameAt(i)); + offsetToLabel[inMapOffset] = labels.size() - 1; + } + for (auto i = 0; i < map.childrenCount(); ++i) { + addLabels( + map.childAt(i), + labels, + offsetToLabel, + offsetToLabel[map.inMapDescriptorAt(i).offset()], + ""); + } + break; + } + case Kind::ArrayWithOffsets: { + const auto& array = node->asArrayWithOffsets(); + const auto offsetsOffset = array.offsetsDescriptor().offset(); + const auto lengthsOffset = array.lengthsDescriptor().offset(); + NIMBLE_DASSERT(labelIndex < labels.size(), "Unexpected label index."); + NIMBLE_DASSERT( + offsetToLabel.size() > offsetsOffset, "Unexpected offset."); + NIMBLE_DASSERT( + offsetToLabel.size() > lengthsOffset, "Unexpected offset."); + labels.push_back(labels[labelIndex] + name); + labelIndex = labels.size() - 1; + offsetToLabel[offsetsOffset] = labelIndex; + offsetToLabel[lengthsOffset] = labelIndex; + addLabels(array.elements(), labels, offsetToLabel, labelIndex, ""); + break; + } + } +} +} // namespace + +StreamLabels::StreamLabels(const std::shared_ptr& root) { + size_t maxOffset = 0; + size_t uniqueLabels = 1; + SchemaReader::traverseSchema( + root, + [&](uint32_t level, + const Type& type, + const SchemaReader::NodeInfo& info) { + switch (type.kind()) { + case Kind::Scalar: { + maxOffset = std::max( + maxOffset, type.asScalar().scalarDescriptor().offset()); + break; + } + case Kind::Array: { + maxOffset = std::max( + maxOffset, type.asArray().lengthsDescriptor().offset()); + break; + } + case Kind::Map: { + maxOffset = std::max( + maxOffset, type.asMap().lengthsDescriptor().offset()); + break; + } + case Kind::Row: { + maxOffset = std::max( + maxOffset, type.asRow().nullsDescriptor().offset()); + uniqueLabels += type.asRow().childrenCount(); + break; + } + case Kind::FlatMap: { + const auto& map = type.asFlatMap(); + maxOffset = + std::max(maxOffset, map.nullsDescriptor().offset()); + for (auto i = 0; i < map.childrenCount(); ++i) { + maxOffset = std::max( + maxOffset, map.inMapDescriptorAt(i).offset()); + } + uniqueLabels += map.childrenCount(); + break; + } + case Kind::ArrayWithOffsets: { + const auto& array = type.asArrayWithOffsets(); + maxOffset = + std::max(maxOffset, array.offsetsDescriptor().offset()); + maxOffset = + std::max(maxOffset, array.lengthsDescriptor().offset()); + break; + } + } + }); + + // Populate labels + labels_.reserve(uniqueLabels + 1); + offsetToLabel_.resize(maxOffset + 1); + + labels_.push_back(""); + addLabels(root, labels_, offsetToLabel_, 0, ""); +} + +std::string_view StreamLabels::streamLabel(offset_size offset) const { + NIMBLE_ASSERT(offset < offsetToLabel_.size(), "Stream offset out of range"); + return labels_[offsetToLabel_[offset]]; +} + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/StreamLabels.h b/dwio/nimble/velox/StreamLabels.h new file mode 100644 index 0000000..53f2278 --- /dev/null +++ b/dwio/nimble/velox/StreamLabels.h @@ -0,0 +1,20 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/velox/SchemaReader.h" + +namespace facebook::nimble { + +class StreamLabels { + public: + explicit StreamLabels(const std::shared_ptr& root); + + std::string_view streamLabel(offset_size offset) const; + + private: + std::vector labels_; + std::vector offsetToLabel_; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/TabletSections.h b/dwio/nimble/velox/TabletSections.h new file mode 100644 index 0000000..6f0fec9 --- /dev/null +++ b/dwio/nimble/velox/TabletSections.h @@ -0,0 +1,10 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#pragma once + +namespace facebook::nimble { + +constexpr std::string_view kSchemaSection = "columnar.schema"; +constexpr std::string_view kMetadataSection = "columnar.metadata"; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/VeloxReader.cpp b/dwio/nimble/velox/VeloxReader.cpp new file mode 100644 index 0000000..f2c9982 --- /dev/null +++ b/dwio/nimble/velox/VeloxReader.cpp @@ -0,0 +1,410 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include "dwio/nimble/velox/VeloxReader.h" +#include +#include +#include +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/velox/ChunkedStreamDecoder.h" +#include "dwio/nimble/velox/MetadataGenerated.h" +#include "dwio/nimble/velox/SchemaReader.h" +#include "dwio/nimble/velox/SchemaSerialization.h" +#include "dwio/nimble/velox/SchemaTypes.h" +#include "dwio/nimble/velox/SchemaUtils.h" +#include "dwio/nimble/velox/TabletSections.h" +#include "velox/common/time/CpuWallTimer.h" +#include "velox/type/Type.h" + +namespace facebook::nimble { + +namespace { + +std::shared_ptr loadSchema(const Tablet& tablet) { + auto section = tablet.loadOptionalSection(std::string(kSchemaSection)); + NIMBLE_CHECK(section.has_value(), "Schema not found."); + return SchemaDeserializer::deserialize(section->content().data()); +} + +std::map loadMetadata(const Tablet& tablet) { + std::map result; + auto section = tablet.loadOptionalSection(std::string(kMetadataSection)); + + if (!section.has_value()) { + return result; + } + + auto metadata = + flatbuffers::GetRoot(section->content().data()); + auto entryCount = metadata->entries()->size(); + for (auto i = 0; i < entryCount; ++i) { + auto* entry = metadata->entries()->Get(i); + result.insert({entry->key()->str(), entry->value()->str()}); + } + + return result; +} + +std::shared_ptr createFlatType( + const std::vector& selectedFeatures, + const velox::TypePtr& veloxType) { + NIMBLE_ASSERT( + !selectedFeatures.empty(), + "Empty feature selection not allowed for struct encoding."); + + auto& valueType = veloxType->asMap().valueType(); + return velox::ROW( + std::vector(selectedFeatures), + std::vector>( + selectedFeatures.size(), valueType)); +} + +} // namespace + +const std::vector& VeloxReader::preloadedOptionalSections() { + static std::vector sections{std::string(kSchemaSection)}; + return sections; +} + +VeloxReader::VeloxReader( + velox::memory::MemoryPool& pool, + velox::ReadFile* file, + std::shared_ptr selector, + VeloxReadParams params) + : VeloxReader( + pool, + std::make_shared( + pool, + file, /* preloadOptionalSections */ + preloadedOptionalSections()), + std::move(selector), + std::move(params)) {} + +VeloxReader::VeloxReader( + velox::memory::MemoryPool& pool, + std::shared_ptr file, + std::shared_ptr selector, + VeloxReadParams params) + : VeloxReader( + pool, + std::make_shared( + pool, + std::move(file), /* preloadOptionalSections */ + preloadedOptionalSections()), + std::move(selector), + std::move(params)) {} + +VeloxReader::VeloxReader( + velox::memory::MemoryPool& pool, + std::shared_ptr tablet, + std::shared_ptr selector, + VeloxReadParams params) + : pool_{pool}, + tablet_{std::move(tablet)}, + parameters_{std::move(params)}, + schema_{loadSchema(*tablet_)}, + streamLabels_{schema_}, + type_{ + selector ? selector->getSchema() + : std::dynamic_pointer_cast( + convertToVeloxType(*schema_))}, + barrier_{ + params.decodingExecutor + ? std::make_unique( + params.decodingExecutor) + : nullptr}, + logger_{parameters_.metricsLogger} { + static_assert(std::is_same_v); + + if (!selector) { + selector = std::make_shared(type_); + } + auto schemaWithId = selector->getSchemaWithId(); + rootFieldReaderFactory_ = FieldReaderFactory::create( + parameters_, + pool_, + schema_, + schemaWithId, + offsets_, + [selector](auto nodeId) { return selector->shouldReadNode(nodeId); }, + barrier_.get()); + + // We scope down the allowed stripes based on the passed in offset ranges. + // These ranges represent file splits. + // File splits contain a file path and a range of bytes (not rows) withing + // that file. It is possible that multiple file splits map to the same file + // (but each covers a diffrent range within this file). + // File splits are guaranteed to not overlap and also to cover the entire file + // range. + // We want to guarantee that each row in the file is processed by + // exactly one split. When a reader is created, we only know about a single + // split (the current range passed in), and we have no additional context + // about other splits being processed by other readers. Therefore, we apply + // the following heuristics to guarantee uniqueness across splits: + // 1. We transpose the passed in range to match stripe boundaries. + // 2. We consider a stripe to be part of this range, only if the stripe + // beginning offset falls inside the range. + // NOTE: With these heuristics, it is possible that a file split will map to + // zero rows in a file (for example, if the file split is falls completely + // inside a single stripe, without covering byte 0). This is perfectly ok, as + // usually the caller will then fetch another file split to process and other + // file splits will cover the rest of the file. + + firstStripe_ = tablet_->stripeCount(); + lastStripe_ = 0; + firstRow_ = 0; + lastRow_ = 0; + uint64_t rows = 0; + for (auto i = 0; i < tablet_->stripeCount(); ++i) { + auto stripeOffset = tablet_->stripeOffset(i); + if ((stripeOffset >= parameters_.fileRangeStartOffset) && + (stripeOffset < parameters_.fileRangeEndOffset)) { + if (i < firstStripe_) { + firstStripe_ = i; + firstRow_ = rows; + } + if (i >= lastStripe_) { + lastStripe_ = i + 1; + lastRow_ = rows + tablet_->stripeRowCount(i); + } + } + + rows += tablet_->stripeRowCount(i); + } + + nextStripe_ = firstStripe_; + + if (parameters_.stripeCountCallback) { + parameters_.stripeCountCallback(lastStripe_ - firstStripe_); + } + + VLOG(1) << "Tablet handling stripes: " << firstStripe_ << " -> " + << lastStripe_ << " (rows " << firstRow_ << " -> " << lastRow_ + << "). Total stripes: " << tablet_->stripeCount() + << ". Total rows: " << tablet_->tabletRowCount(); + + if (!logger_) { + logger_ = std::make_shared(); + } +} + +void VeloxReader::loadStripeIfAny() { + if (nextStripe_ < lastStripe_) { + loadStripe(); + } +} + +void VeloxReader::loadStripe() { + try { + if (loadedStripe_ != std::numeric_limits::max() && + loadedStripe_ == nextStripe_) { + // We are not reloading the current stripe, but we expect all + // decoders/readers to be reset after calling loadStripe(), therefore, we + // need to explicitly reset all decoders and readers. + rootReader_->reset(); + + rowsRemainingInStripe_ = tablet_->stripeRowCount(nextStripe_); + ++nextStripe_; + return; + } + + StripeLoadMetrics metrics{}; + velox::CpuWallTiming timing{}; + { + velox::CpuWallTimer timer{timing}; + // LoadAll returns all the stream available in a stripe. + // The streams returned might be a subset of the total streams available + // in the file, as the current stripe might have captured/encountered less + // streams than later stripes. + // In the extreme case, a stripe can return zero streams (for example, if + // all the streams in that stripe were contained all nulls). + auto streams = + tablet_->load(nextStripe_, offsets_, [this](offset_size offset) { + return streamLabels_.streamLabel(offset); + }); + for (uint32_t i = 0; i < streams.size(); ++i) { + if (!streams[i]) { + // As this stream is not present in current stripe (might be present + // in previous one) we set to nullptr, One of the case is where you + // are projecting more fields in FlatMap than the stripe actually + // has. + decoders_[offsets_[i]] = nullptr; + } else { + metrics.totalStreamSize += streams[i]->getStream().size(); + decoders_[offsets_[i]] = std::make_unique( + pool_, + std::make_unique( + pool_, std::move(streams[i])), + *logger_); + ++metrics.streamCount; + } + } + rowsRemainingInStripe_ = tablet_->stripeRowCount(nextStripe_); + loadedStripe_ = nextStripe_; + ++nextStripe_; + rootReader_ = rootFieldReaderFactory_->createReader(decoders_); + } + metrics.stripeIndex = loadedStripe_; + metrics.rowsInStripe = rowsRemainingInStripe_; + metrics.cpuUsec = timing.cpuNanos / 1000; + metrics.wallTimeUsec = timing.wallNanos / 1000; + if (parameters_.blockedOnIoMsCallback) { + parameters_.blockedOnIoMsCallback(timing.wallNanos / 1000000); + } + logger_->logStripeLoad(metrics); + } catch (const std::exception& e) { + logger_->logException(MetricsLogger::kStripeLoadOperation, e.what()); + throw; + } catch (...) { + logger_->logException( + MetricsLogger::kStripeLoadOperation, + folly::to(folly::exceptionStr(std::current_exception()))); + throw; + } +} + +bool VeloxReader::next(uint64_t rowCount, velox::VectorPtr& result) { + if (rowsRemainingInStripe_ == 0) { + if (nextStripe_ < lastStripe_) { + loadStripe(); + } else { + return false; + } + } + uint64_t rowsToRead = std::min(rowsRemainingInStripe_, rowCount); + std::optional startTime; + if (parameters_.decodingTimeUsCallback) { + startTime = std::chrono::steady_clock::now(); + } + rootReader_->next(rowsToRead, result, nullptr /*scatterBitmap*/); + if (barrier_) { + // Wait for all reader tasks to complete. + barrier_->waitAll(); + } + if (startTime.has_value()) { + parameters_.decodingTimeUsCallback( + std::chrono::duration_cast( + std::chrono::steady_clock::now() - startTime.value()) + .count()); + } + + // Update reader state + rowsRemainingInStripe_ -= rowsToRead; + return true; +} + +const Tablet& VeloxReader::getTabletView() const { + return *tablet_; +} + +const std::shared_ptr& VeloxReader::getType() const { + return type_; +} + +const std::shared_ptr& VeloxReader::schema() const { + return schema_; +} + +const std::map& VeloxReader::metadata() const { + if (!metadata_.has_value()) { + metadata_ = loadMetadata(*tablet_); + } + + return metadata_.value(); +} + +uint64_t VeloxReader::seekToRow(uint64_t rowNumber) { + if (isEmptyFile()) { + return 0; + } + + if (rowNumber < firstRow_) { + LOG(INFO) << "Trying to seek to row " << rowNumber + << " which is outside of the allowed range [" << firstRow_ << ", " + << lastRow_ << ")."; + + nextStripe_ = firstStripe_; + rowsRemainingInStripe_ = 0; + return firstRow_; + } + + if (rowNumber >= lastRow_) { + LOG(INFO) << "Trying to seek to row " << rowNumber + << " which is outside of the allowed range [" << firstRow_ << ", " + << lastRow_ << ")."; + + nextStripe_ = lastStripe_; + rowsRemainingInStripe_ = 0; + return lastRow_; + } + + auto rowsSkipped = skipStripes(0, rowNumber); + loadStripe(); + skipInCurrentStripe(rowNumber - rowsSkipped); + return rowNumber; +} + +uint64_t VeloxReader::skipRows(uint64_t numberOfRowsToSkip) { + if (isEmptyFile() || numberOfRowsToSkip == 0) { + LOG(INFO) << "Nothing to skip!"; + return 0; + } + + // When we skipped or exhausted the whole file we can return 0 + if (rowsRemainingInStripe_ == 0 && nextStripe_ == lastStripe_) { + LOG(INFO) << "Current index is beyond EOF. Nothing to skip."; + return 0; + } + + // Skips remaining rows in stripe + if (rowsRemainingInStripe_ >= numberOfRowsToSkip) { + skipInCurrentStripe(numberOfRowsToSkip); + return numberOfRowsToSkip; + } + + uint64_t rowsSkipped = rowsRemainingInStripe_; + auto rowsToSkip = numberOfRowsToSkip; + // Skip the leftover rows from currently loaded stripe + rowsToSkip -= rowsRemainingInStripe_; + rowsSkipped += skipStripes(nextStripe_, rowsToSkip); + if (nextStripe_ >= lastStripe_) { + LOG(INFO) << "Skipped to last allowed row in the file."; + return rowsSkipped; + } + + loadStripe(); + skipInCurrentStripe(numberOfRowsToSkip - rowsSkipped); + return numberOfRowsToSkip; +} + +uint64_t VeloxReader::skipStripes( + uint32_t startStripeIndex, + uint64_t rowsToSkip) { + NIMBLE_DCHECK( + startStripeIndex <= lastStripe_, + fmt::format("Invalid stripe {}.", startStripeIndex)); + + uint64_t totalRowsToSkip = rowsToSkip; + while (startStripeIndex < lastStripe_ && + rowsToSkip >= tablet_->stripeRowCount(startStripeIndex)) { + rowsToSkip -= tablet_->stripeRowCount(startStripeIndex); + ++startStripeIndex; + } + + nextStripe_ = startStripeIndex; + rowsRemainingInStripe_ = + nextStripe_ >= lastStripe_ ? 0 : tablet_->stripeRowCount(nextStripe_); + + return totalRowsToSkip - rowsToSkip; +} + +void VeloxReader::skipInCurrentStripe(uint64_t rowsToSkip) { + NIMBLE_DCHECK( + rowsToSkip <= rowsRemainingInStripe_, + "Not Enough rows to skip in stripe!"); + rowsRemainingInStripe_ -= rowsToSkip; + rootReader_->skip(rowsToSkip); +} + +VeloxReader::~VeloxReader() = default; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/VeloxReader.h b/dwio/nimble/velox/VeloxReader.h new file mode 100644 index 0000000..fe262df --- /dev/null +++ b/dwio/nimble/velox/VeloxReader.h @@ -0,0 +1,165 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include +#include "dwio/nimble/common/MetricsLogger.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/tablet/Tablet.h" +#include "dwio/nimble/velox/FieldReader.h" +#include "dwio/nimble/velox/SchemaReader.h" +#include "dwio/nimble/velox/StreamLabels.h" +#include "folly/container/F14Set.h" +#include "velox/common/file/File.h" +#include "velox/common/memory/Memory.h" +#include "velox/dwio/common/ColumnSelector.h" +#include "velox/dwio/common/ExecutorBarrier.h" +#include "velox/dwio/common/FlatMapHelper.h" +#include "velox/type/Type.h" +#include "velox/vector/BaseVector.h" + +// The VeloxReader reads (a projection from) a file into a velox VectorPtr. +// The current implementation only uses FlatVector rather than any of the +// fancier vectors (dictionary, etc), and only supports the types needed for ML. + +namespace facebook::nimble { + +struct VeloxReadParams : public FieldReaderParams { + uint64_t fileRangeStartOffset = 0; + uint64_t fileRangeEndOffset = std::numeric_limits::max(); + + // Optional reader decoding executor. When supplied, decoding into a Velox + // vector will be parallelized by this executor, if the column type supports + // parallel decoding. + std::shared_ptr decodingExecutor; + + // Metric logger with pro-populated access info. + std::shared_ptr metricsLogger; + + // Report the number of stripes that will be read (consider the given range). + std::function stripeCountCallback; + + // Report the Wall time (ms) that we're blocked waiting on IO. + std::function blockedOnIoMsCallback; + + // Report the Wall time (us) that we spend decoding. + std::function decodingTimeUsCallback; +}; + +class VeloxReader { + public: + VeloxReader( + velox::memory::MemoryPool& pool, + velox::ReadFile* file, + std::shared_ptr selector = + nullptr, + VeloxReadParams params = {}); + + VeloxReader( + velox::memory::MemoryPool& pool, + std::shared_ptr file, + std::shared_ptr selector = + nullptr, + VeloxReadParams params = {}); + + VeloxReader( + velox::memory::MemoryPool& pool, + std::shared_ptr tablet, + std::shared_ptr selector = + nullptr, + VeloxReadParams params = {}); + + ~VeloxReader(); + + // Fills |result| with up to rowCount new rows, returning whether any new rows + // were read. |result| may be nullptr, in which case it will be allocated via + // pool_. If it is not nullptr its type must match type_. + bool next(uint64_t rowCount, velox::VectorPtr& result); + + const Tablet& getTabletView() const; + + const std::shared_ptr& getType() const; + + const std::shared_ptr& schema() const; + + const std::map& metadata() const; + + velox::memory::MemoryPool& getMemoryPool() const { + return pool_; + } + + // Seeks to |rowNumber| from the beginning of the file (row 0). + // If |rowNumber| is greater than the number of rows in the file, the + // seek will stop at the end of file and following reads will return + // false. + // Returns the current row number the reader is pointing to. If + // |rowNumber| is greater than the number of rows in the file, this will + // return the last row number. + uint64_t seekToRow(uint64_t rowNumber); + + // Skips |numberOfRowsToSkip| rows from current row index. + // If |numberOfRowsToSkip| is greater than the remaining rows in the + // file, skip will stop at the end of file, and following reads will + // return false. + // Returns the number of rows skipped. If |numberOfRowsToSkip| is + // greater than the remaining rows in the file, this will return the + // total number of rows skipped, until reaching the end of file. + uint64_t skipRows(uint64_t numberOfRowsToSkip); + + // Loads the next stripe if any + void loadStripeIfAny(); + + private: + // Loads the next stripe's streams. + void loadStripe(); + + // True if the file contain zero rows. + bool isEmptyFile() const { + return ((lastRow_ - firstRow_) == 0); + } + + // Skips |rowsToSkip| in the currently loaded stripe. |rowsToSkip| must be + // less than rowsRemaining_. + void skipInCurrentStripe(uint64_t rowsToSkip); + + // Skips over multiple stripes, starting from a stripe with index + // |startStripeIndex|. Keeps skipping stripes for as long as remaining rows to + // skip is greater than the next stripe's row count. + // This method does not load the last stripe, or skips inside the last stripe. + // Returns the total number of rows skipped. + uint64_t skipStripes(uint32_t startStripeIndex, uint64_t rowsToSkip); + + static const std::vector& preloadedOptionalSections(); + + velox::memory::MemoryPool& pool_; + std::shared_ptr tablet_; + const VeloxReadParams parameters_; + std::shared_ptr schema_; + StreamLabels streamLabels_; + std::shared_ptr type_; + std::vector offsets_; + folly::F14FastMap> decoders_; + std::unique_ptr rootFieldReaderFactory_; + std::unique_ptr rootReader_; + mutable std::optional> metadata_; + uint32_t firstStripe_; + uint32_t lastStripe_; + uint64_t firstRow_; + uint64_t lastRow_; + + // Reading state for reader + uint32_t nextStripe_ = 0; + uint64_t rowsRemainingInStripe_ = 0; + // stripe currently loaded. Initially state is no stripe loaded (INT_MAX) + uint32_t loadedStripe_ = std::numeric_limits::max(); + + std::unique_ptr barrier_; + // Right now each reader is considered its own session if not passed from + // writer option. + std::shared_ptr logger_; + + friend class VeloxReaderHelper; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/VeloxWriter.cpp b/dwio/nimble/velox/VeloxWriter.cpp new file mode 100644 index 0000000..dca479c --- /dev/null +++ b/dwio/nimble/velox/VeloxWriter.cpp @@ -0,0 +1,814 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/VeloxWriter.h" + +#include +#include + +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/encodings/Encoding.h" +#include "dwio/nimble/encodings/EncodingSelectionPolicy.h" +#include "dwio/nimble/encodings/SentinelEncoding.h" +#include "dwio/nimble/tablet/Tablet.h" +#include "dwio/nimble/velox/BufferGrowthPolicy.h" +#include "dwio/nimble/velox/ChunkedStreamWriter.h" +#include "dwio/nimble/velox/EncodingLayoutTree.h" +#include "dwio/nimble/velox/FieldWriter.h" +#include "dwio/nimble/velox/FlatMapLayoutPlanner.h" +#include "dwio/nimble/velox/FlushPolicy.h" +#include "dwio/nimble/velox/MetadataGenerated.h" +#include "dwio/nimble/velox/SchemaBuilder.h" +#include "dwio/nimble/velox/SchemaSerialization.h" +#include "dwio/nimble/velox/SchemaTypes.h" +#include "dwio/nimble/velox/TabletSections.h" +#include "folly/ScopeGuard.h" +#include "velox/common/time/CpuWallTimer.h" +#include "velox/dwio/common/ExecutorBarrier.h" +#include "velox/type/Type.h" + +namespace facebook::nimble { + +namespace detail { + +class WriterContext : public FieldWriterContext { + public: + const VeloxWriterOptions options; + std::unique_ptr flushPolicy; + velox::CpuWallTiming totalFlushTiming; + velox::CpuWallTiming stripeFlushTiming; + velox::CpuWallTiming encodingSelectionTiming; + // Right now each writer is considered its own session if not passed from + // writer option. + std::shared_ptr logger; + + uint64_t memoryUsed{0}; + uint64_t bytesWritten{0}; + uint64_t rowsInFile{0}; + uint64_t rowsInStripe{0}; + uint64_t stripeSize{0}; + std::vector rowsPerStripe; + + WriterContext(MemoryPool& memoryPool, VeloxWriterOptions options) + : FieldWriterContext{memoryPool, options.reclaimerFactory()}, + options{std::move(options)}, + logger{this->options.metricsLogger} { + flushPolicy = this->options.flushPolicyFactory(); + inputBufferGrowthPolicy = this->options.lowMemoryMode + ? std::make_unique() + : this->options.inputGrowthPolicyFactory(); + if (!logger) { + logger = std::make_shared(); + } + } + + void nextStripe() { + totalFlushTiming.add(stripeFlushTiming); + stripeFlushTiming.clear(); + rowsPerStripe.push_back(rowsInStripe); + memoryUsed = 0; + rowsInStripe = 0; + stripeSize = 0; + ++stripeIndex_; + } + + size_t getStripeIndex() const { + return stripeIndex_; + } + + private: + size_t stripeIndex_{0}; +}; + +} // namespace detail + +namespace { + +constexpr uint32_t kInitialSchemaSectionSize = 1 << 20; // 1MB + +class WriterStreamContext : public StreamContext { + public: + bool isNullStream = false; + const EncodingLayout* encoding; +}; + +class FlatmapEncodingLayoutContext : public TypeBuilderContext { + public: + explicit FlatmapEncodingLayoutContext( + folly::F14FastMap + keyEncodings) + : keyEncodings{std::move(keyEncodings)} {} + + const folly::F14FastMap + keyEncodings; +}; + +template +std::string_view encode( + std::optional encodingLayout, + detail::WriterContext& context, + Buffer& buffer, + const StreamData& streamData) { + NIMBLE_DASSERT( + streamData.data().size() % sizeof(T) == 0, + fmt::format("Unexpected size {}", streamData.data().size())); + std::span data{ + reinterpret_cast(streamData.data().data()), + streamData.data().size() / sizeof(T)}; + + std::unique_ptr> policy; + if (encodingLayout.has_value()) { + policy = std::make_unique>( + encodingLayout.value(), + context.options.compressionOptions, + context.options.encodingSelectionPolicyFactory); + + } else { + policy = std::unique_ptr>( + static_cast*>( + context.options + .encodingSelectionPolicyFactory(TypeTraits::dataType) + .release())); + } + + if (streamData.hasNulls()) { + std::span notNulls = streamData.nonNulls(); + return EncodingFactory::encodeNullable( + std::move(policy), data, notNulls, buffer); + } else { + return EncodingFactory::encode(std::move(policy), data, buffer); + } +} + +template +std::string_view encodeStreamTyped( + detail::WriterContext& context, + Buffer& buffer, + const StreamData& streamData) { + const auto* streamContext = + streamData.descriptor().context(); + + std::optional encodingLayout; + if (streamContext && streamContext->encoding) { + encodingLayout.emplace(*streamContext->encoding); + } + + try { + return encode(encodingLayout, context, buffer, streamData); + } catch (const NimbleUserError& e) { + if (e.errorCode() != error_code::IncompatibleEncoding || + !encodingLayout.has_value()) { + throw; + } + + // Incompatible captured encoding.Try again without a captured encoding. + return encode(std::nullopt, context, buffer, streamData); + } +} + +std::string_view encodeStream( + detail::WriterContext& context, + Buffer& buffer, + const StreamData& streamData) { + auto scalarKind = streamData.descriptor().scalarKind(); + switch (scalarKind) { + case ScalarKind::Bool: + return encodeStreamTyped(context, buffer, streamData); + case ScalarKind::Int8: + return encodeStreamTyped(context, buffer, streamData); + case ScalarKind::Int16: + return encodeStreamTyped(context, buffer, streamData); + case ScalarKind::Int32: + return encodeStreamTyped(context, buffer, streamData); + case ScalarKind::UInt32: + return encodeStreamTyped(context, buffer, streamData); + case ScalarKind::Int64: + return encodeStreamTyped(context, buffer, streamData); + case ScalarKind::Float: + return encodeStreamTyped(context, buffer, streamData); + case ScalarKind::Double: + return encodeStreamTyped(context, buffer, streamData); + case ScalarKind::String: + case ScalarKind::Binary: + return encodeStreamTyped(context, buffer, streamData); + default: + NIMBLE_UNREACHABLE(fmt::format("Unsupported scalar kind {}", scalarKind)); + } +} + +template +void findNodeIds( + const velox::dwio::common::TypeWithId& typeWithId, + Set& output, + std::function predicate) { + if (predicate(typeWithId)) { + output.insert(typeWithId.id()); + } + + for (const auto& child : typeWithId.getChildren()) { + findNodeIds(*child, output, predicate); + } +} + +WriterStreamContext& getStreamContext( + const StreamDescriptorBuilder& descriptor) { + auto* context = descriptor.context(); + if (context) { + return *context; + } + + descriptor.setContext(std::make_unique()); + return *descriptor.context(); +} + +// NOTE: This is a temporary method. We currently use TypeWithId to assing +// node ids to each node in the schema tree. Using TypeWithId is not ideal, as +// it is not very intuitive to users to figure out node ids. In the future, +// we'll design a new way to identify nodes in the tree (probably based on +// multi-level ordinals). But until then, we keep using a "simple" (yet +// restrictive) external configuration and perform internal conversion to node +// ids. Once the new language is ready, we'll switch to using it instead and +// this translation logic will be removed. +std::unique_ptr createRootField( + detail::WriterContext& context, + const std::shared_ptr& type) { + if (!context.options.flatMapColumns.empty()) { + context.flatMapNodeIds.clear(); + context.flatMapNodeIds.reserve(context.options.flatMapColumns.size()); + for (const auto& column : context.options.flatMapColumns) { + context.flatMapNodeIds.insert(type->childByName(column)->id()); + } + } + + if (!context.options.dictionaryArrayColumns.empty()) { + context.dictionaryArrayNodeIds.clear(); + context.dictionaryArrayNodeIds.reserve( + context.options.dictionaryArrayColumns.size()); + for (const auto& column : context.options.dictionaryArrayColumns) { + findNodeIds( + *type->childByName(column), + context.dictionaryArrayNodeIds, + [](const velox::dwio::common::TypeWithId& type) { + return type.type()->kind() == velox::TypeKind::ARRAY; + }); + } + } + + return FieldWriter::create(context, type, [&](const TypeBuilder& type) { + switch (type.kind()) { + case Kind::Row: { + getStreamContext(type.asRow().nullsDescriptor()).isNullStream = true; + break; + } + case Kind::FlatMap: { + getStreamContext(type.asFlatMap().nullsDescriptor()).isNullStream = + true; + break; + } + default: + break; + } + }); +} + +void initializeEncodingLayouts( + const TypeBuilder& typeBuilder, + const EncodingLayoutTree& encodingLayoutTree) { + { +#define _SET_STREAM_CONTEXT(builder, descriptor, identifier) \ + if (auto* encodingLayout = encodingLayoutTree.encodingLayout( \ + EncodingLayoutTree::StreamIdentifiers::identifier)) { \ + auto& streamContext = getStreamContext(builder.descriptor()); \ + streamContext.encoding = encodingLayout; \ + } + + if (typeBuilder.kind() == Kind::FlatMap) { + if (encodingLayoutTree.schemaKind() == Kind::Map) { + // Schema evolution - If a map is converted to flatmap, we should not + // fail, but also not try to replay captured encodings. + return; + } + NIMBLE_CHECK( + encodingLayoutTree.schemaKind() == Kind::FlatMap, + "Incompatible encoding layout node. Expecting flatmap node."); + folly::F14FastMap + keyEncodings; + keyEncodings.reserve(encodingLayoutTree.childrenCount()); + for (auto i = 0; i < encodingLayoutTree.childrenCount(); ++i) { + auto& child = encodingLayoutTree.child(i); + keyEncodings.emplace(child.name(), child); + } + const auto& mapBuilder = typeBuilder.asFlatMap(); + mapBuilder.setContext(std::make_unique( + std::move(keyEncodings))); + + _SET_STREAM_CONTEXT(mapBuilder, nullsDescriptor, FlatMap::NullsStream); + } else { + switch (typeBuilder.kind()) { + case Kind::Scalar: { + NIMBLE_CHECK( + encodingLayoutTree.schemaKind() == Kind::Scalar, + "Incompatible encoding layout node. Expecting scalar node."); + _SET_STREAM_CONTEXT( + typeBuilder.asScalar(), scalarDescriptor, Scalar::ScalarStream); + break; + } + case Kind::Row: { + NIMBLE_CHECK( + encodingLayoutTree.schemaKind() == Kind::Row, + "Incompatible encoding layout node. Expecting row node."); + auto& rowBuilder = typeBuilder.asRow(); + _SET_STREAM_CONTEXT(rowBuilder, nullsDescriptor, Row::NullsStream); + for (auto i = 0; i < rowBuilder.childrenCount() && + i < encodingLayoutTree.childrenCount(); + ++i) { + initializeEncodingLayouts( + rowBuilder.childAt(i), encodingLayoutTree.child(i)); + } + break; + } + case Kind::Array: { + NIMBLE_CHECK( + encodingLayoutTree.schemaKind() == Kind::Array, + "Incompatible encoding layout node. Expecting array node."); + auto& arrayBuilder = typeBuilder.asArray(); + _SET_STREAM_CONTEXT( + arrayBuilder, lengthsDescriptor, Array::LengthsStream); + if (encodingLayoutTree.childrenCount() > 0) { + NIMBLE_CHECK( + encodingLayoutTree.childrenCount() == 1, + "Invalid encoding layout tree. Array node should have exactly one child."); + initializeEncodingLayouts( + arrayBuilder.elements(), encodingLayoutTree.child(0)); + } + break; + } + case Kind::Map: { + if (encodingLayoutTree.schemaKind() == Kind::FlatMap) { + // Schema evolution - If a flatmap is converted to map, we should + // not fail, but also not try to replay captured encodings. + return; + } + NIMBLE_CHECK( + encodingLayoutTree.schemaKind() == Kind::Map, + "Incompatible encoding layout node. Expecting map node."); + auto& mapBuilder = typeBuilder.asMap(); + + _SET_STREAM_CONTEXT( + mapBuilder, lengthsDescriptor, Map::LengthsStream); + if (encodingLayoutTree.childrenCount() > 0) { + NIMBLE_CHECK( + encodingLayoutTree.childrenCount() == 2, + "Invalid encoding layout tree. Map node should have exactly two children."); + initializeEncodingLayouts( + mapBuilder.keys(), encodingLayoutTree.child(0)); + initializeEncodingLayouts( + mapBuilder.values(), encodingLayoutTree.child(1)); + } + + break; + } + case Kind::ArrayWithOffsets: { + NIMBLE_CHECK( + encodingLayoutTree.schemaKind() == Kind::ArrayWithOffsets, + "Incompatible encoding layout node. Expecting offset array node."); + auto& arrayBuilder = typeBuilder.asArrayWithOffsets(); + _SET_STREAM_CONTEXT( + arrayBuilder, offsetsDescriptor, ArrayWithOffsets::OffsetsStream); + _SET_STREAM_CONTEXT( + arrayBuilder, lengthsDescriptor, ArrayWithOffsets::LengthsStream); + if (encodingLayoutTree.childrenCount() > 0) { + NIMBLE_CHECK( + encodingLayoutTree.childrenCount() == 2, + "Invalid encoding layout tree. ArrayWithOffset node should have exactly two children."); + initializeEncodingLayouts( + arrayBuilder.elements(), encodingLayoutTree.child(0)); + } + break; + } + case Kind::FlatMap: { + NIMBLE_UNREACHABLE("Flatmap handled already"); + } + } + } +#undef _SET_STREAM_CONTEXT + } +} + +} // namespace + +VeloxWriter::VeloxWriter( + MemoryPool& memoryPool, + const velox::TypePtr& schema, + std::unique_ptr file, + VeloxWriterOptions options) + : schema_{velox::dwio::common::TypeWithId::create(schema)}, + file_{std::move(file)}, + writerMemoryPool_{memoryPool.addAggregateChild( + fmt::format("nimble_writer_{}", folly::Random::rand64()), + options.reclaimerFactory())}, + encodingMemoryPool_{writerMemoryPool_->addLeafChild( + "encoding", + true, + options.reclaimerFactory())}, + context_{std::make_unique( + *writerMemoryPool_, + std::move(options))}, + writer_{ + *encodingMemoryPool_, + file_.get(), + {.layoutPlanner = context_->options.featureReordering.has_value() + ? std::make_unique( + [&sb = context_->schemaBuilder]() { return sb.getRoot(); }, + context_->options.featureReordering.value()) + : nullptr}}, + root_{createRootField(*context_, schema_)}, + spillConfig_{options.spillConfig} { + NIMBLE_CHECK(file_, "File is null"); + if (context_->options.encodingLayoutTree.has_value()) { + context_->flatmapFieldAddedEventHandler = + [&](const TypeBuilder& flatmap, + std::string_view fieldKey, + const TypeBuilder& fieldType) { + auto* context = flatmap.context(); + if (context) { + auto it = context->keyEncodings.find(fieldKey); + if (it != context->keyEncodings.end()) { + initializeEncodingLayouts(fieldType, it->second); + } + } + }; + initializeEncodingLayouts( + *root_->typeBuilder(), context_->options.encodingLayoutTree.value()); + } +} + +VeloxWriter::~VeloxWriter() {} + +bool VeloxWriter::write(const velox::VectorPtr& vector) { + if (lastException_) { + std::rethrow_exception(lastException_); + } + + NIMBLE_CHECK(file_, "Writer is already closed"); + try { + auto size = vector->size(); + root_->write(vector, OrderedRanges::of(0, size)); + + uint64_t memoryUsed = 0; + for (const auto& stream : context_->streams()) { + memoryUsed += stream->memoryUsed(); + } + + context_->memoryUsed = memoryUsed; + context_->rowsInFile += size; + context_->rowsInStripe += size; + + return tryWriteStripe(); + } catch (...) { + lastException_ = std::current_exception(); + throw; + } +} + +void VeloxWriter::close() { + if (lastException_) { + std::rethrow_exception(lastException_); + } + + if (file_) { + try { + auto exitGuard = + folly::makeGuard([this]() { context_->flushPolicy->onClose(); }); + flush(); + root_->close(); + + if (!context_->options.metadata.empty()) { + auto& metadata = context_->options.metadata; + auto it = metadata.cbegin(); + flatbuffers::FlatBufferBuilder builder(kInitialSchemaSectionSize); + auto entries = builder.CreateVector< + flatbuffers::Offset>( + metadata.size(), [&builder, &it](size_t /* i */) { + auto entry = serialization::CreateMetadataEntry( + builder, + builder.CreateString(it->first), + builder.CreateString(it->second)); + ++it; + return entry; + }); + + builder.Finish(serialization::CreateMetadata(builder, entries)); + writer_.writeOptionalSection( + std::string(kMetadataSection), + {reinterpret_cast(builder.GetBufferPointer()), + builder.GetSize()}); + } + + { + SchemaSerializer serializer; + writer_.writeOptionalSection( + std::string(kSchemaSection), + serializer.serialize(context_->schemaBuilder)); + } + + writer_.close(); + file_->close(); + context_->bytesWritten = file_->size(); + + auto runStats = getRunStats(); + // TODO: compute and populate input size. + FileCloseMetrics metrics{ + .rowCount = context_->rowsInFile, + .stripeCount = context_->getStripeIndex(), + .fileSize = context_->bytesWritten, + .totalFlushCpuUsec = runStats.flushCpuTimeUsec, + .totalFlushWallTimeUsec = runStats.flushWallTimeUsec}; + context_->logger->logFileClose(metrics); + file_ = nullptr; + } catch (const std::exception& e) { + lastException_ = std::current_exception(); + context_->logger->logException( + MetricsLogger::kFileCloseOperation, e.what()); + file_ = nullptr; + throw; + } catch (...) { + lastException_ = std::current_exception(); + context_->logger->logException( + MetricsLogger::kFileCloseOperation, + folly::to( + folly::exceptionStr(std::current_exception()))); + file_ = nullptr; + throw; + } + } +} + +void VeloxWriter::flush() { + if (lastException_) { + std::rethrow_exception(lastException_); + } + + try { + tryWriteStripe(true); + } catch (...) { + lastException_ = std::current_exception(); + throw; + } +} + +void VeloxWriter::writeChunk(bool lastChunk) { + uint64_t previousFlushWallTime = context_->stripeFlushTiming.wallNanos; + std::atomic chunkSize = 0; + { + LoggingScope scope{*context_->logger}; + velox::CpuWallTimer veloxTimer{context_->stripeFlushTiming}; + + if (!encodingBuffer_) { + encodingBuffer_ = std::make_unique(*encodingMemoryPool_); + } + streams_.resize(context_->schemaBuilder.nodeCount()); + + // When writing null streams, we write the nulls as data, and the stream + // itself is non-nullable. This adpater class is how we expose the nulls as + // values. + class NullsAsDataStreamData : public StreamData { + public: + explicit NullsAsDataStreamData(StreamData& streamData) + : StreamData(streamData.descriptor()), streamData_{streamData} { + streamData_.materialize(); + } + + inline virtual std::string_view data() const override { + return { + reinterpret_cast(streamData_.nonNulls().data()), + streamData_.nonNulls().size()}; + } + + inline virtual std::span nonNulls() const override { + return {}; + } + + inline virtual bool hasNulls() const override { + return false; + } + + inline virtual bool empty() const override { + return streamData_.empty(); + } + inline virtual uint64_t memoryUsed() const override { + return streamData_.memoryUsed(); + } + + inline virtual void reset() override { + streamData_.reset(); + } + + private: + StreamData& streamData_; + }; + + auto encode = [&](StreamData& streamData) { + const auto offset = streamData.descriptor().offset(); + auto encoded = encodeStream(*context_, *encodingBuffer_, streamData); + if (!encoded.empty()) { + ChunkedStreamWriter chunkWriter{*encodingBuffer_}; + NIMBLE_DASSERT(offset < streams_.size(), "Stream offset out of range."); + auto& stream = streams_[offset]; + for (auto& buffer : chunkWriter.encode(encoded)) { + chunkSize += buffer.size(); + stream.content.push_back(std::move(buffer)); + } + } + streamData.reset(); + }; + + auto processStream = [&](StreamData& streamData, + std::function encoder) { + const auto offset = streamData.descriptor().offset(); + const auto* context = + streamData.descriptor().context(); + + const auto minStreamSize = + lastChunk ? 0 : context_->options.minStreamChunkRawSize; + + if (context && context->isNullStream) { + // For null streams we promote the null values to be written as + // boolean data. + // We still apply the same null logic, where if all values are + // non-nulls, we omit the entire stream. + if ((streamData.hasNulls() && + streamData.nonNulls().size() > minStreamSize) || + (lastChunk && !streamData.empty() && + !streams_[offset].content.empty())) { + encoder(streamData, true); + } + } else { + if (streamData.data().size() > minStreamSize || + (lastChunk && streamData.nonNulls().size() > 0 && + !streams_[offset].content.empty())) { + encoder(streamData, false); + } + } + }; + + if (context_->options.encodingExecutor) { + velox::dwio::common::ExecutorBarrier barrier{ + context_->options.encodingExecutor}; + for (auto& streamData : context_->streams()) { + processStream( + *streamData, [&](StreamData& innerStreamData, bool isNullStream) { + barrier.add([&innerStreamData, isNullStream, &encode]() { + if (isNullStream) { + NullsAsDataStreamData nullsStreamData{innerStreamData}; + encode(nullsStreamData); + } else { + encode(innerStreamData); + } + }); + }); + } + barrier.waitAll(); + } else { + for (auto& streamData : context_->streams()) { + processStream( + *streamData, + [&encode](StreamData& innerStreamData, bool isNullStream) { + if (isNullStream) { + NullsAsDataStreamData nullsStreamData{innerStreamData}; + encode(nullsStreamData); + } else { + encode(innerStreamData); + } + }); + } + } + + if (lastChunk) { + root_->reset(); + } + + context_->stripeSize += chunkSize; + } + + // Consider getting this from flush timing. + auto flushWallTimeMs = + (context_->stripeFlushTiming.wallNanos - previousFlushWallTime) / + 1'000'000; + VLOG(1) << "writeChunk milliseconds: " << flushWallTimeMs + << ", chunk bytes: " << chunkSize; +} + +uint32_t VeloxWriter::writeStripe() { + writeChunk(true); + + uint64_t previousFlushWallTime = context_->stripeFlushTiming.wallNanos; + uint64_t stripeSize = 0; + { + LoggingScope scope{*context_->logger}; + velox::CpuWallTimer veloxTimer{context_->stripeFlushTiming}; + + size_t nonEmptyCount = 0; + for (auto i = 0; i < streams_.size(); ++i) { + auto& source = streams_[i]; + if (!source.content.empty()) { + source.offset = i; + if (nonEmptyCount != i) { + streams_[nonEmptyCount] = std::move(source); + } + ++nonEmptyCount; + } + } + streams_.resize(nonEmptyCount); + + uint64_t startSize = writer_.size(); + writer_.writeStripe(context_->rowsInStripe, std::move(streams_)); + stripeSize = writer_.size() - startSize; + encodingBuffer_.reset(); + // TODO: once chunked string fields are supported, move string buffer + // reset to writeChunk() + context_->resetStringBuffer(); + } + + NIMBLE_ASSERT( + stripeSize < std::numeric_limits::max(), + fmt::format("unexpected stripe size {}", stripeSize)); + + // Consider getting this from flush timing. + auto flushWallTimeMs = + (context_->stripeFlushTiming.wallNanos - previousFlushWallTime) / + 1'000'000; + + VLOG(1) << "writeStripe milliseconds: " << flushWallTimeMs + << ", on disk stripe bytes: " << stripeSize; + + return static_cast(stripeSize); +} + +bool VeloxWriter::tryWriteStripe(bool force) { + if (context_->rowsInStripe == 0) { + return false; + } + + auto shouldFlush = [&]() { + return context_->flushPolicy->shouldFlush(StripeProgress{ + .rawStripeSize = context_->memoryUsed, + .stripeSize = context_->stripeSize, + .bufferSize = + static_cast(context_->bufferMemoryPool->currentBytes()), + }); + }; + + auto decision = force ? FlushDecision::Stripe : shouldFlush(); + if (decision == FlushDecision::None) { + return false; + } + + try { + // TODO: we can improve merge the last chunk write with stripe + if (decision == FlushDecision::Chunk && context_->options.enableChunking) { + writeChunk(false); + decision = shouldFlush(); + } + + if (decision != FlushDecision::Stripe) { + return false; + } + + StripeFlushMetrics metrics{ + .inputSize = context_->stripeSize, + .rowCount = context_->rowsInStripe, + .trackedMemory = context_->memoryUsed, + }; + + metrics.stripeSize = writeStripe(); + context_->logger->logStripeFlush(metrics); + + context_->nextStripe(); + return true; + } catch (const std::exception& e) { + context_->logger->logException( + MetricsLogger::kStripeFlushOperation, e.what()); + throw; + } catch (...) { + context_->logger->logException( + MetricsLogger::kStripeFlushOperation, + folly::to(folly::exceptionStr(std::current_exception()))); + throw; + } +} + +VeloxWriter::RunStats VeloxWriter::getRunStats() const { + return RunStats{ + .bytesWritten = context_->bytesWritten, + .stripeCount = folly::to(context_->getStripeIndex()), + .rowsPerStripe = context_->rowsPerStripe, + .flushCpuTimeUsec = context_->totalFlushTiming.cpuNanos / 1000, + .flushWallTimeUsec = context_->totalFlushTiming.wallNanos / 1000, + .encodingSelectionCpuTimeUsec = + context_->encodingSelectionTiming.cpuNanos / 1000, + .inputBufferReallocCount = context_->inputBufferGrowthStats.count, + .inputBufferReallocItemCount = + context_->inputBufferGrowthStats.itemCount}; +} +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/VeloxWriter.h b/dwio/nimble/velox/VeloxWriter.h new file mode 100644 index 0000000..22889ca --- /dev/null +++ b/dwio/nimble/velox/VeloxWriter.h @@ -0,0 +1,78 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/tablet/Tablet.h" +#include "dwio/nimble/velox/FieldWriter.h" +#include "dwio/nimble/velox/VeloxWriterOptions.h" +#include "velox/common/file/File.h" +#include "velox/dwio/common/TypeWithId.h" +#include "velox/vector/BaseVector.h" +#include "velox/vector/DecodedVector.h" + +// The VeloxWriter takes a velox VectorPtr and writes it to an Nimble file +// format. + +namespace facebook::nimble { + +namespace detail { + +class WriterContext; + +} // namespace detail + +// Writer that takes velox vector as input and produces nimble file. +class VeloxWriter { + public: + struct RunStats { + uint64_t bytesWritten; + uint32_t stripeCount; + std::vector rowsPerStripe; + uint64_t flushCpuTimeUsec; + uint64_t flushWallTimeUsec; + uint64_t encodingSelectionCpuTimeUsec; + // These 2 stats should be from memory pool and have better + // coverage in the future. + uint64_t inputBufferReallocCount; + uint64_t inputBufferReallocItemCount; + }; + + VeloxWriter( + velox::memory::MemoryPool& memoryPool, + const velox::TypePtr& schema, + std::unique_ptr file, + VeloxWriterOptions options); + + ~VeloxWriter(); + + // Return value of 'true' means this write ended with a flush. + bool write(const velox::VectorPtr& vector); + + void close(); + + void flush(); + + RunStats getRunStats() const; + + private: + std::shared_ptr schema_; + std::unique_ptr file_; + std::shared_ptr writerMemoryPool_; + std::shared_ptr encodingMemoryPool_; + std::unique_ptr context_; + TabletWriter writer_; + std::unique_ptr root_; + + std::unique_ptr encodingBuffer_; + std::vector streams_; + std::exception_ptr lastException_; + const velox::common::SpillConfig* const spillConfig_; + + // Returning 'true' if stripe was written. + bool tryWriteStripe(bool force = false); + void writeChunk(bool lastChunk = true); + uint32_t writeStripe(); +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/VeloxWriterDefaultMetadataOSS.cpp b/dwio/nimble/velox/VeloxWriterDefaultMetadataOSS.cpp new file mode 100644 index 0000000..0797b9e --- /dev/null +++ b/dwio/nimble/velox/VeloxWriterDefaultMetadataOSS.cpp @@ -0,0 +1,11 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include "dwio/nimble/velox/VeloxWriterOptions.h" + +namespace facebook::nimble::detail { + +std::unordered_map defaultMetadata() { + return {}; +} + +} // namespace facebook::nimble::detail diff --git a/dwio/nimble/velox/VeloxWriterOptions.h b/dwio/nimble/velox/VeloxWriterOptions.h new file mode 100644 index 0000000..88fa5ab --- /dev/null +++ b/dwio/nimble/velox/VeloxWriterOptions.h @@ -0,0 +1,115 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include "dwio/nimble/common/MetricsLogger.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/encodings/EncodingSelectionPolicy.h" +#include "dwio/nimble/velox/BufferGrowthPolicy.h" +#include "dwio/nimble/velox/EncodingLayoutTree.h" +#include "dwio/nimble/velox/FlushPolicy.h" +#include "folly/container/F14Set.h" +#include "velox/common/base/SpillConfig.h" +#include "velox/type/Type.h" + +// Options used by Velox writer that affect the output file format + +namespace facebook::nimble { + +namespace detail { +std::unordered_map defaultMetadata(); +} + +// NOTE: the object could be large when encodingOverrides are +// supplied. It's strongly advised to move instead of copying +// it. +struct VeloxWriterOptions { + // Property bag for storing user metadata in the file. + std::unordered_map metadata = + detail::defaultMetadata(); + + // Columns that should be encoded as flat maps + folly::F14FastSet flatMapColumns; + + // Columns that should be encoded as dictionary arrays + // NOTE: For each column, ALL the arrays inside this column will be encoded + // using dictionary arrays. In the future we'll have finer control on + // individual arrays within a column. + folly::F14FastSet dictionaryArrayColumns; + + // The metric logger would come populated with access descriptor information, + // application generated query id or specific sampling configs. + std::shared_ptr metricsLogger; + + // Optional feature reordering config. + // The key for this config is a (top-level) flat map column ordinal and + // the value is an ordered collection of feature ids. When provided, the + // writer will make sure that flat map features are grouped together and + // ordered based on this config. + std::optional>>> + featureReordering; + + // Optional captured encoding layout tree. + // Encoding layout tree is overlayed on the writer tree and the captured + // encodings are attempted to be used first, before resolving to perform an + // encoding selection. + // Captured encodings can be used to speed up writes (as no encoding selection + // is needed at runtime) and cal also provide better selected encodings, based + // on history data. + std::optional encodingLayoutTree; + + // Compression settings to be used when encoding and compressing data streams + CompressionOptions compressionOptions; + + // In low-memory mode, the writer is trying to perform smaller (and more + // precise) buffer allocations. This means that overall, the writer will + // consume less memory, but will come with an additional cost, of more + // reallocations and extra data copy. + // TODO: This options should be removed and integrated into the + // inputGrowthPolicyFactory option (e.g. allow the caller to set an + // ExactGrowthPolicy, as defined here: dwio/nimble/velox/BufferGrowthPolicy.h) + bool lowMemoryMode = false; + + // When flushing data streams into chunks, streams with raw data size smaller + // than this threshold will not be flushed. + // Note: this threshold is ignored when it is time to flush a stripe. + uint64_t minStreamChunkRawSize = 1024; + + // The factory function that produces the root encoding selection policy. + // Encoding selection policy is the way to balance the tradeoffs of + // different performance factors (at both read and write times). Heuristics + // based, ML based or specialized policies can be specified. + EncodingSelectionPolicyFactory encodingSelectionPolicyFactory = + [encodingFactory = ManualEncodingSelectionPolicyFactory{}]( + DataType dataType) -> std::unique_ptr { + return encodingFactory.createPolicy(dataType); + }; + + // Provides policy that controls stripe sizes and memory footprint. + std::function()> flushPolicyFactory = []() { + // Buffering 256MB data before encoding stripes. + return std::make_unique(256 << 20); + }; + + // When the writer needs to buffer data, and internal buffers don't have + // enough capacity, the writer is using this policy to claculate the the new + // capacity for the vuffers. + std::function()> + inputGrowthPolicyFactory = + []() -> std::unique_ptr { + return DefaultInputBufferGrowthPolicy::withDefaultRanges(); + }; + + std::function()> + reclaimerFactory = []() { return nullptr; }; + + const velox::common::SpillConfig* spillConfig{nullptr}; + + // If provided, internal encoding operations will happen in parallel using + // this executor. + std::shared_ptr encodingExecutor; + + bool enableChunking = false; +}; + +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/tests/BufferGrowthPolicyTest.cpp b/dwio/nimble/velox/tests/BufferGrowthPolicyTest.cpp new file mode 100644 index 0000000..40b2363 --- /dev/null +++ b/dwio/nimble/velox/tests/BufferGrowthPolicyTest.cpp @@ -0,0 +1,127 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include +#include + +#include "dwio/nimble/velox/BufferGrowthPolicy.h" + +namespace facebook::nimble { + +struct DefaultInputBufferGrowthPolicyTestCase { + std::map rangedConfigs; + uint64_t size; + uint64_t capacity; + uint64_t expectedNewCapacity; +}; + +class DefaultInputBufferGrowthPolicyTest + : public ::testing::Test, + public ::testing::WithParamInterface< + DefaultInputBufferGrowthPolicyTestCase> {}; + +TEST_P(DefaultInputBufferGrowthPolicyTest, GetExtendedCapacity) { + const auto& testCase = GetParam(); + DefaultInputBufferGrowthPolicy policy{testCase.rangedConfigs}; + + ASSERT_EQ( + policy.getExtendedCapacity(testCase.size, testCase.capacity), + testCase.expectedNewCapacity); +} + +INSTANTIATE_TEST_CASE_P( + DefaultInputBufferGrowthPolicyMinCapacityTestSuite, + DefaultInputBufferGrowthPolicyTest, + testing::Values( + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}}, + .size = 8, + .capacity = 0, + .expectedNewCapacity = 16}, + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}}, + .size = 16, + .capacity = 0, + .expectedNewCapacity = 16}, + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}}, + .size = 8, + .capacity = 4, + .expectedNewCapacity = 16}, + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}}, + .size = 16, + .capacity = 10, + .expectedNewCapacity = 16}, + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{4, 4.0f}, {16, 2.0f}, {1024, 1.5f}}, + .size = 2, + .capacity = 0, + .expectedNewCapacity = 4})); + +INSTANTIATE_TEST_CASE_P( + DefaultInputBufferGrowthPolicyInRangeGrowthTestSuite, + DefaultInputBufferGrowthPolicyTest, + testing::Values( + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}}, + .size = 8, + .capacity = 8, + .expectedNewCapacity = 8}, + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}}, + .size = 20, + .capacity = 16, + .expectedNewCapacity = 32}, + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}}, + .size = 24, + .capacity = 20, + .expectedNewCapacity = 40}, + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}}, + .size = 60, + .capacity = 20, + .expectedNewCapacity = 80}, + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}, {1024, 1.5f}}, + .size = 24, + .capacity = 20, + .expectedNewCapacity = 40}, + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}, {1024, 1.5f}}, + .size = 1028, + .capacity = 1024, + .expectedNewCapacity = 1536}, + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}, {1024, 1.5f}}, + .size = 1028, + .capacity = 1000, + .expectedNewCapacity = 1500}, + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}, {1024, 1.5f}}, + .size = 512, + .capacity = 16, + .expectedNewCapacity = 512})); + +// We determine the growth factor only once and grow the +// capacity until it suffices. +INSTANTIATE_TEST_CASE_P( + DefaultInputBufferGrowthPolicyCrossRangeTestSuite, + DefaultInputBufferGrowthPolicyTest, + testing::Values( + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}}, + .size = 20, + .capacity = 0, + .expectedNewCapacity = 32}, + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}, {1024, 1.5f}}, + .size = 2048, + .capacity = 1000, + .expectedNewCapacity = 2250}, + DefaultInputBufferGrowthPolicyTestCase{ + .rangedConfigs = {{16, 2.0f}, {1024, 1.5f}, {2048, 1.2f}}, + .size = 2048, + .capacity = 1000, + .expectedNewCapacity = 2073})); +} // namespace facebook::nimble diff --git a/dwio/nimble/velox/tests/CMakeLists.txt b/dwio/nimble/velox/tests/CMakeLists.txt new file mode 100644 index 0000000..ae965f9 --- /dev/null +++ b/dwio/nimble/velox/tests/CMakeLists.txt @@ -0,0 +1,38 @@ +add_library(nimble_velox_schema_utils SchemaUtils.cpp) +target_link_libraries( + nimble_velox_schema_utils nimble_velox_schema_fb nimble_velox_schema_builder + nimble_velox_schema_reader gtest gtest_main) + +add_executable( + nimble_velox_tests + BufferGrowthPolicyTest.cpp + EncodingLayoutTreeTests.cpp + FlatMapLayoutPlannerTests.cpp + OrderedRangesTests.cpp + SchemaTests.cpp + SerializationTests.cpp + TypeTests.cpp + VeloxReaderTests.cpp + VeloxWriterTests.cpp) +add_test(nimble_velox_tests nimble_velox_tests) + +target_link_libraries( + nimble_velox_tests + nimble_velox_common + nimble_velox_schema_utils + nimble_velox_reader + nimble_velox_writer + nimble_velox_serializer + nimble_velox_deserializer + nimble_velox_field_writer + nimble_velox_flatmap_layout_planner + nimble_common_file_writer + nimble_common + nimble_encodings + velox_vector + velox_vector_fuzzer + velox_vector_test_lib + gmock + gtest + gtest_main + Folly::folly) diff --git a/dwio/nimble/velox/tests/ChunkedStreamDecoderTests.cpp b/dwio/nimble/velox/tests/ChunkedStreamDecoderTests.cpp new file mode 100644 index 0000000..6023b69 --- /dev/null +++ b/dwio/nimble/velox/tests/ChunkedStreamDecoderTests.cpp @@ -0,0 +1,333 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include + +#include "dwio/nimble/encodings/tests/TestUtils.h" +#include "dwio/nimble/velox/ChunkedStreamDecoder.h" +#include "dwio/nimble/velox/ChunkedStreamWriter.h" + +using namespace ::facebook; + +namespace { + +DEFINE_uint32(seed, 0, "Override test seed."); +DEFINE_uint32( + min_chunk_count, + 2, + "Minumum chunks to use in multi-chunk tests."); +DEFINE_uint32( + max_chunk_count, + 10, + "Maximum chunks to use in multi-chunk tests."); +DEFINE_uint32(max_chunk_size, 10, "Maximum items to store in each chunk."); + +DEFINE_uint32( + reset_iterations, + 2, + "How many iterations to run, resetting the decoder in between."); + +size_t paddedBitmapByteCount(uint32_t count) { + return velox::bits::nbytes(count) + velox::simd::kPadding; +} + +template +nimble::Vector +generateData(RNG& rng, velox::memory::MemoryPool& memoryPool, size_t size) { + nimble::Vector data{&memoryPool, size}; + for (auto i = 0; i < size; ++i) { + uint64_t value = + folly::Random::rand64(std::numeric_limits::max(), rng); + data[i] = *reinterpret_cast(&value); + LOG(INFO) << "Data[" << i << "]: " << data[i]; + } + return data; +} + +template +nimble::Vector +generateNullData(RNG& rng, velox::memory::MemoryPool& memoryPool, size_t size) { + nimble::Vector data{&memoryPool, size}; + for (auto i = 0; i < size; ++i) { + data[i] = static_cast(folly::Random::oneIn(2, rng)); + LOG(INFO) << "Null[" << i << "]: " << data[i]; + } + return data; +} + +class TestStreamLoader : public nimble::StreamLoader { + public: + explicit TestStreamLoader(std::string stream) : stream_{std::move(stream)} {} + + const std::string_view getStream() const override { + return stream_; + } + + private: + const std::string stream_; +}; + +template +std::unique_ptr createStream( + velox::memory::MemoryPool& memoryPool, + const std::vector>& values, + const std::vector>>& nulls, + nimble::CompressionParams compressionParams = { + .type = nimble::CompressionType::Uncompressed}) { + NIMBLE_ASSERT(values.size() == nulls.size(), "Data and nulls size mismatch."); + + std::string stream; + + for (auto i = 0; i < values.size(); ++i) { + nimble::Buffer buffer{memoryPool}; + nimble::ChunkedStreamWriter writer{buffer, compressionParams}; + std::vector segments; + + if (nulls[i].has_value()) { + // Remove null entries from data + nimble::Vector data( + &memoryPool, + std::accumulate( + nulls[i]->data(), nulls[i]->data() + nulls[i]->size(), 0UL)); + uint32_t offset = 0; + for (auto j = 0; j < nulls[i]->size(); ++j) { + if (nulls[i].value()[j]) { + data[offset++] = values[i][j]; + } + } + segments = writer.encode(nimble::test::Encoder::encodeNullable( + buffer, data, nulls[i].value())); + } else { + segments = + writer.encode(nimble::test::Encoder::encode(buffer, values[i])); + } + + for (const auto& segment : segments) { + stream += segment; + } + } + + return std::make_unique(std::move(stream)); +} + +template +T getValue(const std::vector>& data, size_t offset) { + size_t i = 0; + while (data[i].size() <= offset) { + offset -= data[i].size(); + ++i; + } + + return data[i][offset]; +} + +template +bool getNullValue( + const std::vector>& data, + const std::vector>>& nulls, + size_t offset) { + size_t i = 0; + while (data[i].size() <= offset) { + offset -= data[i].size(); + ++i; + } + + if (!nulls[i].has_value()) { + return true; + } + + return nulls[i].value()[offset]; +} + +template +void test( + bool multipleChunks, + bool hasNulls, + bool skips, + bool scatter, + bool compress) { + uint32_t seed = FLAGS_seed == 0 ? folly::Random::rand32() : FLAGS_seed; + LOG(INFO) << "seed: " << seed; + std::mt19937 rng{seed}; + + auto memoryPool = velox::memory::deprecatedAddDefaultLeafMemoryPool(); + + auto chunkCount = multipleChunks + ? folly::Random::rand32(FLAGS_min_chunk_count, FLAGS_max_chunk_count, rng) + : 1; + std::vector> data; + std::vector>> nulls(chunkCount); + + size_t totalSize = 0; + for (auto i = 0; i < chunkCount; ++i) { + auto size = folly::Random::rand32(1, FLAGS_max_chunk_size, rng); + LOG(INFO) << "Chunk: " << i << ", Size: " << size; + data.push_back(generateData(rng, *memoryPool, size)); + if (hasNulls && folly::Random::oneIn(2, rng)) { + nulls[i] = generateNullData(rng, *memoryPool, size); + } + totalSize += size; + } + + auto streamLoader = createStream>( + *memoryPool, + data, + nulls, + compress + ? nimble::CompressionParams{.type = nimble::CompressionType::Zstd} + : nimble::CompressionParams{ + .type = nimble::CompressionType::Uncompressed}); + + nimble::ChunkedStreamDecoder decoder{ + *memoryPool, + std::make_unique( + *memoryPool, std::move(streamLoader)), + /* metricLogger */ {}}; + + for (auto batchSize = 1; batchSize <= totalSize; ++batchSize) { + for (auto iteration = 0; iteration < FLAGS_reset_iterations; ++iteration) { + LOG(INFO) << "batchSize: " << batchSize; + auto offset = 0; + while (offset < totalSize) { + const auto outputSize = folly::Random::oneIn(5, rng) + ? 0 + : std::min(batchSize, totalSize - offset); + if (skips && folly::Random::oneIn(2, rng)) { + LOG(INFO) << "Skipping " << outputSize; + decoder.skip(outputSize); + } else { + std::vector scatterMap; + std::optional scatterBitmap; + auto scatterSize = outputSize; + if (scatter) { + scatterSize = + folly::Random::rand32(outputSize, outputSize * 2, rng); + scatterMap.resize(paddedBitmapByteCount(scatterSize)); + std::vector indices(scatterSize); + std::iota(indices.begin(), indices.end(), 0); + std::shuffle(indices.begin(), indices.end(), rng); + for (auto i = 0; i < outputSize; ++i) { + nimble::bits::setBit( + indices[i], reinterpret_cast(scatterMap.data())); + } + for (auto i = 0; i < scatterSize; ++i) { + LOG(INFO) << "Scatter[" << i << "]: " + << nimble::bits::getBit( + i, reinterpret_cast(scatterMap.data())); + } + scatterBitmap.emplace(scatterMap.data(), scatterSize); + } + + if (hasNulls || scatter) { + std::vector output(scatterSize); + std::vector outputNulls( + paddedBitmapByteCount(scatterSize)); + uint32_t count = 0; + const auto nonNullCount = decoder.next( + outputSize, + output.data(), + [&]() { + ++count; + return outputNulls.data(); + }, + scatterBitmap.has_value() ? &scatterBitmap.value() : nullptr); + + LOG(INFO) << "offset: " << offset << ", batchSize: " << batchSize + << ", outputSize: " << outputSize + << ", scatterSize: " << scatterSize + << ", nonNullCount: " << nonNullCount; + if (nonNullCount != scatterSize || (outputSize == 0 && scatter)) { + EXPECT_EQ(1, count); + EXPECT_EQ( + nimble::bits::countSetBits( + 0, + scatterSize, + reinterpret_cast(outputNulls.data())), + nonNullCount); + } else { + EXPECT_EQ(0, count); + } + + if (nonNullCount == scatterSize) { + for (auto i = 0; i < scatterSize; ++i) { + EXPECT_EQ(getValue(data, offset + i), output[i]) + << "Index: " << i << ", Offset: " << offset; + } + } else { + if (scatter) { + size_t scatterOffset = 0; + for (auto i = 0; i < scatterSize; ++i) { + if (!nimble::bits::getBit( + i, + reinterpret_cast(scatterMap.data()))) { + EXPECT_FALSE(nimble::bits::getBit( + i, reinterpret_cast(outputNulls.data()))); + } else { + const auto isNotNull = nimble::bits::getBit( + i, reinterpret_cast(outputNulls.data())); + EXPECT_EQ( + getNullValue(data, nulls, offset + scatterOffset), + isNotNull) + << "Index: " << i << ", Offset: " << offset + << ", scatterOffset: " << scatterOffset; + if (isNotNull) { + EXPECT_EQ( + getValue(data, offset + scatterOffset), output[i]) + << "Index: " << i << ", Offset: " << offset + << ", scatterOffset: " << scatterOffset; + } + ++scatterOffset; + } + } + } else { + for (auto i = 0; i < scatterSize; ++i) { + const auto isNotNull = nimble::bits::getBit( + i, reinterpret_cast(outputNulls.data())); + EXPECT_EQ(getNullValue(data, nulls, offset + i), isNotNull) + << "Index: " << i << ", Offset: " << offset; + if (isNotNull) { + EXPECT_EQ(getValue(data, offset + i), output[i]) + << "Index: " << i << ", Offset: " << offset; + } + } + } + } + } else { + std::vector output(outputSize); + EXPECT_EQ(outputSize, decoder.next(outputSize, output.data())); + + for (auto i = 0; i < outputSize; ++i) { + EXPECT_EQ(getValue(data, offset + i), output[i]) + << "Index: " << i << ", Offset: " << offset; + } + } + } + + offset += outputSize; + } + + decoder.reset(); + } + } +} + +} // namespace + +TEST(ChunkedStreamDecoderTests, Decode) { + for (auto multipleChunks : {false, true}) { + for (auto hasNulls : {false, true}) { + for (auto skip : {false, true}) { + for (auto scatter : {false, true}) { + for (auto compress : {false, true}) { + for (int i = 0; i < 3; ++i) { + LOG(INFO) << "Interation: " << i + << ", Multiple Chunks: " << multipleChunks + << ", Has Nulls: " << hasNulls << ", Skips: " << skip + << ", Scatter: " << scatter; + test(multipleChunks, hasNulls, skip, scatter, compress); + } + } + } + } + } + } +} diff --git a/dwio/nimble/velox/tests/ChunkedStreamTests.cpp b/dwio/nimble/velox/tests/ChunkedStreamTests.cpp new file mode 100644 index 0000000..f17732c --- /dev/null +++ b/dwio/nimble/velox/tests/ChunkedStreamTests.cpp @@ -0,0 +1,176 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include + +#include "dwio/nimble/tablet/Tablet.h" +#include "dwio/nimble/velox/ChunkedStream.h" +#include "dwio/nimble/velox/ChunkedStreamWriter.h" + +using namespace ::facebook; + +namespace { +template +std::string randomString(RNG rng, uint32_t length) { + std::string random; + random.resize(folly::Random::rand32(length, rng)); + for (auto i = 0; i < random.size(); ++i) { + random[i] = folly::Random::rand32(256, rng); + } + + return random; +} + +class TestStreamLoader : public nimble::StreamLoader { + public: + explicit TestStreamLoader(std::string stream) : stream_{std::move(stream)} {} + const std::string_view getStream() const override { + return stream_; + } + + private: + const std::string stream_; +}; + +template +std::tuple, std::unique_ptr> +createChunkedStream( + RNG rng, + nimble::Buffer& buffer, + size_t chunkCount, + bool compress = false) { + std::vector data; + data.resize(chunkCount); + for (auto i = 0; i < data.size(); ++i) { + data[i] = randomString(rng, 100); + if (compress) { + data[i] += std::string(200, 'a'); + } + } + std::string result; + for (auto i = 0; i < data.size(); ++i) { + std::vector segments; + { + nimble::CompressionParams compressionParams{ + .type = nimble::CompressionType::Uncompressed}; + if (compress) { + compressionParams.type = nimble::CompressionType::Zstd; + compressionParams.zstdLevel = 3; + } + nimble::ChunkedStreamWriter writer{buffer, std::move(compressionParams)}; + segments = writer.encode(data[i]); + } + for (const auto& segment : segments) { + result += segment; + } + } + + return {data, std::make_unique(result)}; +} + +} // namespace + +TEST(ChunkedStreamTests, SingleChunkNoCompression) { + uint32_t seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng{seed}; + + auto memoryPool = velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*memoryPool}; + auto [data, result] = createChunkedStream(rng, buffer, /* chunkCount */ 1); + ASSERT_GT(result->getStream().size(), 0); + EXPECT_EQ(data.size(), 1); + + nimble::InMemoryChunkedStream reader{*memoryPool, std::move(result)}; + // Run multiple times to verify that reset() is working + for (auto i = 0; i < 3; ++i) { + ASSERT_TRUE(reader.hasNext()); + EXPECT_EQ( + nimble::CompressionType::Uncompressed, reader.peekCompressionType()); + auto chunk = reader.nextChunk(); + EXPECT_EQ(data[0], chunk); + EXPECT_FALSE(reader.hasNext()); + reader.reset(); + } +} + +TEST(ChunkedStreamTests, MultiChunkNoCompression) { + uint32_t seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng{seed}; + + auto memoryPool = velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*memoryPool}; + auto [data, result] = createChunkedStream( + rng, + buffer, + /* chunkCount */ std::max(2U, folly::Random::rand32(20, rng))); + ASSERT_GT(result->getStream().size(), 0); + EXPECT_GE(data.size(), 2); + + nimble::InMemoryChunkedStream reader{*memoryPool, std::move(result)}; + // Run multiple times to verify that reset() is working + for (auto i = 0; i < 3; ++i) { + for (auto j = 0; j < data.size(); ++j) { + ASSERT_TRUE(reader.hasNext()); + EXPECT_EQ( + nimble::CompressionType::Uncompressed, reader.peekCompressionType()); + auto chunk = reader.nextChunk(); + EXPECT_EQ(data[j], chunk); + } + EXPECT_FALSE(reader.hasNext()); + reader.reset(); + } +} + +TEST(ChunkedStreamTests, SingleChunkWithCompression) { + uint32_t seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng{seed}; + + auto memoryPool = velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*memoryPool}; + auto [data, result] = + createChunkedStream(rng, buffer, /* chunkCount */ 1, /* compress */ true); + ASSERT_GT(result->getStream().size(), 0); + EXPECT_EQ(data.size(), 1); + + nimble::InMemoryChunkedStream reader{*memoryPool, std::move(result)}; + // Run multiple times to verify that reset() is working + for (auto i = 0; i < 3; ++i) { + ASSERT_TRUE(reader.hasNext()); + EXPECT_EQ(nimble::CompressionType::Zstd, reader.peekCompressionType()); + auto chunk = reader.nextChunk(); + EXPECT_EQ(data[0], chunk); + EXPECT_FALSE(reader.hasNext()); + reader.reset(); + } +} + +TEST(ChunkedStreamTests, MultiChunkWithCompression) { + uint32_t seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng{seed}; + + auto memoryPool = velox::memory::deprecatedAddDefaultLeafMemoryPool(); + nimble::Buffer buffer{*memoryPool}; + auto [data, result] = createChunkedStream( + rng, + buffer, + /* chunkCount */ std::max(2U, folly::Random::rand32(20, rng)), + /* compress */ true); + ASSERT_GT(result->getStream().size(), 0); + EXPECT_GE(data.size(), 2); + + nimble::InMemoryChunkedStream reader{*memoryPool, std::move(result)}; + // Run multiple times to verify that reset() is working + for (auto i = 0; i < 3; ++i) { + for (auto j = 0; j < data.size(); ++j) { + ASSERT_TRUE(reader.hasNext()); + EXPECT_EQ(nimble::CompressionType::Zstd, reader.peekCompressionType()); + auto chunk = reader.nextChunk(); + EXPECT_EQ(data[j], chunk); + } + EXPECT_FALSE(reader.hasNext()); + reader.reset(); + } +} diff --git a/dwio/nimble/velox/tests/EncodingLayoutTreeTests.cpp b/dwio/nimble/velox/tests/EncodingLayoutTreeTests.cpp new file mode 100644 index 0000000..ea520b1 --- /dev/null +++ b/dwio/nimble/velox/tests/EncodingLayoutTreeTests.cpp @@ -0,0 +1,201 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include + +#include "dwio/nimble/velox/EncodingLayoutTree.h" + +using namespace ::facebook; + +namespace { + +std::optional cloneAsOptional( + const nimble::EncodingLayout* encodingLayout) { + if (!encodingLayout) { + return std::nullopt; + } + + std::string output; + output.resize(1024); + auto size = encodingLayout->serialize(output); + return { + nimble::EncodingLayout::create({output.data(), static_cast(size)}) + .first}; +} + +void verifyEncodingLayout( + const std::optional& expected, + const std::optional& actual) { + ASSERT_EQ(expected.has_value(), actual.has_value()); + if (!expected.has_value()) { + return; + } + + ASSERT_EQ(expected->encodingType(), actual->encodingType()); + ASSERT_EQ(expected->compressionType(), actual->compressionType()); + ASSERT_EQ(expected->childrenCount(), actual->childrenCount()); + + for (auto i = 0; i < expected->childrenCount(); ++i) { + verifyEncodingLayout(expected->child(i), actual->child(i)); + } +} + +void verifyEncodingLayoutTree( + const nimble::EncodingLayoutTree& expected, + const nimble::EncodingLayoutTree& actual) { + ASSERT_EQ(expected.schemaKind(), actual.schemaKind()); + ASSERT_EQ(expected.name(), actual.name()); + ASSERT_EQ(expected.childrenCount(), actual.childrenCount()); + + for (uint8_t i = 0; i < std::numeric_limits::max(); ++i) { + verifyEncodingLayout( + cloneAsOptional(expected.encodingLayout(i)), + cloneAsOptional(actual.encodingLayout(i))); + } + + for (auto i = 0; i < expected.childrenCount(); ++i) { + verifyEncodingLayoutTree(expected.child(i), actual.child(i)); + } +} + +void test(const nimble::EncodingLayoutTree& expected) { + std::string output; + output.resize(2048); + auto size = expected.serialize(output); + + auto actual = nimble::EncodingLayoutTree::create({output.data(), size}); + + verifyEncodingLayoutTree(expected, actual); +} + +} // namespace + +TEST(EncodingLayoutTreeTests, SingleNode) { + nimble::EncodingLayoutTree expected{ + nimble::Kind::Row, + {{ + nimble::EncodingLayoutTree::StreamIdentifiers::Row::NullsStream, + nimble::EncodingLayout{ + nimble::EncodingType::SparseBool, + nimble::CompressionType::Zstrong, + { + nimble::EncodingLayout{ + nimble::EncodingType::FixedBitWidth, + nimble::CompressionType::Uncompressed}, + }}, + }}, + " abc ", + }; + + test(expected); +} + +TEST(EncodingLayoutTreeTests, SingleNodeMultipleStreams) { + nimble::EncodingLayoutTree expected{ + nimble::Kind::Row, + { + { + 2, + nimble::EncodingLayout{ + nimble::EncodingType::SparseBool, + nimble::CompressionType::Zstrong, + { + nimble::EncodingLayout{ + nimble::EncodingType::FixedBitWidth, + nimble::CompressionType::Uncompressed}, + }}, + }, + { + 4, + nimble::EncodingLayout{ + nimble::EncodingType::Dictionary, + nimble::CompressionType::Zstd, + { + nimble::EncodingLayout{ + nimble::EncodingType::Constant, + nimble::CompressionType::Uncompressed}, + }}, + }, + }, + " abcd ", + }; + + test(expected); +} + +TEST(EncodingLayoutTreeTests, WithChildren) { + nimble::EncodingLayoutTree expected{ + nimble::Kind::Row, + { + { + 1, + nimble::EncodingLayout{ + nimble::EncodingType::SparseBool, + nimble::CompressionType::Zstrong, + { + nimble::EncodingLayout{ + nimble::EncodingType::FixedBitWidth, + nimble::CompressionType::Uncompressed}, + }}, + }, + }, + " abc ", + { + { + nimble::Kind::Scalar, + { + { + 0, + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Zstd, + { + nimble::EncodingLayout{ + nimble::EncodingType::Constant, + nimble::CompressionType::Uncompressed}, + }}, + }, + }, + " abc1 ", + }, + { + nimble::Kind::Array, + {}, + "", + }, + }, + }; + + test(expected); +} + +TEST(EncodingLayoutTreeTests, SingleNodeNoEncoding) { + nimble::EncodingLayoutTree expected{ + nimble::Kind::Row, + {}, + " abc ", + }; + + test(expected); +} + +TEST(EncodingLayoutTreeTests, SingleNodeEmptyName) { + nimble::EncodingLayoutTree expected{ + nimble::Kind::Row, + { + { + 9, + nimble::EncodingLayout{ + nimble::EncodingType::SparseBool, + nimble::CompressionType::Zstrong, + { + nimble::EncodingLayout{ + nimble::EncodingType::FixedBitWidth, + nimble::CompressionType::Uncompressed}, + }}, + }, + }, + "", + }; + + test(expected); +} diff --git a/dwio/nimble/velox/tests/FlatMapLayoutPlannerTests.cpp b/dwio/nimble/velox/tests/FlatMapLayoutPlannerTests.cpp new file mode 100644 index 0000000..5fa74ec --- /dev/null +++ b/dwio/nimble/velox/tests/FlatMapLayoutPlannerTests.cpp @@ -0,0 +1,416 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include +#include + +#include "dwio/nimble/velox/FlatMapLayoutPlanner.h" +#include "dwio/nimble/velox/tests/SchemaUtils.h" +#include "folly/Random.h" + +using namespace ::facebook; + +void addNamedTypes( + const nimble::TypeBuilder& node, + std::string prefix, + std::vector>& result) { + switch (node.kind()) { + case nimble::Kind::Scalar: { + result.emplace_back( + node.asScalar().scalarDescriptor().offset(), prefix + "s"); + break; + } + case nimble::Kind::Row: { + auto& row = node.asRow(); + result.emplace_back(row.nullsDescriptor().offset(), prefix + "r"); + for (auto i = 0; i < row.childrenCount(); ++i) { + addNamedTypes( + row.childAt(i), + fmt::format("{}r.{}({}).", prefix, row.nameAt(i), i), + result); + } + break; + } + case nimble::Kind::Array: { + auto& array = node.asArray(); + result.emplace_back(array.lengthsDescriptor().offset(), prefix + "a"); + addNamedTypes( + array.elements(), folly::to(prefix, "a."), result); + break; + } + case nimble::Kind::ArrayWithOffsets: { + auto& arrayWithOffsets = node.asArrayWithOffsets(); + result.emplace_back( + arrayWithOffsets.offsetsDescriptor().offset(), prefix + "da.o"); + result.emplace_back( + arrayWithOffsets.lengthsDescriptor().offset(), prefix + "da.l"); + addNamedTypes( + arrayWithOffsets.elements(), + folly::to(prefix, "da.e:"), + result); + + break; + } + case nimble::Kind::Map: { + auto& map = node.asMap(); + result.emplace_back(map.lengthsDescriptor().offset(), prefix + "m"); + addNamedTypes(map.keys(), folly::to(prefix, "m.k:"), result); + addNamedTypes( + map.values(), folly::to(prefix, "m.v:"), result); + break; + } + case nimble::Kind::FlatMap: { + auto& flatmap = node.asFlatMap(); + result.emplace_back(flatmap.nullsDescriptor().offset(), prefix + "f"); + for (auto i = 0; i < flatmap.childrenCount(); ++i) { + result.emplace_back( + flatmap.inMapDescriptorAt(i).offset(), + fmt::format("{}f.{}({}).im", prefix, flatmap.nameAt(i), i)); + addNamedTypes( + flatmap.childAt(i), + fmt::format("{}f.{}({}).", prefix, flatmap.nameAt(i), i), + result); + } + break; + } + } +} + +std::vector> getNamedTypes( + const nimble::TypeBuilder& root) { + std::vector> namedTypes; + namedTypes.reserve(0); // To silence CLANGTIDY + addNamedTypes(root, "", namedTypes); + return namedTypes; +} + +void testStreamLayout( + std::mt19937& rng, + nimble::FlatMapLayoutPlanner& planner, + nimble::SchemaBuilder& builder, + std::vector&& streams, + std::vector&& expected) { + std::random_shuffle(streams.begin(), streams.end()); + + ASSERT_EQ(expected.size(), streams.size()); + streams = planner.getLayout(std::move(streams)); + ASSERT_EQ(expected.size(), streams.size()); + + for (auto i = 0; i < expected.size(); ++i) { + EXPECT_EQ(expected[i], streams[i].content.front()) << "i = " << i; + } + + // Now we test that planner can handle the case where less streams are + // provided than the actual nodes in the schema. + std::vector streamSubset; + std::vector expectedSubset; + streamSubset.reserve(streams.size()); + expectedSubset.reserve(streams.size()); + for (auto i = 0; i < streams.size(); ++i) { + if (folly::Random::oneIn(2, rng)) { + streamSubset.push_back(streams[i]); + for (const auto& e : expected) { + if (streamSubset.back().content.front() == e) { + expectedSubset.push_back(e); + break; + } + } + } + } + + std::random_shuffle(streamSubset.begin(), streamSubset.end()); + + ASSERT_EQ(expectedSubset.size(), streamSubset.size()); + streamSubset = planner.getLayout(std::move(streamSubset)); + ASSERT_EQ(expectedSubset.size(), streamSubset.size()); + + for (auto i = 0; i < expectedSubset.size(); ++i) { + EXPECT_EQ(expectedSubset[i], streamSubset[i].content.front()); + } +} + +TEST(FlatMapLayoutPlannerTests, ReorderFlatMap) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + nimble::SchemaBuilder builder; + + nimble::test::FlatMapChildAdder fm1; + nimble::test::FlatMapChildAdder fm2; + nimble::test::FlatMapChildAdder fm3; + + SCHEMA( + builder, + ROW({ + {"c1", TINYINT()}, + {"c2", FLATMAP(Int8, TINYINT(), fm1)}, + {"c3", ARRAY(TINYINT())}, + {"c4", FLATMAP(Int8, TINYINT(), fm2)}, + {"c5", TINYINT()}, + {"c6", FLATMAP(Int8, ARRAY(TINYINT()), fm3)}, + })); + + fm1.addChild("2"); + fm1.addChild("5"); + fm1.addChild("42"); + fm1.addChild("7"); + fm2.addChild("2"); + fm2.addChild("5"); + fm2.addChild("42"); + fm2.addChild("7"); + fm3.addChild("2"); + fm3.addChild("5"); + fm3.addChild("42"); + fm3.addChild("7"); + + auto namedTypes = getNamedTypes(*builder.getRoot()); + + nimble::FlatMapLayoutPlanner planner{ + [&]() { return builder.getRoot(); }, + {{1, {3, 42, 9, 2, 21}}, {5, {3, 2, 42, 21}}}}; + + std::vector streams; + streams = planner.getLayout(std::move(streams)); + ASSERT_EQ(0, streams.size()); + + streams.reserve(namedTypes.size()); + for (auto i = 0; i < namedTypes.size(); ++i) { + streams.push_back(nimble::Stream{ + std::get<0>(namedTypes[i]), {std::get<1>(namedTypes[i])}}); + } + + std::vector expected{ + // Row should always be first + "r", + // Followed by feature order + "r.c2(1).f", + "r.c2(1).f.42(2).im", + "r.c2(1).f.42(2).s", + "r.c2(1).f.2(0).im", + "r.c2(1).f.2(0).s", + "r.c6(5).f", + "r.c6(5).f.2(0).im", + "r.c6(5).f.2(0).a", + "r.c6(5).f.2(0).a.s", + "r.c6(5).f.42(2).im", + "r.c6(5).f.42(2).a", + "r.c6(5).f.42(2).a.s", + // From here, streams follow schema order + "r.c1(0).s", + "r.c2(1).f.5(1).im", + "r.c2(1).f.5(1).s", + "r.c2(1).f.7(3).im", + "r.c2(1).f.7(3).s", + "r.c3(2).a", + "r.c3(2).a.s", + "r.c4(3).f", + "r.c4(3).f.2(0).im", + "r.c4(3).f.2(0).s", + "r.c4(3).f.5(1).im", + "r.c4(3).f.5(1).s", + "r.c4(3).f.42(2).im", + "r.c4(3).f.42(2).s", + "r.c4(3).f.7(3).im", + "r.c4(3).f.7(3).s", + "r.c5(4).s", + "r.c6(5).f.5(1).im", + "r.c6(5).f.5(1).a", + "r.c6(5).f.5(1).a.s", + "r.c6(5).f.7(3).im", + "r.c6(5).f.7(3).a", + "r.c6(5).f.7(3).a.s", + }; + + testStreamLayout( + rng, planner, builder, std::move(streams), std::move(expected)); +} + +TEST(FlatMapLayoutPlannerTests, ReorderFlatMapDynamicFeatures) { + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng(seed); + + nimble::SchemaBuilder builder; + + nimble::test::FlatMapChildAdder fm; + + SCHEMA( + builder, + ROW({ + {"c1", TINYINT()}, + {"c2", FLATMAP(Int8, TINYINT(), fm)}, + {"c3", ARRAY(TINYINT())}, + })); + + fm.addChild("2"); + fm.addChild("5"); + fm.addChild("42"); + fm.addChild("7"); + + auto namedTypes = getNamedTypes(*builder.getRoot()); + + nimble::FlatMapLayoutPlanner planner{ + [&]() { return builder.getRoot(); }, {{1, {3, 42, 9, 2, 21}}}}; + + std::vector streams; + streams.reserve(namedTypes.size()); + for (auto i = 0; i < namedTypes.size(); ++i) { + streams.push_back(nimble::Stream{ + std::get<0>(namedTypes[i]), {std::get<1>(namedTypes[i])}}); + } + + std::vector expected{ + // Row should always be first + "r", + // Followed by feature order + "r.c2(1).f", + "r.c2(1).f.42(2).im", + "r.c2(1).f.42(2).s", + "r.c2(1).f.2(0).im", + "r.c2(1).f.2(0).s", + // From here, streams follow schema order + "r.c1(0).s", + "r.c2(1).f.5(1).im", + "r.c2(1).f.5(1).s", + "r.c2(1).f.7(3).im", + "r.c2(1).f.7(3).s", + "r.c3(2).a", + "r.c3(2).a.s", + }; + + testStreamLayout( + rng, planner, builder, std::move(streams), std::move(expected)); + + fm.addChild("21"); + fm.addChild("3"); + fm.addChild("57"); + + namedTypes = getNamedTypes(*builder.getRoot()); + + streams.clear(); + streams.reserve(namedTypes.size()); + for (auto i = 0; i < namedTypes.size(); ++i) { + streams.push_back(nimble::Stream{ + std::get<0>(namedTypes[i]), {std::get<1>(namedTypes[i])}}); + } + + expected = { + // Row should always be first + "r", + // Followed by feature order + "r.c2(1).f", + "r.c2(1).f.3(5).im", + "r.c2(1).f.3(5).s", + "r.c2(1).f.42(2).im", + "r.c2(1).f.42(2).s", + "r.c2(1).f.2(0).im", + "r.c2(1).f.2(0).s", + "r.c2(1).f.21(4).im", + "r.c2(1).f.21(4).s", + // From here, streams follow schema order + "r.c1(0).s", + "r.c2(1).f.5(1).im", + "r.c2(1).f.5(1).s", + "r.c2(1).f.7(3).im", + "r.c2(1).f.7(3).s", + "r.c2(1).f.57(6).im", + "r.c2(1).f.57(6).s", + "r.c3(2).a", + "r.c3(2).a.s", + }; + + testStreamLayout( + rng, planner, builder, std::move(streams), std::move(expected)); +} + +TEST(FlatMapLayoutPlannerTests, IncompatibleOrdinals) { + nimble::SchemaBuilder builder; + + nimble::test::FlatMapChildAdder fm1; + nimble::test::FlatMapChildAdder fm2; + nimble::test::FlatMapChildAdder fm3; + + SCHEMA( + builder, + ROW({ + {"c1", TINYINT()}, + {"c2", FLATMAP(Int8, TINYINT(), fm1)}, + {"c3", ARRAY(TINYINT())}, + {"c4", FLATMAP(Int8, TINYINT(), fm2)}, + {"c5", TINYINT()}, + {"c6", FLATMAP(Int8, ARRAY(TINYINT()), fm3)}, + })); + + fm1.addChild("2"); + fm1.addChild("5"); + fm1.addChild("42"); + fm1.addChild("7"); + fm2.addChild("2"); + fm2.addChild("5"); + + nimble::FlatMapLayoutPlanner planner{ + [&]() { return builder.getRoot(); }, + {{1, {3, 42, 9, 2, 21}}, {2, {3, 2, 42, 21}}}}; + try { + planner.getLayout({}); + FAIL() << "Factory should have failed."; + } catch (const nimble::NimbleUserError& e) { + EXPECT_THAT( + e.what(), testing::HasSubstr("for feature ordering is not a flat map")); + } +} + +TEST(FlatMapLayoutPlannerTests, OrdinalOutOfRange) { + nimble::SchemaBuilder builder; + + nimble::test::FlatMapChildAdder fm1; + nimble::test::FlatMapChildAdder fm2; + nimble::test::FlatMapChildAdder fm3; + + SCHEMA( + builder, + ROW({ + {"c1", TINYINT()}, + {"c2", FLATMAP(Int8, TINYINT(), fm1)}, + {"c3", ARRAY(TINYINT())}, + {"c4", FLATMAP(Int8, TINYINT(), fm2)}, + {"c5", TINYINT()}, + {"c6", FLATMAP(Int8, ARRAY(TINYINT()), fm3)}, + })); + + fm1.addChild("2"); + fm1.addChild("5"); + fm1.addChild("42"); + fm1.addChild("7"); + fm2.addChild("2"); + fm2.addChild("5"); + + nimble::FlatMapLayoutPlanner planner{ + [&]() { return builder.getRoot(); }, + {{6, {3, 42, 9, 2, 21}}, {3, {3, 2, 42, 21}}}}; + try { + planner.getLayout({}); + FAIL() << "Factory should have failed."; + } catch (const nimble::NimbleUserError& e) { + EXPECT_THAT( + e.what(), testing::HasSubstr("for feature ordering is out of range")); + } +} + +TEST(FlatMapLayoutPlannerTests, IncompatibleSchema) { + nimble::SchemaBuilder builder; + + SCHEMA(builder, MAP(TINYINT(), STRING())); + + nimble::FlatMapLayoutPlanner planner{ + [&]() { return builder.getRoot(); }, {{3, {3, 2, 42, 21}}}}; + try { + planner.getLayout({}); + FAIL() << "Factory should have failed."; + } catch (const nimble::NimbleInternalError& e) { + EXPECT_THAT( + e.what(), + testing::HasSubstr( + "Flat map layout planner requires row as the schema root")); + } +} diff --git a/dwio/nimble/velox/tests/OrderedRangesTests.cpp b/dwio/nimble/velox/tests/OrderedRangesTests.cpp new file mode 100644 index 0000000..49d7916 --- /dev/null +++ b/dwio/nimble/velox/tests/OrderedRangesTests.cpp @@ -0,0 +1,54 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include +#include +#include +#include "dwio/nimble/velox/OrderedRanges.h" + +namespace facebook::nimble::tests { +namespace { + +using namespace testing; +using range_helper::OrderedRanges; + +template +std::vector> collectRange(const T& t) { + std::vector> ret; + t.apply([&](auto begin, auto size) { ret.emplace_back(begin, size); }); + return ret; +} + +template +std::vector collectEach(const T& t) { + std::vector ret; + t.applyEach([&](auto val) { ret.emplace_back(val); }); + return ret; +} + +TEST(OrderedRanges, apply) { + OrderedRanges ranges; + ranges.add(0, 1); + ranges.add(1, 1); + ranges.add(100, 1); + ranges.add(50, 50); + EXPECT_THAT( + collectRange(ranges), + ElementsAre( + std::make_tuple(0, 2), + std::make_tuple(100, 1), + std::make_tuple(50, 50))); +} + +TEST(OrderedRanges, applyEach) { + OrderedRanges ranges; + ranges.add(0, 5); + ranges.add(5, 30); + ranges.add(35, 15); + + std::vector expected(50); + std::iota(expected.begin(), expected.end(), 0); + EXPECT_THAT(collectEach(ranges), ElementsAreArray(expected)); +} + +} // namespace +} // namespace facebook::nimble::tests diff --git a/dwio/nimble/velox/tests/SchemaTests.cpp b/dwio/nimble/velox/tests/SchemaTests.cpp new file mode 100644 index 0000000..5ccaaa8 --- /dev/null +++ b/dwio/nimble/velox/tests/SchemaTests.cpp @@ -0,0 +1,367 @@ +#include +#include + +#include "dwio/nimble/common/Exceptions.h" +#include "dwio/nimble/velox/SchemaBuilder.h" +#include "dwio/nimble/velox/SchemaReader.h" +#include "dwio/nimble/velox/SchemaTypes.h" +#include "dwio/nimble/velox/StreamLabels.h" +#include "dwio/nimble/velox/tests/SchemaUtils.h" + +using namespace ::facebook; + +namespace { +void verifyLabels( + const std::vector>& schemaNodes, + std::vector expected) { + nimble::StreamLabels streamLabels{ + nimble::SchemaReader::getSchema(schemaNodes)}; + std::vector actual; + actual.reserve(schemaNodes.size()); + for (size_t i = 0, end = schemaNodes.size(); i < end; ++i) { + actual.push_back(streamLabels.streamLabel(schemaNodes[i]->offset())); + } + + EXPECT_EQ(actual, expected); +} + +} // namespace + +TEST(SchemaTests, SchemaUtils) { + nimble::SchemaBuilder builder; + + nimble::test::FlatMapChildAdder fm1; + nimble::test::FlatMapChildAdder fm2; + + SCHEMA( + builder, + ROW({ + {"c1", TINYINT()}, + {"c2", ARRAY(TINYINT())}, + {"c3", FLATMAP(Int8, TINYINT(), fm1)}, + {"c4", MAP(TINYINT(), TINYINT())}, + {"c5", FLATMAP(Float, ARRAY(BIGINT()), fm2)}, + {"c6", SMALLINT()}, + {"c7", INTEGER()}, + {"c8", BIGINT()}, + {"c9", REAL()}, + {"c10", DOUBLE()}, + {"c11", BOOLEAN()}, + {"c12", STRING()}, + {"c13", BINARY()}, + {"c14", OFFSETARRAY(INTEGER())}, + })); + + auto nodes = builder.getSchemaNodes(); + nimble::test::verifySchemaNodes( + nodes, + { + {nimble::Kind::Row, 19, nimble::ScalarKind::Bool, std::nullopt, 14}, + {nimble::Kind::Scalar, 0, nimble::ScalarKind::Int8, "c1"}, + {nimble::Kind::Array, 2, nimble::ScalarKind::UInt32, "c2"}, + {nimble::Kind::Scalar, 1, nimble::ScalarKind::Int8}, + {nimble::Kind::FlatMap, 3, nimble::ScalarKind::Int8, "c3", 0}, + {nimble::Kind::Map, 6, nimble::ScalarKind::UInt32, "c4"}, + {nimble::Kind::Scalar, 4, nimble::ScalarKind::Int8}, + {nimble::Kind::Scalar, 5, nimble::ScalarKind::Int8}, + {nimble::Kind::FlatMap, 7, nimble::ScalarKind::Float, "c5", 0}, + {nimble::Kind::Scalar, 8, nimble::ScalarKind::Int16, "c6"}, + {nimble::Kind::Scalar, 9, nimble::ScalarKind::Int32, "c7"}, + {nimble::Kind::Scalar, 10, nimble::ScalarKind::Int64, "c8"}, + {nimble::Kind::Scalar, 11, nimble::ScalarKind::Float, "c9"}, + {nimble::Kind::Scalar, 12, nimble::ScalarKind::Double, "c10"}, + {nimble::Kind::Scalar, 13, nimble::ScalarKind::Bool, "c11"}, + {nimble::Kind::Scalar, 14, nimble::ScalarKind::String, "c12"}, + {nimble::Kind::Scalar, 15, nimble::ScalarKind::Binary, "c13"}, + {nimble::Kind::ArrayWithOffsets, + 18, + nimble::ScalarKind::UInt32, + "c14"}, + {nimble::Kind::Scalar, 17, nimble::ScalarKind::UInt32}, + {nimble::Kind::Scalar, 16, nimble::ScalarKind::Int32}, + }); + + verifyLabels(nodes, {"/", "/0", "/1", "/1", "/2", "/3", "/3", + "/3", "/4", "/5", "/6", "/7", "/8", "/9", + "/10", "/11", "/12", "/13", "/13", "/13"}); + + fm2.addChild("f1"); + + nodes = builder.getSchemaNodes(); + nimble::test::verifySchemaNodes( + nodes, + { + {nimble::Kind::Row, 19, nimble::ScalarKind::Bool, std::nullopt, 14}, + {nimble::Kind::Scalar, 0, nimble::ScalarKind::Int8, "c1"}, + {nimble::Kind::Array, 2, nimble::ScalarKind::UInt32, "c2"}, + {nimble::Kind::Scalar, 1, nimble::ScalarKind::Int8}, + {nimble::Kind::FlatMap, 3, nimble::ScalarKind::Int8, "c3", 0}, + {nimble::Kind::Map, 6, nimble::ScalarKind::UInt32, "c4"}, + {nimble::Kind::Scalar, 4, nimble::ScalarKind::Int8}, + {nimble::Kind::Scalar, 5, nimble::ScalarKind::Int8}, + {nimble::Kind::FlatMap, 7, nimble::ScalarKind::Float, "c5", 1}, + {nimble::Kind::Scalar, 22, nimble::ScalarKind::Bool, "f1"}, + {nimble::Kind::Array, 21, nimble::ScalarKind::UInt32}, + {nimble::Kind::Scalar, 20, nimble::ScalarKind::Int64}, + {nimble::Kind::Scalar, 8, nimble::ScalarKind::Int16, "c6"}, + {nimble::Kind::Scalar, 9, nimble::ScalarKind::Int32, "c7"}, + {nimble::Kind::Scalar, 10, nimble::ScalarKind::Int64, "c8"}, + {nimble::Kind::Scalar, 11, nimble::ScalarKind::Float, "c9"}, + {nimble::Kind::Scalar, 12, nimble::ScalarKind::Double, "c10"}, + {nimble::Kind::Scalar, 13, nimble::ScalarKind::Bool, "c11"}, + {nimble::Kind::Scalar, 14, nimble::ScalarKind::String, "c12"}, + {nimble::Kind::Scalar, 15, nimble::ScalarKind::Binary, "c13"}, + {nimble::Kind::ArrayWithOffsets, + 18, + nimble::ScalarKind::UInt32, + "c14"}, + {nimble::Kind::Scalar, 17, nimble::ScalarKind::UInt32}, + {nimble::Kind::Scalar, 16, nimble::ScalarKind::Int32}, + }); + + verifyLabels( + nodes, {"/", "/0", "/1", "/1", "/2", "/3", "/3", "/3", + "/4", "/4/f1", "/4/f1", "/4/f1", "/5", "/6", "/7", "/8", + "/9", "/10", "/11", "/12", "/13", "/13", "/13"}); + + fm1.addChild("f1"); + fm1.addChild("f2"); + fm2.addChild("f2"); + fm2.addChild("f3"); + + nodes = builder.getSchemaNodes(); + nimble::test::verifySchemaNodes( + nodes, + { + {nimble::Kind::Row, 19, nimble::ScalarKind::Bool, std::nullopt, 14}, + {nimble::Kind::Scalar, 0, nimble::ScalarKind::Int8, "c1"}, + {nimble::Kind::Array, 2, nimble::ScalarKind::UInt32, "c2"}, + {nimble::Kind::Scalar, 1, nimble::ScalarKind::Int8}, + {nimble::Kind::FlatMap, 3, nimble::ScalarKind::Int8, "c3", 2}, + {nimble::Kind::Scalar, 24, nimble::ScalarKind::Bool, "f1"}, + {nimble::Kind::Scalar, 23, nimble::ScalarKind::Int8}, + {nimble::Kind::Scalar, 26, nimble::ScalarKind::Bool, "f2"}, + {nimble::Kind::Scalar, 25, nimble::ScalarKind::Int8}, + {nimble::Kind::Map, 6, nimble::ScalarKind::UInt32, "c4"}, + {nimble::Kind::Scalar, 4, nimble::ScalarKind::Int8}, + {nimble::Kind::Scalar, 5, nimble::ScalarKind::Int8}, + {nimble::Kind::FlatMap, 7, nimble::ScalarKind::Float, "c5", 3}, + {nimble::Kind::Scalar, 22, nimble::ScalarKind::Bool, "f1"}, + {nimble::Kind::Array, 21, nimble::ScalarKind::UInt32}, + {nimble::Kind::Scalar, 20, nimble::ScalarKind::Int64}, + {nimble::Kind::Scalar, 29, nimble::ScalarKind::Bool, "f2"}, + {nimble::Kind::Array, 28, nimble::ScalarKind::UInt32}, + {nimble::Kind::Scalar, 27, nimble::ScalarKind::Int64}, + {nimble::Kind::Scalar, 32, nimble::ScalarKind::Bool, "f3"}, + {nimble::Kind::Array, 31, nimble::ScalarKind::UInt32}, + {nimble::Kind::Scalar, 30, nimble::ScalarKind::Int64}, + {nimble::Kind::Scalar, 8, nimble::ScalarKind::Int16, "c6"}, + {nimble::Kind::Scalar, 9, nimble::ScalarKind::Int32, "c7"}, + {nimble::Kind::Scalar, 10, nimble::ScalarKind::Int64, "c8"}, + {nimble::Kind::Scalar, 11, nimble::ScalarKind::Float, "c9"}, + {nimble::Kind::Scalar, 12, nimble::ScalarKind::Double, "c10"}, + {nimble::Kind::Scalar, 13, nimble::ScalarKind::Bool, "c11"}, + {nimble::Kind::Scalar, 14, nimble::ScalarKind::String, "c12"}, + {nimble::Kind::Scalar, 15, nimble::ScalarKind::Binary, "c13"}, + {nimble::Kind::ArrayWithOffsets, + 18, + nimble::ScalarKind::UInt32, + "c14"}, + {nimble::Kind::Scalar, 17, nimble::ScalarKind::UInt32}, + {nimble::Kind::Scalar, 16, nimble::ScalarKind::Int32}, + }); + + verifyLabels( + nodes, {"/", "/0", "/1", "/1", "/2", "/2/f1", "/2/f1", + "/2/f2", "/2/f2", "/3", "/3", "/3", "/4", "/4/f1", + "/4/f1", "/4/f1", "/4/f2", "/4/f2", "/4/f2", "/4/f3", "/4/f3", + "/4/f3", "/5", "/6", "/7", "/8", "/9", "/10", + "/11", "/12", "/13", "/13", "/13"}); +} + +TEST(SchemaTests, RoundTrip) { + nimble::SchemaBuilder builder; + // ROW( + // c1:INT, + // c2:FLATMAP>, + // c3:MAP, + // c4:FLATMAP, + // c5:BOOL, + // c6:OFFSETARRAY) + + auto row = builder.createRowTypeBuilder(6); + { + auto scalar = builder.createScalarTypeBuilder(nimble::ScalarKind::Int32); + row->addChild("c1", scalar); + } + + auto flatMapCol2 = builder.createFlatMapTypeBuilder(nimble::ScalarKind::Int8); + row->addChild("c2", flatMapCol2); + + { + auto map = builder.createMapTypeBuilder(); + auto keys = builder.createScalarTypeBuilder(nimble::ScalarKind::String); + auto values = builder.createScalarTypeBuilder(nimble::ScalarKind::Float); + map->setChildren(std::move(keys), std::move(values)); + row->addChild("c3", map); + } + + auto flatMapCol4 = + builder.createFlatMapTypeBuilder(nimble::ScalarKind::Int64); + row->addChild("c4", flatMapCol4); + + { + auto scalar = builder.createScalarTypeBuilder(nimble::ScalarKind::Bool); + row->addChild("c5", scalar); + } + + { + auto arrayWithOffsets = builder.createArrayWithOffsetsTypeBuilder(); + auto elements = builder.createScalarTypeBuilder(nimble::ScalarKind::Float); + arrayWithOffsets->setChildren(std::move(elements)); + row->addChild("c6", arrayWithOffsets); + } + + auto nodes = builder.getSchemaNodes(); + nimble::test::verifySchemaNodes( + nodes, + { + {nimble::Kind::Row, 0, nimble::ScalarKind::Bool, std::nullopt, 6}, + {nimble::Kind::Scalar, 1, nimble::ScalarKind::Int32, "c1", 0}, + {nimble::Kind::FlatMap, 2, nimble::ScalarKind::Int8, "c2", 0}, + {nimble::Kind::Map, 3, nimble::ScalarKind::UInt32, "c3"}, + {nimble::Kind::Scalar, 4, nimble::ScalarKind::String, std::nullopt}, + {nimble::Kind::Scalar, 5, nimble::ScalarKind::Float, std::nullopt}, + {nimble::Kind::FlatMap, 6, nimble::ScalarKind::Int64, "c4", 0}, + {nimble::Kind::Scalar, 7, nimble::ScalarKind::Bool, "c5"}, + {nimble::Kind::ArrayWithOffsets, 9, nimble::ScalarKind::UInt32, "c6"}, + { + nimble::Kind::Scalar, + 8, + nimble::ScalarKind::UInt32, + std::nullopt, + }, + {nimble::Kind::Scalar, 10, nimble::ScalarKind::Float, std::nullopt}, + }); + + verifyLabels( + nodes, {"/", "/0", "/1", "/2", "/2", "/2", "/3", "/4", "/5", "/5", "/5"}); + + { + auto array = builder.createArrayTypeBuilder(); + auto elements = builder.createScalarTypeBuilder(nimble::ScalarKind::Double); + array->setChildren(elements); + flatMapCol2->addChild("f1", array); + } + + { + auto array = builder.createArrayTypeBuilder(); + auto elements = builder.createScalarTypeBuilder(nimble::ScalarKind::Double); + array->setChildren(elements); + flatMapCol2->addChild("f2", array); + } + + { + auto scalar = builder.createScalarTypeBuilder(nimble::ScalarKind::Int32); + flatMapCol4->addChild("f1", scalar); + } + + nodes = builder.getSchemaNodes(); + + nimble::test::verifySchemaNodes( + nodes, + { + {nimble::Kind::Row, 0, nimble::ScalarKind::Bool, std::nullopt, 6}, + {nimble::Kind::Scalar, 1, nimble::ScalarKind::Int32, "c1", 0}, + {nimble::Kind::FlatMap, 2, nimble::ScalarKind::Int8, "c2", 2}, + {nimble::Kind::Scalar, 13, nimble::ScalarKind::Bool, "f1"}, + {nimble::Kind::Array, 11, nimble::ScalarKind::UInt32}, + {nimble::Kind::Scalar, 12, nimble::ScalarKind::Double}, + {nimble::Kind::Scalar, 16, nimble::ScalarKind::Bool, "f2"}, + {nimble::Kind::Array, 14, nimble::ScalarKind::UInt32}, + {nimble::Kind::Scalar, 15, nimble::ScalarKind::Double}, + {nimble::Kind::Map, 3, nimble::ScalarKind::UInt32, "c3"}, + {nimble::Kind::Scalar, 4, nimble::ScalarKind::String}, + {nimble::Kind::Scalar, 5, nimble::ScalarKind::Float}, + {nimble::Kind::FlatMap, 6, nimble::ScalarKind::Int64, "c4", 1}, + {nimble::Kind::Scalar, 18, nimble::ScalarKind::Bool, "f1"}, + {nimble::Kind::Scalar, 17, nimble::ScalarKind::Int32}, + {nimble::Kind::Scalar, 7, nimble::ScalarKind::Bool, "c5"}, + {nimble::Kind::ArrayWithOffsets, 9, nimble::ScalarKind::UInt32, "c6"}, + {nimble::Kind::Scalar, 8, nimble::ScalarKind::UInt32}, + {nimble::Kind::Scalar, 10, nimble::ScalarKind::Float}, + }); + + verifyLabels( + nodes, + {"/", + "/0", + "/1", + "/1/f1", + "/1/f1", + "/1/f1", + "/1/f2", + "/1/f2", + "/1/f2", + "/2", + "/2", + "/2", + "/3", + "/3/f1", + "/3/f1", + "/4", + "/5", + "/5", + "/5"}); + + { + auto array = builder.createArrayTypeBuilder(); + auto elements = builder.createScalarTypeBuilder(nimble::ScalarKind::Double); + array->setChildren(elements); + flatMapCol2->addChild("f3", array); + } + + { + auto scalar = builder.createScalarTypeBuilder(nimble::ScalarKind::Int32); + flatMapCol4->addChild("f2", scalar); + } + + nodes = builder.getSchemaNodes(); + + nimble::test::verifySchemaNodes( + nodes, + { + {nimble::Kind::Row, 0, nimble::ScalarKind::Bool, std::nullopt, 6}, + {nimble::Kind::Scalar, 1, nimble::ScalarKind::Int32, "c1", 0}, + {nimble::Kind::FlatMap, 2, nimble::ScalarKind::Int8, "c2", 3}, + {nimble::Kind::Scalar, 13, nimble::ScalarKind::Bool, "f1"}, + {nimble::Kind::Array, 11, nimble::ScalarKind::UInt32}, + {nimble::Kind::Scalar, 12, nimble::ScalarKind::Double}, + {nimble::Kind::Scalar, 16, nimble::ScalarKind::Bool, "f2"}, + {nimble::Kind::Array, 14, nimble::ScalarKind::UInt32}, + {nimble::Kind::Scalar, 15, nimble::ScalarKind::Double}, + {nimble::Kind::Scalar, 21, nimble::ScalarKind::Bool, "f3"}, + {nimble::Kind::Array, 19, nimble::ScalarKind::UInt32}, + {nimble::Kind::Scalar, 20, nimble::ScalarKind::Double}, + {nimble::Kind::Map, 3, nimble::ScalarKind::UInt32, "c3"}, + {nimble::Kind::Scalar, 4, nimble::ScalarKind::String}, + {nimble::Kind::Scalar, 5, nimble::ScalarKind::Float}, + {nimble::Kind::FlatMap, 6, nimble::ScalarKind::Int64, "c4", 2}, + {nimble::Kind::Scalar, 18, nimble::ScalarKind::Bool, "f1"}, + {nimble::Kind::Scalar, 17, nimble::ScalarKind::Int32}, + {nimble::Kind::Scalar, 23, nimble::ScalarKind::Bool, "f2"}, + {nimble::Kind::Scalar, 22, nimble::ScalarKind::Int32}, + {nimble::Kind::Scalar, 7, nimble::ScalarKind::Bool, "c5"}, + {nimble::Kind::ArrayWithOffsets, 9, nimble::ScalarKind::UInt32, "c6"}, + {nimble::Kind::Scalar, 8, nimble::ScalarKind::UInt32}, + {nimble::Kind::Scalar, 10, nimble::ScalarKind::Float}, + }); + + verifyLabels(nodes, {"/", "/0", "/1", "/1/f1", "/1/f1", "/1/f1", + "/1/f2", "/1/f2", "/1/f2", "/1/f3", "/1/f3", "/1/f3", + "/2", "/2", "/2", "/3", "/3/f1", "/3/f1", + "/3/f2", "/3/f2", "/4", "/5", "/5", "/5"}); + + auto result = nimble::SchemaReader::getSchema(nodes); + nimble::test::compareSchema(nodes, result); +} diff --git a/dwio/nimble/velox/tests/SchemaUtils.cpp b/dwio/nimble/velox/tests/SchemaUtils.cpp new file mode 100644 index 0000000..09da96c --- /dev/null +++ b/dwio/nimble/velox/tests/SchemaUtils.cpp @@ -0,0 +1,182 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#include + +#include "dwio/nimble/velox/SchemaReader.h" +#include "dwio/nimble/velox/tests/SchemaUtils.h" + +namespace facebook::nimble::test { +void verifySchemaNodes( + const std::vector>& nodes, + std::vector expected) { + ASSERT_EQ(expected.size(), nodes.size()); + for (auto i = 0; i < expected.size(); ++i) { + ASSERT_TRUE(nodes[i]) << "i = " << i; + EXPECT_EQ(expected[i].kind(), nodes[i]->kind()) << "i = " << i; + EXPECT_EQ(expected[i].offset(), nodes[i]->offset()) << "i = " << i; + EXPECT_EQ(expected[i].name(), nodes[i]->name()) << "i = " << i; + EXPECT_EQ(expected[i].childrenCount(), nodes[i]->childrenCount()) + << "i = " << i; + EXPECT_EQ(expected[i].scalarKind(), nodes[i]->scalarKind()) << "i = " << i; + } +} + +void compareSchema( + uint32_t& index, + const std::vector>& nodes, + const std::shared_ptr& type, + std::optional name = std::nullopt) { + auto& node = nodes[index++]; + EXPECT_EQ(node->name().has_value(), name.has_value()); + if (name.has_value()) { + EXPECT_EQ(name.value(), node->name().value()); + } + + EXPECT_EQ(type->kind(), node->kind()); + + switch (type->kind()) { + case nimble::Kind::Scalar: { + auto& scalar = type->asScalar(); + EXPECT_EQ(scalar.scalarDescriptor().offset(), node->offset()); + EXPECT_EQ(scalar.scalarDescriptor().scalarKind(), node->scalarKind()); + EXPECT_EQ(0, node->childrenCount()); + break; + } + case nimble::Kind::Row: { + auto& row = type->asRow(); + EXPECT_EQ(row.nullsDescriptor().offset(), node->offset()); + EXPECT_EQ(nimble::ScalarKind::Bool, node->scalarKind()); + EXPECT_EQ(row.childrenCount(), node->childrenCount()); + + for (auto i = 0; i < row.childrenCount(); ++i) { + compareSchema(index, nodes, row.childAt(i), row.nameAt(i)); + } + + break; + } + case nimble::Kind::Array: { + auto& array = type->asArray(); + EXPECT_EQ(array.lengthsDescriptor().offset(), node->offset()); + EXPECT_EQ(nimble::ScalarKind::UInt32, node->scalarKind()); + EXPECT_EQ(0, node->childrenCount()); + + compareSchema(index, nodes, array.elements()); + + break; + } + case nimble::Kind::ArrayWithOffsets: { + auto& arrayWithOffsets = type->asArrayWithOffsets(); + EXPECT_EQ(arrayWithOffsets.lengthsDescriptor().offset(), node->offset()); + EXPECT_EQ(nimble::ScalarKind::UInt32, node->scalarKind()); + EXPECT_EQ(0, node->childrenCount()); + + auto& offsetNode = nodes[index++]; + EXPECT_FALSE(offsetNode->name().has_value()); + EXPECT_EQ(Kind::Scalar, offsetNode->kind()); + EXPECT_EQ(ScalarKind::UInt32, offsetNode->scalarKind()); + EXPECT_EQ( + arrayWithOffsets.offsetsDescriptor().offset(), offsetNode->offset()); + + compareSchema(index, nodes, arrayWithOffsets.elements()); + + break; + } + case nimble::Kind::Map: { + auto& map = type->asMap(); + EXPECT_EQ(map.lengthsDescriptor().offset(), node->offset()); + EXPECT_EQ(nimble::ScalarKind::UInt32, node->scalarKind()); + EXPECT_EQ(0, node->childrenCount()); + + compareSchema(index, nodes, map.keys()); + compareSchema(index, nodes, map.values()); + + break; + } + case nimble::Kind::FlatMap: { + auto& map = type->asFlatMap(); + EXPECT_EQ(map.nullsDescriptor().offset(), node->offset()); + EXPECT_EQ(map.keyScalarKind(), node->scalarKind()); + EXPECT_EQ(map.childrenCount(), node->childrenCount()); + + for (auto i = 0; i < map.childrenCount(); ++i) { + auto& inMapNode = nodes[index++]; + ASSERT_TRUE(inMapNode->name().has_value()); + EXPECT_EQ(map.nameAt(i), inMapNode->name().value()); + EXPECT_EQ(Kind::Scalar, inMapNode->kind()); + EXPECT_EQ(ScalarKind::Bool, inMapNode->scalarKind()); + EXPECT_EQ(map.inMapDescriptorAt(i).offset(), inMapNode->offset()); + compareSchema(index, nodes, map.childAt(i)); + } + + break; + } + default: + FAIL() << "Unknown type kind: " << (int)type->kind(); + } +} + +void compareSchema( + const std::vector>& nodes, + const std::shared_ptr& root) { + uint32_t index = 0; + compareSchema(index, nodes, root); + EXPECT_EQ(nodes.size(), index); +} + +std::shared_ptr row( + nimble::SchemaBuilder& builder, + std::vector>> + children) { + auto row = builder.createRowTypeBuilder(children.size()); + for (const auto& pair : children) { + row->addChild(pair.first, pair.second); + } + return row; +} + +std::shared_ptr array( + nimble::SchemaBuilder& builder, + std::shared_ptr elements) { + auto array = builder.createArrayTypeBuilder(); + array->setChildren(std::move(elements)); + + return array; +} + +std::shared_ptr arrayWithOffsets( + nimble::SchemaBuilder& builder, + std::shared_ptr elements) { + auto arrayWithOffsets = builder.createArrayWithOffsetsTypeBuilder(); + arrayWithOffsets->setChildren(std::move(elements)); + + return arrayWithOffsets; +} + +std::shared_ptr map( + nimble::SchemaBuilder& builder, + std::shared_ptr keys, + std::shared_ptr values) { + auto map = builder.createMapTypeBuilder(); + map->setChildren(std::move(keys), std::move(values)); + + return map; +} + +std::shared_ptr flatMap( + nimble::SchemaBuilder& builder, + nimble::ScalarKind keyScalarKind, + std::function(nimble::SchemaBuilder&)> + valueFactory, + FlatMapChildAdder& childAdder) { + auto map = builder.createFlatMapTypeBuilder(keyScalarKind); + childAdder.initialize(builder, *map, valueFactory); + return map; +} + +void schema( + nimble::SchemaBuilder& builder, + std::function factory) { + factory(builder); +} + +} // namespace facebook::nimble::test diff --git a/dwio/nimble/velox/tests/SchemaUtils.h b/dwio/nimble/velox/tests/SchemaUtils.h new file mode 100644 index 0000000..ffd3512 --- /dev/null +++ b/dwio/nimble/velox/tests/SchemaUtils.h @@ -0,0 +1,106 @@ +// (c) Facebook, Inc. and its affiliates. Confidential and proprietary. + +#pragma once + +#include + +#include "dwio/nimble/velox/SchemaBuilder.h" +#include "dwio/nimble/velox/SchemaReader.h" +#include "dwio/nimble/velox/SchemaTypes.h" + +namespace facebook::nimble::test { + +void verifySchemaNodes( + const std::vector>& nodes, + std::vector expected); + +void compareSchema( + const std::vector>& nodes, + const std::shared_ptr& root); + +std::shared_ptr row( + nimble::SchemaBuilder& builder, + std::vector>> + children); + +std::shared_ptr array( + nimble::SchemaBuilder& builder, + std::shared_ptr elements); + +std::shared_ptr arrayWithOffsets( + nimble::SchemaBuilder& builder, + std::shared_ptr elements); + +std::shared_ptr map( + nimble::SchemaBuilder& builder, + std::shared_ptr keys, + std::shared_ptr values); + +class FlatMapChildAdder { + public: + void addChild(std::string name) { + NIMBLE_CHECK(schemaBuilder_, "Flat map child adder is not intialized."); + typeBuilder_->addChild(name, valueFactory_(*schemaBuilder_)); + } + + private: + void initialize( + nimble::SchemaBuilder& schemaBuilder, + nimble::FlatMapTypeBuilder& typeBuilder, + std::function( + nimble::SchemaBuilder&)> valueFactory) { + schemaBuilder_ = &schemaBuilder; + typeBuilder_ = &typeBuilder; + valueFactory_ = std::move(valueFactory); + } + + nimble::SchemaBuilder* schemaBuilder_; + nimble::FlatMapTypeBuilder* typeBuilder_; + std::function(nimble::SchemaBuilder&)> + valueFactory_; + + friend std::shared_ptr flatMap( + nimble::SchemaBuilder& builder, + nimble::ScalarKind keyScalarKind, + std::function( + nimble::SchemaBuilder&)> valueFactory, + FlatMapChildAdder& childAdder); +}; + +std::shared_ptr flatMap( + nimble::SchemaBuilder& builder, + nimble::ScalarKind keyScalarKind, + std::function(nimble::SchemaBuilder&)> + valueFactory, + FlatMapChildAdder& childAdder); + +void schema( + nimble::SchemaBuilder& builder, + std::function factory); + +#define SCHEMA(schemaBuilder, types) \ + facebook::nimble::test::schema( \ + schemaBuilder, [&](nimble::SchemaBuilder& builder) { types; }) + +#define ROW(...) facebook::nimble::test::row(builder, __VA_ARGS__) +#define TINYINT() builder.createScalarTypeBuilder(nimble::ScalarKind::Int8) +#define SMALLINT() builder.createScalarTypeBuilder(nimble::ScalarKind::Int16) +#define INTEGER() builder.createScalarTypeBuilder(nimble::ScalarKind::Int32) +#define BIGINT() builder.createScalarTypeBuilder(nimble::ScalarKind::Int64) +#define REAL() builder.createScalarTypeBuilder(nimble::ScalarKind::Float) +#define DOUBLE() builder.createScalarTypeBuilder(nimble::ScalarKind::Double) +#define BOOLEAN() builder.createScalarTypeBuilder(nimble::ScalarKind::Bool) +#define STRING() builder.createScalarTypeBuilder(nimble::ScalarKind::String) +#define BINARY() builder.createScalarTypeBuilder(nimble::ScalarKind::Binary) +#define ARRAY(elements) facebook::nimble::test::array(builder, elements) +#define OFFSETARRAY(elements) \ + facebook::nimble::test::arrayWithOffsets(builder, elements) +#define MAP(keys, values) facebook::nimble::test::map(builder, keys, values) +#define FLATMAP(keyKind, values, adder) \ + facebook::nimble::test::flatMap( \ + builder, \ + nimble::ScalarKind::keyKind, \ + [&](nimble::SchemaBuilder& builder) { return values; }, \ + adder) + +} // namespace facebook::nimble::test diff --git a/dwio/nimble/velox/tests/SerializationTests.cpp b/dwio/nimble/velox/tests/SerializationTests.cpp new file mode 100644 index 0000000..7156ba6 --- /dev/null +++ b/dwio/nimble/velox/tests/SerializationTests.cpp @@ -0,0 +1,196 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include + +#include "dwio/nimble/velox/Deserializer.h" +#include "dwio/nimble/velox/Serializer.h" +#include "velox/vector/BaseVector.h" +#include "velox/vector/ComplexVector.h" +#include "velox/vector/DecodedVector.h" +#include "velox/vector/NullsBuilder.h" +#include "velox/vector/SelectivityVector.h" +#include "velox/vector/fuzzer/VectorFuzzer.h" +#include "velox/vector/tests/utils/VectorMaker.h" + +using namespace facebook; +using namespace facebook::nimble; + +class SerializationTests : public ::testing::Test { + protected: + static void SetUpTestCase() { + velox::memory::MemoryManager::testingSetInstance({}); + } + + void SetUp() override { + rootPool_ = velox::memory::memoryManager()->addRootPool("default_root"); + pool_ = velox::memory::memoryManager()->addLeafPool("default_leaf"); + } + + std::shared_ptr rootPool_; + std::shared_ptr pool_; + + static bool vectorEquals( + const velox::VectorPtr& expected, + const velox::VectorPtr& actual, + velox::vector_size_t index) { + return expected->equalValueAt(actual.get(), index, index); + } + + template + void writeAndVerify( + velox::memory::MemoryPool& pool, + const velox::TypePtr& type, + std::function generator, + std::function validator, + size_t count) { + SerializerOptions options{ + .compressionType = CompressionType::Zstd, + .compressionThreshold = 32, + .compressionLevel = 3, + }; + Serializer serializer{options, *rootPool_, type}; + Deserializer deserializer{ + pool, + SchemaReader::getSchema(serializer.schemaBuilder().getSchemaNodes())}; + + velox::VectorPtr output; + for (auto i = 0; i < count; ++i) { + auto input = generator(type); + auto serialized = + serializer.serialize(input, OrderedRanges::of(0, input->size())); + deserializer.deserialize(serialized, output); + + ASSERT_EQ(output->size(), input->size()); + for (auto j = 0; j < input->size(); ++j) { + ASSERT_TRUE(validator(output, input, j)) + << "Content mismatch at index " << j + << "\nReference: " << input->toString(j) + << "\nResult: " << output->toString(j); + } + } + } +}; + +TEST_F(SerializationTests, FuzzSimple) { + auto type = velox::ROW({ + {"bool_val", velox::BOOLEAN()}, + {"byte_val", velox::TINYINT()}, + {"short_val", velox::SMALLINT()}, + {"int_val", velox::INTEGER()}, + {"long_val", velox::BIGINT()}, + {"float_val", velox::REAL()}, + {"double_val", velox::DOUBLE()}, + {"string_val", velox::VARCHAR()}, + {"binary_val", velox::VARBINARY()}, + // {"ts_val", velox::TIMESTAMP()}, + }); + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + + // Small batches creates more edge cases. + size_t batchSize = 10; + velox::VectorFuzzer noNulls( + { + .vectorSize = batchSize, + .nullRatio = 0, + .stringLength = 20, + .stringVariableLength = true, + }, + pool_.get(), + seed); + + auto iterations = 20; + auto batches = 20; + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + *pool_, + type, + [&](auto& type) { + return noNulls.fuzzInputRow( + std::dynamic_pointer_cast(type)); + }, + vectorEquals, + batches); + } +} + +TEST_F(SerializationTests, FuzzComplex) { + auto type = velox::ROW({ + {"array", velox::ARRAY(velox::REAL())}, + {"map", velox::MAP(velox::INTEGER(), velox::DOUBLE())}, + {"row", + velox::ROW({ + {"a", velox::REAL()}, + {"b", velox::INTEGER()}, + })}, + {"nested", + velox::ARRAY(velox::ROW({ + {"a", velox::INTEGER()}, + {"b", velox::MAP(velox::REAL(), velox::REAL())}, + }))}, + }); + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + + // Small batches creates more edge cases. + size_t batchSize = 10; + velox::VectorFuzzer noNulls( + { + .vectorSize = batchSize, + .nullRatio = 0, + .stringLength = 20, + .stringVariableLength = true, + .containerLength = 5, + .containerVariableLength = true, + }, + pool_.get(), + seed); + + auto iterations = 20; + auto batches = 20; + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + *pool_.get(), + type, + [&](auto& type) { + return noNulls.fuzzInputRow( + std::dynamic_pointer_cast(type)); + }, + vectorEquals, + batches); + } +} + +TEST_F(SerializationTests, RootNotRow) { + auto type = velox::MAP(velox::INTEGER(), velox::ARRAY(velox::DOUBLE())); + auto seed = folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + + // Small batches creates more edge cases. + size_t batchSize = 10; + velox::VectorFuzzer noNulls( + { + .vectorSize = batchSize, + .nullRatio = 0, + .stringLength = 20, + .stringVariableLength = true, + .containerLength = 5, + .containerVariableLength = true, + }, + pool_.get(), + seed); + + auto iterations = 20; + auto batches = 20; + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + *pool_.get(), + type, + [&](auto& type) { return noNulls.fuzz(type); }, + vectorEquals, + batches); + } +} diff --git a/dwio/nimble/velox/tests/TypeTests.cpp b/dwio/nimble/velox/tests/TypeTests.cpp new file mode 100644 index 0000000..91de766 --- /dev/null +++ b/dwio/nimble/velox/tests/TypeTests.cpp @@ -0,0 +1,415 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include + +#include "dwio/nimble/common/tests/NimbleFileWriter.h" +#include "dwio/nimble/velox/VeloxReader.h" +#include "dwio/nimble/velox/VeloxWriterOptions.h" +#include "velox/type/Type.h" +#include "velox/vector/ComplexVector.h" +#include "velox/vector/FlatVector.h" +#include "velox/vector/fuzzer/VectorFuzzer.h" + +using namespace ::facebook; + +class TypeTests : public testing::Test { + protected: + static void SetUpTestCase() { + velox::memory::MemoryManager::testingSetInstance({}); + } + + void SetUp() override { + rootPool_ = velox::memory::memoryManager()->addRootPool("default_root"); + leafPool_ = rootPool_->addLeafChild("default_leaf"); + } + + std::shared_ptr rootPool_; + std::shared_ptr leafPool_; +}; + +TEST_F(TypeTests, MatchingSchema) { + const uint32_t batchSize = 10; + + auto type = velox::ROW({ + {"simple", velox::TINYINT()}, + {"array", velox::ARRAY(velox::BIGINT())}, + {"map", velox::MAP(velox::INTEGER(), velox::BIGINT())}, + {"struct", + velox::ROW({ + {"a", velox::REAL()}, + {"b", velox::DOUBLE()}, + })}, + {"nested", + velox::MAP( + velox::INTEGER(), + velox::ROW({ + {"a", velox::REAL()}, + {"b", + velox::ARRAY(velox::MAP(velox::INTEGER(), velox::BIGINT()))}, + }))}, + {"arraywithoffsets", velox::ARRAY(velox::BIGINT())}, + }); + + velox::VectorFuzzer fuzzer({.vectorSize = batchSize}, leafPool_.get()); + auto vector = fuzzer.fuzzInputFlatRow(type); + auto file = nimble::test::createNimbleFile(*rootPool_, vector); + + velox::InMemoryReadFile readFile(file); + auto selector = std::make_shared( + std::dynamic_pointer_cast(vector->type())); + nimble::VeloxReader reader(*leafPool_, &readFile, std::move(selector)); + + velox::VectorPtr result; + ASSERT_TRUE(reader.next(batchSize, result)); + ASSERT_EQ(result->type()->kind(), velox::TypeKind::ROW); + ASSERT_TRUE(result->type()->kindEquals(vector->type())); + ASSERT_EQ(batchSize, result->size()); + + for (auto i = 0; i < result->size(); ++i) { + ASSERT_TRUE(vector->equalValueAt(result.get(), i, i)) + << "Content mismatch row " << i << "\nExpected: " << vector->toString(i) + << "\nActual: " << result->toString(i); + } +} + +TEST_F(TypeTests, ExtraColumnWithRename) { + const uint32_t batchSize = 10; + auto fileType = velox::ROW({ + {"simple", velox::TINYINT()}, + {"array", velox::ARRAY(velox::BIGINT())}, + {"map", velox::MAP(velox::INTEGER(), velox::BIGINT())}, + {"struct", + velox::ROW({ + {"a", velox::REAL()}, + {"b", velox::DOUBLE()}, + })}, + {"nested", + velox::MAP( + velox::INTEGER(), + velox::ROW({ + {"a", velox::REAL()}, + {"b", + velox::ARRAY(velox::MAP(velox::INTEGER(), velox::BIGINT()))}, + }))}, + {"arraywithoffsets", velox::ARRAY(velox::INTEGER())}, + }); + + auto newType = velox::ROW({ + {"simple", velox::TINYINT()}, + {"array", velox::ARRAY(velox::BIGINT())}, + {"map", velox::MAP(velox::INTEGER(), velox::BIGINT())}, + {"struct_rename", + velox::ROW({ + {"a", velox::REAL()}, + {"b", velox::DOUBLE()}, + })}, + {"nested", + velox::MAP( + velox::INTEGER(), + velox::ROW({ + {"a", velox::REAL()}, + {"b", + velox::ARRAY(velox::MAP(velox::INTEGER(), velox::BIGINT()))}, + }))}, + {"arraywithoffsets", velox::ARRAY(velox::INTEGER())}, + {"new", velox::TINYINT()}, + }); + + velox::VectorFuzzer fuzzer({.vectorSize = batchSize}, leafPool_.get()); + auto vector = fuzzer.fuzzInputFlatRow(fileType); + auto file = nimble::test::createNimbleFile(*rootPool_, vector); + + velox::InMemoryReadFile readFile(file); + auto selector = std::make_shared( + std::dynamic_pointer_cast(newType)); + nimble::VeloxReader reader(*leafPool_, &readFile, std::move(selector)); + + velox::VectorPtr result; + ASSERT_TRUE(reader.next(batchSize, result)); + ASSERT_EQ(result->type()->kind(), velox::TypeKind::ROW); + ASSERT_TRUE(result->type()->kindEquals(newType)); + + for (auto i = 0; i < fileType->size(); ++i) { + auto& expectedChild = vector->as()->childAt(i); + auto& actualChild = result->as()->childAt(i); + ASSERT_EQ(batchSize, expectedChild->size()); + ASSERT_EQ(batchSize, actualChild->size()); + + for (auto j = 0; j < batchSize; ++j) { + ASSERT_TRUE(expectedChild->equalValueAt(actualChild.get(), j, j)) + << "Content mismatch at column " << i << ", row " << j + << "\nExpected: " << expectedChild->toString(j) + << "\nActual: " << actualChild->toString(j); + } + } + + auto& extraChild = + result->as()->childAt(newType->size() - 1); + ASSERT_EQ(batchSize, extraChild->size()); + for (auto i = 0; i < batchSize; ++i) { + ASSERT_TRUE(extraChild->isNullAt(i)); + } +} + +TEST_F(TypeTests, SameTypeWithProjection) { + const uint32_t batchSize = 10; + auto type = velox::ROW({ + {"simple", velox::TINYINT()}, + {"array", velox::ARRAY(velox::BIGINT())}, + {"map", velox::MAP(velox::INTEGER(), velox::BIGINT())}, + {"struct", + velox::ROW({ + {"a", velox::REAL()}, + {"b", velox::DOUBLE()}, + })}, + {"nested", + velox::MAP( + velox::INTEGER(), + velox::ROW({ + {"a", velox::REAL()}, + {"b", + velox::ARRAY(velox::MAP(velox::INTEGER(), velox::BIGINT()))}, + }))}, + {"arraywithoffsets", velox::ARRAY(velox::BIGINT())}, + }); + + velox::VectorFuzzer fuzzer({.vectorSize = batchSize}, leafPool_.get()); + auto vector = fuzzer.fuzzInputFlatRow(type); + auto file = nimble::test::createNimbleFile(*rootPool_, vector); + + velox::InMemoryReadFile readFile(file); + auto selector = std::make_shared( + std::dynamic_pointer_cast(type), + std::vector{"array", "nested", "arraywithoffsets"}); + nimble::VeloxReader reader(*leafPool_, &readFile, std::move(selector)); + + velox::VectorPtr result; + ASSERT_TRUE(reader.next(batchSize, result)); + ASSERT_EQ(result->type()->kind(), velox::TypeKind::ROW); + ASSERT_TRUE(result->type()->kindEquals(type)); + + for (auto i : {1, 4, 5}) { + auto& expectedChild = vector->as()->childAt(i); + auto& actualChild = result->as()->childAt(i); + ASSERT_EQ(batchSize, expectedChild->size()); + ASSERT_EQ(batchSize, actualChild->size()); + + for (auto j = 0; j < batchSize; ++j) { + ASSERT_TRUE(expectedChild->equalValueAt(actualChild.get(), j, j)) + << "Content mismatch at column " << i << ", row " << j + << "\nExpected: " << expectedChild->toString(j) + << "\nActual: " << actualChild->toString(j); + } + } + + // Not projected + for (auto i : {0, 2, 3}) { + ASSERT_EQ(result->as()->childAt(i), nullptr); + } +} + +TEST_F(TypeTests, ProjectingNewColumn) { + const uint32_t batchSize = 10; + auto fileType = velox::ROW({ + {"simple", velox::TINYINT()}, + {"array", velox::ARRAY(velox::BIGINT())}, + {"map", velox::MAP(velox::INTEGER(), velox::BIGINT())}, + {"struct", + velox::ROW({ + {"a", velox::REAL()}, + {"b", velox::DOUBLE()}, + })}, + {"nested", + velox::MAP( + velox::INTEGER(), + velox::ROW({ + {"a", velox::REAL()}, + {"b", + velox::ARRAY(velox::MAP(velox::INTEGER(), velox::BIGINT()))}, + }))}, + }); + + auto newType = velox::ROW({ + {"simple", velox::TINYINT()}, + {"array", velox::ARRAY(velox::BIGINT())}, + {"map", velox::MAP(velox::INTEGER(), velox::BIGINT())}, + {"struct_rename", + velox::ROW({ + {"a", velox::REAL()}, + {"b", velox::DOUBLE()}, + })}, + {"nested", + velox::MAP( + velox::INTEGER(), + velox::ROW({ + {"a", velox::REAL()}, + {"b", + velox::ARRAY(velox::MAP(velox::INTEGER(), velox::BIGINT()))}, + }))}, + {"new", velox::TINYINT()}, + }); + + velox::VectorFuzzer fuzzer({.vectorSize = batchSize}, leafPool_.get()); + auto vector = fuzzer.fuzzInputFlatRow(fileType); + auto file = nimble::test::createNimbleFile(*rootPool_, vector); + + velox::InMemoryReadFile readFile(file); + auto selector = std::make_shared( + std::dynamic_pointer_cast(newType), + std::vector{"struct_rename", "new"}); + nimble::VeloxReader reader(*leafPool_, &readFile, std::move(selector)); + + velox::VectorPtr result; + ASSERT_TRUE(reader.next(batchSize, result)); + ASSERT_EQ(result->type()->kind(), velox::TypeKind::ROW); + ASSERT_TRUE(result->type()->kindEquals(newType)); + + for (auto i : {3}) { + auto& expectedChild = vector->as()->childAt(i); + auto& actualChild = result->as()->childAt(i); + ASSERT_EQ(batchSize, expectedChild->size()); + ASSERT_EQ(batchSize, actualChild->size()); + + for (auto j = 0; j < batchSize; ++j) { + ASSERT_TRUE(expectedChild->equalValueAt(actualChild.get(), j, j)) + << "Content mismatch at column " << i << ", row " << j + << "\nExpected: " << expectedChild->toString(j) + << "\nActual: " << actualChild->toString(j); + } + } + + // Projected, but missing + for (auto i : {5}) { + auto& nullChild = result->as()->childAt(i); + ASSERT_EQ(batchSize, nullChild->size()); + for (auto j = 0; j < batchSize; ++j) { + ASSERT_TRUE(nullChild->isNullAt(j)) + << "Expecting null value at column " << i << ", row " << j; + } + } + + // Not projected + for (auto i : {0, 1, 2, 4}) { + ASSERT_EQ(result->as()->childAt(i), nullptr); + } +} + +TEST_F(TypeTests, FlatMapFeatureSelection) { + const uint32_t batchSize = 200; + auto type = velox::ROW({ + {"map", velox::MAP(velox::INTEGER(), velox::BIGINT())}, + }); + + velox::RowVectorPtr vector = nullptr; + std::unordered_set uniqueKeys; + velox::VectorFuzzer fuzzer( + {.vectorSize = batchSize, .nullRatio = 0.1, .containerLength = 2}, + leafPool_.get()); + + while (uniqueKeys.empty()) { + vector = fuzzer.fuzzInputFlatRow(type); + auto map = vector->childAt(0)->as(); + auto keys = map->mapKeys()->asFlatVector(); + auto values = map->mapValues()->asFlatVector(); + + for (auto i = 0; i < batchSize; ++i) { + for (auto j = 0; j < map->sizeAt(i); ++j) { + size_t idx = map->offsetAt(i) + j; + if (!values->isNullAt(i)) { + uniqueKeys.emplace(keys->valueAt(i)); + } + } + } + } + + nimble::VeloxWriterOptions options{ + .flatMapColumns = {"map"}, + }; + auto file = + nimble::test::createNimbleFile(*rootPool_, vector, std::move(options)); + + velox::InMemoryReadFile readFile(file); + auto selector = std::make_shared( + std::dynamic_pointer_cast(type)); + + auto findKeyNotInMap = [&uniqueKeys](int32_t start) { + while (uniqueKeys.find(start) != uniqueKeys.end()) { + ++start; + } + return start; + }; + + int32_t existingKey = *uniqueKeys.begin(); + int32_t nonExistingKey1 = findKeyNotInMap(0); + int32_t nonExistingKey2 = findKeyNotInMap(nonExistingKey1 + 1); + + auto expectedInnerType = + velox::ROW({velox::BIGINT(), velox::BIGINT(), velox::BIGINT()}); + + nimble::VeloxReadParams params; + params.readFlatMapFieldAsStruct = {"map"}, + params.flatMapFeatureSelector = { + {"map", + {{folly::to(nonExistingKey1), + folly::to(existingKey), + folly::to(nonExistingKey2)}}}}; + + nimble::VeloxReader reader( + *leafPool_, &readFile, std::move(selector), std::move(params)); + + velox::VectorPtr result; + ASSERT_TRUE(reader.next(batchSize, result)); + ASSERT_EQ(result->type()->kind(), velox::TypeKind::ROW); + ASSERT_EQ(1, result->type()->as().size()); + ASSERT_EQ( + result->as()->childAt(0)->type()->kind(), + velox::TypeKind::ROW); + ASSERT_TRUE(result->as()->childAt(0)->type()->kindEquals( + expectedInnerType)); + + auto resultStruct = + result->as()->childAt(0)->as(); + ASSERT_EQ(3, resultStruct->childrenSize()); + ASSERT_EQ( + folly::to(nonExistingKey1), + resultStruct->type()->as().nameOf(0)); + ASSERT_EQ( + folly::to(existingKey), + resultStruct->type()->as().nameOf(1)); + ASSERT_EQ( + folly::to(nonExistingKey2), + resultStruct->type()->as().nameOf(2)); + ASSERT_EQ( + velox::TypeKind::BIGINT, + resultStruct->as()->childAt(0)->typeKind()); + ASSERT_EQ( + velox::TypeKind::BIGINT, + resultStruct->as()->childAt(1)->typeKind()); + ASSERT_EQ( + velox::TypeKind::BIGINT, + resultStruct->as()->childAt(2)->typeKind()); + + for (auto i : {1}) { + ASSERT_EQ(batchSize, resultStruct->size()); + auto featureVector = resultStruct->childAt(i); + bool foundNotNullValue = false; + + for (auto j = 0; j < batchSize; ++j) { + if (!featureVector->isNullAt(j)) { + foundNotNullValue = true; + break; + } + } + + ASSERT_TRUE(foundNotNullValue); + } + + for (auto i : {0, 2}) { + auto& nullChild = resultStruct->childAt(i); + ASSERT_EQ(batchSize, nullChild->size()); + for (auto j = 0; j < batchSize; ++j) { + ASSERT_TRUE(nullChild->isNullAt(j)) + << "Expecting null value at column " << i << ", row " << j; + } + } +} diff --git a/dwio/nimble/velox/tests/VeloxReaderTests.cpp b/dwio/nimble/velox/tests/VeloxReaderTests.cpp new file mode 100644 index 0000000..717f568 --- /dev/null +++ b/dwio/nimble/velox/tests/VeloxReaderTests.cpp @@ -0,0 +1,4117 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include +#include +#include + +#include "dwio/nimble/common/Buffer.h" +#include "dwio/nimble/common/Types.h" +#include "dwio/nimble/common/Vector.h" +#include "dwio/nimble/common/tests/NimbleFileWriter.h" +#include "dwio/nimble/common/tests/TestUtils.h" +#include "dwio/nimble/velox/SchemaUtils.h" +#include "dwio/nimble/velox/VeloxReader.h" +#include "dwio/nimble/velox/VeloxWriter.h" +#include "folly/FileUtil.h" +#include "folly/Random.h" +#include "folly/executors/CPUThreadPoolExecutor.h" +#include "velox/dwio/common/ColumnSelector.h" +#include "velox/type/Type.h" +#include "velox/vector/BaseVector.h" +#include "velox/vector/ComplexVector.h" +#include "velox/vector/DecodedVector.h" +#include "velox/vector/NullsBuilder.h" +#include "velox/vector/SelectivityVector.h" +#include "velox/vector/fuzzer/VectorFuzzer.h" +#include "velox/vector/tests/utils/VectorMaker.h" + +using namespace ::facebook; + +DEFINE_string( + output_test_file_path, + "", + "If provided, files created during tests will be writtern to this path. " + "Each test will overwrite the previous file, so this is mainly useful when a single test is executed."); + +DEFINE_uint32( + reader_tests_seed, + 0, + "If provided, this seed will be used when executing tests. " + "Otherwise, a random seed will be used."); + +namespace { +struct VeloxMapGeneratorConfig { + std::shared_ptr rowType; + velox::TypeKind keyType; + std::string stringKeyPrefix = "test_"; + uint32_t maxSizeForMap = 10; + unsigned long seed = FLAGS_reader_tests_seed > 0 ? FLAGS_reader_tests_seed + : folly::Random::rand32(); + bool hasNulls = true; +}; + +// Generates a batch of MpaVector Data +class VeloxMapGenerator { + public: + VeloxMapGenerator( + velox::memory::MemoryPool* pool, + VeloxMapGeneratorConfig config) + : leafPool_{pool}, config_{config}, rng_(config_.seed), buffer_(*pool) { + LOG(INFO) << "seed: " << config_.seed; + } + + velox::VectorPtr generateBatch(velox::vector_size_t batchSize) { + auto offsets = velox::allocateOffsets(batchSize, leafPool_); + auto rawOffsets = offsets->template asMutable(); + auto sizes = velox::allocateSizes(batchSize, leafPool_); + auto rawSizes = sizes->template asMutable(); + velox::vector_size_t childSize = 0; + for (auto i = 0; i < batchSize; ++i) { + rawOffsets[i] = childSize; + auto length = folly::Random::rand32(rng_) % (config_.maxSizeForMap + 1); + rawSizes[i] = length; + childSize += length; + } + + // create keys + auto keys = generateKeys(batchSize, childSize, rawSizes); + auto offset = 0; + // encode keys + if (folly::Random::oneIn(2, rng_)) { + offset = 0; + auto indices = velox::AlignedBuffer::allocate( + childSize, leafPool_); + auto rawIndices = indices->asMutable(); + for (auto i = 0; i < batchSize; ++i) { + auto mapSize = rawSizes[i]; + for (auto j = 0; j < mapSize; ++j) { + rawIndices[offset + j] = offset + mapSize - j - 1; + } + offset += mapSize; + } + keys = velox::BaseVector::wrapInDictionary( + nullptr, indices, childSize, keys); + } + + velox::VectorFuzzer fuzzer( + { + .vectorSize = static_cast(childSize), + .nullRatio = 0.1, + .stringLength = 20, + .stringVariableLength = true, + .containerLength = 5, + .containerVariableLength = true, + .dictionaryHasNulls = config_.hasNulls, + }, + leafPool_, + config_.seed); + + // Generate a random null vector. + velox::NullsBuilder builder{batchSize, leafPool_}; + if (config_.hasNulls) { + for (auto i = 0; i < batchSize; ++i) { + if (folly::Random::oneIn(10, rng_)) { + builder.setNull(i); + } + } + } + auto nulls = builder.build(); + std::vector children; + for (auto& featureColumn : config_.rowType->children()) { + velox::VectorPtr map = std::make_shared( + leafPool_, + featureColumn, + nulls, + batchSize, + offsets, + sizes, + keys, + fuzzer.fuzz(featureColumn->asMap().valueType())); + // Encode map + if (folly::Random::oneIn(2, rng_)) { + map = fuzzer.fuzzDictionary(map); + } + children.push_back(map); + } + + return std::make_shared( + leafPool_, config_.rowType, nullptr, batchSize, std::move(children)); + } + + std::mt19937& rng() { + return rng_; + } + + private: + std::shared_ptr generateKeys( + velox::vector_size_t batchSize, + velox::vector_size_t childSize, + velox::vector_size_t* rawSizes) { + switch (config_.keyType) { +#define SCALAR_CASE(veloxKind, cppType) \ + case velox::TypeKind::veloxKind: { \ + auto keys = \ + velox::BaseVector::create(velox::veloxKind(), childSize, leafPool_); \ + auto rawKeyValues = keys->asFlatVector()->mutableRawValues(); \ + auto offset = 0; \ + for (auto i = 0; i < batchSize; ++i) { \ + for (auto j = 0; j < rawSizes[i]; ++j) { \ + rawKeyValues[offset++] = folly::to(j); \ + } \ + } \ + return keys; \ + } + SCALAR_CASE(TINYINT, int8_t) + SCALAR_CASE(SMALLINT, int16_t) + SCALAR_CASE(INTEGER, int32_t) + SCALAR_CASE(BIGINT, int64_t) + +#undef SCALAR_CASE + case velox::TypeKind::VARCHAR: { + auto keys = + velox::BaseVector::create(velox::VARCHAR(), childSize, leafPool_); + auto flatVector = keys->asFlatVector(); + auto offset = 0; + for (auto i = 0; i < batchSize; ++i) { + for (auto j = 0; j < rawSizes[i]; ++j) { + auto key = config_.stringKeyPrefix + folly::to(j); + flatVector->set( + offset++, {key.data(), static_cast(key.size())}); + } + } + return keys; + } + default: + NIMBLE_NOT_SUPPORTED("Unsupported Key Type"); + } + } + velox::memory::MemoryPool* leafPool_; + VeloxMapGeneratorConfig config_; + std::mt19937 rng_; + nimble::Buffer buffer_; +}; + +template +void fillKeysVector( + velox::VectorPtr& vector, + velox::vector_size_t offset, + T& key) { + auto flatVectorMutable = + static_cast&>(*vector).mutableRawValues(); + flatVectorMutable[offset] = key; +} + +template +std::string getStringKey(T key) { + return folly::to(key); +} + +template <> +std::string getStringKey(velox::StringView key) { + return std::string(key); +} + +// utility function to convert an input Map velox::VectorPtr to outVector if +// isKeyPresent +template +void filterFlatMap( + const velox::VectorPtr& vector, + velox::VectorPtr& outVector, + std::function isKeyPresent) { + auto mapVector = vector->as(); + auto offsets = mapVector->rawOffsets(); + auto sizes = mapVector->rawSizes(); + auto keysVector = mapVector->mapKeys()->asFlatVector(); + auto valuesVector = mapVector->mapValues(); + + if (outVector == nullptr) { + outVector = velox::BaseVector::create( + vector->type(), vector->size(), vector->pool()); + } + auto resultVector = outVector->as(); + auto newKeysVector = resultVector->mapKeys(); + velox::VectorPtr newValuesVector = velox::BaseVector::create( + mapVector->mapValues()->type(), 0, mapVector->pool()); + auto* offsetsPtr = resultVector->mutableOffsets(vector->size()) + ->asMutable(); + auto* lengthsPtr = resultVector->mutableSizes(vector->size()) + ->asMutable(); + newKeysVector->resize(keysVector->size()); + newValuesVector->resize(valuesVector->size()); + resultVector->setNullCount(vector->size()); + + velox::vector_size_t offset = 0; + for (velox::vector_size_t index = 0; index < mapVector->size(); ++index) { + offsetsPtr[index] = offset; + if (!mapVector->isNullAt(index)) { + resultVector->setNull(index, false); + for (velox::vector_size_t i = offsets[index]; + i < offsets[index] + sizes[index]; + ++i) { + auto keyValue = keysVector->valueAtFast(i); + auto&& stringKeyValue = getStringKey(keyValue); + if (isKeyPresent(stringKeyValue)) { + fillKeysVector(newKeysVector, offset, keyValue); + newValuesVector->copy(valuesVector.get(), offset, i, 1); + ++offset; + } + } + } else { + resultVector->setNull(index, true); + } + lengthsPtr[index] = offset - offsetsPtr[index]; + } + + newKeysVector->resize(offset, false); + newValuesVector->resize(offset, false); + resultVector->setKeysAndValues( + std::move(newKeysVector), std::move(newValuesVector)); +} + +// compare two map vector, where expected map will be converted a new vector +// based on isKeyPresent Functor +template +void compareFlatMapAsFilteredMap( + velox::VectorPtr expected, + velox::VectorPtr actual, + std::function isKeyPresent) { + auto flat = velox::BaseVector::create( + expected->type(), expected->size(), expected->pool()); + flat->copy(expected.get(), 0, 0, expected->size()); + auto expectedRow = flat->as(); + auto actualRow = actual->as(); + EXPECT_EQ(expectedRow->childrenSize(), actualRow->childrenSize()); + for (auto i = 0; i < expectedRow->childrenSize(); ++i) { + velox::VectorPtr outVector; + filterFlatMap(expectedRow->childAt(i), outVector, isKeyPresent); + for (int j = 0; j < outVector->size(); j++) { + ASSERT_TRUE(outVector->equalValueAt(actualRow->childAt(i).get(), j, j)) + << "Content mismatch at index " << j + << "\nReference: " << outVector->toString(j) + << "\nResult: " << actualRow->childAt(i)->toString(j); + } + } +} + +template +void verifyUpcastedScalars( + const velox::VectorPtr& expected, + uint32_t& idxInExpected, + const velox::VectorPtr& result, + uint32_t readSize) { + ASSERT_TRUE(expected->isScalar() && result->isScalar()); + auto flatExpected = expected->asFlatVector(); + auto flatResult = result->asFlatVector(); + for (uint32_t i = 0; i < result->size(); ++i) { + EXPECT_EQ(expected->isNullAt(idxInExpected), result->isNullAt(i)) + << "Unexpected null status. index: " << i << ", readSize: " << readSize; + if (!result->isNullAt(i)) { + if constexpr ( + nimble::isIntegralType() || nimble::isBoolType()) { + EXPECT_EQ( + static_cast(flatExpected->valueAtFast(idxInExpected)), + flatResult->valueAtFast(i)) + << "Unexpected value. index: " << i << ", readSize: " << readSize; + } else { + EXPECT_DOUBLE_EQ( + static_cast(flatExpected->valueAtFast(idxInExpected)), + flatResult->valueAtFast(i)) + << "Unexpected value. index: " << i << ", readSize: " << readSize; + } + } + ++idxInExpected; + } +} + +size_t streamsReadCount( + velox::memory::MemoryPool& pool, + velox::ReadFile* readFile, + const std::vector& chunks) { + // Assumed for the algorithm + VELOX_CHECK_EQ(false, readFile->shouldCoalesce()); + nimble::Tablet tablet(pool, readFile); + VELOX_CHECK_GE(tablet.stripeCount(), 1); + auto offsets = tablet.streamOffsets(0); + std::unordered_set streamOffsets; + LOG(INFO) << "Number of streams: " << offsets.size(); + for (auto offset : offsets) { + LOG(INFO) << "Stream offset: " << offset; + streamOffsets.insert(offset); + } + size_t readCount = 0; + auto fileSize = readFile->size(); + for (const auto [offset, size] : chunks) { + // This is to prevent the case when the file is too small, then entire file + // is read from 0 to the end. It can also happen that we don't read from 0 + // to the end, but just the last N bytes (a big block at the end). If that + // read coincidently starts at the beginning of a stream, I may think that + // I'm reading a stream. So I'm also guarding against it. + if (streamOffsets.contains(offset) && (offset + size) != fileSize) { + ++readCount; + } + } + return readCount; +} + +} // namespace + +class VeloxReaderTests : public ::testing::Test { + protected: + static void SetUpTestCase() { + velox::memory::MemoryManager::testingSetInstance({}); + } + + void SetUp() override { + rootPool_ = velox::memory::memoryManager()->addRootPool("default_root"); + leafPool_ = rootPool_->addLeafChild("default_leaf"); + } + + static bool vectorEquals( + const velox::VectorPtr& expected, + const velox::VectorPtr& actual, + velox::vector_size_t index) { + return expected->equalValueAt(actual.get(), index, index); + } + + void verifyReadersEqual( + std::unique_ptr lhs, + std::unique_ptr rhs, + int32_t expectedNumTotalArrays, + int32_t expectedNumUniqueArrays) { + velox::VectorPtr leftResult; + velox::VectorPtr rightResult; + ASSERT_TRUE(lhs->next(expectedNumTotalArrays, leftResult)); + ASSERT_TRUE(rhs->next(expectedNumTotalArrays, rightResult)); + ASSERT_EQ( + leftResult->wrappedVector() + ->as() + ->childAt(0) + ->loadedVector() + ->wrappedVector() + ->size(), + expectedNumUniqueArrays); + ASSERT_EQ( + rightResult->wrappedVector() + ->as() + ->childAt(0) + ->loadedVector() + ->wrappedVector() + ->size(), + expectedNumUniqueArrays); + + for (int i = 0; i < expectedNumTotalArrays; ++i) { + vectorEquals(leftResult, rightResult, i); + } + + // Ensure no extra data floating in left/right + ASSERT_FALSE(lhs->next(1, leftResult)); + ASSERT_FALSE(rhs->next(1, rightResult)); + } + + void testVeloxTypeFromNimbleSchema( + velox::memory::MemoryPool& memoryPool, + nimble::VeloxWriterOptions writerOptions, + const velox::RowVectorPtr& vector) { + const auto& veloxRowType = + std::dynamic_pointer_cast(vector->type()); + auto file = + nimble::test::createNimbleFile(*rootPool_, vector, writerOptions); + auto inMemFile = velox::InMemoryReadFile(file); + + nimble::VeloxReader veloxReader( + memoryPool, + &inMemFile, + std::make_shared(veloxRowType)); + const auto& veloxTypeResult = convertToVeloxType(*veloxReader.schema()); + + EXPECT_EQ(*veloxRowType, *veloxTypeResult) + << "Expected: " << veloxRowType->toString() + << ", actual: " << veloxTypeResult->toString(); + } + + template + void getFieldDefaultValue(nimble::Vector& input, uint32_t index) { + static_assert( + T() == 0, "Default Constructor value is not zero initialized"); + input[index] = T(); + } + + template <> + void getFieldDefaultValue( + nimble::Vector& input, + uint32_t index) { + input[index] = std::string(); + } + + template + void + verifyDefaultValue(T valueToBeUpdatedWith, T defaultValue, int32_t size) { + nimble::Vector testData(leafPool_.get(), size); + for (int i = 0; i < testData.size(); ++i) { + getFieldDefaultValue(testData, i); + ASSERT_EQ(testData[i], defaultValue); + testData[i] = valueToBeUpdatedWith; + getFieldDefaultValue(testData, i); + ASSERT_EQ(testData[i], defaultValue); + } + } + + template + void testFlatMapNullValues() { + auto type = velox::ROW( + {{"fld", velox::MAP(velox::INTEGER(), velox::CppToType::create())}}); + + std::string file; + auto writeFile = std::make_unique(&file); + + facebook::nimble::VeloxWriterOptions writerOptions; + writerOptions.flatMapColumns.insert("fld"); + + nimble::VeloxWriter writer( + *rootPool_, type, std::move(writeFile), std::move(writerOptions)); + + facebook::velox::test::VectorMaker vectorMaker(leafPool_.get()); + auto values = vectorMaker.flatVectorNullable( + {std::nullopt, std::nullopt, std::nullopt}); + auto keys = vectorMaker.flatVector({1, 2, 3}); + auto vector = vectorMaker.rowVector( + {"fld"}, {vectorMaker.mapVector({0, 1, 2}, keys, values)}); + + writer.write(vector); + writer.close(); + + nimble::VeloxReadParams readParams; + velox::InMemoryReadFile readFile(file); + auto selector = std::make_shared(type); + nimble::VeloxReader reader( + *leafPool_, &readFile, std::move(selector), readParams); + + velox::VectorPtr output; + auto size = 3; + reader.next(size, output); + for (auto i = 0; i < size; ++i) { + EXPECT_TRUE(vectorEquals(vector, output, i)); + } + } + + template + void writeAndVerify( + std::mt19937& rng, + velox::memory::MemoryPool& pool, + const velox::RowTypePtr& type, + std::function generator, + std::function validator, + size_t count, + nimble::VeloxWriterOptions writerOptions = {}, + nimble::VeloxReadParams readParams = {}, + std::function isKeyPresent = nullptr, + std::function comparator = nullptr, + bool multiSkip = false, + bool checkMemoryLeak = false) { + std::string file; + auto writeFile = std::make_unique(&file); + nimble::FlushDecision decision; + writerOptions.enableChunking = true; + writerOptions.minStreamChunkRawSize = folly::Random::rand32(30, rng); + writerOptions.flushPolicyFactory = [&]() { + return std::make_unique( + [&](auto&) { return decision; }); + }; + + std::vector expected; + nimble::VeloxWriter writer( + *rootPool_, type, std::move(writeFile), std::move(writerOptions)); + bool perBatchFlush = folly::Random::oneIn(2, rng); + for (auto i = 0; i < count; ++i) { + auto vector = generator(type); + int32_t rowIndex = 0; + while (rowIndex < vector->size()) { + decision = nimble::FlushDecision::None; + auto batchSize = vector->size() - rowIndex; + // Randomly produce chunks + if (folly::Random::oneIn(2, rng)) { + batchSize = folly::Random::rand32(0, batchSize, rng) + 1; + decision = nimble::FlushDecision::Chunk; + } + if ((perBatchFlush || folly::Random::oneIn(5, rng)) && + (rowIndex + batchSize == vector->size())) { + decision = nimble::FlushDecision::Stripe; + } + writer.write(vector->slice(rowIndex, batchSize)); + rowIndex += batchSize; + } + expected.push_back(vector); + } + writer.close(); + + if (!FLAGS_output_test_file_path.empty()) { + folly::writeFile(file, FLAGS_output_test_file_path.c_str()); + } + + velox::InMemoryReadFile readFile(file); + auto selector = std::make_shared(type); + // new pool with to limit already used memory and with tracking enabled + auto leakDetectPool = + facebook::velox::memory::deprecatedDefaultMemoryManager().addRootPool( + "memory_leak_detect"); + auto readerPool = leakDetectPool->addLeafChild("reader_pool"); + + nimble::VeloxReader reader( + *readerPool.get(), &readFile, selector, readParams); + if (folly::Random::oneIn(2, rng)) { + LOG(INFO) << "using executor"; + readParams.decodingExecutor = + std::make_shared(1); + } + + auto rootTypeFromSchema = convertToVeloxType(*reader.schema()); + EXPECT_EQ(*type, *rootTypeFromSchema) + << "Expected: " << type->toString() + << ", actual: " << rootTypeFromSchema->toString(); + + velox::VectorPtr result; + velox::vector_size_t numIncrements = 0, prevMemory = 0; + for (auto i = 0; i < expected.size(); ++i) { + auto& current = expected.at(i); + ASSERT_TRUE(reader.next(current->size(), result)); + ASSERT_EQ(result->size(), current->size()); + if (comparator) { + comparator(result); + } + if (isKeyPresent) { + compareFlatMapAsFilteredMap(current, result, isKeyPresent); + } else { + for (auto j = 0; j < result->size(); ++j) { + ASSERT_TRUE(validator(current, result, j)) + << "Content mismatch at batch " << i << " at row " << j + << "\nReference: " << current->toString(j) + << "\nResult: " << result->toString(j); + } + } + + // validate skip + if (i % 2 == 0) { + nimble::VeloxReader reader1(pool, &readFile, selector, readParams); + nimble::VeloxReader reader2(pool, &readFile, selector, readParams); + auto rowCount = expected.at(0)->size(); + velox::vector_size_t remaining = rowCount; + uint32_t skipCount = 0; + do { + auto toSkip = folly::Random::rand32(1, remaining, rng); + velox::VectorPtr result1; + velox::VectorPtr result2; + reader1.next(toSkip, result1); + reader2.skipRows(toSkip); + remaining -= toSkip; + + if (remaining > 0) { + auto toRead = folly::Random::rand32(1, remaining, rng); + reader1.next(toRead, result1); + reader2.next(toRead, result2); + + ASSERT_EQ(result1->size(), result2->size()); + + for (auto j = 0; j < result1->size(); ++j) { + ASSERT_TRUE(vectorEquals(result1, result2, j)) + << "Content mismatch at index " << j + << " skipCount = " << skipCount + << " remaining = " << remaining << " to read = " << toRead + << "\nReference: " << result1->toString(j) + << "\nResult: " << result2->toString(j); + } + + remaining -= toRead; + } + skipCount += 1; + } while (multiSkip && remaining > 0); + } + + // validate memory usage + if (readerPool->currentBytes() > prevMemory) { + numIncrements++; + } + prevMemory = readerPool->currentBytes(); + } + ASSERT_FALSE(reader.next(1, result)); + if (checkMemoryLeak) { + EXPECT_LE(numIncrements, 3 * expected.size() / 4); + } + } + + std::unique_ptr getReaderForLifeCycleTest( + const std::shared_ptr schema, + size_t batchSize, + std::mt19937& rng, + nimble::VeloxWriterOptions writerOptions = {}, + nimble::VeloxReadParams readParams = {}) { + velox::VectorFuzzer fuzzer( + {.vectorSize = batchSize, .nullRatio = 0.5}, + leafPool_.get(), + folly::Random::rand32(rng)); + auto vector = fuzzer.fuzzInputFlatRow(schema); + auto file = + nimble::test::createNimbleFile(*rootPool_, vector, writerOptions); + + std::unique_ptr readFile = + std::make_unique(file); + + std::shared_ptr tablet = + std::make_shared(*leafPool_, std::move(readFile)); + auto selector = + std::make_shared(schema); + std::unique_ptr reader = + std::make_unique( + *leafPool_, tablet, std::move(selector), readParams); + return reader; + } + + std::unique_ptr getReaderForWrite( + velox::memory::MemoryPool& pool, + const velox::RowTypePtr& type, + std::vector> + generators, + size_t batchSize, + bool multiSkip = false, + bool checkMemoryLeak = false) { + nimble::VeloxWriterOptions writerOptions = {}; + nimble::VeloxReadParams readParams = {}; + + std::string file; + auto writeFile = std::make_unique(&file); + writerOptions.enableChunking = true; + writerOptions.flushPolicyFactory = [&]() { + return std::make_unique( + [&](auto&) { return nimble::FlushDecision::None; }); + }; + writerOptions.dictionaryArrayColumns.insert("dictionaryArray"); + + nimble::VeloxWriter writer( + *rootPool_, type, std::move(writeFile), std::move(writerOptions)); + + for (int i = 0; i < generators.size(); ++i) { + auto vectorGenerator = generators.at(i); + auto vector = vectorGenerator(type); + // In the future, we can take in batchSize as param and try every size + writer.write(vector); + } + writer.close(); + + auto readFile = std::make_unique(file); + auto selector = std::make_shared(type); + + return std::make_unique( + *leafPool_, std::move(readFile), std::move(selector), readParams); + } + + std::shared_ptr rootPool_; + std::shared_ptr leafPool_; +}; + +TEST_F(VeloxReaderTests, DontReadUnselectedColumnsFromFile) { + auto type = velox::ROW({ + {"tinyint_val", velox::TINYINT()}, + {"smallint_val", velox::SMALLINT()}, + {"int_val", velox::INTEGER()}, + {"long_val", velox::BIGINT()}, + {"float_val", velox::REAL()}, + {"double_val", velox::DOUBLE()}, + {"string_val", velox::VARCHAR()}, + {"array_val", velox::ARRAY(velox::BIGINT())}, + {"map_val", velox::MAP(velox::INTEGER(), velox::BIGINT())}, + }); + + size_t batchSize = 100; + auto selectedColumnNames = + std::vector{"tinyint_val", "double_val"}; + + velox::VectorFuzzer fuzzer({.vectorSize = batchSize}, leafPool_.get()); + auto vector = fuzzer.fuzzInputFlatRow(type); + + auto file = nimble::test::createNimbleFile(*rootPool_, vector); + + uint32_t readSize = 1; + for (auto useChaniedBuffers : {false, true}) { + nimble::testing::InMemoryTrackableReadFile readFile( + file, useChaniedBuffers); + // We want to check stream by stream if they are being read + readFile.setShouldCoalesce(false); + + auto selector = std::make_shared( + std::dynamic_pointer_cast(vector->type()), + selectedColumnNames); + nimble::VeloxReader reader(*leafPool_, &readFile, std::move(selector)); + + velox::VectorPtr result; + reader.next(readSize, result); + + auto chunks = readFile.chunks(); + + for (auto [offset, size] : chunks) { + LOG(INFO) << "Stream read: " << offset; + } + + EXPECT_EQ( + streamsReadCount(*leafPool_, &readFile, chunks), + selectedColumnNames.size()); + } +} + +TEST_F(VeloxReaderTests, DontReadUnprojectedFeaturesFromFile) { + auto type = velox::ROW({ + {"float_features", velox::MAP(velox::INTEGER(), velox::REAL())}, + }); + auto rowType = std::dynamic_pointer_cast(type); + + int batchSize = 500; + uint32_t seed = FLAGS_reader_tests_seed > 0 ? FLAGS_reader_tests_seed + : folly::Random::rand32(); + + VeloxMapGeneratorConfig generatorConfig{ + .rowType = rowType, + .keyType = velox::TypeKind::INTEGER, + .maxSizeForMap = 10, + .seed = seed, + .hasNulls = false, + }; + + VeloxMapGenerator generator(leafPool_.get(), generatorConfig); + auto vector = generator.generateBatch(batchSize); + + nimble::VeloxWriterOptions writerOptions; + writerOptions.flatMapColumns.insert("float_features"); + + auto file = nimble::test::createNimbleFile( + *rootPool_, vector, std::move(writerOptions)); + + for (auto useChaniedBuffers : {false, true}) { + facebook::nimble::testing::InMemoryTrackableReadFile readFile( + file, useChaniedBuffers); + // We want to check stream by stream if they are being read + readFile.setShouldCoalesce(false); + + auto selector = std::make_shared( + std::dynamic_pointer_cast(vector->type())); + + nimble::VeloxReadParams params; + params.readFlatMapFieldAsStruct.insert("float_features"); + auto& selectedFeatures = + params.flatMapFeatureSelector["float_features"].features; + std::mt19937 rng(seed); + for (int i = 0; i < generatorConfig.maxSizeForMap; ++i) { + if (folly::Random::oneIn(2, rng)) { + selectedFeatures.push_back(folly::to(i)); + } + } + // Features list can't be empty. + if (selectedFeatures.empty()) { + selectedFeatures = {folly::to( + folly::Random::rand32(generatorConfig.maxSizeForMap))}; + } + + LOG(INFO) << "Selected features (" << selectedFeatures.size() + << ") :" << folly::join(", ", selectedFeatures); + + nimble::VeloxReader reader( + *leafPool_, &readFile, std::move(selector), params); + + uint32_t readSize = 1000; + velox::VectorPtr result; + reader.next(readSize, result); + + auto selectedFeaturesSet = std::unordered_set( + selectedFeatures.cbegin(), selectedFeatures.cend()); + + // We have those streams: Row, FlatMap, N*(Values + inMap) + // Row: Empty stream. Not read. + // FlatMap: Empty if !hasNulls + // N: Number of features + // Values: Empty if all rows are null (if inMap all false) + // inMap: Non-empty + // + // Therefore the formula is: 0 + 0 + N*(Values*any(inMap) + inMap) + ASSERT_FALSE(generatorConfig.hasNulls); + int expectedNonEmptyStreamsCount = 0; // 0 if !hasNulls + auto rowResult = result->as(); + ASSERT_EQ(rowResult->childrenSize(), 1); // FlatMap + auto flatMap = rowResult->childAt(0)->as(); + + for (int feature = 0; feature < flatMap->childrenSize(); ++feature) { + // Each feature will have at least inMap stream + ++expectedNonEmptyStreamsCount; + if (selectedFeaturesSet.contains( + flatMap->type()->asRow().nameOf(feature))) { + auto columnResult = flatMap->childAt(feature); + for (int row = 0; row < columnResult->size(); ++row) { + // Values stream for this column will only exist if there's at least + // one element inMap in this column (if not all rows are null at + // either row level or element level) + if (!flatMap->isNullAt(row) && !columnResult->isNullAt(row)) { + ++expectedNonEmptyStreamsCount; + // exit row iteration, we know that there's at least one element + break; + } + } + } + } + + auto chunks = readFile.chunks(); + + LOG(INFO) << "Total streams read: " << chunks.size(); + for (auto [offset, size] : chunks) { + LOG(INFO) << "Stream read: " << offset; + } + + EXPECT_EQ( + streamsReadCount(*leafPool_, &readFile, chunks), + expectedNonEmptyStreamsCount); + } +} + +TEST_F(VeloxReaderTests, ReadComplexData) { + auto type = velox::ROW({ + {"tinyint_val", velox::TINYINT()}, + {"smallint_val", velox::SMALLINT()}, + {"int_val", velox::INTEGER()}, + {"long_val", velox::BIGINT()}, + {"float_val", velox::REAL()}, + {"double_val", velox::DOUBLE()}, + {"bool_val", velox::BOOLEAN()}, + {"string_val", velox::VARCHAR()}, + {"array_val", velox::ARRAY(velox::BIGINT())}, + {"map_val", velox::MAP(velox::INTEGER(), velox::BIGINT())}, + {"struct_val", + velox::ROW({ + {"float_val", velox::REAL()}, + {"double_val", velox::DOUBLE()}, + })}, + {"nested_val", + velox::MAP( + velox::INTEGER(), + velox::ROW({ + {"float_val", velox::REAL()}, + {"array_val", + velox::ARRAY(velox::MAP(velox::INTEGER(), velox::BIGINT()))}, + }))}, + }); + + auto typeUpcast = velox::ROW({ + {"tinyint_val", velox::SMALLINT()}, + {"smallint_val", velox::INTEGER()}, + {"int_val", velox::BIGINT()}, + {"long_val", velox::BIGINT()}, + {"float_val", velox::DOUBLE()}, + {"double_val", velox::DOUBLE()}, + {"bool_val", velox::INTEGER()}, + {"string_val", velox::VARCHAR()}, + {"array_val", velox::ARRAY(velox::BIGINT())}, + {"map_val", velox::MAP(velox::INTEGER(), velox::BIGINT())}, + {"struct_val", + velox::ROW({ + {"float_val", velox::REAL()}, + {"double_val", velox::DOUBLE()}, + })}, + {"nested_val", + velox::MAP( + velox::INTEGER(), + velox::ROW({ + {"float_val", velox::REAL()}, + {"array_val", + velox::ARRAY(velox::MAP(velox::INTEGER(), velox::BIGINT()))}, + }))}, + }); + + for (size_t batchSize : {5, 1234}) { + velox::VectorFuzzer fuzzer({.vectorSize = batchSize}, leafPool_.get()); + auto vector = fuzzer.fuzzInputFlatRow(type); + auto file = nimble::test::createNimbleFile(*rootPool_, vector); + + for (bool upcast : {false, true}) { + for (uint32_t readSize : {1, 2, 5, 7, 20, 100, 555, 2000}) { + velox::InMemoryReadFile readFile(file); + auto selector = std::make_shared( + std::dynamic_pointer_cast( + upcast ? typeUpcast : vector->type())); + nimble::VeloxReader reader(*leafPool_, &readFile, std::move(selector)); + + velox::vector_size_t rowIndex = 0; + std::vector childRowIndices( + vector->as()->childrenSize(), 0); + velox::VectorPtr result; + while (reader.next(readSize, result)) { + ASSERT_EQ(result->type()->kind(), velox::TypeKind::ROW); + if (upcast) { + verifyUpcastedScalars( + vector->as()->childAt(0), + childRowIndices[0], + result->as()->childAt(0), + readSize); + verifyUpcastedScalars( + vector->as()->childAt(1), + childRowIndices[1], + result->as()->childAt(1), + readSize); + verifyUpcastedScalars( + vector->as()->childAt(2), + childRowIndices[2], + result->as()->childAt(2), + readSize); + verifyUpcastedScalars( + vector->as()->childAt(3), + childRowIndices[3], + result->as()->childAt(3), + readSize); + verifyUpcastedScalars( + vector->as()->childAt(4), + childRowIndices[4], + result->as()->childAt(4), + readSize); + verifyUpcastedScalars( + vector->as()->childAt(5), + childRowIndices[5], + result->as()->childAt(5), + readSize); + verifyUpcastedScalars( + vector->as()->childAt(6), + childRowIndices[6], + result->as()->childAt(6), + readSize); + } else { + for (velox::vector_size_t i = 0; i < result->size(); ++i) { + ASSERT_TRUE(vector->equalValueAt(result.get(), rowIndex, i)) + << "Content mismatch at index " << rowIndex + << "\nReference: " << vector->toString(rowIndex) + << "\nResult: " << result->toString(i); + + ++rowIndex; + } + } + } + } + } + } +} + +TEST_F(VeloxReaderTests, Lifetime) { + velox::StringView s{"012345678901234567890123456789"}; + std::vector strings{s, s, s, s, s}; + std::vector> stringsOfStrings{ + strings, strings, strings, strings, strings}; + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + auto vector = vectorMaker.rowVector( + {vectorMaker.flatVector({1, 2, 3, 4, 5}), + vectorMaker.flatVector(strings), + vectorMaker.arrayVector(stringsOfStrings), + vectorMaker.mapVector( + 5, + /*sizeAt*/ [](auto row) { return row; }, + /*keyAt*/ [](auto row) { return row; }, + /*valueAt*/ + [&s](auto /* row */) { return s; }), + vectorMaker.rowVector( + /* childNames */ {"a", "b"}, + /* children */ + {vectorMaker.flatVector({1., 2., 3., 4., 5.}), + vectorMaker.flatVector(strings)})}); + + velox::VectorPtr result; + { + auto file = nimble::test::createNimbleFile(*rootPool_, vector); + velox::InMemoryReadFile readFile(file); + auto selector = std::make_shared( + std::dynamic_pointer_cast(vector->type())); + nimble::VeloxReader reader(*leafPool_, &readFile, std::move(selector)); + + ASSERT_TRUE(reader.next(vector->size(), result)); + ASSERT_FALSE(reader.next(vector->size(), result)); + } + + // At this point, the reader is dropped, so the vector should be + // self-contained and doesn't rely on the reader state. + + ASSERT_EQ(vector->size(), result->size()); + ASSERT_EQ(result->type()->kind(), velox::TypeKind::ROW); + + for (int32_t i = 0; i < result->size(); ++i) { + ASSERT_TRUE(vector->equalValueAt(result.get(), i, i)) + << "Content mismatch at index " << i + << "\nReference: " << vector->toString(i) + << "\nResult: " << result->toString(i); + } +} + +TEST_F(VeloxReaderTests, AllValuesNulls) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + auto vector = vectorMaker.rowVector( + {vectorMaker.flatVectorNullable( + {std::nullopt, std::nullopt, std::nullopt}), + vectorMaker.flatVectorNullable( + {std::nullopt, std::nullopt, std::nullopt}), + velox::BaseVector::createNullConstant( + velox::ROW({{"foo", velox::INTEGER()}}), 3, leafPool_.get()), + velox::BaseVector::createNullConstant( + velox::MAP(velox::INTEGER(), velox::BIGINT()), 3, leafPool_.get()), + velox::BaseVector::createNullConstant( + velox::ARRAY(velox::INTEGER()), 3, leafPool_.get())}); + + auto projectedType = velox::ROW({ + {"c0", velox::INTEGER()}, + {"c1", velox::DOUBLE()}, + {"c2", velox::ROW({{"foo", velox::INTEGER()}})}, + {"c3", velox::MAP(velox::INTEGER(), velox::BIGINT())}, + {"c4", velox::ARRAY(velox::INTEGER())}, + }); + velox::VectorPtr result; + { + nimble::VeloxWriterOptions options; + options.flatMapColumns.insert("c3"); + options.dictionaryArrayColumns.insert("c4"); + auto file = nimble::test::createNimbleFile(*rootPool_, vector, options); + velox::InMemoryReadFile readFile(file); + + nimble::VeloxReadParams params; + params.readFlatMapFieldAsStruct.insert("c3"); + params.flatMapFeatureSelector.insert({"c3", {{"1"}}}); + auto selector = + std::make_shared(projectedType); + nimble::VeloxReader reader(*leafPool_, &readFile, selector, params); + + ASSERT_TRUE(reader.next(vector->size(), result)); + ASSERT_FALSE(reader.next(vector->size(), result)); + } + + // At this point, the reader is dropped, so the vector should be + // self-contained and doesn't rely on the reader state. + + ASSERT_EQ(vector->size(), result->size()); + auto& vectorType = result->type(); + ASSERT_EQ(vectorType->kind(), velox::TypeKind::ROW); + ASSERT_EQ(vectorType->size(), projectedType->size()); + ASSERT_EQ(vectorType->childAt(3)->kind(), velox::TypeKind::ROW); + ASSERT_EQ(vectorType->childAt(4)->kind(), velox::TypeKind::ARRAY); + + auto resultRow = result->as(); + for (int32_t i = 0; i < result->size(); ++i) { + for (auto j = 0; j < vectorType->size(); ++j) { + ASSERT_TRUE(resultRow->childAt(j)->isNullAt(i)); + } + } +} + +TEST_F(VeloxReaderTests, StringBuffers) { + // Creating a string column with 10 identical strings. + // We will perform 2 reads of 5 rows each, and compare the string buffers + // generated. + // Note: all strings are long enough to force Velox to store them in string + // buffers instead of inlining them. + std::string s{"012345678901234567890123456789"}; + std::vector column{s, s, s, s, s, s, s, s, s, s}; + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + auto vector = vectorMaker.rowVector({vectorMaker.flatVector(column)}); + + velox::VectorPtr result; + auto file = nimble::test::createNimbleFile(*rootPool_, vector); + velox::InMemoryReadFile readFile(file); + auto selector = std::make_shared( + std::dynamic_pointer_cast(vector->type())); + nimble::VeloxReader reader(*leafPool_, &readFile, std::move(selector)); + + ASSERT_TRUE(reader.next(5, result)); + + ASSERT_EQ(5, result->size()); + ASSERT_EQ(result->type()->kind(), velox::TypeKind::ROW); + auto rowVector = result->as(); + ASSERT_EQ(1, rowVector->childrenSize()); + const auto& buffers1 = rowVector->childAt(0) + ->as>() + ->stringBuffers(); + EXPECT_LE( + 1, + rowVector->childAt(0) + ->as>() + ->stringBuffers() + .size()); + + // Capture string buffer size after first batch read + auto bufferSizeFirst = std::accumulate( + buffers1.begin(), buffers1.end(), 0, [](int sum, const auto& buffer) { + return sum + buffer->size(); + }); + + ASSERT_TRUE(reader.next(5, result)); + rowVector = result->as(); + ASSERT_EQ(1, rowVector->childrenSize()); + const auto& buffers2 = rowVector->childAt(0) + ->as>() + ->stringBuffers(); + + ASSERT_EQ(5, result->size()); + EXPECT_LE( + 1, + rowVector->childAt(0) + ->as>() + ->stringBuffers() + .size()); + + // Capture string buffer size after second batch read. Since both batched + // contain exactly the same strings ,batch sizes should match. + auto bufferSizeSecond = std::accumulate( + buffers2.begin(), buffers2.end(), 0, [](int sum, const auto& buffer) { + return sum + buffer->size(); + }); + + EXPECT_EQ(bufferSizeFirst, bufferSizeSecond); +} + +TEST_F(VeloxReaderTests, NullVectors) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + // In the following table, the first 5 rows contain nulls and the last 5 + // rows don't. + auto vector = vectorMaker.rowVector( + {vectorMaker.flatVectorNullable( + {1, 2, std::nullopt, 4, 5, 6, 7, 8, 9, 10}), + vectorMaker.flatVectorNullable( + {"1", std::nullopt, "3", "4", "5", "6", "7", "8", "9", "10"}), + vectorMaker.arrayVectorNullable( + {std::vector>{1.0, 2.2, std::nullopt}, + {}, + std::nullopt, + std::vector>{1.1, 2.0}, + {}, + std::vector>{6.1}, + std::vector>{7.1}, + std::vector>{8.1}, + std::vector>{9.1}, + std::vector>{10.1}}), + vectorMaker.mapVector( + 10, + /*sizeAt*/ [](auto row) { return row; }, + /*keyAt*/ [](auto row) { return row; }, + /*valueAt*/ + [](auto row) { return row; }, + /*isNullAt*/ [](auto row) { return row < 5 && row % 2 == 0; }), + vectorMaker.rowVector( + /* childNames */ {"a", "b"}, + /* children */ + {vectorMaker.flatVector({1, 2, 3, 4, 5, 6, 7, 8, 9, 10}), + vectorMaker.flatVector( + {1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.10})})}); + vector->childAt(4)->setNull(2, true); // Set null in row vector + + auto file = nimble::test::createNimbleFile(*rootPool_, vector); + velox::InMemoryReadFile readFile(file); + auto selector = std::make_shared( + std::dynamic_pointer_cast(vector->type())); + + nimble::VeloxReader reader(*leafPool_, &readFile, std::move(selector)); + + velox::VectorPtr result; + + // When reader is reading the first 5 rows, it should find null entries and + // vectors should indicate that nulls exist. + ASSERT_TRUE(reader.next(5, result)); + ASSERT_EQ(5, result->size()); + ASSERT_EQ(velox::TypeKind::ROW, result->type()->kind()); + + auto rowVector = result->as(); + + ASSERT_EQ(5, rowVector->childrenSize()); + EXPECT_TRUE(rowVector->childAt(0)->mayHaveNulls()); + EXPECT_TRUE(rowVector->childAt(1)->mayHaveNulls()); + EXPECT_TRUE(rowVector->childAt(2)->mayHaveNulls()); + EXPECT_TRUE(rowVector->childAt(3)->mayHaveNulls()); + EXPECT_TRUE(rowVector->childAt(4)->mayHaveNulls()); + + for (int32_t i = 0; i < result->size(); ++i) { + ASSERT_TRUE(vector->equalValueAt(result.get(), i, i)) + << "Content mismatch at index " << i + << "\nReference: " << vector->toString(i) + << "\nResult: " << result->toString(i); + } + + // When reader is reading the last 5 rows, it should identify that no null + // exist and optimize vectors to efficiently indicate that. + ASSERT_TRUE(reader.next(5, result)); + rowVector = result->as(); + + EXPECT_FALSE(rowVector->childAt(0)->mayHaveNulls()); + EXPECT_FALSE(rowVector->childAt(1)->mayHaveNulls()); + EXPECT_FALSE(rowVector->childAt(2)->mayHaveNulls()); + EXPECT_FALSE(rowVector->childAt(3)->mayHaveNulls()); + EXPECT_FALSE(rowVector->childAt(4)->mayHaveNulls()); + + for (int32_t i = 0; i < result->size(); ++i) { + ASSERT_TRUE(vector->equalValueAt(result.get(), i + 5, i)) + << "Content mismatch at index " << i + 5 + << "\nReference: " << vector->toString(i + 5) + << "\nResult: " << result->toString(i); + } + + ASSERT_FALSE(reader.next(1, result)); +} + +TEST_F(VeloxReaderTests, ArrayWithOffsetsCaching) { + auto type = velox::ROW({ + {"dictionaryArray", velox::ARRAY(velox::INTEGER())}, + }); + auto rowType = std::dynamic_pointer_cast(type); + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + // Test cache hit on second write + auto generator = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector( + {{1, 2, 3}, + {1, 2, 3}, + {1, 2, 3, 4, 5}, + {1, 2, 3, 4, 5}, + {1, 2, 3, 4, 5}, + {1, 2, 3, 4, 5}}), + }); + }; + auto halfGenerator = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector( + {{1, 2, 3}, {1, 2, 3}, {1, 2, 3, 4, 5}}), + }); + }; + auto otherHalfGenerator = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector( + {{1, 2, 3, 4, 5}, {1, 2, 3, 4, 5}, {1, 2, 3, 4, 5}}), + }); + }; + + auto leftGenerators = + std::vector>( + {generator}); + auto rightGenerators = + std::vector>( + {halfGenerator, otherHalfGenerator}); + auto expectedTotalArrays = 6; + auto expectedUniqueArrays = 2; + auto left = getReaderForWrite(*leafPool_, rowType, leftGenerators, 1); + auto right = getReaderForWrite(*leafPool_, rowType, rightGenerators, 1); + + verifyReadersEqual( + std::move(left), + std::move(right), + expectedTotalArrays, + expectedUniqueArrays); + + // Test cache miss on second write + // We can re-use the totalGenerator since the total output is unchanged + auto halfGeneratorCacheMiss = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector({{1, 2, 3}, {1, 2, 3}}), + }); + }; + auto otherHalfGeneratorCacheMiss = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector( + {{1, 2, 3, 4, 5}, + {1, 2, 3, 4, 5}, + {1, 2, 3, 4, 5}, + {1, 2, 3, 4, 5}}), + }); + }; + + rightGenerators = + std::vector>( + {halfGeneratorCacheMiss, otherHalfGeneratorCacheMiss}); + expectedTotalArrays = 6; + expectedUniqueArrays = 2; + left = getReaderForWrite(*leafPool_, rowType, leftGenerators, 1); + right = getReaderForWrite(*leafPool_, rowType, rightGenerators, 1); + + verifyReadersEqual( + std::move(left), + std::move(right), + expectedTotalArrays, + expectedUniqueArrays); + + // Check cached value against null on next write. + auto fullGeneratorWithNull = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVectorNullable( + {std::vector>{1, 2, 3}, + std::vector>{1, 2, 3}, + std::vector>{1, 2, 3, 4, 5}, + std::nullopt, + std::vector>{1, 2, 3, 4, 5}, + std::vector>{1, 2, 3, 4, 5}}), + }); + }; + auto halfGeneratorNoNull = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVectorNullable( + {std::vector>{1, 2, 3}, + std::vector>{1, 2, 3}, + std::vector>{1, 2, 3, 4, 5}}), + }); + }; + auto otherHalfGeneratorWithNull = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVectorNullable( + {std::nullopt, + std::vector>{1, 2, 3, 4, 5}, + std::vector>{1, 2, 3, 4, 5}}), + }); + }; + + leftGenerators = + std::vector>( + {fullGeneratorWithNull}); + rightGenerators = + std::vector>( + {halfGeneratorNoNull, otherHalfGeneratorWithNull}); + + expectedTotalArrays = 6; // null array included in count + expectedUniqueArrays = 2; // {{1, 2, 3}, {1, 2, 3, 4, 5}} + left = getReaderForWrite(*leafPool_, rowType, leftGenerators, 1); + right = getReaderForWrite(*leafPool_, rowType, rightGenerators, 1); + + verifyReadersEqual( + std::move(left), + std::move(right), + expectedTotalArrays, + expectedUniqueArrays); + + // Check cache stores last valid value, and not null. + // Re-use total generator from previous test case + auto halfGeneratorWithNull = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVectorNullable( + {std::vector>{1, 2, 3}, + std::vector>{1, 2, 3}, + std::vector>{1, 2, 3, 4, 5}, + std::nullopt}), + }); + }; + auto otherHalfGeneratorNoNull = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVectorNullable( + {std::vector>{1, 2, 3, 4, 5}, + std::vector>{1, 2, 3, 4, 5}}), + }); + }; + rightGenerators = + std::vector>( + {halfGeneratorWithNull, otherHalfGeneratorNoNull}); + + expectedTotalArrays = 6; // null array included in count + expectedUniqueArrays = 2; // {{1, 2, 3}, {1, 2, 3, 4, 5}} + left = getReaderForWrite(*leafPool_, rowType, leftGenerators, 1); + right = getReaderForWrite(*leafPool_, rowType, rightGenerators, 1); + + verifyReadersEqual( + std::move(left), + std::move(right), + expectedTotalArrays, + expectedUniqueArrays); + + // Check cached value against empty vector + auto fullGeneratorWithEmpty = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector( + {{1, 2, 3}, + {1, 2, 3}, + {1, 2, 3, 4, 5}, + {}, + {1, 2, 3, 4, 5}, + {1, 2, 3, 4, 5}}), + }); + }; + + auto halfGeneratorNoEmpty = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector( + {{1, 2, 3}, {1, 2, 3}, {1, 2, 3, 4, 5}}), + }); + }; + auto otherHalfGeneratorWithEmpty = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector( + {{}, {1, 2, 3, 4, 5}, {1, 2, 3, 4, 5}}), + }); + }; + + leftGenerators = + std::vector>( + {fullGeneratorWithEmpty}); + rightGenerators = + std::vector>( + {halfGeneratorNoEmpty, otherHalfGeneratorWithEmpty}); + expectedTotalArrays = 6; // empty array included in count + expectedUniqueArrays = 4; // {{1, 2, 3}, {1, 2, 3, 4, 5}, {}, {1, 2, 3, 4, 5}} + left = getReaderForWrite(*leafPool_, rowType, leftGenerators, 1); + right = getReaderForWrite(*leafPool_, rowType, rightGenerators, 1); + + verifyReadersEqual( + std::move(left), + std::move(right), + expectedTotalArrays, + expectedUniqueArrays); + + // Test empty vector stored in cache from first write + // Re-use total generator from last test case + auto halfGeneratorWithEmpty = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector( + {{1, 2, 3}, {1, 2, 3}, {1, 2, 3, 4, 5}, {}}), + }); + }; + auto otherHalfGeneratorNoEmpty = [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector( + {{1, 2, 3, 4, 5}, {1, 2, 3, 4, 5}}), + }); + }; + + leftGenerators = + std::vector>( + {fullGeneratorWithEmpty}); + rightGenerators = + std::vector>( + {halfGeneratorWithEmpty, otherHalfGeneratorNoEmpty}); + expectedTotalArrays = 6; // empty array included in count + expectedUniqueArrays = 4; // {{1, 2, 3}, {1, 2, 3, 4, 5}, {}, {1, 2, 3, 4, 5}} + left = getReaderForWrite(*leafPool_, rowType, leftGenerators, 1); + right = getReaderForWrite(*leafPool_, rowType, rightGenerators, 1); + + verifyReadersEqual( + std::move(left), + std::move(right), + expectedTotalArrays, + expectedUniqueArrays); +} + +TEST_F(VeloxReaderTests, FuzzSimple) { + auto type = velox::ROW({ + {"bool_val", velox::BOOLEAN()}, + {"byte_val", velox::TINYINT()}, + {"short_val", velox::SMALLINT()}, + {"int_val", velox::INTEGER()}, + {"long_val", velox::BIGINT()}, + {"float_val", velox::REAL()}, + {"double_val", velox::DOUBLE()}, + {"string_val", velox::VARCHAR()}, + {"binary_val", velox::VARBINARY()}, + // {"ts_val", velox::TIMESTAMP()}, + }); + auto rowType = std::dynamic_pointer_cast(type); + uint32_t seed = FLAGS_reader_tests_seed > 0 ? FLAGS_reader_tests_seed + : folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + + // Small batches creates more edge cases. + size_t batchSize = 10; + velox::VectorFuzzer noNulls( + { + .vectorSize = batchSize, + .nullRatio = 0, + .stringLength = 20, + .stringVariableLength = true, + }, + leafPool_.get(), + seed); + + velox::VectorFuzzer hasNulls{ + { + .vectorSize = batchSize, + .nullRatio = 0.05, + .stringLength = 10, + .stringVariableLength = true, + }, + leafPool_.get(), + seed}; + + auto iterations = 20; + auto batches = 20; + std::mt19937 rng{seed}; + for (auto parallelismFactor : {0U, 1U, std::thread::hardware_concurrency()}) { + LOG(INFO) << "Parallelism Factor: " << parallelismFactor; + nimble::VeloxWriterOptions writerOptions; + if (parallelismFactor > 0) { + writerOptions.encodingExecutor = + std::make_shared(parallelismFactor); + } + + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + rng, + *leafPool_, + rowType, + [&](auto& type) { return noNulls.fuzzInputRow(type); }, + vectorEquals, + batches, + writerOptions); + writeAndVerify( + rng, + *leafPool_, + rowType, + [&](auto& type) { return hasNulls.fuzzInputRow(type); }, + vectorEquals, + batches, + writerOptions); + } + } +} + +TEST_F(VeloxReaderTests, FuzzComplex) { + auto type = velox::ROW({ + {"array", velox::ARRAY(velox::REAL())}, + {"dict_array", velox::ARRAY(velox::REAL())}, + {"map", velox::MAP(velox::INTEGER(), velox::DOUBLE())}, + {"row", + velox::ROW({ + {"a", velox::REAL()}, + {"b", velox::INTEGER()}, + })}, + {"nested", + velox::ARRAY(velox::ROW({ + {"a", velox::INTEGER()}, + {"b", velox::MAP(velox::REAL(), velox::REAL())}, + }))}, + {"nested_map_array1", + velox::MAP(velox::INTEGER(), velox::ARRAY(velox::REAL()))}, + {"nested_map_array2", + velox::MAP(velox::INTEGER(), velox::ARRAY(velox::INTEGER()))}, + }); + auto rowType = std::dynamic_pointer_cast(type); + uint32_t seed = FLAGS_reader_tests_seed > 0 ? FLAGS_reader_tests_seed + : folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + + nimble::VeloxWriterOptions writerOptions; + writerOptions.dictionaryArrayColumns.insert("nested_map_array1"); + writerOptions.dictionaryArrayColumns.insert("nested_map_array2"); + writerOptions.dictionaryArrayColumns.insert("dict_array"); + + // Small batches creates more edge cases. + size_t batchSize = 10; + velox::VectorFuzzer noNulls( + { + .vectorSize = batchSize, + .nullRatio = 0, + .stringLength = 20, + .stringVariableLength = true, + .containerLength = 5, + .containerVariableLength = true, + }, + leafPool_.get(), + seed); + + velox::VectorFuzzer hasNulls{ + { + .vectorSize = batchSize, + .nullRatio = 0.05, + .stringLength = 10, + .stringVariableLength = true, + .containerLength = 5, + .containerVariableLength = true, + }, + leafPool_.get(), + seed}; + + auto iterations = 20; + auto batches = 20; + std::mt19937 rng{seed}; + + for (auto parallelismFactor : {0U, 1U, std::thread::hardware_concurrency()}) { + LOG(INFO) << "Parallelism Factor: " << parallelismFactor; + writerOptions.encodingExecutor = parallelismFactor > 0 + ? std::make_shared(parallelismFactor) + : nullptr; + + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& type) { return noNulls.fuzzInputRow(type); }, + vectorEquals, + batches, + writerOptions); + writeAndVerify( + rng, + *leafPool_, + rowType, + [&](auto& type) { return hasNulls.fuzzInputRow(type); }, + vectorEquals, + batches, + writerOptions); + } + } +} + +TEST_F(VeloxReaderTests, ArrayWithOffsets) { + auto type = velox::ROW({ + {"dictionaryArray", velox::ARRAY(velox::INTEGER())}, + }); + auto rowType = std::dynamic_pointer_cast(type); + uint32_t seed = FLAGS_reader_tests_seed > 0 ? FLAGS_reader_tests_seed + : folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + + nimble::VeloxWriterOptions writerOptions; + writerOptions.dictionaryArrayColumns.insert("dictionaryArray"); + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto iterations = 20; + auto batches = 20; + std::mt19937 rng{seed}; + int expectedNumArrays = 0; + bool checkMemoryLeak = true; + + auto compare = [&](const velox::VectorPtr& vector) { + ASSERT_EQ( + vector->wrappedVector() + ->as() + ->childAt(0) + ->loadedVector() + ->wrappedVector() + ->size(), + expectedNumArrays); + }; + for (auto i = 0; i < iterations; ++i) { + expectedNumArrays = 1; + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector({{1, 2}}), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector({{1, 2}, {1, 2}}), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector({{}}), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector({{}, {}, {}}), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + + expectedNumArrays = 3; + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector( + {{1, 2}, {1, 2}, {2, 3}, {5, 6, 7}, {5, 6, 7}}), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector( + {{1, 2}, {1, 2}, {2, 3}, {}, {}}), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector( + {{}, {}, {2, 3}, {5, 6, 7}, {5, 6, 7}}), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector( + {{1, 2}, {1, 2}, {}, {5, 6, 7}, {5, 6, 7}}), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + + expectedNumArrays = 4; + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVector( + {{1, 3}, {1, 2}, {}, {5, 6, 7}, {5, 6, 7}}), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + + expectedNumArrays = 5; + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + // The middle element is a 0 length element and not null + vectorMaker.arrayVector( + {{1, 3}, {1, 2}, {}, {1, 2}, {5, 6, 7}}), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + // The middle element is a 0 length element and not null + vectorMaker.arrayVector( + {{1, 3}, {1, 2}, {}, {1, 2}, {5, 6, 7}}), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + } +} + +TEST_F(VeloxReaderTests, ArrayWithOffsetsNullable) { + auto type = velox::ROW({ + {"dictionaryArray", velox::ARRAY(velox::INTEGER())}, + }); + auto rowType = std::dynamic_pointer_cast(type); + uint32_t seed = FLAGS_reader_tests_seed > 0 ? FLAGS_reader_tests_seed + : folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + + nimble::VeloxWriterOptions writerOptions; + writerOptions.dictionaryArrayColumns.insert("dictionaryArray"); + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto iterations = 20; + auto batches = 20; + std::mt19937 rng{seed}; + int expectedNumArrays = 0; + bool checkMemoryLeak = true; + + auto compare = [&](const velox::VectorPtr& vector) { + ASSERT_EQ( + vector->wrappedVector() + ->as() + ->childAt(0) + ->loadedVector() + ->wrappedVector() + ->size(), + expectedNumArrays); + }; + for (auto i = 0; i < iterations; ++i) { + expectedNumArrays = 1; + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVectorNullable({{}, std::nullopt}), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVectorNullable({std::nullopt}), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + + expectedNumArrays = 2; + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVectorNullable({ + std::vector>{1, 2, std::nullopt}, + {}, + std::vector>{1, 2, std::nullopt}, + std::nullopt, + std::vector>{1, 2, std::nullopt}, + std::vector>{1, 2, std::nullopt}, + std::vector>{1, 2}, + }), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + + expectedNumArrays = 2; + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVectorNullable({ + std::vector>{1, 3}, + std::vector>{1, 2}, + {}, + std::vector>{1, 2}, + std::nullopt, + std::vector>{1, 2}, + }), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + + expectedNumArrays = 1; + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVectorNullable({ + std::vector>{1, 2}, + std::vector>{1, 2}, + {}, + std::nullopt, + std::vector>{1, 2}, + }), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + + expectedNumArrays = 1; + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVectorNullable({ + std::vector>{1, 2}, + std::vector>{1, 2}, + std::vector>{1, 2}, + {}, + std::nullopt, + }), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + + expectedNumArrays = 1; + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVectorNullable({ + {}, + std::nullopt, + std::vector>{1, 2}, + std::vector>{1, 2}, + std::vector>{1, 2}, + }), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + false, + checkMemoryLeak); + } +} + +TEST_F(VeloxReaderTests, ArrayWithOffsetsMultiskips) { + auto type = velox::ROW({ + {"dictionaryArray", velox::ARRAY(velox::INTEGER())}, + }); + auto rowType = std::dynamic_pointer_cast(type); + uint32_t seed = FLAGS_reader_tests_seed > 0 ? FLAGS_reader_tests_seed + : folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + + nimble::VeloxWriterOptions writerOptions; + writerOptions.dictionaryArrayColumns.insert("dictionaryArray"); + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto iterations = 50; + auto batches = 20; + std::mt19937 rng{seed}; + int expectedNumArrays = 0; + bool checkMemoryLeak = true; + + auto compare = [&](const velox::VectorPtr& vector) { + ASSERT_EQ( + vector->wrappedVector() + ->as() + ->childAt(0) + ->loadedVector() + ->wrappedVector() + ->size(), + expectedNumArrays); + }; + + auto strideVector = [&](const std::vector>& vector) { + std::vector> stridedVector; + + for (auto const& vec : vector) { + for (uint32_t idx = 0; idx < folly::Random::rand32(1, 5, rng); ++idx) { + stridedVector.push_back(vec); + } + } + return stridedVector; + }; + + for (auto i = 0; i < iterations; ++i) { + expectedNumArrays = 6; + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + // The middle element is a 0 length element and not null + vectorMaker.arrayVector(strideVector( + {{1, 2}, {1, 2, 3}, {}, {1, 2, 3}, {}, {4, 5, 6, 7}})), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + true, + checkMemoryLeak); + + expectedNumArrays = 3; + writeAndVerify( + rng, + *leafPool_.get(), + rowType, + [&](auto& /*type*/) { + return vectorMaker.rowVector( + {"dictionaryArray"}, + { + vectorMaker.arrayVectorNullable({ + std::vector>{1, 2}, + std::vector>{1, 2, 3}, + std::nullopt, + std::vector>{1, 2, 3}, + std::nullopt, + std::vector>{4, 5, 6, 7}, + }), + }); + }, + vectorEquals, + batches, + writerOptions, + {}, + nullptr, + compare, + true, + checkMemoryLeak); + } +} + +// convert map to struct +template +bool compareFlatMap( + const velox::VectorPtr& expected, + const velox::VectorPtr& actual, + velox::vector_size_t index) { + auto mapVector = expected->as(); + auto offsets = mapVector->rawOffsets(); + auto sizes = mapVector->rawSizes(); + auto keysVector = mapVector->mapKeys()->asFlatVector(); + auto valuesVector = mapVector->mapValues(); + + auto structVector = actual->as(); + folly::F14FastMap columnOffsets( + structVector->childrenSize()); + for (auto i = 0; i < structVector->childrenSize(); ++i) { + columnOffsets[structVector->type()->asRow().nameOf(i)] = i; + } + + std::unordered_set keys; + if (!mapVector->isNullAt(index)) { + for (auto i = offsets[index]; i < offsets[index] + sizes[index]; ++i) { + auto key = keysVector->valueAtFast(i); + keys.insert(folly::to(key)); + if (!valuesVector->equalValueAt( + structVector->childAt(columnOffsets[folly::to(key)]) + .get(), + i, + index)) { + return false; + } + } + } + // missing keys should be null + for (const auto& columnOffset : columnOffsets) { + if (keys.count(folly::to(columnOffset.first)) == 0 && + !structVector->childAt(columnOffset.second)->isNullAt(index)) { + return false; + } + } + + return true; +} + +template +bool compareFlatMaps( + const velox::VectorPtr& expected, + const velox::VectorPtr& actual, + velox::vector_size_t index) { + auto flat = velox::BaseVector::create( + expected->type(), expected->size(), expected->pool()); + flat->copy(expected.get(), 0, 0, expected->size()); + auto expectedRow = flat->as(); + auto actualRow = actual->as(); + EXPECT_EQ(expectedRow->childrenSize(), actualRow->childrenSize()); + for (auto i = 0; i < expectedRow->childrenSize(); ++i) { + auto columnType = actualRow->type()->childAt(i); + if (columnType->kind() != velox::TypeKind::ROW) { + return false; + } + if (!compareFlatMap( + expectedRow->childAt(i), actualRow->childAt(i), index)) { + return false; + } + } + return true; +} + +TEST_F(VeloxReaderTests, FlatMapNullValues) { + testFlatMapNullValues(); + testFlatMapNullValues(); + testFlatMapNullValues(); + testFlatMapNullValues(); + testFlatMapNullValues(); + testFlatMapNullValues(); + testFlatMapNullValues(); +} + +TEST_F(VeloxReaderTests, FlatMapToStruct) { + auto floatFeatures = velox::MAP(velox::INTEGER(), velox::REAL()); + auto idListFeatures = + velox::MAP(velox::INTEGER(), velox::ARRAY(velox::BIGINT())); + auto idScoreListFeatures = + velox::MAP(velox::INTEGER(), velox::MAP(velox::BIGINT(), velox::REAL())); + auto rowColumn = velox::MAP( + velox::INTEGER(), + velox::ROW({{"a", velox::INTEGER()}, {"b", velox::REAL()}})); + + auto type = velox::ROW({ + {"float_features", floatFeatures}, + {"id_list_features", idListFeatures}, + {"id_score_list_features", idScoreListFeatures}, + {"row_column", rowColumn}, + }); + auto rowType = std::dynamic_pointer_cast(type); + + VeloxMapGeneratorConfig generatorConfig{ + .rowType = rowType, + .keyType = velox::TypeKind::INTEGER, + .maxSizeForMap = 10}; + VeloxMapGenerator generator(leafPool_.get(), generatorConfig); + + nimble::VeloxWriterOptions writerOptions; + writerOptions.flatMapColumns.insert("float_features"); + writerOptions.flatMapColumns.insert("id_list_features"); + writerOptions.flatMapColumns.insert("id_score_list_features"); + writerOptions.flatMapColumns.insert("row_column"); + + nimble::VeloxReadParams params; + params.readFlatMapFieldAsStruct.insert("float_features"); + params.readFlatMapFieldAsStruct.insert("id_list_features"); + params.readFlatMapFieldAsStruct.insert("id_score_list_features"); + params.readFlatMapFieldAsStruct.insert("row_column"); + for (auto i = 0; i < 10; ++i) { + params.flatMapFeatureSelector["float_features"].features.push_back( + folly::to(i)); + params.flatMapFeatureSelector["id_list_features"].features.push_back( + folly::to(i)); + params.flatMapFeatureSelector["id_score_list_features"].features.push_back( + folly::to(i)); + params.flatMapFeatureSelector["row_column"].features.push_back( + folly::to(i)); + } + + auto iterations = 20; + auto batches = 10; + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + generator.rng(), + *leafPool_, + rowType, + [&](auto&) { return generator.generateBatch(10); }, + compareFlatMaps, + batches, + writerOptions, + params); + } +} + +TEST_F(VeloxReaderTests, FlatMapToStructForComplexType) { + auto rowColumn = velox::MAP( + velox::INTEGER(), + velox::ROW( + {{"a", velox::INTEGER()}, + {"b", velox::MAP(velox::INTEGER(), velox::ARRAY(velox::REAL()))}})); + + auto type = velox::ROW({ + {"row_column", rowColumn}, + }); + auto rowType = std::dynamic_pointer_cast(type); + + VeloxMapGeneratorConfig generatorConfig{ + .rowType = rowType, + .keyType = velox::TypeKind::INTEGER, + .maxSizeForMap = 10}; + VeloxMapGenerator generator(leafPool_.get(), generatorConfig); + + nimble::VeloxWriterOptions writerOptions; + writerOptions.flatMapColumns.insert("row_column"); + + nimble::VeloxReadParams params; + params.readFlatMapFieldAsStruct.insert("row_column"); + for (auto i = 0; i < 10; ++i) { + params.flatMapFeatureSelector["row_column"].features.push_back( + folly::to(i)); + } + + auto iterations = 20; + auto batches = 10; + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + generator.rng(), + *leafPool_, + rowType, + [&](auto&) { return generator.generateBatch(10); }, + compareFlatMaps, + batches, + writerOptions, + params); + } +} + +TEST_F(VeloxReaderTests, StringKeyFlatMapAsStruct) { + auto stringKeyFeatures = velox::MAP(velox::VARCHAR(), velox::REAL()); + auto type = velox::ROW({ + {"string_key_feature", stringKeyFeatures}, + }); + auto rowType = std::dynamic_pointer_cast(type); + + nimble::VeloxWriterOptions writerOptions; + writerOptions.flatMapColumns.insert("string_key_feature"); + + VeloxMapGeneratorConfig generatorConfig{ + .rowType = rowType, + .keyType = velox::TypeKind::VARCHAR, + .maxSizeForMap = 10, + .stringKeyPrefix = "testKeyString_", + }; + VeloxMapGenerator generator(leafPool_.get(), generatorConfig); + + nimble::VeloxReadParams params; + params.readFlatMapFieldAsStruct.emplace("string_key_feature"); + for (auto i = 0; i < 10; ++i) { + params.flatMapFeatureSelector["string_key_feature"].features.push_back( + "testKeyString_" + folly::to(i)); + } + + auto iterations = 10; + auto batches = 1; + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + generator.rng(), + *leafPool_, + rowType, + [&](auto&) { return generator.generateBatch(10); }, + compareFlatMaps, + batches, + writerOptions, + params); + } + + iterations = 20; + batches = 10; + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + generator.rng(), + *leafPool_, + rowType, + [&](auto&) { return generator.generateBatch(10); }, + compareFlatMaps, + batches, + writerOptions, + params); + } +} + +TEST_F(VeloxReaderTests, FlatMapAsMapEncoding) { + auto floatFeatures = velox::MAP(velox::INTEGER(), velox::REAL()); + auto idListFeatures = + velox::MAP(velox::INTEGER(), velox::ARRAY(velox::BIGINT())); + auto idScoreListFeatures = + velox::MAP(velox::INTEGER(), velox::MAP(velox::BIGINT(), velox::REAL())); + auto type = velox::ROW({ + {"float_features", floatFeatures}, + {"id_list_features", idListFeatures}, + {"id_score_list_features", idScoreListFeatures}, + }); + auto rowType = std::dynamic_pointer_cast(type); + VeloxMapGeneratorConfig generatorConfig{ + .rowType = rowType, + .keyType = velox::TypeKind::INTEGER, + }; + VeloxMapGenerator generator(leafPool_.get(), generatorConfig); + + nimble::VeloxWriterOptions writerOptions; + writerOptions.flatMapColumns.emplace("float_features"); + writerOptions.flatMapColumns.emplace("id_list_features"); + writerOptions.flatMapColumns.emplace("id_score_list_features"); + + // Verify the flatmap read without feature selection they are read as + // MapEncoding + nimble::VeloxReadParams params; + auto iterations = 10; + auto batches = 10; + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + generator.rng(), + *leafPool_, + rowType, + [&](auto&) { return generator.generateBatch(10); }, + vectorEquals, + batches, + writerOptions, + params); + } + + for (auto i = 0; i < 10; ++i) { + params.flatMapFeatureSelector["float_features"].features.push_back( + folly::to(i)); + params.flatMapFeatureSelector["id_list_features"].features.push_back( + folly::to(i)); + params.flatMapFeatureSelector["id_score_list_features"].features.push_back( + folly::to(i)); + } + + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + generator.rng(), + *leafPool_, + rowType, + [&](auto&) { return generator.generateBatch(10); }, + vectorEquals, + batches, + writerOptions, + params); + } + + { + // Selecting only odd values column from flat map + params.flatMapFeatureSelector.clear(); + for (auto i = 0; i < 10; ++i) { + if (i % 2 == 1) { + params.flatMapFeatureSelector["float_features"].features.push_back( + folly::to(i)); + params.flatMapFeatureSelector["id_list_features"].features.push_back( + folly::to(i)); + params.flatMapFeatureSelector["id_score_list_features"] + .features.push_back(folly::to(i)); + } + } + + std::unordered_set floatFeaturesLookup{ + params.flatMapFeatureSelector["float_features"].features.begin(), + params.flatMapFeatureSelector["float_features"].features.end()}; + std::unordered_set idListFeaturesLookup{ + params.flatMapFeatureSelector["id_list_features"].features.begin(), + params.flatMapFeatureSelector["id_list_features"].features.end()}; + std::unordered_set idScoreListFeaturesLookup{ + params.flatMapFeatureSelector["id_score_list_features"] + .features.begin(), + params.flatMapFeatureSelector["id_score_list_features"].features.end()}; + auto isKeyPresent = [&](std::string& key) { + return floatFeaturesLookup.find(key) != floatFeaturesLookup.end() || + idListFeaturesLookup.find(key) != idListFeaturesLookup.end() || + idScoreListFeaturesLookup.find(key) != + idScoreListFeaturesLookup.end(); + }; + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + generator.rng(), + *leafPool_, + rowType, + [&](auto&) { return generator.generateBatch(10); }, + vectorEquals, + batches, + writerOptions, + params, + isKeyPresent); + } + } + + { + // Exclude odd values column from flat map + params.flatMapFeatureSelector.clear(); + std::unordered_set floatFeaturesLookup; + std::unordered_set idListFeaturesLookup; + std::unordered_set idScoreListFeaturesLookup; + + params.flatMapFeatureSelector["float_features"].mode = + nimble::SelectionMode::Exclude; + params.flatMapFeatureSelector["id_list_features"].mode = + nimble::SelectionMode::Exclude; + params.flatMapFeatureSelector["id_score_list_features"].mode = + nimble::SelectionMode::Exclude; + for (auto i = 0; i < 10; ++i) { + std::string iStr = folly::to(i); + if (i % 2 == 1) { + params.flatMapFeatureSelector["float_features"].features.push_back( + iStr); + params.flatMapFeatureSelector["id_list_features"].features.push_back( + iStr); + params.flatMapFeatureSelector["id_score_list_features"] + .features.push_back(iStr); + } else { + floatFeaturesLookup.insert(iStr); + idListFeaturesLookup.insert(iStr); + idScoreListFeaturesLookup.insert(iStr); + } + } + + auto isKeyPresent = [&](std::string& key) { + return floatFeaturesLookup.find(key) != floatFeaturesLookup.end() || + idListFeaturesLookup.find(key) != idListFeaturesLookup.end() || + idScoreListFeaturesLookup.find(key) != + idScoreListFeaturesLookup.end(); + }; + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + generator.rng(), + *leafPool_, + rowType, + [&](auto&) { return generator.generateBatch(10); }, + vectorEquals, + batches, + writerOptions, + params, + isKeyPresent); + } + } +} + +TEST_F(VeloxReaderTests, StringKeyFlatMapAsMapEncoding) { + auto stringKeyFeatures = velox::MAP(velox::VARCHAR(), velox::REAL()); + auto type = velox::ROW({ + {"string_key_feature", stringKeyFeatures}, + }); + auto rowType = std::dynamic_pointer_cast(type); + + VeloxMapGeneratorConfig generatorConfig{ + .rowType = rowType, + .keyType = velox::TypeKind::VARCHAR, + .stringKeyPrefix = "testKeyString_", + }; + VeloxMapGenerator generator(leafPool_.get(), generatorConfig); + + nimble::VeloxWriterOptions writerOptions; + writerOptions.flatMapColumns.insert("string_key_feature"); + + nimble::VeloxReadParams params; + // Selecting only keys with even index to it + for (auto i = 0; i < 10; ++i) { + if (i % 2 == 0) { + params.flatMapFeatureSelector["string_key_feature"].features.push_back( + "testKeyString_" + folly::to(i)); + } + } + + std::unordered_set stringKeyFeature{ + params.flatMapFeatureSelector["string_key_feature"].features.begin(), + params.flatMapFeatureSelector["string_key_feature"].features.end()}; + + auto isKeyPresent = [&](std::string& key) { + return stringKeyFeature.find(key) != stringKeyFeature.end(); + }; + + auto iterations = 10; + + // Keeping the batchCount 1 produces the case where flatmap readers nulls + // column is empty, as decodedmap produce the mayHavenulls as false + auto batches = 1; + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + generator.rng(), + *leafPool_, + rowType, + [&](auto&) { return generator.generateBatch(10); }, + nullptr, /* for key present use a fix function */ + batches, + writerOptions, + params, + isKeyPresent); + } + + iterations = 20; + batches = 10; + for (auto i = 0; i < iterations; ++i) { + writeAndVerify( + generator.rng(), + *leafPool_, + rowType, + [&](auto&) { return generator.generateBatch(10); }, + nullptr, /* for key present use a fix function */ + batches, + writerOptions, + params, + isKeyPresent); + } +} + +class TestNimbleReaderFactory { + public: + TestNimbleReaderFactory( + velox::memory::MemoryPool& leafPool, + velox::memory::MemoryPool& rootPool, + std::vector vectors, + const nimble::VeloxWriterOptions& writerOptions = {}) + : memoryPool_(leafPool) { + file_ = std::make_unique( + nimble::test::createNimbleFile(rootPool, vectors, writerOptions)); + type_ = std::dynamic_pointer_cast(vectors[0]->type()); + } + + nimble::VeloxReader createReader(nimble::VeloxReadParams params = {}) { + auto selector = + std::make_shared(type_); + return nimble::VeloxReader( + this->memoryPool_, file_.get(), std::move(selector), params); + } + + nimble::Tablet createTablet() { + return nimble::Tablet(memoryPool_, file_.get()); + } + + private: + std::unique_ptr file_; + std::shared_ptr type_; + velox::memory::MemoryPool& memoryPool_; +}; + +std::vector createSkipSeekVectors( + velox::memory::MemoryPool& pool, + const std::vector& rowsPerStripe) { + uint32_t seed = FLAGS_reader_tests_seed > 0 ? FLAGS_reader_tests_seed + : folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng{seed}; + + std::vector vectors(rowsPerStripe.size()); + velox::test::VectorMaker vectorMaker{&pool}; + + for (auto i = 0; i < rowsPerStripe.size(); ++i) { + std::string s; + vectors[i] = vectorMaker.rowVector( + {"a", "b", "dictionaryArray"}, + {vectorMaker.flatVector( + rowsPerStripe[i], + /* valueAt */ + [&rng](auto /* row */) { return folly::Random::rand32(rng); }, + /* isNullAt */ [](auto row) { return row % 2 == 1; }), + vectorMaker.flatVector( + rowsPerStripe[i], + /* valueAt */ + [&s, &rng](auto /* row */) { + s = "arf_" + folly::to(folly::Random::rand32(rng)); + return velox::StringView(s.data(), s.size()); + }, + /* isNullAt */ [](auto row) { return row % 2 == 1; }), + vectorMaker.arrayVector( + rowsPerStripe[i], + /* sizeAt */ + [](auto /* row */) { return 1; }, + /* valueAt */ + [](auto row) { + /* duplicated values to check cache usage */ + return row / 4; + })}); + } + + return vectors; +} + +void readAndVerifyContent( + nimble::VeloxReader& reader, + std::vector expectedVectors, + uint32_t rowsToRead, + uint32_t expectedNumberOfRows, + uint32_t expectedStripe, + velox::vector_size_t expectedRowInStripe) { + velox::VectorPtr result; + EXPECT_TRUE(reader.next(rowsToRead, result)); + ASSERT_EQ(result->type()->kind(), velox::TypeKind::ROW); + velox::RowVector* rowVec = result->as(); + ASSERT_EQ(rowVec->childAt(0)->type()->kind(), velox::TypeKind::INTEGER); + ASSERT_EQ(rowVec->childAt(1)->type()->kind(), velox::TypeKind::VARCHAR); + ASSERT_EQ(rowVec->childAt(2)->type()->kind(), velox::TypeKind::ARRAY); + const int curRows = result->size(); + ASSERT_EQ(curRows, expectedNumberOfRows); + ASSERT_LT(expectedStripe, expectedVectors.size()); + auto& expected = expectedVectors[expectedStripe]; + + for (velox::vector_size_t i = 0; i < curRows; ++i) { + if (!expected->equalValueAt(result.get(), i + expectedRowInStripe, i)) { + ASSERT_TRUE( + expected->equalValueAt(result.get(), i + expectedRowInStripe, i)) + << "Content mismatch at index " << i + << "\nReference: " << expected->toString(i + expectedRowInStripe) + << "\nResult: " << result->toString(i); + } + } +} + +TEST_F(VeloxReaderTests, ReaderSeekTest) { + // Generate an Nimble file with 3 stripes and 10 rows each + auto vectors = createSkipSeekVectors(*leafPool_, {10, 10, 10}); + nimble::VeloxWriterOptions writerOptions; + writerOptions.dictionaryArrayColumns.insert("dictionaryArray"); + + TestNimbleReaderFactory readerFactory( + *leafPool_, *rootPool_, vectors, writerOptions); + auto reader = readerFactory.createReader(); + + auto rowResult = reader.skipRows(0); + EXPECT_EQ(0, rowResult); + rowResult = reader.seekToRow(0); + EXPECT_EQ(0, rowResult); + + // [Stripe# 0, Current Pos: 0] seek to 1 position + rowResult = reader.seekToRow(1); + EXPECT_EQ(rowResult, 1); + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 1, + /* expectedNumberOfRows */ 1, + /* expectedStripe */ 0, + /* expectedRowInStripe */ 1); + + // [Stripe# 0, Current Pos: 2] seek to 5 position + rowResult = reader.seekToRow(5); + // [Stripe# 0, Current Pos: 5] seeks start from rowIdx 0 + EXPECT_EQ(rowResult, 5); + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 6, + /* expectedNumberOfRows */ 5, + /* expectedStripe */ 0, + /* expectedRowInStripe */ 5); + + // [Stripe# 0, Current Pos: 10] seek to 10 position + rowResult = reader.seekToRow(10); + // [Stripe# 1, Current Pos: 0] + EXPECT_EQ(rowResult, 10); + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 10, + /* expectedNumberOfRows */ 10, + /* expectedStripe */ 1, + /* expectedRowInStripe */ 0); + + // [Stripe# 2, Current Pos: 0] + rowResult = reader.seekToRow(29); + // [Stripe# 2, Current Pos: 9] + EXPECT_EQ(rowResult, 29); + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 2, + /* expectedNumberOfRows */ 1, + /* expectedStripe */ 2, + /* expectedRowInStripe */ 9); + + // seek past + { + rowResult = reader.seekToRow(32); + // Seeks with rows >= totalRows in Nimble file, seeks to lastRow + EXPECT_EQ(rowResult, 30); + velox::VectorPtr result; + EXPECT_FALSE(reader.next(1, result)); + } +} + +TEST_F(VeloxReaderTests, ReaderSkipTest) { + // Generate an Nimble file with 3 stripes and 10 rows each + auto vectors = createSkipSeekVectors(*leafPool_, {10, 10, 10}); + nimble::VeloxWriterOptions writerOptions; + writerOptions.dictionaryArrayColumns.insert("dictionaryArray"); + + TestNimbleReaderFactory readerFactory( + *leafPool_, *rootPool_, vectors, writerOptions); + auto reader = readerFactory.createReader(); + + // Current position in Comments below represent the position in stripe + // [Stripe# 0, Current Pos: 0], After skip [Stripe# 0, Current Pos: 1] + auto rowResult = reader.skipRows(1); + EXPECT_EQ(rowResult, 1); + // readAndVerifyContent() moves the rowPosition in reader + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 1, + /* expectedNumberOfRows */ 1, + /* expectedStripe */ 0, + /* expectedRowInStripe */ 1); + + // [Stripe# 0, Current Pos: 2], After skip [Stripe# 0, Current Pos: 7] + rowResult = reader.skipRows(5); + EXPECT_EQ(rowResult, 5); + // reader don't read across stripe so expectedRow is 3 + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 4, + /* expectedNumberOfRows */ 3, + /* expectedStripe */ 0, + /* expectedRowInStripe */ 7); + + // [Stripe# 1, Current Pos: 0], After skip [Stripe# 2, Current Pos: 0] + rowResult = reader.skipRows(10); + EXPECT_EQ(rowResult, 10); + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 1, + /* expectedNumberOfRows */ 1, + /* expectedStripe */ 2, + /* expectedRowInStripe */ 0); + + // [Stripe# 2, Current Pos: 1], After skip [Stripe# 2, Current Pos: 9] + rowResult = reader.skipRows(8); + EXPECT_EQ(rowResult, 8); + // reader don't read across stripe so expectedRow is 3 + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 2, + /* expectedNumberOfRows */ 1, + /* expectedStripe */ 2, + /* expectedRowInStripe */ 9); + + { + // [Stripe# 3, Current Pos: 0], Reached EOF + rowResult = reader.skipRows(5); + EXPECT_EQ(rowResult, 0); + velox::VectorPtr result; + EXPECT_FALSE(reader.next(1, result)); + } + + // Try to seek to start and test skip + rowResult = reader.seekToRow(0); + EXPECT_EQ(0, rowResult); + // [Stripe# 0, Current Pos: 0], After skip [Stripe# 1, Current Pos: 2] + rowResult = reader.skipRows(12); + EXPECT_EQ(rowResult, 12); + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 10, + /* expectedNumberOfRows */ 8, + /* expectedStripe */ 1, + /* expectedRowInStripe */ 2); + + // Test continuous skip calls and then readandVerify + reader.seekToRow(0); + // [Stripe# 0, Current Pos: 0], After skip [Stripe# 1, Current Pos: 0] + for (int i = 0; i < 10; ++i) { + rowResult = reader.skipRows(1); + EXPECT_EQ(rowResult, 1); + } + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 1, + /* expectedNumberOfRows */ 1, + /* expectedStripe */ 1, + /* expectedRowInStripe */ 0); + + // Continuous skip calls across stripe + // [Stripe# 1, Current Pos: 1], After skip [Stripe# 2, Current Pos: 9] + for (int i = 0; i < 6; ++i) { + rowResult = reader.skipRows(3); + } + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 2, + /* expectedNumberOfRows */ 1, + /* expectedStripe */ 2, + /* expectedRowInStripe */ 9); + + { + // Current position: EOF + velox::VectorPtr result; + EXPECT_FALSE(reader.next(1, result)); + } + + // Read the data(This also move the reader state) follow by skips and verify + reader.seekToRow(0); + for (int i = 0; i < 11; ++i) { + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 1, + /* expectedNumberOfRows */ 1, + /* expectedStripe */ (i / 10), + /* expectedRowInStripe */ (i % 10)); + } + // [Stripe# 1, Current Pos: 1], After skip [Stripe# 1, Current Pos: 6] + rowResult = reader.skipRows(5); + EXPECT_EQ(rowResult, 5); + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 5, + /* expectedNumberOfRows */ 4, + /* expectedStripe */ 1, + /* expectedRowInStripe */ 6); + + { + // verify the skip to more rows then file have + reader.seekToRow(0); + // [Stripe# 0, Current Pos: 0], After skip EOF + rowResult = reader.skipRows(32); + EXPECT_EQ(30, rowResult); + velox::VectorPtr result; + EXPECT_FALSE(reader.next(1, result)); + + reader.seekToRow(0); + // [Stripe# 0, Current Pos: 0], After skip [Stripe# 2, Current Pos: 2] + rowResult = reader.skipRows(22); + EXPECT_EQ(22, rowResult); + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 9, + /* expectedNumberOfRows */ 8, + /* expectedStripe */ 2, + /* expectedRowInStripe */ 2); + EXPECT_FALSE(reader.next(1, result)); + } +} + +TEST_F(VeloxReaderTests, ReaderSkipSingleStripeTest) { + // Generate an Nimble file with 1 stripe and 12 rows + auto vectors = createSkipSeekVectors(*leafPool_, {12}); + nimble::VeloxWriterOptions writerOptions; + writerOptions.dictionaryArrayColumns.insert("dictionaryArray"); + + TestNimbleReaderFactory readerFactory( + *leafPool_, *rootPool_, vectors, writerOptions); + auto reader = readerFactory.createReader(); + + // Current position in Comments below represent the position in stripe + // [Stripe# 0, Current Pos: 0], After skip [Stripe# 0, Current Pos: 1] + auto rowResult = reader.skipRows(1); + EXPECT_EQ(rowResult, 1); + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 12, + /* expectedNumberOfRows */ 11, + /* expectedStripe */ 0, + /* expectedRowInStripe */ 1); + + // Current pos : EOF, try to read skip past it + { + rowResult = reader.skipRows(13); + EXPECT_EQ(rowResult, 0); + rowResult = reader.skipRows(1); + EXPECT_EQ(rowResult, 0); + velox::VectorPtr result; + EXPECT_FALSE(reader.next(1, result)); + } + + // Seek to position 2 and then skip 11 rows to reach EOF + rowResult = reader.seekToRow(2); + EXPECT_EQ(rowResult, 2); + rowResult = reader.skipRows(11); + EXPECT_EQ(rowResult, 10); + { + velox::VectorPtr result; + EXPECT_FALSE(reader.next(1, result)); + } + + // Seek to 0 and skip 13 rows + rowResult = reader.seekToRow(0); + EXPECT_EQ(rowResult, 0); + rowResult = reader.skipRows(13); + EXPECT_EQ(rowResult, 12); + { + velox::VectorPtr result; + EXPECT_FALSE(reader.next(1, result)); + } +} + +TEST_F(VeloxReaderTests, ReaderSeekSingleStripeTest) { + // Generate an Nimble file with 1 stripes and 11 rows + auto vectors = createSkipSeekVectors(*leafPool_, {11}); + nimble::VeloxWriterOptions writerOptions; + writerOptions.dictionaryArrayColumns.insert("dictionaryArray"); + + TestNimbleReaderFactory readerFactory( + *leafPool_, *rootPool_, vectors, writerOptions); + auto reader = readerFactory.createReader(); + + // Current position in Comments below represent the position in stripe + // [Stripe# 0, Current Pos: 0], After skip [Stripe# 0, Current Pos: 5] + auto rowResult = reader.seekToRow(5); + EXPECT_EQ(rowResult, 5); + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 12, + /* expectedNumberOfRows */ 6, + /* expectedStripe */ 0, + /* expectedRowInStripe */ 5); + + // Current pos : EOF, try to read skip past it + { + rowResult = reader.seekToRow(15); + EXPECT_EQ(rowResult, 11); + velox::VectorPtr result; + EXPECT_FALSE(reader.next(1, result)); + rowResult = reader.seekToRow(10000); + EXPECT_EQ(rowResult, 11); + EXPECT_FALSE(reader.next(1, result)); + } +} + +TEST_F(VeloxReaderTests, ReaderSkipUnevenStripesTest) { + // Generate an Nimble file with 4 stripes + auto vectors = createSkipSeekVectors(*leafPool_, {12, 15, 25, 18}); + nimble::VeloxWriterOptions writerOptions; + writerOptions.dictionaryArrayColumns.insert("dictionaryArray"); + + TestNimbleReaderFactory readerFactory( + *leafPool_, *rootPool_, vectors, writerOptions); + auto reader = readerFactory.createReader(); + + // Current position in Comments below represent the position in stripe + // [Stripe# 0, Current Pos: 0], After skip [Stripe# 2, Current Pos: 8] + auto rowResult = reader.skipRows(35); + EXPECT_EQ(rowResult, 35); + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ 12, + /* expectedNumberOfRows */ 12, + /* expectedStripe */ 2, + /* expectedRowInStripe */ 8); + + // [Stripe# 2, Current Pos: 20], After skip EOF + { + rowResult = reader.skipRows(25); + EXPECT_EQ(rowResult, 23); + velox::VectorPtr result; + EXPECT_FALSE(reader.next(1, result)); + } +} + +// this test is created to keep an eye on the default value for T() for +// primitive type. Recently it came to our notice that T() does zero +// initialize the value for optimized builds. T() we have used a bit in the +// code to zero out the result. This is a dummy test to fail fast if it is not +// zero initialized for primitive types +TEST_F(VeloxReaderTests, TestPrimitiveFieldDefaultValue) { + verifyDefaultValue(2, 0, 10); + verifyDefaultValue(2, 0, 30); + verifyDefaultValue(2, 0, 30); + verifyDefaultValue(2, 0, 30); + verifyDefaultValue(2, 0, 30); + verifyDefaultValue(2, 0, 30); + verifyDefaultValue(2, 0, 30); + verifyDefaultValue(2, 0.0, 30); + verifyDefaultValue(2.1, 0.0, 30); + verifyDefaultValue(true, false, 30); + verifyDefaultValue(3.2, 0.0, 30); + verifyDefaultValue("test", "", 30); +} + +struct RangeTestParams { + uint64_t rangeStart; + uint64_t rangeEnd; + + // Tuple arguments: rowsToRead, expectedNumberOfRows, expectedStripe, + // expectedRowInStripe + std::vector> expectedReads; + + // Tuple arguments: seekToRow, expectedSeekResult + std::vector> expectedSeeks; + + // Tuple arguments: skipRows, expectedSkipResult + std::vector> expectedSkips; +}; + +TEST_F(VeloxReaderTests, RangeReads) { + // Generate an Nimble file with 4 stripes + auto vectors = createSkipSeekVectors(*leafPool_, {10, 15, 25, 9}); + nimble::VeloxWriterOptions writerOptions; + writerOptions.dictionaryArrayColumns.insert("dictionaryArray"); + + TestNimbleReaderFactory readerFactory( + *leafPool_, *rootPool_, vectors, writerOptions); + + auto test = [&readerFactory, &vectors](RangeTestParams params) { + auto reader = readerFactory.createReader(nimble::VeloxReadParams{ + .fileRangeStartOffset = params.rangeStart, + .fileRangeEndOffset = params.rangeEnd}); + + for (const auto& expectedRead : params.expectedReads) { + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ std::get<0>(expectedRead), + /* expectedNumberOfRows */ std::get<1>(expectedRead), + /* expectedStripe */ std::get<2>(expectedRead), + /* expectedRowInStripe */ std::get<3>(expectedRead)); + } + + { + velox::VectorPtr result; + EXPECT_FALSE(reader.next(1, result)); + } + + for (const auto& expectedSeek : params.expectedSeeks) { + EXPECT_EQ( + std::get<1>(expectedSeek), + reader.seekToRow(std::get<0>(expectedSeek))); + } + + reader.seekToRow(0); + for (const auto& expectedSkip : params.expectedSkips) { + EXPECT_EQ( + std::get<1>(expectedSkip), + reader.skipRows(std::get<0>(expectedSkip))); + } + + reader.seekToRow(0); + for (const auto& expectedRead : params.expectedReads) { + readAndVerifyContent( + reader, + vectors, + /* rowsToRead */ std::get<0>(expectedRead), + /* expectedNumberOfRows */ std::get<1>(expectedRead), + /* expectedStripe */ std::get<2>(expectedRead), + /* expectedRowInStripe */ std::get<3>(expectedRead)); + } + + { + velox::VectorPtr result; + EXPECT_FALSE(reader.next(1, result)); + } + }; + + // Try to read all data in the file. Since we cover the entire file (end is + // bigger than the file size), we expect to be able to read all lines. + LOG(INFO) << "--> Range covers the entire file"; + LOG(INFO) << "File: |--s0--|--s1--|--s2--|--s3--|"; + LOG(INFO) << "Range: |---------------------------------|"; + LOG(INFO) << "Expected: |--s0--|--s1--|--s2--|--s3--|"; + test({ + .rangeStart = 0, + .rangeEnd = 100'000'000, + + // Reads stop at stripe boundaries, so we need to invoke several reads + // to read the entire file. + .expectedReads = + {{30, 10, 0, 0}, {30, 15, 1, 0}, {30, 25, 2, 0}, {30, 9, 3, 0}}, + + // Seeks should be allowed to anywhere in this file (rows 0 to 59) + .expectedSeeks = + {{0, 0}, + {5, 5}, + {10, 10}, + {15, 15}, + {25, 25}, + {30, 30}, + {45, 45}, + {50, 50}, + {55, 55}, + {59, 59}, + {60, 59}}, + + // Skips should cover the entire file (59 rows) + .expectedSkips = {{0, 0}, {10, 10}, {20, 20}, {30, 29}, {1, 0}}, + }); + + // Test a range covering only the first stripe. + // Using range starting at 0 guarantees we cover the first stripe. + // Since first stripe is much greater than 1 byte, using range ending at 1, + // guarantees we don't cover any other stripe other than the fisrt stripe. + LOG(INFO) << "--> Range covers beginning of first stripe"; + LOG(INFO) << "File: |--s0--|--s1--|--s2--|--s3--|"; + LOG(INFO) << "Range: |---|"; + LOG(INFO) << "Expected: |--s0--|"; + test({ + .rangeStart = 0, + .rangeEnd = 1, + + // Reads should only find rows in stripe 0. + .expectedReads = {{5, 5, 0, 0}, {10, 5, 0, 5}}, + + // Seeks should be allowed to access rows in first stripe only (rows 0 + // to 10) + .expectedSeeks = + {{0, 0}, {5, 5}, {10, 10}, {15, 10}, {30, 10}, {59, 10}, {60, 10}}, + + // Skips should cover first stripe only (59 rows) + .expectedSkips = {{0, 0}, {5, 5}, {10, 5}, {1, 0}}, + }); + + auto tablet = readerFactory.createTablet(); + + // Test a range starting somewhere in the first stripe (but not at zero + // offset) to exactly the end of the first stripe. This should be resolved + // to zero stripes. + LOG(INFO) << "--> Range covers end of stripe 0"; + LOG(INFO) << "File: |--s0--|--s1--|--s2--|--s3--|"; + LOG(INFO) << "Range: |---|"; + LOG(INFO) << "Expected: >"; + test({ + .rangeStart = 1, + .rangeEnd = tablet.stripeOffset(1), + + // No read should succeed, as we have zero stripes to read from + .expectedReads = {}, + + // All seeks should be ignored + .expectedSeeks = + {{0, 0}, {5, 0}, {10, 0}, {15, 0}, {30, 0}, {59, 0}, {60, 0}}, + + // All skips should be ignored + .expectedSkips = {{0, 0}, {5, 0}, {59, 0}}, + }); + + // Test a range starting somewhere in stripe 0 (but not at zero) to somwhere + // in stripe 1. This should resolve to only stripe 1. + LOG(INFO) << "--> Range covers beginning of stripe 1"; + LOG(INFO) << "File: |--s0--|--s1--|--s2--|--s3--|"; + LOG(INFO) << "Range: |------|"; + LOG(INFO) << "Expected: |--s1--|"; + + test({ + .rangeStart = 1, + .rangeEnd = tablet.stripeOffset(1) + 1, + + // Reads should all resolve to stripe 1 + .expectedReads = {{5, 5, 1, 0}, {20, 10, 1, 5}}, + + // Seeks should succeed if they are in range [10, 25). Otherwise, they + // should return the edges of stripe 1. + .expectedSeeks = + {{0, 10}, + {5, 10}, + {10, 10}, + {15, 15}, + {25, 25}, + {26, 25}, + {59, 25}, + {60, 25}}, + + // Skips should allow skipping only 15 rows (number of rows in stripe 1) + .expectedSkips = {{0, 0}, {5, 5}, {11, 10}, {1, 0}}, + }); + + // Test a range starting exactly on stripe 2 to somwhere + // in stripe 2. This should resolve to only stripe 2. + LOG(INFO) << "--> Range starts at beginning of stripe 2"; + LOG(INFO) << "File: |--s0--|--s1--|--s2--|--s3--|"; + LOG(INFO) << "Range: |---|"; + LOG(INFO) << "Expected: |--s1--|"; + test({ + .rangeStart = tablet.stripeOffset(1), + .rangeEnd = tablet.stripeOffset(1) + 1, + + // Reads should all resolve to stripe 1 + .expectedReads = {{5, 5, 1, 0}, {20, 10, 1, 5}}, + + // Seeks should succeed if they are in range [10, 25). Otherwise, they + // should return the edges of stripe 1. + .expectedSeeks = + {{0, 10}, + {5, 10}, + {10, 10}, + {15, 15}, + {25, 25}, + {26, 25}, + {59, 25}, + {60, 25}}, + + // Skips should allow skipping only 15 rows (number of rows in stripe 1) + .expectedSkips = {{0, 0}, {5, 5}, {11, 10}, {1, 0}}, + }); + + // Test a range spanning multiple stripes. We'll start somewhere in stripe 0 + // and end somewhere in stripe 2. This should resolve to stripes 1 and 2. + LOG(INFO) << "--> Range spans stripes (0, 1 ,2)"; + LOG(INFO) << "File: |--s0--|--s1--|--s2--|--s3--|"; + LOG(INFO) << "Range: |------------|"; + LOG(INFO) << "Expected: |--s1--|--s2--|"; + test({ + .rangeStart = tablet.stripeOffset(1) - 1, + .rangeEnd = tablet.stripeOffset(2) + 1, + + // Reads should all resolve to stripes 1 and 2 (rows [15 to 50)). + // Reads stop at stripe boundaries, so we need to invoke several reads + // to continue to the next stripe. + .expectedReads = + {{5, 5, 1, 0}, {20, 10, 1, 5}, {20, 20, 2, 0}, {20, 5, 2, 20}}, + + // Seeks should succeed if they are in range [10, 50). Otherwise, they + // should return the edges of stripe 1 and 2. + .expectedSeeks = + {{0, 10}, + {5, 10}, + {10, 10}, + {15, 15}, + {25, 25}, + {26, 26}, + {49, 49}, + {50, 50}, + {59, 50}, + {60, 50}}, + + // Skips should allow skipping only 40 rows (number of rows in stripes 1 + // and 2) + .expectedSkips = {{0, 0}, {5, 5}, {11, 11}, {23, 23}, {2, 1}, {1, 0}}, + }); + + // Test a range spanning multiple stripes. We'll start at the beginning of + // stripe 1 and end somewhere in stripe 3. This should resolve to stripes 1, + // 2 and 3. + LOG(INFO) << "--> Range spans stripes (1 ,2, 3)"; + LOG(INFO) << "File: |--s0--|--s1--|--s2--|--s3--|"; + LOG(INFO) << "Range: |----------------|"; + LOG(INFO) << "Expected: |--s1--|--s2--|--s3--|"; + test({ + .rangeStart = tablet.stripeOffset(1), + .rangeEnd = tablet.stripeOffset(3) + 1, + + // Reads should all resolve to stripes 1, 2 and 3 (rows [15 to 59)). + // Reads stop at stripe boundaries, so we need to invoke several reads + // to continue to the next stripe. + .expectedReads = + {{5, 5, 1, 0}, + {20, 10, 1, 5}, + {20, 20, 2, 0}, + {20, 5, 2, 20}, + {20, 9, 3, 0}}, + + // Seeks should succeed if they are in range [10, 59). Otherwise, they + // should return the edges of stripe 1 and 3. + .expectedSeeks = + {{0, 10}, + {5, 10}, + {10, 10}, + {15, 15}, + {25, 25}, + {26, 26}, + {49, 49}, + {50, 50}, + {59, 59}, + {60, 59}}, + + // Skips should allow skipping only 49 rows (number of rows in stripes 1 + // to 3) + .expectedSkips = {{0, 0}, {5, 5}, {11, 11}, {32, 32}, {2, 1}, {1, 0}}, + }); + + // Test last stripe. + LOG(INFO) << "--> Range covers stripe 3"; + LOG(INFO) << "File: |--s0--|--s1--|--s2--|--s3--|"; + LOG(INFO) << "Range: |----------|"; + LOG(INFO) << "Expected: |--s3--|"; + test({ + .rangeStart = tablet.stripeOffset(3), + .rangeEnd = 100'000'000, + + // Reads should all resolve to stripe 3 (rows 50 to 59). + .expectedReads = {{5, 5, 3, 0}, {5, 4, 3, 5}}, + + // Seeks should succeed if they are in range [50, 59). Otherwise, they + // should return the edges of stripe 3. + .expectedSeeks = + {{0, 50}, + {10, 50}, + {15, 50}, + {26, 50}, + {49, 50}, + {50, 50}, + {59, 59}, + {60, 59}}, + + // Skips should allow skipping only 9 rows (number of rows in stripe 3) + .expectedSkips = {{0, 0}, {5, 5}, {5, 4}, {1, 0}}, + }); +} + +TEST_F(VeloxReaderTests, TestScalarFieldLifeCycle) { + auto testScalarFieldLifeCycle = + [&](const std::shared_ptr schema, + int32_t batchSize, + std::mt19937& rng) { + velox::VectorPtr result; + auto reader = getReaderForLifeCycleTest(schema, 4 * batchSize, rng); + EXPECT_TRUE(reader->next(batchSize, result)); + // Hold the reference to values Buffers + auto child = result->as()->childAt(0); + velox::BaseVector* rowPtr = result.get(); + velox::Buffer* rawNulls = child->nulls().get(); + velox::BufferPtr values = child->values(); + // Reset the child so that it can be reused + child.reset(); + EXPECT_TRUE(reader->next(batchSize, result)); + child = result->as()->childAt(0); + EXPECT_EQ(rawNulls, child->nulls().get()); + EXPECT_NE(values.get(), child->values().get()); + EXPECT_EQ(rowPtr, result.get()); + + // Hold the reference to NULL buffer + velox::BufferPtr nulls = child->nulls(); + velox::Buffer* rawValues = child->values().get(); + child.reset(); + EXPECT_TRUE(reader->next(batchSize, result)); + child = result->as()->childAt(0); + EXPECT_NE(nulls.get(), child->nulls().get()); + EXPECT_EQ(rawValues, child->values().get()); + EXPECT_EQ(rowPtr, result.get()); + + rawNulls = nulls.get(); + // Hold reference to child ScalarVector and it should use another + // ScalarVector along with childBuffers + EXPECT_TRUE(reader->next(batchSize, result)); + auto child1 = result->as()->childAt(0); + EXPECT_NE(child, child1); + EXPECT_EQ(rowPtr, result.get()); + // after VectorPtr is reset its buffer also reset + EXPECT_NE(rawNulls, child1->nulls().get()); + EXPECT_NE(rawValues, child1->values().get()); + }; + + uint32_t seed = FLAGS_reader_tests_seed > 0 ? FLAGS_reader_tests_seed + : folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng{seed}; + std::vector> types = { + velox::ROW({{"tinyInt", velox::TINYINT()}}), + velox::ROW({{"smallInt", velox::SMALLINT()}}), + velox::ROW({{"int", velox::INTEGER()}}), + velox::ROW({{"bigInt", velox::BIGINT()}}), + velox::ROW({{"Real", velox::REAL()}}), + velox::ROW({{"Double", velox::DOUBLE()}}), + velox::ROW({{"VARCHAR", velox::VARCHAR()}}), + }; + for (auto& type : types) { + LOG(INFO) << "Field Type: " << type->nameOf(0); + for (int i = 0; i < 10; ++i) { + testScalarFieldLifeCycle(type, 10, rng); + } + } +} + +TEST_F(VeloxReaderTests, TestArrayFieldLifeCycle) { + uint32_t seed = FLAGS_reader_tests_seed > 0 ? FLAGS_reader_tests_seed + : folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng{seed}; + auto type = velox::ROW({{"arr_val", velox::ARRAY(velox::BIGINT())}}); + auto testArrayFieldLifeCycle = + [&](const std::shared_ptr type, + int32_t batchSize, + std::mt19937& rng) { + velox::VectorPtr result; + auto reader = getReaderForLifeCycleTest(type, 4 * batchSize, rng); + EXPECT_TRUE(reader->next(batchSize, result)); + // Hold the reference to internal Buffers and element doesn't change + auto child = std::dynamic_pointer_cast( + result->as()->childAt(0)); + velox::BaseVector* childPtr = child.get(); + velox::BaseVector* rowPtr = result.get(); + velox::Buffer* rawNulls = child->nulls().get(); + velox::Buffer* rawSizes = child->sizes().get(); + velox::BufferPtr offsets = child->offsets(); + auto elementsPtr = child->elements().get(); + child.reset(); + EXPECT_TRUE(reader->next(batchSize, result)); + child = std::dynamic_pointer_cast( + result->as()->childAt(0)); + + EXPECT_EQ(rawNulls, child->nulls().get()); + EXPECT_EQ(rawSizes, child->sizes().get()); + EXPECT_NE(offsets, child->offsets()); + EXPECT_EQ(elementsPtr, child->elements().get()); + EXPECT_EQ(rowPtr, result.get()); + + // Hold the reference to Elements vector, other buffer should be + // reused + auto elements = child->elements(); + velox::Buffer* rawOffsets = child->offsets().get(); + childPtr = child.get(); + child.reset(); + EXPECT_TRUE(reader->next(batchSize, result)); + child = std::dynamic_pointer_cast( + result->as()->childAt(0)); + EXPECT_EQ(rawNulls, child->nulls().get()); + EXPECT_EQ(rawSizes, child->sizes().get()); + EXPECT_EQ(rawOffsets, child->offsets().get()); + EXPECT_NE(elements, child->elements()); + EXPECT_EQ(childPtr, child.get()); + EXPECT_EQ(rowPtr, result.get()); + + // Don't release the Child Array vector to row vector, all the buffers + // in array should not be resused. + elementsPtr = child->elements().get(); + EXPECT_TRUE(reader->next(batchSize, result)); + auto child1 = std::dynamic_pointer_cast( + result->as()->childAt(0)); + EXPECT_NE(rawNulls, child1->nulls().get()); + EXPECT_NE(rawSizes, child1->sizes().get()); + EXPECT_NE(rawOffsets, child1->offsets().get()); + EXPECT_NE(elementsPtr, child1->elements().get()); + EXPECT_NE(childPtr, child1.get()); + EXPECT_EQ(rowPtr, result.get()); + }; + for (int i = 0; i < 10; ++i) { + testArrayFieldLifeCycle(type, 10, rng); + } +} + +TEST_F(VeloxReaderTests, TestMapFieldLifeCycle) { + auto testMapFieldLifeCycle = + [&](const std::shared_ptr type, + int32_t batchSize, + std::mt19937& rng) { + velox::VectorPtr result; + auto reader = getReaderForLifeCycleTest(type, 5 * batchSize, rng); + EXPECT_TRUE(reader->next(batchSize, result)); + // Hold the reference to internal Buffers and element doesn't change + auto child = std::dynamic_pointer_cast( + result->as()->childAt(0)); + velox::BaseVector* childPtr = child.get(); + velox::BaseVector* rowPtr = result.get(); + velox::Buffer* rawNulls = child->nulls().get(); + velox::BufferPtr sizes = child->sizes(); + velox::Buffer* rawOffsets = child->offsets().get(); + auto keysPtr = child->mapKeys().get(); + auto valuesPtr = child->mapValues().get(); + child.reset(); + EXPECT_TRUE(reader->next(batchSize, result)); + child = std::dynamic_pointer_cast( + result->as()->childAt(0)); + + EXPECT_EQ(rawNulls, child->nulls().get()); + EXPECT_NE(sizes, child->sizes()); + EXPECT_EQ(rawOffsets, child->offsets().get()); + EXPECT_EQ(keysPtr, child->mapKeys().get()); + EXPECT_EQ(valuesPtr, child->mapValues().get()); + EXPECT_EQ(rowPtr, result.get()); + + // Hold the reference to keys vector, other buffer should be reused + auto mapKeys = child->mapKeys(); + velox::Buffer* rawSizes = child->sizes().get(); + childPtr = child.get(); + child.reset(); + EXPECT_TRUE(reader->next(batchSize, result)); + child = std::dynamic_pointer_cast( + result->as()->childAt(0)); + EXPECT_EQ(rawNulls, child->nulls().get()); + EXPECT_EQ(rawSizes, child->sizes().get()); + EXPECT_EQ(rawOffsets, child->offsets().get()); + EXPECT_NE(mapKeys, child->mapKeys()); + EXPECT_EQ(valuesPtr, child->mapValues().get()); + EXPECT_EQ(childPtr, child.get()); + EXPECT_EQ(rowPtr, result.get()); + + // Hold the reference to values vector, other buffer should be reused + keysPtr = child->mapKeys().get(); + auto mapValues = child->mapValues(); + childPtr = child.get(); + child.reset(); + EXPECT_TRUE(reader->next(batchSize, result)); + child = std::dynamic_pointer_cast( + result->as()->childAt(0)); + EXPECT_EQ(rawNulls, child->nulls().get()); + EXPECT_EQ(rawSizes, child->sizes().get()); + EXPECT_EQ(rawOffsets, child->offsets().get()); + EXPECT_EQ(keysPtr, child->mapKeys().get()); + EXPECT_NE(mapValues, child->mapValues()); + EXPECT_EQ(childPtr, child.get()); + EXPECT_EQ(rowPtr, result.get()); + + // Don't release the Child map vector to row vector, all the buffers + // in array should not be resused. + valuesPtr = child->mapValues().get(); + EXPECT_TRUE(reader->next(batchSize, result)); + auto child1 = std::dynamic_pointer_cast( + result->as()->childAt(0)); + EXPECT_NE(rawNulls, child1->nulls().get()); + EXPECT_NE(rawSizes, child1->sizes().get()); + EXPECT_NE(rawOffsets, child1->offsets().get()); + EXPECT_NE(keysPtr, child1->mapKeys().get()); + EXPECT_NE(valuesPtr, child1->mapValues().get()); + EXPECT_NE(childPtr, child1.get()); + EXPECT_EQ(rowPtr, result.get()); + }; + uint32_t seed = FLAGS_reader_tests_seed > 0 ? FLAGS_reader_tests_seed + : folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng{seed}; + auto type = + velox::ROW({{"map_val", velox::MAP(velox::INTEGER(), velox::REAL())}}); + for (int i = 0; i < 10; ++i) { + if (i == 4) { + LOG(INFO) << i; + } + testMapFieldLifeCycle(type, 10, rng); + testMapFieldLifeCycle(type, 10, rng); + } +} + +TEST_F(VeloxReaderTests, TestFlatMapAsMapFieldLifeCycle) { + auto testFlatMapFieldLifeCycle = + [&](const std::shared_ptr type, + int32_t batchSize, + std::mt19937& rng) { + velox::VectorPtr result; + nimble::VeloxWriterOptions writeOptions; + writeOptions.flatMapColumns.insert("flat_map"); + auto reader = + getReaderForLifeCycleTest(type, 5 * batchSize, rng, writeOptions); + EXPECT_TRUE(reader->next(batchSize, result)); + // Hold the reference to internal Buffers and element doesn't change + auto child = std::dynamic_pointer_cast( + result->as()->childAt(0)); + velox::BaseVector* childPtr = child.get(); + velox::BaseVector* rowPtr = result.get(); + velox::Buffer* rawNulls = child->nulls().get(); + velox::BufferPtr sizes = child->sizes(); + velox::Buffer* rawOffsets = child->offsets().get(); + auto keysPtr = child->mapKeys().get(); + child.reset(); + EXPECT_TRUE(reader->next(batchSize, result)); + child = std::dynamic_pointer_cast( + result->as()->childAt(0)); + + EXPECT_EQ(rawNulls, child->nulls().get()); + EXPECT_NE(sizes, child->sizes()); + EXPECT_EQ(rawOffsets, child->offsets().get()); + EXPECT_EQ(keysPtr, child->mapKeys().get()); + EXPECT_EQ(rowPtr, result.get()); + + // Hold the reference to keys vector, other buffer should be reused + auto mapKeys = child->mapKeys(); + velox::Buffer* rawSizes = child->sizes().get(); + childPtr = child.get(); + child.reset(); + EXPECT_TRUE(reader->next(batchSize, result)); + child = std::dynamic_pointer_cast( + result->as()->childAt(0)); + EXPECT_EQ(rawNulls, child->nulls().get()); + EXPECT_EQ(rawSizes, child->sizes().get()); + EXPECT_EQ(rawOffsets, child->offsets().get()); + EXPECT_NE(mapKeys, child->mapKeys()); + EXPECT_EQ(childPtr, child.get()); + EXPECT_EQ(rowPtr, result.get()); + + // Don't release the Child map vector to row vector, all the buffers + // in array should not be resused. + EXPECT_TRUE(reader->next(batchSize, result)); + auto child1 = std::dynamic_pointer_cast( + result->as()->childAt(0)); + EXPECT_NE(rawNulls, child1->nulls().get()); + EXPECT_NE(rawSizes, child1->sizes().get()); + EXPECT_NE(rawOffsets, child1->offsets().get()); + EXPECT_NE(keysPtr, child1->mapKeys().get()); + EXPECT_NE(childPtr, child1.get()); + EXPECT_EQ(rowPtr, result.get()); + }; + uint32_t seed = FLAGS_reader_tests_seed > 0 ? FLAGS_reader_tests_seed + : folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng{seed}; + auto type = + velox::ROW({{"flat_map", velox::MAP(velox::INTEGER(), velox::REAL())}}); + for (int i = 0; i < 10; ++i) { + testFlatMapFieldLifeCycle(type, 10, rng); + testFlatMapFieldLifeCycle(type, 10, rng); + } +} + +TEST_F(VeloxReaderTests, TestRowFieldLifeCycle) { + auto testRowFieldLifeCycle = + [&](const std::shared_ptr type, + int32_t batchSize, + std::mt19937& rng) { + velox::VectorPtr result; + auto reader = getReaderForLifeCycleTest(type, 5 * batchSize, rng); + EXPECT_TRUE(reader->next(batchSize, result)); + // Hold the reference to internal Buffers and element doesn't change + auto child = std::dynamic_pointer_cast( + result->as()->childAt(0)); + velox::BaseVector* childPtr = child.get(); + velox::BaseVector* rowPtr = result.get(); + velox::BufferPtr nulls = child->nulls(); + velox::BaseVector* childPtrAtIdx0 = child->childAt(0).get(); + velox::BaseVector* childPtrAtIdx1 = child->childAt(1).get(); + child.reset(); + EXPECT_TRUE(reader->next(batchSize, result)); + child = std::dynamic_pointer_cast( + result->as()->childAt(0)); + + EXPECT_NE(nulls, child->nulls()); + EXPECT_EQ(childPtrAtIdx0, child->childAt(0).get()); + EXPECT_EQ(childPtrAtIdx1, child->childAt(1).get()); + EXPECT_EQ(rowPtr, result.get()); + + // Hold the reference to one of child vector, sibling should not + // change + auto childAtIdx0 = child->childAt(0); + velox::Buffer* rawNulls = child->nulls().get(); + childPtr = child.get(); + child.reset(); + EXPECT_TRUE(reader->next(batchSize, result)); + child = std::dynamic_pointer_cast( + result->as()->childAt(0)); + EXPECT_EQ(rawNulls, child->nulls().get()); + EXPECT_NE(childAtIdx0, child->childAt(0)); + EXPECT_EQ(childPtrAtIdx1, child->childAt(1).get()); + EXPECT_EQ(childPtr, child.get()); + EXPECT_EQ(rowPtr, result.get()); + + // Don't release the Child-row vector to row vector, all the buffers + // in array should not be resused. + EXPECT_TRUE(reader->next(batchSize, result)); + auto child1 = std::dynamic_pointer_cast( + result->as()->childAt(0)); + EXPECT_NE(rawNulls, child1->nulls().get()); + EXPECT_NE(child->childAt(0), child1->childAt(0)); + EXPECT_NE(child->childAt(1), child1->childAt(1)); + EXPECT_NE(childPtr, child1.get()); + EXPECT_EQ(rowPtr, result.get()); + }; + + auto type = velox::ROW( + {{"row_val", + velox::ROW( + {{"a", velox::INTEGER()}, {"b", velox::ARRAY(velox::BIGINT())}})}}); + uint32_t seed = FLAGS_reader_tests_seed > 0 ? FLAGS_reader_tests_seed + : folly::Random::rand32(); + LOG(INFO) << "seed: " << seed; + std::mt19937 rng{seed}; + for (int i = 0; i < 20; ++i) { + testRowFieldLifeCycle(type, 10, rng); + } +} + +TEST_F(VeloxReaderTests, VeloxTypeFromNimbleSchema) { + auto type = velox::ROW({ + {"tinyint_val", velox::TINYINT()}, + {"smallint_val", velox::SMALLINT()}, + {"int_val", velox::INTEGER()}, + {"long_val", velox::BIGINT()}, + {"float_val", velox::REAL()}, + {"double_val", velox::DOUBLE()}, + {"binary_val", velox::VARBINARY()}, + {"string_val", velox::VARCHAR()}, + {"array_val", velox::ARRAY(velox::BIGINT())}, + {"map_val", velox::MAP(velox::INTEGER(), velox::BIGINT())}, + {"struct_val", + velox::ROW({ + {"float_val", velox::REAL()}, + {"double_val", velox::DOUBLE()}, + })}, + {"nested_map_row_val", + velox::MAP( + velox::INTEGER(), + velox::ROW({ + {"float_val", velox::REAL()}, + {"array_val", + velox::ARRAY(velox::MAP(velox::INTEGER(), velox::BIGINT()))}, + }))}, + {"dictionary_array_val", velox::ARRAY(velox::BIGINT())}, + }); + + velox::VectorFuzzer fuzzer({.vectorSize = 100}, leafPool_.get()); + auto vector = fuzzer.fuzzInputFlatRow(type); + + nimble::VeloxWriterOptions writerOptions; + writerOptions.flatMapColumns.insert("nested_map_row_val"); + writerOptions.dictionaryArrayColumns.insert("dictionary_array_val"); + testVeloxTypeFromNimbleSchema(*leafPool_, writerOptions, vector); +} + +TEST_F(VeloxReaderTests, VeloxTypeFromNimbleSchemaEmptyFlatMap) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + uint32_t numRows = 5; + auto vector = vectorMaker.rowVector( + {"col_0", "col_1"}, + { + vectorMaker.flatVector( + numRows, + [](velox::vector_size_t row) { return 1000 + row; }, + [](velox::vector_size_t row) { return row == 1; }), + vectorMaker.mapVector( + numRows, + /*sizeAt*/ + [](velox::vector_size_t /* mapRow */) { return 0; }, + /*keyAt*/ + [](velox::vector_size_t /* mapRow */, + velox::vector_size_t /* row */) { return ""; }, + /*valueAt*/ + [](velox::vector_size_t /* mapRow */, + velox::vector_size_t /* row */) { return 0; }, + /*isNullAt*/ + [](velox::vector_size_t /* mapRow */) { return true; }), + }); + nimble::VeloxWriterOptions writerOptions; + writerOptions.flatMapColumns.insert("col_1"); + testVeloxTypeFromNimbleSchema(*leafPool_, writerOptions, vector); +} + +TEST_F(VeloxReaderTests, MissingMetadata) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + auto vector = + vectorMaker.rowVector({vectorMaker.flatVector({1, 2, 3})}); + + nimble::VeloxWriterOptions options; + auto file = nimble::test::createNimbleFile(*rootPool_, vector, options); + for (auto useChaniedBuffers : {false, true}) { + nimble::testing::InMemoryTrackableReadFile readFile( + file, useChaniedBuffers); + + nimble::VeloxReader reader(*leafPool_, &readFile); + { + readFile.resetChunks(); + const auto& metadata = reader.metadata(); + // Default metadata injects at least one entry + ASSERT_LE(1, metadata.size()); + EXPECT_EQ(1, readFile.chunks().size()); + } + + { + // Metadata is loaded lazily, so reading again just to be sure all is + // well. + readFile.resetChunks(); + const auto& metadata = reader.metadata(); + ASSERT_LE(1, metadata.size()); + EXPECT_EQ(0, readFile.chunks().size()); + } + } +} + +TEST_F(VeloxReaderTests, WithMetadata) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + auto vector = + vectorMaker.rowVector({vectorMaker.flatVector({1, 2, 3})}); + + nimble::VeloxWriterOptions options{ + .metadata = {{"key 1", "value 1"}, {"key 2", "value 2"}}, + }; + auto file = nimble::test::createNimbleFile(*rootPool_, vector, options); + for (auto useChaniedBuffers : {false, true}) { + nimble::testing::InMemoryTrackableReadFile readFile( + file, useChaniedBuffers); + + nimble::VeloxReader reader(*leafPool_, &readFile); + + { + readFile.resetChunks(); + auto metadata = reader.metadata(); + ASSERT_EQ(2, metadata.size()); + ASSERT_TRUE(metadata.contains("key 1")); + ASSERT_TRUE(metadata.contains("key 2")); + ASSERT_EQ("value 1", metadata["key 1"]); + ASSERT_EQ("value 2", metadata["key 2"]); + + EXPECT_EQ(1, readFile.chunks().size()); + } + + { + // Metadata is loaded lazily, so reading again just to be sure all is + // well. + readFile.resetChunks(); + auto metadata = reader.metadata(); + ASSERT_EQ(2, metadata.size()); + ASSERT_TRUE(metadata.contains("key 1")); + ASSERT_TRUE(metadata.contains("key 2")); + ASSERT_EQ("value 1", metadata["key 1"]); + ASSERT_EQ("value 2", metadata["key 2"]); + + EXPECT_EQ(0, readFile.chunks().size()); + } + } +} + +TEST_F(VeloxReaderTests, InaccurateSchemaWithSelection) { + // Some compute engines (e.g. Presto) sometimes don't have the full schema + // to pass into the reader (if column projection is used). The reader needs + // the schema in order to correctly construct the output vector. However, + // for unprojected columns, the reader just need to put a placeholder null + // column (so ordinals will work as expected), and the actual column type + // doesn't matter. In this case, we expect the compute engine to construct a + // column selector, with dummy nodes in the schema for the unprojected + // columns. This test verifies that the reader handles this correctly. + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + auto vector = vectorMaker.rowVector( + {"int1", "int2", "string", "double", "row1", "row2", "int3", "int4"}, + {vectorMaker.flatVector({11, 12, 13, 14, 15}), + vectorMaker.flatVector({21, 22, 23, 24, 25}), + vectorMaker.flatVector({"s1", "s2", "s3", "s4", "s5"}), + vectorMaker.flatVector({1.1, 2.2, 3.3, 4.4, 5.5}), + vectorMaker.rowVector( + /* childNames */ {"a1", "b1"}, + /* children */ + {vectorMaker.flatVector({111, 112, 113, 114, 115}), + vectorMaker.flatVector({"s111", "s112", "s113", "s114", "s115"})}), + vectorMaker.rowVector( + /* childNames */ {"a2", "b2"}, + /* children */ + {vectorMaker.flatVector({211, 212, 213, 214, 215}), + vectorMaker.flatVector({"s211", "s212", "s213", "s214", "s215"})}), + vectorMaker.flatVector({31, 32, 33, 34, 35}), + vectorMaker.flatVector({41, 42, 43, 44, 45})}); + + velox::VectorPtr result; + { + auto file = nimble::test::createNimbleFile(*rootPool_, vector); + velox::InMemoryReadFile readFile(file); + auto inaccurateType = velox::ROW({ + {"c1", velox::VARCHAR()}, + {"c2", velox::INTEGER()}, + {"c3", velox::VARCHAR()}, + {"c4", velox::VARCHAR()}, + {"c5", velox::VARCHAR()}, + {"c6", velox::ROW({velox::INTEGER(), velox::VARCHAR()})}, + {"c7", velox::INTEGER()}, + // We didn't add the last column on purpose, to test that the reader + // can handle smaller schemas. + }); + + std::unordered_set projected{1, 2, 5, 6}; + auto selector = std::make_shared( + inaccurateType, + std::vector{projected.begin(), projected.end()}); + nimble::VeloxReader reader(*leafPool_, &readFile, std::move(selector)); + + ASSERT_TRUE(reader.next(vector->size(), result)); + const auto& rowResult = result->as(); + ASSERT_EQ(inaccurateType->size(), rowResult->childrenSize()); + for (auto i = 0; i < rowResult->childrenSize(); ++i) { + const auto& child = rowResult->childAt(i); + if (projected.count(i) == 0) { + ASSERT_EQ(rowResult->childAt(i), nullptr); + } else { + ASSERT_EQ(5, child->size()); + for (auto j = 0; j < child->size(); ++j) { + ASSERT_FALSE(child->isNullAt(j)); + ASSERT_TRUE(child->equalValueAt(vector->childAt(i).get(), j, j)); + } + } + } + ASSERT_FALSE(reader.next(vector->size(), result)); + } +} + +TEST_F(VeloxReaderTests, ChunkStreamsWithNulls) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + std::vector vectors{ + vectorMaker.rowVector({ + vectorMaker.flatVectorNullable( + {std::nullopt, std::nullopt, std::nullopt}), + vectorMaker.flatVectorNullable( + {2, std::nullopt, std::nullopt}), + vectorMaker.flatVectorNullable( + {std::nullopt, std::nullopt, std::nullopt}), + }), + vectorMaker.rowVector({ + vectorMaker.flatVectorNullable( + {std::nullopt, std::nullopt, std::nullopt}), + vectorMaker.flatVectorNullable( + {std::nullopt, std::nullopt, std::nullopt}), + vectorMaker.flatVectorNullable( + {std::nullopt, std::nullopt, std::nullopt}), + }), + vectorMaker.rowVector({ + vectorMaker.flatVectorNullable( + {std::nullopt, std::nullopt, std::nullopt}), + vectorMaker.flatVectorNullable( + {2, std::nullopt, std::nullopt}), + vectorMaker.flatVectorNullable( + {std::nullopt, 3, std::nullopt}), + })}; + + for (auto enableChunking : {false, true}) { + nimble::VeloxWriterOptions options{ + .flushPolicyFactory = + [&]() { + return std::make_unique( + [&](auto&) { return nimble::FlushDecision::Chunk; }); + }, + .enableChunking = enableChunking}; + auto file = nimble::test::createNimbleFile( + *rootPool_, vectors, options, /* flushAfterWrite */ false); + velox::InMemoryReadFile readFile(file); + nimble::VeloxReader reader(*leafPool_, &readFile, /* selector */ nullptr); + + velox::VectorPtr result; + for (const auto& expected : vectors) { + ASSERT_TRUE(reader.next(expected->size(), result)); + ASSERT_EQ(expected->size(), result->size()); + for (auto i = 0; i < expected->size(); ++i) { + LOG(INFO) << expected->toString(i); + ASSERT_TRUE(expected->equalValueAt(result.get(), i, i)) + << "Content mismatch at index " << i + << "\nReference: " << expected->toString(i) + << "\nResult: " << result->toString(i); + } + } + } +} diff --git a/dwio/nimble/velox/tests/VeloxWriterTests.cpp b/dwio/nimble/velox/tests/VeloxWriterTests.cpp new file mode 100644 index 0000000..c770902 --- /dev/null +++ b/dwio/nimble/velox/tests/VeloxWriterTests.cpp @@ -0,0 +1,1789 @@ +// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +#include +#include +#include + +#include "common/strings/UUID.h" +#include "dwio/nimble/common/EncodingPrimitives.h" +#include "dwio/nimble/common/tests/TestUtils.h" +#include "dwio/nimble/encodings/EncodingLayoutCapture.h" +#include "dwio/nimble/velox/ChunkedStream.h" +#include "dwio/nimble/velox/EncodingLayoutTree.h" +#include "dwio/nimble/velox/SchemaSerialization.h" +#include "dwio/nimble/velox/TabletSections.h" +#include "dwio/nimble/velox/VeloxReader.h" +#include "dwio/nimble/velox/VeloxWriter.h" +#include "folly/FileUtil.h" +#include "thrift/lib/cpp2/protocol/DebugProtocol.h" +#include "velox/common/memory/SharedArbitrator.h" +#include "velox/exec/MemoryReclaimer.h" +#include "velox/vector/VectorStream.h" +#include "velox/vector/fuzzer/VectorFuzzer.h" +#include "velox/vector/tests/utils/VectorMaker.h" + +using namespace ::facebook; + +class VeloxWriterTests : public testing::Test { + protected: + static void SetUpTestCase() { + velox::memory::SharedArbitrator::registerFactory(); + velox::memory::MemoryManager::testingSetInstance( + {.arbitratorKind = "SHARED"}); + } + + void SetUp() override { + rootPool_ = velox::memory::memoryManager()->addRootPool("default_root"); + leafPool_ = rootPool_->addLeafChild("default_leaf"); + } + + std::shared_ptr rootPool_; + std::shared_ptr leafPool_; +}; + +TEST_F(VeloxWriterTests, EmptyFile) { + auto type = velox::ROW({{"simple", velox::INTEGER()}}); + + std::string file; + auto writeFile = std::make_unique(&file); + + nimble::VeloxWriter writer(*rootPool_, type, std::move(writeFile), {}); + writer.close(); + + velox::InMemoryReadFile readFile(file); + nimble::VeloxReader reader(*leafPool_, &readFile); + + velox::VectorPtr result; + ASSERT_FALSE(reader.next(1, result)); +} + +TEST_F(VeloxWriterTests, ExceptionOnClose) { + class ThrowingWriteFile final : public velox::WriteFile { + public: + void append(std::string_view /* data */) final { + throw std::runtime_error("error/" + strings::generateUUID()); + } + void flush() final { + throw std::runtime_error("error/" + strings::generateUUID()); + } + void close() final { + throw std::runtime_error("error/" + strings::generateUUID()); + } + uint64_t size() const final { + throw std::runtime_error("error/" + strings::generateUUID()); + } + }; + + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + auto vector = vectorMaker.rowVector( + {"col0"}, {vectorMaker.flatVector({1, 2, 3})}); + + std::string file; + auto writeFile = std::make_unique(); + + nimble::VeloxWriter writer( + *rootPool_, + vector->type(), + std::move(writeFile), + {.flushPolicyFactory = [&]() { + return std::make_unique( + [&](auto&) { return nimble::FlushDecision::Stripe; }); + }}); + std::string error; + try { + writer.write(vector); + FAIL() << "Expecting exception"; + } catch (const std::runtime_error& e) { + EXPECT_TRUE(std::string{e.what()}.starts_with("error/")); + error = e.what(); + } + + try { + writer.write(vector); + FAIL() << "Expecting exception"; + } catch (const std::runtime_error& e) { + EXPECT_EQ(error, e.what()); + } + + try { + writer.flush(); + FAIL() << "Expecting exception"; + } catch (const std::runtime_error& e) { + EXPECT_EQ(error, e.what()); + } + + try { + writer.close(); + FAIL() << "Expecting exception"; + } catch (const std::runtime_error& e) { + EXPECT_EQ(error, e.what()); + } + + try { + writer.close(); + FAIL() << "Expecting exception"; + } catch (const std::runtime_error& e) { + EXPECT_EQ(error, e.what()); + } +} + +TEST_F(VeloxWriterTests, EmptyFileNoSchema) { + const uint32_t batchSize = 10; + auto type = velox::ROW({{"simple", velox::INTEGER()}}); + nimble::VeloxWriterOptions writerOptions; + + std::string file; + auto writeFile = std::make_unique(&file); + + nimble::VeloxWriter writer( + *rootPool_, type, std::move(writeFile), std::move(writerOptions)); + writer.close(); + + velox::InMemoryReadFile readFile(file); + nimble::VeloxReader reader(*leafPool_, &readFile); + + velox::VectorPtr result; + ASSERT_FALSE(reader.next(batchSize, result)); +} + +TEST_F(VeloxWriterTests, RootHasNulls) { + auto batchSize = 5; + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + auto vector = vectorMaker.rowVector( + {"col0"}, {vectorMaker.flatVector(batchSize, [](auto row) { + return row; + })}); + + // add nulls + for (auto i = 0; i < batchSize; ++i) { + vector->setNull(i, i % 2 == 0); + } + + std::string file; + auto writeFile = std::make_unique(&file); + nimble::VeloxWriter writer( + *rootPool_, vector->type(), std::move(writeFile), {}); + writer.write(vector); + writer.close(); + + velox::InMemoryReadFile readFile(file); + nimble::VeloxReader reader(*leafPool_, &readFile); + + velox::VectorPtr result; + ASSERT_TRUE(reader.next(batchSize, result)); + ASSERT_EQ(result->size(), batchSize); + for (auto i = 0; i < batchSize; ++i) { + ASSERT_TRUE(result->equalValueAt(vector.get(), i, i)); + } +} + +TEST_F(VeloxWriterTests, FeatureReorderingNonFlatmapColumn) { + const uint32_t batchSize = 10; + + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + auto vector = vectorMaker.rowVector( + {"map", "flatmap"}, + {vectorMaker.mapVector( + 5, + /* sizeAt */ [](auto row) { return row % 3; }, + /* keyAt */ [](auto /* row */, auto mapIndex) { return mapIndex; }, + /* valueAt */ [](auto row, auto /* mapIndex */) { return row; }, + /* isNullAt */ [](auto /* row */) { return false; }), + vectorMaker.mapVector( + 5, + /* sizeAt */ [](auto row) { return row % 3; }, + /* keyAt */ [](auto /* row */, auto mapIndex) { return mapIndex; }, + /* valueAt */ [](auto row, auto /* mapIndex */) { return row; }, + /* isNullAt */ [](auto /* row */) { return false; })}); + + std::string file; + auto writeFile = std::make_unique(&file); + + nimble::VeloxWriter writer( + *rootPool_, + vector->type(), + std::move(writeFile), + {.flatMapColumns = {"flatmap"}, + .featureReordering = + std::vector>>{ + {0, {1, 2}}, {1, {3, 4}}}}); + + try { + writer.write(vector); + writer.close(); + FAIL(); + } catch (const nimble::NimbleUserError& e) { + EXPECT_EQ("USER", e.errorSource()); + EXPECT_EQ("INVALID_ARGUMENT", e.errorCode()); + EXPECT_EQ( + "Column 'map' for feature ordering is not a flat map.", + e.errorMessage()); + } +} + +namespace { +std::vector generateBatches( + const std::shared_ptr& type, + size_t batchCount, + size_t size, + uint32_t seed, + velox::memory::MemoryPool& pool) { + velox::VectorFuzzer fuzzer( + {.vectorSize = size, .nullRatio = 0.1}, &pool, seed); + std::vector batches; + + for (size_t i = 0; i < batchCount; ++i) { + batches.push_back(fuzzer.fuzzInputFlatRow(type)); + } + return batches; +} +} // namespace + +struct RawStripeSizeFlushPolicyTestCase { + const size_t batchCount; + const uint32_t rawStripeSize; + const uint32_t stripeCount; +}; + +class RawStripeSizeFlushPolicyTest + : public VeloxWriterTests, + public ::testing::WithParamInterface {}; + +TEST_P(RawStripeSizeFlushPolicyTest, RawStripeSizeFlushPolicy) { + auto type = velox::ROW({{"simple", velox::INTEGER()}}); + nimble::VeloxWriterOptions writerOptions{.flushPolicyFactory = []() { + // Buffering 256MB data before encoding stripes. + return std::make_unique( + GetParam().rawStripeSize); + }}; + + std::string file; + auto writeFile = std::make_unique(&file); + + nimble::VeloxWriter writer( + *rootPool_, type, std::move(writeFile), std::move(writerOptions)); + auto batches = + generateBatches(type, GetParam().batchCount, 4000, 20221110, *leafPool_); + + for (const auto& batch : batches) { + writer.write(batch); + } + writer.close(); + + velox::InMemoryReadFile readFile(file); + auto selector = std::make_shared(type); + nimble::VeloxReader reader(*leafPool_, &readFile, std::move(selector)); + + EXPECT_EQ(GetParam().stripeCount, reader.getTabletView().stripeCount()); +} + +namespace { +class MockReclaimer : public velox::memory::MemoryReclaimer { + public: + explicit MockReclaimer() {} + void setEnterArbitrationFunc(std::function&& func) { + enterArbitrationFunc_ = func; + } + void enterArbitration() override { + if (enterArbitrationFunc_) { + enterArbitrationFunc_(); + } + } + + private: + std::function enterArbitrationFunc_; +}; +} // namespace + +TEST_F(VeloxWriterTests, MemoryReclaimPath) { + auto rootPool = velox::memory::memoryManager()->addRootPool( + "root", 4L << 20, velox::exec::MemoryReclaimer::create()); + auto writerPool = rootPool->addAggregateChild( + "writer", velox::exec::MemoryReclaimer::create()); + + auto type = velox::ROW( + {{"simple_int", velox::INTEGER()}, {"simple_double", velox::DOUBLE()}}); + std::string file; + auto writeFile = std::make_unique(&file); + std::atomic_bool reclaimEntered = false; + nimble::VeloxWriterOptions writerOptions{.reclaimerFactory = [&]() { + auto reclaimer = std::make_unique(); + reclaimer->setEnterArbitrationFunc([&]() { reclaimEntered = true; }); + return reclaimer; + }}; + nimble::VeloxWriter writer( + *writerPool, type, std::move(writeFile), std::move(writerOptions)); + auto batches = generateBatches(type, 100, 4000, 20221110, *leafPool_); + + EXPECT_THROW( + { + for (const auto& batch : batches) { + writer.write(batch); + } + }, + velox::VeloxException); + ASSERT_TRUE(reclaimEntered.load()); +} + +TEST_F(VeloxWriterTests, FlushHugeStrings) { + nimble::VeloxWriterOptions writerOptions{.flushPolicyFactory = []() { + return std::make_unique(1 * 1024 * 1024); + }}; + + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + // Each vector contains 99 strings with 36 characters each (36*99=3564) + 100 + // bytes for null vector + 99 string_views (99*16=1584) for a total of 5248 + // bytes, so writing 200 batches should exceed the flush theshold of 1MB + auto vector = vectorMaker.rowVector( + {"string"}, + { + vectorMaker.flatVector( + 100, + [](auto /* row */) { + return std::string("abcdefghijklmnopqrstuvwxyz0123456789"); + }, + [](auto row) { return row == 6; }), + }); + + std::string file; + auto writeFile = std::make_unique(&file); + + nimble::VeloxWriter writer( + *rootPool_, + vector->type(), + std::move(writeFile), + std::move(writerOptions)); + + // Writing 500 batches should produce 3 stripes, as each 200 vectors will + // exceed the flush threshold. + for (auto i = 0; i < 500; ++i) { + writer.write(vector); + } + writer.close(); + + velox::InMemoryReadFile readFile(file); + auto selector = std::make_shared( + std::dynamic_pointer_cast(vector->type())); + nimble::VeloxReader reader(*leafPool_, &readFile, std::move(selector)); + + EXPECT_EQ(3, reader.getTabletView().stripeCount()); +} + +TEST_F(VeloxWriterTests, EncodingLayout) { + nimble::EncodingLayoutTree expected{ + nimble::Kind::Row, + {}, + "", + { + {nimble::Kind::Map, + { + { + 0, + nimble::EncodingLayout{ + nimble::EncodingType::Dictionary, + nimble::CompressionType::Uncompressed, + { + nimble::EncodingLayout{ + nimble::EncodingType::FixedBitWidth, + nimble::CompressionType::Zstrong}, + std::nullopt, + }}, + }, + }, + "", + { + // Map keys + {nimble::Kind::Scalar, {}, ""}, + // Map Values + {nimble::Kind::Scalar, + { + { + 0, + nimble::EncodingLayout{ + nimble::EncodingType::MainlyConstant, + nimble::CompressionType::Uncompressed, + { + std::nullopt, + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Zstrong}, + }}, + }, + }, + ""}, + }}, + {nimble::Kind::FlatMap, + {}, + "", + { + { + nimble::Kind::Scalar, + { + { + 0, + nimble::EncodingLayout{ + nimble::EncodingType::MainlyConstant, + nimble::CompressionType::Uncompressed, + { + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Uncompressed}, + nimble::EncodingLayout{ + nimble::EncodingType::FixedBitWidth, + nimble::CompressionType::Uncompressed}, + }}, + }, + }, + "1", + }, + { + nimble::Kind::Scalar, + { + { + 0, + nimble::EncodingLayout{ + nimble::EncodingType::Constant, + nimble::CompressionType::Uncompressed, + }, + }, + }, + "2", + }, + }}, + }}; + + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + auto vector = vectorMaker.rowVector( + {"map", "flatmap"}, + {vectorMaker.mapVector( + 5, + /* sizeAt */ [](auto row) { return row % 3; }, + /* keyAt */ + [](auto /* row */, auto mapIndex) { return mapIndex; }, + /* valueAt */ [](auto row, auto /* mapIndex */) { return row; }, + /* isNullAt */ [](auto /* row */) { return false; }), + vectorMaker.mapVector( + std::vector>>>>{ + std::vector>>{ + {0, 2}, + {2, 3}, + }, + std::nullopt, + {}, + std::vector>>{ + {1, 4}, + {0, std::nullopt}, + }, + std::vector>>{ + {1, std::nullopt}, + }, + })}); + + std::string file; + auto writeFile = std::make_unique(&file); + + nimble::VeloxWriter writer( + *rootPool_, + vector->type(), + std::move(writeFile), + { + .flatMapColumns = {"flatmap"}, + .encodingLayoutTree = std::move(expected), + // Boosting acceptance ratio by 100x to make sure it is always + // accepted (even if compressed size if bigger than uncompressed size) + .compressionOptions = {.compressionAcceptRatio = 100}, + }); + + writer.write(vector); + writer.close(); + + for (auto useChaniedBuffers : {false, true}) { + nimble::testing::InMemoryTrackableReadFile readFile( + file, useChaniedBuffers); + nimble::Tablet tablet{*leafPool_, &readFile}; + auto section = + tablet.loadOptionalSection(std::string(nimble::kSchemaSection)); + NIMBLE_CHECK(section.has_value(), "Schema not found."); + auto schema = + nimble::SchemaDeserializer::deserialize(section->content().data()); + auto& mapNode = schema->asRow().childAt(0)->asMap(); + auto& mapValuesNode = mapNode.values()->asScalar(); + auto& flatMapNode = schema->asRow().childAt(1)->asFlatMap(); + ASSERT_EQ(3, flatMapNode.childrenCount()); + + auto findChild = + [](const facebook::nimble::FlatMapType& map, + std::string_view key) -> std::shared_ptr { + for (auto i = 0; i < map.childrenCount(); ++i) { + if (map.nameAt(i) == key) { + return map.childAt(i); + } + } + return nullptr; + }; + const auto& flatMapKey1Node = findChild(flatMapNode, "1")->asScalar(); + const auto& flatMapKey2Node = findChild(flatMapNode, "2")->asScalar(); + + for (auto i = 0; i < tablet.stripeCount(); ++i) { + std::vector identifiers{ + mapNode.lengthsDescriptor().offset(), + mapValuesNode.scalarDescriptor().offset(), + flatMapKey1Node.scalarDescriptor().offset(), + flatMapKey2Node.scalarDescriptor().offset()}; + auto streams = tablet.load(i, identifiers); + { + nimble::InMemoryChunkedStream chunkedStream{ + *leafPool_, std::move(streams[0])}; + ASSERT_TRUE(chunkedStream.hasNext()); + // Verify Map stream + auto capture = + nimble::EncodingLayoutCapture::capture(chunkedStream.nextChunk()); + EXPECT_EQ(nimble::EncodingType::Dictionary, capture.encodingType()); + EXPECT_EQ( + nimble::EncodingType::FixedBitWidth, + capture.child(nimble::EncodingIdentifiers::Dictionary::Alphabet) + ->encodingType()); + EXPECT_EQ( + nimble::CompressionType::Zstrong, + capture.child(nimble::EncodingIdentifiers::Dictionary::Alphabet) + ->compressionType()); + } + + { + nimble::InMemoryChunkedStream chunkedStream{ + *leafPool_, std::move(streams[1])}; + ASSERT_TRUE(chunkedStream.hasNext()); + // Verify Map Values stream + auto capture = + nimble::EncodingLayoutCapture::capture(chunkedStream.nextChunk()); + EXPECT_EQ(nimble::EncodingType::MainlyConstant, capture.encodingType()); + EXPECT_EQ( + nimble::EncodingType::Trivial, + capture + .child(nimble::EncodingIdentifiers::MainlyConstant::OtherValues) + ->encodingType()); + EXPECT_EQ( + nimble::CompressionType::Zstrong, + capture + .child(nimble::EncodingIdentifiers::MainlyConstant::OtherValues) + ->compressionType()); + } + + { + nimble::InMemoryChunkedStream chunkedStream{ + *leafPool_, std::move(streams[2])}; + ASSERT_TRUE(chunkedStream.hasNext()); + // Verify FlatMap Kay "1" stream + auto capture = + nimble::EncodingLayoutCapture::capture(chunkedStream.nextChunk()); + EXPECT_EQ(nimble::EncodingType::MainlyConstant, capture.encodingType()); + EXPECT_EQ( + nimble::EncodingType::Trivial, + capture + .child(nimble::EncodingIdentifiers::MainlyConstant::IsCommon) + ->encodingType()); + EXPECT_EQ( + nimble::CompressionType::Uncompressed, + capture + .child(nimble::EncodingIdentifiers::MainlyConstant::IsCommon) + ->compressionType()); + EXPECT_EQ( + nimble::EncodingType::FixedBitWidth, + capture + .child(nimble::EncodingIdentifiers::MainlyConstant::OtherValues) + ->encodingType()); + EXPECT_EQ( + nimble::CompressionType::Uncompressed, + capture + .child(nimble::EncodingIdentifiers::MainlyConstant::OtherValues) + ->compressionType()); + } + + { + nimble::InMemoryChunkedStream chunkedStream{ + *leafPool_, std::move(streams[3])}; + ASSERT_TRUE(chunkedStream.hasNext()); + // Verify FlatMap Kay "2" stream + auto capture = + nimble::EncodingLayoutCapture::capture(chunkedStream.nextChunk()); + EXPECT_EQ(nimble::EncodingType::Constant, capture.encodingType()); + } + } + } +} + +TEST_F(VeloxWriterTests, EncodingLayoutSchemaMismatch) { + nimble::EncodingLayoutTree expected{ + nimble::Kind::Row, + {}, + "", + { + { + nimble::Kind::Scalar, + { + { + 0, + nimble::EncodingLayout{ + nimble::EncodingType::Dictionary, + nimble::CompressionType::Uncompressed, + { + nimble::EncodingLayout{ + nimble::EncodingType::FixedBitWidth, + nimble::CompressionType::Zstrong}, + std::nullopt, + }}, + }, + }, + "", + }, + }}; + + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + auto vector = vectorMaker.rowVector( + {"map"}, + { + vectorMaker.mapVector( + 5, + /* sizeAt */ [](auto row) { return row % 3; }, + /* keyAt */ + [](auto /* row */, auto mapIndex) { return mapIndex; }, + /* valueAt */ [](auto row, auto /* mapIndex */) { return row; }, + /* isNullAt */ [](auto /* row */) { return false; }), + }); + + std::string file; + auto writeFile = std::make_unique(&file); + + try { + nimble::VeloxWriter writer( + *rootPool_, + vector->type(), + std::move(writeFile), + { + .encodingLayoutTree = std::move(expected), + .compressionOptions = {.compressionAcceptRatio = 100}, + }); + FAIL() << "Writer should fail on incompatible encoding layout node"; + } catch (const nimble::NimbleUserError& e) { + EXPECT_NE( + std::string(e.what()).find( + "Incompatible encoding layout node. Expecting map node"), + std::string::npos); + } +} + +TEST_F(VeloxWriterTests, EncodingLayoutSchemaEvolutionMapToFlatmap) { + nimble::EncodingLayoutTree expected{ + nimble::Kind::Row, + {}, + "", + { + {nimble::Kind::Map, + { + { + 0, + nimble::EncodingLayout{ + nimble::EncodingType::Dictionary, + nimble::CompressionType::Uncompressed, + { + nimble::EncodingLayout{ + nimble::EncodingType::FixedBitWidth, + nimble::CompressionType::Zstrong}, + std::nullopt, + }}, + }, + }, + "", + { + // Map keys + {nimble::Kind::Scalar, {}, ""}, + // Map Values + {nimble::Kind::Scalar, + { + { + 0, + nimble::EncodingLayout{ + nimble::EncodingType::MainlyConstant, + nimble::CompressionType::Uncompressed, + { + std::nullopt, + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Zstrong}, + }}, + }, + }, + ""}, + }}, + }}; + + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + auto vector = vectorMaker.rowVector( + {"map"}, + { + vectorMaker.mapVector( + 5, + /* sizeAt */ [](auto row) { return row % 3; }, + /* keyAt */ + [](auto /* row */, auto mapIndex) { return mapIndex; }, + /* valueAt */ [](auto row, auto /* mapIndex */) { return row; }, + /* isNullAt */ [](auto /* row */) { return false; }), + }); + + std::string file; + auto writeFile = std::make_unique(&file); + + nimble::VeloxWriter writer( + *rootPool_, + vector->type(), + std::move(writeFile), + { + .flatMapColumns = {"map"}, + .encodingLayoutTree = std::move(expected), + .compressionOptions = {.compressionAcceptRatio = 100}, + }); + + writer.write(vector); + writer.close(); + + // Getting here is good enough for now (as it means we didn't fail on node + // type mismatch). Once we add metric collection, we can use these to verify + // that no captured encoding was used. +} + +TEST_F(VeloxWriterTests, EncodingLayoutSchemaEvolutionFlamapToMap) { + nimble::EncodingLayoutTree expected{ + nimble::Kind::Row, + {}, + "", + { + {nimble::Kind::FlatMap, + {}, + "", + { + { + nimble::Kind::Scalar, + { + { + 0, + nimble::EncodingLayout{ + nimble::EncodingType::MainlyConstant, + nimble::CompressionType::Uncompressed, + { + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Uncompressed}, + nimble::EncodingLayout{ + nimble::EncodingType::FixedBitWidth, + nimble::CompressionType::Uncompressed}, + }}, + }, + }, + "1", + }, + { + nimble::Kind::Scalar, + { + { + 0, + nimble::EncodingLayout{ + nimble::EncodingType::Constant, + nimble::CompressionType::Uncompressed, + }, + }, + }, + "2", + }, + }}, + }}; + + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + auto vector = vectorMaker.rowVector( + {"flatmap"}, + { + vectorMaker.mapVector( + 5, + /* sizeAt */ [](auto row) { return row % 3; }, + /* keyAt */ + [](auto /* row */, auto mapIndex) { return mapIndex; }, + /* valueAt */ [](auto row, auto /* mapIndex */) { return row; }, + /* isNullAt */ [](auto /* row */) { return false; }), + }); + + std::string file; + auto writeFile = std::make_unique(&file); + + nimble::VeloxWriter writer( + *rootPool_, + vector->type(), + std::move(writeFile), + { + .encodingLayoutTree = std::move(expected), + .compressionOptions = {.compressionAcceptRatio = 100}, + }); + + writer.write(vector); + writer.close(); + + // Getting here is good enough for now (as it means we didn't fail on node + // type mismatch). Once we add metric collection, we can use these to verify + // that no captured encoding was used. +} + +TEST_F(VeloxWriterTests, EncodingLayoutSchemaEvolutionExpandingRow) { + nimble::EncodingLayoutTree expected{ + nimble::Kind::Row, + {}, + "", + { + {nimble::Kind::Row, + { + { + 0, + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Uncompressed}, + }, + }, + "", + { + { + nimble::Kind::Scalar, + { + { + 0, + nimble::EncodingLayout{ + nimble::EncodingType::Trivial, + nimble::CompressionType::Uncompressed}, + }, + }, + "", + }, + }}, + }}; + + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + // We are adding new top level column and also nested column + auto vector = vectorMaker.rowVector( + {"row1", "row2"}, + { + vectorMaker.rowVector({ + vectorMaker.flatVector({1, 2, 3, 4, 5}), + vectorMaker.flatVector({1, 2, 3, 4, 5}), + }), + vectorMaker.rowVector({ + vectorMaker.flatVector({1, 2, 3, 4, 5}), + vectorMaker.flatVector({1, 2, 3, 4, 5}), + }), + }); + + std::string file; + auto writeFile = std::make_unique(&file); + + nimble::VeloxWriter writer( + *rootPool_, + vector->type(), + std::move(writeFile), + { + .encodingLayoutTree = std::move(expected), + .compressionOptions = {.compressionAcceptRatio = 100}, + }); + + writer.write(vector); + writer.close(); + + // Getting here is good enough for now (as it means we didn't fail on node + // type mismatch). Once we add metric collection, we can use these to verify + // that no captured encoding was used. +} + +#define ASSERT_CHUNK_COUNT(count, chunked) \ + for (auto __i = 0; __i < count; ++__i) { \ + ASSERT_TRUE(chunked.hasNext()); \ + auto chunk = chunked.nextChunk(); \ + EXPECT_LT(0, chunk.size()); \ + } \ + ASSERT_FALSE(chunked.hasNext()); + +void testChunks( + velox::memory::MemoryPool& rootPool, + uint32_t minStreamChunkRawSize, + std::vector> vectors, + std::function verifier, + folly::F14FastSet flatMapColumns = {}) { + ASSERT_LT(0, vectors.size()); + auto& type = std::get<0>(vectors[0])->type(); + + auto leafPool = rootPool.addLeafChild("chunk_leaf"); + auto expected = velox::BaseVector::create(type, 0, leafPool.get()); + + std::string file; + auto writeFile = std::make_unique(&file); + + auto flushDecision = nimble::FlushDecision::None; + nimble::VeloxWriter writer( + rootPool, + type, + std::move(writeFile), + { + .flatMapColumns = std::move(flatMapColumns), + .minStreamChunkRawSize = minStreamChunkRawSize, + .flushPolicyFactory = + [&]() { + return std::make_unique( + [&](auto&) { return flushDecision; }); + }, + .enableChunking = true, + }); + + for (const auto& vector : vectors) { + flushDecision = std::get<1>(vector); + writer.write(std::get<0>(vector)); + expected->append(std::get<0>(vector).get()); + } + + writer.close(); + + folly::writeFile(file, "/tmp/afile"); + + nimble::Tablet tablet{ + *leafPool, std::make_unique(file)}; + verifier(tablet); + + nimble::VeloxReader reader( + *leafPool, std::make_shared(file)); + velox::VectorPtr result; + ASSERT_TRUE(reader.next(expected->size(), result)); + ASSERT_EQ(expected->size(), result->size()); + for (auto i = 0; i < expected->size(); ++i) { + ASSERT_TRUE(expected->equalValueAt(result.get(), i, i)); + } + ASSERT_FALSE(reader.next(1, result)); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsRowAllNullsNoChunks) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto vector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVector({}), + }); + vector->appendNulls(5); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 0, + {{vector, nimble::FlushDecision::None}, + {vector, nimble::FlushDecision::None}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + // Logically, there should be two streams in the tablet. + // However, when writing stripes, we do not write empty streams. + // In this case, the integer column is empty, and therefore, omitted. + ASSERT_EQ(1, tablet.streamCount(0)); + EXPECT_LT(0, tablet.streamSizes(0)[0]); + + auto streamLoaders = tablet.load(0, std::array{0}); + ASSERT_EQ(1, streamLoaders.size()); + + // No chunks used, so expecting single chunk + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[0])}; + ASSERT_CHUNK_COUNT(1, chunked); + }); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsRowAllNullsWithChunksMinSizeBig) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto vector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVector({}), + }); + vector->appendNulls(5); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 1024, + {{vector, nimble::FlushDecision::Chunk}, + {vector, nimble::FlushDecision::Chunk}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + // Logically, there should be two streams in the tablet. + // However, when writing stripes, we do not write empty streams. + // In this case, the integer column is empty, and therefore, omitted. + ASSERT_EQ(1, tablet.streamCount(0)); + EXPECT_LT(0, tablet.streamSizes(0)[0]); + + auto streamLoaders = tablet.load(0, std::array{0}); + ASSERT_EQ(1, streamLoaders.size()); + + // Chunks requested, but min chunk size is too big, so expecting one + // merged chunk + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[0])}; + ASSERT_CHUNK_COUNT(1, chunked); + }); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsRowAllNullsWithChunksMinSizeZero) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto vector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVector({}), + }); + vector->appendNulls(5); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 0, + {{vector, nimble::FlushDecision::Chunk}, + {vector, nimble::FlushDecision::Chunk}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + // Logically, there should be two streams in the tablet. + // However, when writing stripes, we do not write empty streams. + // In this case, the integer column is empty, and therefore, omitted. + ASSERT_EQ(1, tablet.streamCount(0)); + EXPECT_LT(0, tablet.streamSizes(0)[0]); + + auto streamLoaders = tablet.load(0, std::array{0}); + ASSERT_EQ(1, streamLoaders.size()); + + // Chunks requested, and min chunk size is zero, so expecting two + // separate chunks. + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[0])}; + ASSERT_CHUNK_COUNT(2, chunked); + }); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsRowSomeNullsNoChunks) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto nullsVector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVector({}), + }); + nullsVector->appendNulls(5); + + auto nonNullsVector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVector({1, 2, 3}), + }); + nonNullsVector->setNull(1, /* isNull */ true); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 0, + {{nullsVector, nimble::FlushDecision::None}, + {nonNullsVector, nimble::FlushDecision::None}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + // We have values in stream 2, so it is not optimized away. + ASSERT_EQ(2, tablet.streamCount(0)); + EXPECT_LT(0, tablet.streamSizes(0)[0]); + EXPECT_LT(0, tablet.streamSizes(0)[1]); + + auto streamLoaders = tablet.load(0, std::array{0, 1}); + ASSERT_EQ(2, streamLoaders.size()); + { + // No chunks requested, so expecting single chunk. + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[0])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + { + // No chunks requested, so expecting single chunk. + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[1])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + }); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsRowSomeNullsWithChunksMinSizeBig) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto nullsVector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVector({}), + }); + nullsVector->appendNulls(5); + + auto nonNullsVector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVector({1, 2, 3}), + }); + nonNullsVector->setNull(1, /* isNull */ true); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 1024, + {{nullsVector, nimble::FlushDecision::Chunk}, + {nonNullsVector, nimble::FlushDecision::Chunk}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + ASSERT_EQ(2, tablet.streamCount(0)); + EXPECT_LT(0, tablet.streamSizes(0)[0]); + EXPECT_LT(0, tablet.streamSizes(0)[1]); + + auto streamLoaders = tablet.load(0, std::array{0, 1}); + ASSERT_EQ(2, streamLoaders.size()); + { + // Chunks requested, but min chunk size is too big, so expecting one + // merged chunk. + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[0])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + { + // Chunks requested, but min chunk size is too big, so expecting one + // merged chunk. + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[1])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + }); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsRowSomeNullsWithChunksMinSizeZero) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto nullsVector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVector({}), + }); + nullsVector->appendNulls(5); + + auto nonNullsVector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVector({1, 2, 3}), + }); + nonNullsVector->setNull(1, /* isNull */ true); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 0, + {{nullsVector, nimble::FlushDecision::Chunk}, + {nonNullsVector, nimble::FlushDecision::Chunk}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + ASSERT_EQ(2, tablet.streamCount(0)); + EXPECT_LT(0, tablet.streamSizes(0)[0]); + EXPECT_LT(0, tablet.streamSizes(0)[1]); + + auto streamLoaders = tablet.load(0, std::array{0, 1}); + ASSERT_EQ(2, streamLoaders.size()); + { + // Chunks requested, and min chunk size is zero, so expecting two + // separate chunks. + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[0])}; + ASSERT_CHUNK_COUNT(2, chunked); + } + { + // Chunks requested, and min chunk size is zero. However, first write + // didn't have any data, so no chunk was written. + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[1])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + }); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsRowNoNullsNoChunks) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto vector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVector({1, 2, 3}), + }); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 0, + {{vector, nimble::FlushDecision::None}, + {vector, nimble::FlushDecision::None}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + ASSERT_EQ(2, tablet.streamCount(0)); + + // When there are no nulls, the nulls stream is omitted. + EXPECT_EQ(0, tablet.streamSizes(0)[0]); + EXPECT_LT(0, tablet.streamSizes(0)[1]); + + auto streamLoaders = tablet.load(0, std::array{0, 1}); + ASSERT_EQ(2, streamLoaders.size()); + { + // Nulls stream should be missing, as all values are non-null + EXPECT_FALSE(streamLoaders[0]); + } + { + // No chunks requested, so expecting one chunk. + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[1])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + }); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsRowNoNullsWithChunksMinSizeBig) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto vector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVector({1, 2, 3}), + }); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 1024, + {{vector, nimble::FlushDecision::Chunk}, + {vector, nimble::FlushDecision::Chunk}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + ASSERT_EQ(2, tablet.streamCount(0)); + + // When there are no nulls, the nulls stream is omitted. + EXPECT_EQ(0, tablet.streamSizes(0)[0]); + EXPECT_LT(0, tablet.streamSizes(0)[1]); + + auto streamLoaders = tablet.load(0, std::array{0, 1}); + ASSERT_EQ(2, streamLoaders.size()); + { + // Nulls stream should be missing, as all values are non-null + EXPECT_FALSE(streamLoaders[0]); + } + { + // Chunks requested, but min size is too big, so expecting one merged + // chunk. + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[1])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + }); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsRowNoNullsWithChunksMinSizeZero) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto vector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVector({1, 2, 3}), + }); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 0, + {{vector, nimble::FlushDecision::Chunk}, + {vector, nimble::FlushDecision::Chunk}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + ASSERT_EQ(2, tablet.streamCount(0)); + + // When there are no nulls, the nulls stream is omitted. + EXPECT_EQ(0, tablet.streamSizes(0)[0]); + EXPECT_LT(0, tablet.streamSizes(0)[1]); + + auto streamLoaders = tablet.load(0, std::array{0, 1}); + ASSERT_EQ(2, streamLoaders.size()); + { + // Nulls stream should be missing, as all values are non-null + EXPECT_FALSE(streamLoaders[0]); + } + { + // Chunks requested, with min size zero, so expecting two chunks. + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[1])}; + ASSERT_CHUNK_COUNT(2, chunked); + } + }); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsChildAllNullsNoChunks) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto vector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVectorNullable( + {std::nullopt, std::nullopt, std::nullopt}), + }); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 0, + {{vector, nimble::FlushDecision::None}, + {vector, nimble::FlushDecision::None}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + // When all rows are not null, the nulls stream is omitted. + // When all values are null, the values stream is omitted. + // Since these are the last two stream, they are optimized away. + ASSERT_EQ(0, tablet.streamCount(0)); + }); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsChildAllNullsWithChunksMinSizeBig) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto vector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVectorNullable( + {std::nullopt, std::nullopt, std::nullopt}), + }); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 1024, + {{vector, nimble::FlushDecision::Chunk}, + {vector, nimble::FlushDecision::Chunk}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + // When all rows are not null, the nulls stream is omitted. + // When all values are null, the values stream is omitted. + // Since these are the last two stream, they are optimized away. + ASSERT_EQ(0, tablet.streamCount(0)); + }); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsChildAllNullsWithChunksMinSizeZero) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto vector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.flatVectorNullable( + {std::nullopt, std::nullopt, std::nullopt}), + }); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 0, + {{vector, nimble::FlushDecision::Chunk}, + {vector, nimble::FlushDecision::Chunk}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + // When all rows are not null, the nulls stream is omitted. + // When all values are null, the values stream is omitted. + // Since these are the last two stream, they are optimized away. + ASSERT_EQ(0, tablet.streamCount(0)); + }); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsFlatmapAllNullsNoChunks) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto vector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.mapVector( + std::vector>>>>{ + std::nullopt, + std::nullopt, + // std::vector>>{ + // {5, 6}}, + std::nullopt, + }), + }); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 0, + {{vector, nimble::FlushDecision::None}, + {vector, nimble::FlushDecision::None}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + // Expected streams: + // 0: Row nulls stream (expected empty, as all values are not null) + // 1: Flatmap nulls stream + ASSERT_EQ(2, tablet.streamCount(0)); + EXPECT_EQ(0, tablet.streamSizes(0)[0]); + EXPECT_LT(0, tablet.streamSizes(0)[1]); + + auto streamLoaders = tablet.load(0, std::array{0, 1}); + ASSERT_EQ(2, streamLoaders.size()); + + // No chunks used, so expecting single chunk + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[1])}; + ASSERT_CHUNK_COUNT(1, chunked); + }, + /* flatmapColumns */ {"c1"}); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsFlatmapAllNullsWithChunksMinSizeBig) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto vector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.mapVector( + std::vector>>>>{ + std::nullopt, + std::nullopt, + std::nullopt, + }), + }); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 1024, + {{vector, nimble::FlushDecision::Chunk}, + {vector, nimble::FlushDecision::Chunk}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + // Expected streams: + // 0: Row nulls stream (expected empty, as all values are not null) + // 1: Flatmap nulls stream + ASSERT_EQ(2, tablet.streamCount(0)); + EXPECT_EQ(0, tablet.streamSizes(0)[0]); + EXPECT_LT(0, tablet.streamSizes(0)[1]); + + auto streamLoaders = tablet.load(0, std::array{0, 1}); + ASSERT_EQ(2, streamLoaders.size()); + + // Chunks requested, but min size is too big, so expecting single merged + // chunk + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[1])}; + ASSERT_CHUNK_COUNT(1, chunked); + }, + /* flatmapColumns */ {"c1"}); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsFlatmapAllNullsWithChunksMinSizeZero) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto vector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.mapVector( + std::vector>>>>{ + std::nullopt, + std::nullopt, + std::nullopt, + }), + }); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 0, + {{vector, nimble::FlushDecision::Chunk}, + {vector, nimble::FlushDecision::Chunk}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + // Expected streams: + // 0: Row nulls stream (expected empty, as all values are not null) + // 1: Flatmap nulls stream + ASSERT_EQ(2, tablet.streamCount(0)); + EXPECT_EQ(0, tablet.streamSizes(0)[0]); + EXPECT_LT(0, tablet.streamSizes(0)[1]); + + auto streamLoaders = tablet.load(0, std::array{0, 1}); + ASSERT_EQ(2, streamLoaders.size()); + + // Chunks requested, with min size zero, so expecting two chunks + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[1])}; + ASSERT_CHUNK_COUNT(2, chunked); + }, + /* flatmapColumns */ {"c1"}); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsFlatmapSomeNullsNoChunks) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto nullsVector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.mapVector( + std::vector>>>>{ + std::nullopt, + std::nullopt, + std::nullopt, + }), + }); + auto nonNullsVector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.mapVector( + std::vector>>>>{ + std::nullopt, + std::vector>>{ + {5, 6}}, + std::nullopt, + }), + }); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 0, + {{nullsVector, nimble::FlushDecision::None}, + {nonNullsVector, nimble::FlushDecision::None}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + // Expected streams: + // 0: Row nulls stream (expected empty, as all values are not null) + // 1: Flatmap nulls stream + // 2: Scalar stream (flatmap value for key 5) + // 3: Scalar stream (flatmap in-map for key 5) + ASSERT_EQ(4, tablet.streamCount(0)); + EXPECT_EQ(0, tablet.streamSizes(0)[0]); + EXPECT_LT(0, tablet.streamSizes(0)[1]); + EXPECT_LT(0, tablet.streamSizes(0)[2]); + EXPECT_LT(0, tablet.streamSizes(0)[3]); + + auto streamLoaders = + tablet.load(0, std::array{0, 1, 2, 3}); + ASSERT_EQ(4, streamLoaders.size()); + + EXPECT_FALSE(streamLoaders[0]); + EXPECT_TRUE(streamLoaders[1]); + EXPECT_TRUE(streamLoaders[2]); + EXPECT_TRUE(streamLoaders[3]); + + { + // No chunks used, so expecting single chunk + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[1])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + { + // No chunks used, so expecting single chunk + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[2])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + { + // No chunks used, so expecting single chunk + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[3])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + }, + /* flatmapColumns */ {"c1"}); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsFlatmapSomeNullsWithChunksMinSizeBig) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto nullsVector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.mapVector( + std::vector>>>>{ + std::nullopt, + std::nullopt, + std::nullopt, + }), + }); + auto nonNullsVector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.mapVector( + std::vector>>>>{ + std::nullopt, + std::vector>>{ + {5, 6}}, + std::nullopt, + }), + }); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 1024, + {{nullsVector, nimble::FlushDecision::Chunk}, + {nonNullsVector, nimble::FlushDecision::Chunk}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + // Expected streams: + // 0: Row nulls stream (expected empty, as all values are not null) + // 1: Flatmap nulls stream + // 2: Scalar stream (flatmap value for key 5) + // 3: Scalar stream (flatmap in-map for key 5) + ASSERT_EQ(4, tablet.streamCount(0)); + EXPECT_EQ(0, tablet.streamSizes(0)[0]); + EXPECT_LT(0, tablet.streamSizes(0)[1]); + EXPECT_LT(0, tablet.streamSizes(0)[2]); + EXPECT_LT(0, tablet.streamSizes(0)[3]); + + auto streamLoaders = + tablet.load(0, std::array{0, 1, 2, 3}); + ASSERT_EQ(4, streamLoaders.size()); + + EXPECT_FALSE(streamLoaders[0]); + EXPECT_TRUE(streamLoaders[1]); + EXPECT_TRUE(streamLoaders[2]); + EXPECT_TRUE(streamLoaders[3]); + + { + // Chunks requested, but min size is big, so expecting single merged + // chunk + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[1])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + { + // Chunks requested, but min size is big, so expecting single merged + // chunk + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[2])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + { + // Chunks requested, but min size is big, so expecting single merged + // chunk + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[3])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + }, + /* flatmapColumns */ {"c1"}); +} + +TEST_F(VeloxWriterTests, ChunkedStreamsFlatmapSomeNullsWithChunksMinSizeZero) { + velox::test::VectorMaker vectorMaker{leafPool_.get()}; + + auto nullsVector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.mapVector( + std::vector>>>>{ + std::nullopt, + std::nullopt, + std::nullopt, + }), + }); + auto nonNullsVector = vectorMaker.rowVector( + {"c1"}, + { + vectorMaker.mapVector( + std::vector>>>>{ + std::nullopt, + std::vector>>{ + {5, 6}}, + std::nullopt, + }), + }); + + testChunks( + *rootPool_, + /* minStreamChunkRawSize */ 0, + {{nullsVector, nimble::FlushDecision::Chunk}, + {nonNullsVector, nimble::FlushDecision::Chunk}}, + [&](const auto& tablet) { + ASSERT_EQ(1, tablet.stripeCount()); + + // Expected streams: + // 0: Row nulls stream (expected empty, as all values are not null) + // 1: Flatmap nulls stream + // 2: Scalar stream (flatmap value for key 5) + // 3: Scalar stream (flatmap in-map for key 5) + ASSERT_EQ(4, tablet.streamCount(0)); + EXPECT_EQ(0, tablet.streamSizes(0)[0]); + EXPECT_LT(0, tablet.streamSizes(0)[1]); + EXPECT_LT(0, tablet.streamSizes(0)[2]); + EXPECT_LT(0, tablet.streamSizes(0)[3]); + + auto streamLoaders = + tablet.load(0, std::array{0, 1, 2, 3}); + ASSERT_EQ(4, streamLoaders.size()); + + EXPECT_FALSE(streamLoaders[0]); + EXPECT_TRUE(streamLoaders[1]); + EXPECT_TRUE(streamLoaders[2]); + EXPECT_TRUE(streamLoaders[3]); + + { + // Chunks requested, with min size zero, so expecting two chunks + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[1])}; + ASSERT_CHUNK_COUNT(2, chunked); + } + { + // Chunks requested, with min size zero, but first write didn't + // contain any values, so expecting single merged chunk + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[2])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + { + // Chunks requested, with min size zero, but first write didn't + // have any items in the map, so expecting single merged chunk + nimble::InMemoryChunkedStream chunked{ + *leafPool_, std::move(streamLoaders[3])}; + ASSERT_CHUNK_COUNT(1, chunked); + } + }, + /* flatmapColumns */ {"c1"}); +} + +INSTANTIATE_TEST_CASE_P( + RawStripeSizeFlushPolicyTestSuite, + RawStripeSizeFlushPolicyTest, + testing::Values( + RawStripeSizeFlushPolicyTestCase{ + .batchCount = 50, + .rawStripeSize = 256 << 10, + .stripeCount = 4}, + RawStripeSizeFlushPolicyTestCase{ + .batchCount = 100, + .rawStripeSize = 256 << 10, + .stripeCount = 7}, + RawStripeSizeFlushPolicyTestCase{ + .batchCount = 100, + .rawStripeSize = 256 << 11, + .stripeCount = 4}, + RawStripeSizeFlushPolicyTestCase{ + .batchCount = 100, + .rawStripeSize = 256 << 12, + .stripeCount = 2}, + RawStripeSizeFlushPolicyTestCase{ + .batchCount = 100, + .rawStripeSize = 256 << 20, + .stripeCount = 1})); diff --git a/license.header b/license.header new file mode 100644 index 0000000..9e2a164 --- /dev/null +++ b/license.header @@ -0,0 +1,13 @@ + Copyright (c) Meta Platforms, Inc. and its affiliates. + + Licensed 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. diff --git a/scripts/format-check.py b/scripts/format-check.py new file mode 100755 index 0000000..59b976f --- /dev/null +++ b/scripts/format-check.py @@ -0,0 +1,284 @@ +#!/usr/bin/env python3 +# Copyright (c) Facebook, Inc. and its affiliates. +# +# Licensed 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. +# + +import argparse +import os +import sys +from collections import OrderedDict + +import regex +import util + +from util import attrdict + +EXTENSIONS = "cpp,h,inc,prolog" +SCRIPTS = util.script_path() + + +def get_diff(file, formatted): + if not formatted.endswith("\n"): + formatted = formatted + "\n" + + status, stdout, stderr = util.run( + f"diff -u {file} --label {file} --label {file} -", input=formatted + ) + if stdout != "": + stdout = f"diff a/{file} b/{file}\n" + stdout + + return status, stdout, stderr + + +class CppFormatter(str): + def diff(self, commit): + if commit == "": + return get_diff(self, util.run(f"clang-format --style=file {self}")[1]) + else: + return util.run( + f"{SCRIPTS}/git-clang-format -q --extensions='{EXTENSIONS}' --diff --style=file {commit} {self}" + ) + + def fix(self, commit): + if commit == "": + return util.run(f"clang-format -i --style=file {self}")[0] == 0 + else: + return ( + util.run( + f"{SCRIPTS}/git-clang-format -q --extensions='{EXTENSIONS}' --style=file {commit} {self}" + )[0] + == 0 + ) + + +class CMakeFormatter(str): + def diff(self, commit): + return get_diff( + self, util.run(f"cmake-format --first-comment-is-literal True {self}")[1] + ) + + def fix(self, commit): + return ( + util.run(f"cmake-format --first-comment-is-literal True -i {self}")[0] == 0 + ) + + +class PythonFormatter(str): + def diff(self, commit): + return util.run(f"black -q --diff {self}") + + def fix(self, commit): + return util.run(f"black -q {self}")[0] == 0 + + +format_file_types = OrderedDict( + { + "CMakeLists.txt": attrdict({"formatter": CMakeFormatter}), + "*.cmake": attrdict({"formatter": CMakeFormatter}), + "*.cpp": attrdict({"formatter": CppFormatter}), + "*.h": attrdict({"formatter": CppFormatter}), + "*.inc": attrdict({"formatter": CppFormatter}), + "*.prolog": attrdict({"formatter": CppFormatter}), + "*.py": attrdict({"formatter": PythonFormatter}), + } +) + + +def get_formatter(filename): + if filename in format_file_types: + return format_file_types[filename] + + return format_file_types.get("*" + util.get_fileextn(filename), None) + + +def format_command(commit, files, fix): + ok = 0 + for filepath in files: + filename = util.get_filename(filepath) + filetype = get_formatter(filename) + + if filetype is None: + print("Skip : " + filepath, file=sys.stderr) + continue + + file = filetype.formatter(filepath) + + if fix == "show": + status, diff, stderr = file.diff(commit) + + if stderr != "": + ok = 1 + print(f"Error: {file}", file=sys.stderr) + continue + + if diff != "" and diff != "no modified files to format": + ok = 1 + print(f"Fix : {file}", file=sys.stderr) + print(diff) + else: + print(f"Ok : {file}", file=sys.stderr) + + else: + print(f"Fix : {file}", file=sys.stderr) + if not file.fix(commit): + ok = 1 + print(f"Error: {file}", file=sys.stderr) + + return ok + + +def header_command(commit, files, fix): + options = "-vk" if fix == "show" else "-i" + + status, stdout, stderr = util.run( + f"{SCRIPTS}/license-header.py {options} -", input=files + ) + + if stdout != "": + print(stdout) + + return status + + +def tidy_command(commit, files, fix): + files = [file for file in files if regex.match(r".*\.cpp$", file)] + + if not files: + return 0 + + commit = f"--commit {commit}" if commit != "" else "" + fix = "--fix" if fix == "fix" else "" + + status, stdout, stderr = util.run( + f"{SCRIPTS}/run-clang-tidy.py {commit} {fix} -", input=files + ) + + if stdout != "": + print(stdout) + + return status + + +def get_commit(files): + if files == "commit": + return "HEAD^" + + if files == "main" or files == "master": + return util.run(f"git merge-base origin/{files} HEAD")[1] + + return "" + + +def get_files(commit, path): + filelist = [] + + if commit != "": + status, stdout, stderr = util.run( + f"git diff --relative --name-only --diff-filter='ACM' {commit}" + ) + filelist = stdout.splitlines() + else: + for root, _dirs, files in os.walk(path): + for name in files: + filelist.append(os.path.join(root, name)) + + return [ + file + for file in filelist + if "/data/" not in file + and "velox/external/" not in file + and "build/fbcode_builder" not in file + and "build/deps" not in file + and "cmake-build-debug" not in file + ] + + +def help(args): + parser.print_help() + return 0 + + +def add_check_options(subparser, name): + parser = subparser.add_parser(name) + parser.add_argument("--fix", action="store_const", default="show", const="fix") + return parser + + +def add_options(parser): + files = parser.add_subparsers(dest="files") + + tree_parser = add_check_options(files, "tree") + tree_parser.add_argument("path", default="") + + branch_parser = add_check_options(files, "main") + branch_parser = add_check_options(files, "master") + commit_parser = add_check_options(files, "commit") + + +def add_check_command(parser, name): + subparser = parser.add_parser(name) + add_options(subparser) + + return subparser + + +def parse_args(): + global parser + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description="""Check format/header/tidy + + check.py {format,header,tidy} {commit,branch} [--fix] + check.py {format,header,tidy} {tree} [--fix] PATH +""", + ) + command = parser.add_subparsers(dest="command") + command.add_parser("help") + + format_command_parser = add_check_command(command, "format") + header_command_parser = add_check_command(command, "header") + tidy_command_parser = add_check_command(command, "tidy") + + parser.set_defaults(path="") + parser.set_defaults(command="help") + + return parser.parse_args() + + +def run_command(args, command): + commit = get_commit(args.files) + files = get_files(commit, args.path) + + return command(commit, files, args.fix) + + +def format(args): + return run_command(args, format_command) + + +def header(args): + return run_command(args, header_command) + + +def tidy(args): + return run_command(args, tidy_command) + + +def main(): + args = parse_args() + return globals()[args.command](args) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/git-clang-format b/scripts/git-clang-format new file mode 100755 index 0000000..46e7f5c --- /dev/null +++ b/scripts/git-clang-format @@ -0,0 +1,622 @@ +#!/usr/bin/env python3 +# +#===- git-clang-format - ClangFormat Git Integration ---------*- python -*--===# +# +# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +# +#============================================================================== +#LLVM Release License +#============================================================================== +#University of Illinois/NCSA +#Open Source License +# +#Copyright (c) 2003-2010 University of Illinois at Urbana-Champaign. +#All rights reserved. +# +#Developed by: +# +# LLVM Team +# +# University of Illinois at Urbana-Champaign +# +# http://llvm.org +# +#Permission is hereby granted, free of charge, to any person obtaining a copy of +#this software and associated documentation files (the "Software"), to deal with +#the Software without restriction, including without limitation the rights to +#use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies +#of the Software, and to permit persons to whom the Software is furnished to do +#so, subject to the following conditions: +# +# * Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimers. +# +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimers in the +# documentation and/or other materials provided with the distribution. +# +# * Neither the names of the LLVM Team, University of Illinois at +# Urbana-Champaign, nor the names of its contributors may be used to +# endorse or promote products derived from this Software without specific +# prior written permission. +# +#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +#FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +#CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +#LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +#OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS WITH THE +#SOFTWARE. +#===------------------------------------------------------------------------===# + +r""" +clang-format git integration +============================ + +This file provides a clang-format integration for git. Put it somewhere in your +path and ensure that it is executable. Then, "git clang-format" will invoke +clang-format on the changes in current files or a specific commit. + +For further details, run: +git clang-format -h + +Requires Python 2.7 or Python 3 +""" + +from __future__ import absolute_import, division, print_function +import argparse +import collections +import contextlib +import errno +import os +import re +import subprocess +import sys + +usage = 'git clang-format [OPTIONS] [] [] [--] [...]' + +desc = ''' +If zero or one commits are given, run clang-format on all lines that differ +between the working directory and , which defaults to HEAD. Changes are +only applied to the working directory. + +If two commits are given (requires --diff), run clang-format on all lines in the +second that differ from the first . + +The following git-config settings set the default of the corresponding option: + clangFormat.binary + clangFormat.commit + clangFormat.extension + clangFormat.style +''' + +# Name of the temporary index file in which save the output of clang-format. +# This file is created within the .git directory. +temp_index_basename = 'clang-format-index' + + +Range = collections.namedtuple('Range', 'start, count') + + +def main(): + config = load_git_config() + + # In order to keep '--' yet allow options after positionals, we need to + # check for '--' ourselves. (Setting nargs='*' throws away the '--', while + # nargs=argparse.REMAINDER disallows options after positionals.) + argv = sys.argv[1:] + try: + idx = argv.index('--') + except ValueError: + dash_dash = [] + else: + dash_dash = argv[idx:] + argv = argv[:idx] + + default_extensions = ','.join([ + # From clang/lib/Frontend/FrontendOptions.cpp, all lower case + 'c', 'h', # C + 'm', # ObjC + 'mm', # ObjC++ + 'cc', 'cp', 'cpp', 'c++', 'cxx', 'hpp', # C++ + 'cu', # CUDA + # Other languages that clang-format supports + 'proto', 'protodevel', # Protocol Buffers + 'java', # Java + 'js', # JavaScript + 'ts', # TypeScript + ]) + + p = argparse.ArgumentParser( + usage=usage, formatter_class=argparse.RawDescriptionHelpFormatter, + description=desc) + p.add_argument('--binary', + default=config.get('clangformat.binary', 'clang-format'), + help='path to clang-format'), + p.add_argument('--commit', + default=config.get('clangformat.commit', 'HEAD'), + help='default commit to use if none is specified'), + p.add_argument('--diff', action='store_true', + help='print a diff instead of applying the changes') + p.add_argument('--extensions', + default=config.get('clangformat.extensions', + default_extensions), + help=('comma-separated list of file extensions to format, ' + 'excluding the period and case-insensitive')), + p.add_argument('-f', '--force', action='store_true', + help='allow changes to unstaged files') + p.add_argument('-p', '--patch', action='store_true', + help='select hunks interactively') + p.add_argument('-q', '--quiet', action='count', default=0, + help='print less information') + p.add_argument('--style', + default=config.get('clangformat.style', None), + help='passed to clang-format'), + p.add_argument('-v', '--verbose', action='count', default=0, + help='print extra information') + # We gather all the remaining positional arguments into 'args' since we need + # to use some heuristics to determine whether or not was present. + # However, to print pretty messages, we make use of metavar and help. + p.add_argument('args', nargs='*', metavar='', + help='revision from which to compute the diff') + p.add_argument('ignored', nargs='*', metavar='...', + help='if specified, only consider differences in these files') + opts = p.parse_args(argv) + + opts.verbose -= opts.quiet + del opts.quiet + + commits, files = interpret_args(opts.args, dash_dash, opts.commit) + if len(commits) > 1: + if not opts.diff: + die('--diff is required when two commits are given') + else: + if len(commits) > 2: + die('at most two commits allowed; %d given' % len(commits)) + changed_lines = compute_diff_and_extract_lines(commits, files) + if opts.verbose >= 1: + ignored_files = set(changed_lines) + filter_by_extension(changed_lines, opts.extensions.lower().split(',')) + if opts.verbose >= 1: + ignored_files.difference_update(changed_lines) + if ignored_files: + print('Ignoring changes in the following files (wrong extension):') + for filename in ignored_files: + print(' %s' % filename) + if changed_lines: + print('Running clang-format on the following files:') + for filename in changed_lines: + print(' %s' % filename) + if not changed_lines: + print('no modified files to format') + return + # The computed diff outputs absolute paths, so we must cd before accessing + # those files. + cd_to_toplevel() + if len(commits) > 1: + old_tree = commits[1] + new_tree = run_clang_format_and_save_to_tree(changed_lines, + revision=commits[1], + binary=opts.binary, + style=opts.style) + else: + old_tree = create_tree_from_workdir(changed_lines) + new_tree = run_clang_format_and_save_to_tree(changed_lines, + binary=opts.binary, + style=opts.style) + if opts.verbose >= 1: + print('old tree: %s' % old_tree) + print('new tree: %s' % new_tree) + if old_tree == new_tree: + if opts.verbose >= 0: + print('clang-format did not modify any files') + elif opts.diff: + print_diff(old_tree, new_tree) + else: + changed_files = apply_changes(old_tree, new_tree, force=opts.force, + patch_mode=opts.patch) + if (opts.verbose >= 0 and not opts.patch) or opts.verbose >= 1: + print('changed files:') + for filename in changed_files: + print(' %s' % filename) + + +def load_git_config(non_string_options=None): + """Return the git configuration as a dictionary. + + All options are assumed to be strings unless in `non_string_options`, in which + is a dictionary mapping option name (in lower case) to either "--bool" or + "--int".""" + if non_string_options is None: + non_string_options = {} + out = {} + for entry in run('git', 'config', '--list', '--null').split('\0'): + if entry: + name, value = entry.split('\n', 1) + if name in non_string_options: + value = run('git', 'config', non_string_options[name], name) + out[name] = value + return out + + +def interpret_args(args, dash_dash, default_commit): + """Interpret `args` as "[commits] [--] [files]" and return (commits, files). + + It is assumed that "--" and everything that follows has been removed from + args and placed in `dash_dash`. + + If "--" is present (i.e., `dash_dash` is non-empty), the arguments to its + left (if present) are taken as commits. Otherwise, the arguments are checked + from left to right if they are commits or files. If commits are not given, + a list with `default_commit` is used.""" + if dash_dash: + if len(args) == 0: + commits = [default_commit] + else: + commits = args + for commit in commits: + object_type = get_object_type(commit) + if object_type not in ('commit', 'tag'): + if object_type is None: + die("'%s' is not a commit" % commit) + else: + die("'%s' is a %s, but a commit was expected" % (commit, object_type)) + files = dash_dash[1:] + elif args: + commits = [] + while args: + if not disambiguate_revision(args[0]): + break + commits.append(args.pop(0)) + if not commits: + commits = [default_commit] + files = args + else: + commits = [default_commit] + files = [] + return commits, files + + +def disambiguate_revision(value): + """Returns True if `value` is a revision, False if it is a file, or dies.""" + # If `value` is ambiguous (neither a commit nor a file), the following + # command will die with an appropriate error message. + run('git', 'rev-parse', value, verbose=False) + object_type = get_object_type(value) + if object_type is None: + return False + if object_type in ('commit', 'tag'): + return True + die('`%s` is a %s, but a commit or filename was expected' % + (value, object_type)) + + +def get_object_type(value): + """Returns a string description of an object's type, or None if it is not + a valid git object.""" + cmd = ['git', 'cat-file', '-t', value] + p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + stdout, stderr = p.communicate() + if p.returncode != 0: + return None + return convert_string(stdout.strip()) + + +def compute_diff_and_extract_lines(commits, files): + """Calls compute_diff() followed by extract_lines().""" + diff_process = compute_diff(commits, files) + changed_lines = extract_lines(diff_process.stdout) + diff_process.stdout.close() + diff_process.wait() + if diff_process.returncode != 0: + # Assume error was already printed to stderr. + sys.exit(2) + return changed_lines + + +def compute_diff(commits, files): + """Return a subprocess object producing the diff from `commits`. + + The return value's `stdin` file object will produce a patch with the + differences between the working directory and the first commit if a single + one was specified, or the difference between both specified commits, filtered + on `files` (if non-empty). Zero context lines are used in the patch.""" + git_tool = 'diff-index' + if len(commits) > 1: + git_tool = 'diff-tree' + cmd = ['git', git_tool, '-p', '-U0'] + commits + ['--'] + cmd.extend(files) + p = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE) + p.stdin.close() + return p + + +def extract_lines(patch_file): + """Extract the changed lines in `patch_file`. + + The return value is a dictionary mapping filename to a list of (start_line, + line_count) pairs. + + The input must have been produced with ``-U0``, meaning unidiff format with + zero lines of context. The return value is a dict mapping filename to a + list of line `Range`s.""" + matches = {} + for line in patch_file: + line = convert_string(line) + match = re.search(r'^\+\+\+\ [^/]+/(.*)', line) + if match: + filename = match.group(1).rstrip('\r\n') + match = re.search(r'^@@ -[0-9,]+ \+(\d+)(,(\d+))?', line) + if match: + start_line = int(match.group(1)) + line_count = 1 + if match.group(3): + line_count = int(match.group(3)) + if line_count > 0: + matches.setdefault(filename, []).append(Range(start_line, line_count)) + return matches + + +def filter_by_extension(dictionary, allowed_extensions): + """Delete every key in `dictionary` that doesn't have an allowed extension. + + `allowed_extensions` must be a collection of lowercase file extensions, + excluding the period.""" + allowed_extensions = frozenset(allowed_extensions) + for filename in list(dictionary.keys()): + base_ext = filename.rsplit('.', 1) + if len(base_ext) == 1 and '' in allowed_extensions: + continue + if len(base_ext) == 1 or base_ext[1].lower() not in allowed_extensions: + del dictionary[filename] + + +def cd_to_toplevel(): + """Change to the top level of the git repository.""" + toplevel = run('git', 'rev-parse', '--show-toplevel') + os.chdir(toplevel) + + +def create_tree_from_workdir(filenames): + """Create a new git tree with the given files from the working directory. + + Returns the object ID (SHA-1) of the created tree.""" + return create_tree(filenames, '--stdin') + + +def run_clang_format_and_save_to_tree(changed_lines, revision=None, + binary='clang-format', style=None): + """Run clang-format on each file and save the result to a git tree. + + Returns the object ID (SHA-1) of the created tree.""" + def iteritems(container): + try: + return container.iteritems() # Python 2 + except AttributeError: + return container.items() # Python 3 + def index_info_generator(): + for filename, line_ranges in iteritems(changed_lines): + if revision: + git_metadata_cmd = ['git', 'ls-tree', + '%s:%s' % (revision, os.path.dirname(filename)), + os.path.basename(filename)] + git_metadata = subprocess.Popen(git_metadata_cmd, stdin=subprocess.PIPE, + stdout=subprocess.PIPE) + stdout = git_metadata.communicate()[0] + mode = oct(int(stdout.split()[0], 8)) + else: + mode = oct(os.stat(filename).st_mode) + # Adjust python3 octal format so that it matches what git expects + if mode.startswith('0o'): + mode = '0' + mode[2:] + blob_id = clang_format_to_blob(filename, line_ranges, + revision=revision, + binary=binary, + style=style) + yield '%s %s\t%s' % (mode, blob_id, filename) + return create_tree(index_info_generator(), '--index-info') + + +def create_tree(input_lines, mode): + """Create a tree object from the given input. + + If mode is '--stdin', it must be a list of filenames. If mode is + '--index-info' is must be a list of values suitable for "git update-index + --index-info", such as " ". Any other mode + is invalid.""" + assert mode in ('--stdin', '--index-info') + cmd = ['git', 'update-index', '--add', '-z', mode] + with temporary_index_file(): + p = subprocess.Popen(cmd, stdin=subprocess.PIPE) + for line in input_lines: + p.stdin.write(to_bytes('%s\0' % line)) + p.stdin.close() + if p.wait() != 0: + die('`%s` failed' % ' '.join(cmd)) + tree_id = run('git', 'write-tree') + return tree_id + + +def clang_format_to_blob(filename, line_ranges, revision=None, + binary='clang-format', style=None): + """Run clang-format on the given file and save the result to a git blob. + + Runs on the file in `revision` if not None, or on the file in the working + directory if `revision` is None. + + Returns the object ID (SHA-1) of the created blob.""" + clang_format_cmd = [binary] + if style: + clang_format_cmd.extend(['-style='+style]) + clang_format_cmd.extend([ + '-lines=%s:%s' % (start_line, start_line+line_count-1) + for start_line, line_count in line_ranges]) + if revision: + clang_format_cmd.extend(['-assume-filename='+filename]) + git_show_cmd = ['git', 'cat-file', 'blob', '%s:%s' % (revision, filename)] + git_show = subprocess.Popen(git_show_cmd, stdin=subprocess.PIPE, + stdout=subprocess.PIPE) + git_show.stdin.close() + clang_format_stdin = git_show.stdout + else: + clang_format_cmd.extend([filename]) + git_show = None + clang_format_stdin = subprocess.PIPE + try: + clang_format = subprocess.Popen(clang_format_cmd, stdin=clang_format_stdin, + stdout=subprocess.PIPE) + if clang_format_stdin == subprocess.PIPE: + clang_format_stdin = clang_format.stdin + except OSError as e: + if e.errno == errno.ENOENT: + die('cannot find executable "%s"' % binary) + else: + raise + clang_format_stdin.close() + hash_object_cmd = ['git', 'hash-object', '-w', '--path='+filename, '--stdin'] + hash_object = subprocess.Popen(hash_object_cmd, stdin=clang_format.stdout, + stdout=subprocess.PIPE) + clang_format.stdout.close() + stdout = hash_object.communicate()[0] + if hash_object.returncode != 0: + die('`%s` failed' % ' '.join(hash_object_cmd)) + if clang_format.wait() != 0: + die('`%s` failed' % ' '.join(clang_format_cmd)) + if git_show and git_show.wait() != 0: + die('`%s` failed' % ' '.join(git_show_cmd)) + return convert_string(stdout).rstrip('\r\n') + + +@contextlib.contextmanager +def temporary_index_file(tree=None): + """Context manager for setting GIT_INDEX_FILE to a temporary file and deleting + the file afterward.""" + index_path = create_temporary_index(tree) + old_index_path = os.environ.get('GIT_INDEX_FILE') + os.environ['GIT_INDEX_FILE'] = index_path + try: + yield + finally: + if old_index_path is None: + del os.environ['GIT_INDEX_FILE'] + else: + os.environ['GIT_INDEX_FILE'] = old_index_path + os.remove(index_path) + + +def create_temporary_index(tree=None): + """Create a temporary index file and return the created file's path. + + If `tree` is not None, use that as the tree to read in. Otherwise, an + empty index is created.""" + gitdir = run('git', 'rev-parse', '--git-dir') + path = os.path.join(gitdir, temp_index_basename) + if tree is None: + tree = '--empty' + run('git', 'read-tree', '--index-output='+path, tree) + return path + + +def print_diff(old_tree, new_tree): + """Print the diff between the two trees to stdout.""" + # We use the porcelain 'diff' and not plumbing 'diff-tree' because the output + # is expected to be viewed by the user, and only the former does nice things + # like color and pagination. + # + # We also only print modified files since `new_tree` only contains the files + # that were modified, so unmodified files would show as deleted without the + # filter. + subprocess.check_call(['git', 'diff', '--diff-filter=M', old_tree, new_tree, + '--']) + + +def apply_changes(old_tree, new_tree, force=False, patch_mode=False): + """Apply the changes in `new_tree` to the working directory. + + Bails if there are local changes in those files and not `force`. If + `patch_mode`, runs `git checkout --patch` to select hunks interactively.""" + changed_files = run('git', 'diff-tree', '--diff-filter=M', '-r', '-z', + '--name-only', old_tree, + new_tree).rstrip('\0').split('\0') + if not force: + unstaged_files = run('git', 'diff-files', '--name-status', *changed_files) + if unstaged_files: + print('The following files would be modified but ' + 'have unstaged changes:', file=sys.stderr) + print(unstaged_files, file=sys.stderr) + print('Please commit, stage, or stash them first.', file=sys.stderr) + sys.exit(2) + if patch_mode: + # In patch mode, we could just as well create an index from the new tree + # and checkout from that, but then the user will be presented with a + # message saying "Discard ... from worktree". Instead, we use the old + # tree as the index and checkout from new_tree, which gives the slightly + # better message, "Apply ... to index and worktree". This is not quite + # right, since it won't be applied to the user's index, but oh well. + with temporary_index_file(old_tree): + subprocess.check_call(['git', 'checkout', '--patch', new_tree]) + index_tree = old_tree + else: + with temporary_index_file(new_tree): + run('git', 'checkout-index', '-a', '-f') + return changed_files + + +def run(*args, **kwargs): + stdin = kwargs.pop('stdin', '') + verbose = kwargs.pop('verbose', True) + strip = kwargs.pop('strip', True) + for name in kwargs: + raise TypeError("run() got an unexpected keyword argument '%s'" % name) + p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, + stdin=subprocess.PIPE) + stdout, stderr = p.communicate(input=stdin) + + stdout = convert_string(stdout) + stderr = convert_string(stderr) + + if p.returncode == 0: + if stderr: + if verbose: + print('`%s` printed to stderr:' % ' '.join(args), file=sys.stderr) + print(stderr.rstrip(), file=sys.stderr) + if strip: + stdout = stdout.rstrip('\r\n') + return stdout + if verbose: + print('`%s` returned %s' % (' '.join(args), p.returncode), file=sys.stderr) + if stderr: + print(stderr.rstrip(), file=sys.stderr) + sys.exit(2) + + +def die(message): + print('error:', message, file=sys.stderr) + sys.exit(2) + + +def to_bytes(str_input): + # Encode to UTF-8 to get binary data. + if isinstance(str_input, bytes): + return str_input + return str_input.encode('utf-8') + + +def to_string(bytes_input): + if isinstance(bytes_input, str): + return bytes_input + return bytes_input.encode('utf-8') + + +def convert_string(bytes_input): + try: + return to_string(bytes_input.decode('utf-8')) + except AttributeError: # 'str' object has no attribute 'decode'. + return str(bytes_input) + except UnicodeError: + return str(bytes_input) + +if __name__ == '__main__': + main() diff --git a/scripts/license-header.py b/scripts/license-header.py new file mode 100755 index 0000000..8ddf825 --- /dev/null +++ b/scripts/license-header.py @@ -0,0 +1,259 @@ +#!/usr/bin/env python3 +# Copyright (c) Facebook, Inc. and its affiliates. +# +# Licensed 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. + +import argparse +import fnmatch +import os +import sys +from collections import OrderedDict + +import regex + + +class attrdict(dict): + __getattr__ = dict.__getitem__ + __setattr__ = dict.__setitem__ + + +def parse_args(): + parser = argparse.ArgumentParser(description="Update license headers") + parser.add_argument("--header", default="license.header", help="header file") + parser.add_argument( + "--extra", + default=80, + help="extra characters past beginning of file to look for header", + ) + parser.add_argument( + "--editdist", default=12, type=int, help="max edit distance between headers" + ) + parser.add_argument( + "--remove", default=False, action="store_true", help="remove the header" + ) + parser.add_argument( + "--cslash", + default=False, + action="store_true", + help='use C slash "//" style comments', + ) + parser.add_argument( + "-v", default=False, action="store_true", dest="verbose", help="verbose output" + ) + + group = parser.add_mutually_exclusive_group() + group.add_argument( + "-k", default=False, action="store_true", dest="check", help="check headers" + ) + group.add_argument( + "-i", + default=False, + action="store_true", + dest="inplace", + help="edit file inplace", + ) + + parser.add_argument("files", metavar="FILES", nargs="+", help="files to process") + + return parser.parse_args() + + +def file_read(filename): + with open(filename) as file: + return file.read() + + +def file_lines(filename): + return file_read(filename).rstrip().split("\n") + + +def wrapper(prefix, leader, suffix, header): + return prefix + "\n".join([leader + line for line in header]) + suffix + + +def wrapper_chpp(header, args): + if args.cslash: + return wrapper("", "//", "\n", header) + else: + return wrapper("/*\n", " *", "\n */\n", header) + + +def wrapper_hash(header, args): + return wrapper("", "#", "\n", header) + + +file_types = OrderedDict( + { + "CMakeLists.txt": attrdict({"wrapper": wrapper_hash, "hashbang": False}), + "Makefile": attrdict({"wrapper": wrapper_hash, "hashbang": False}), + "*.cpp": attrdict({"wrapper": wrapper_chpp, "hashbang": False}), + "*.dockfile": attrdict({"wrapper": wrapper_hash, "hashbang": False}), + "*.h": attrdict({"wrapper": wrapper_chpp, "hashbang": False}), + "*.inc": attrdict({"wrapper": wrapper_chpp, "hashbang": False}), + "*.java": attrdict({"wrapper": wrapper_chpp, "hashbang": False}), + "*.prolog": attrdict({"wrapper": wrapper_chpp, "hashbang": False}), + "*.py": attrdict({"wrapper": wrapper_hash, "hashbang": True}), + "*.sh": attrdict({"wrapper": wrapper_hash, "hashbang": True}), + "*.thrift": attrdict({"wrapper": wrapper_chpp, "hashbang": False}), + "*.txt": attrdict({"wrapper": wrapper_hash, "hashbang": True}), + "*.yml": attrdict({"wrapper": wrapper_hash, "hashbang": False}), + } +) + +file_pattern = regex.compile( + "|".join(["^" + fnmatch.translate(type) + "$" for type in file_types.keys()]) +) + + +def get_filename(filename): + return os.path.basename(filename) + + +def get_fileextn(filename): + split = os.path.splitext(filename) + if len(split) <= 1: + return "" + + return split[-1] + + +def get_wrapper(filename): + if filename in file_types: + return file_types[filename] + + return file_types["*" + get_fileextn(filename)] + + +def message(file, string): + if file: + print(string, file=file) + + +def main(): + fail = False + log_to = None + + args = parse_args() + + if args.verbose: + log_to = sys.stderr + + if args.check: + log_to = None + + if args.verbose: + log_to = sys.stdout + + header_text = file_lines(args.header) + + if len(args.files) == 1 and args.files[0] == "-": + files = [file.strip() for file in sys.stdin.readlines()] + else: + files = args.files + + for filepath in files: + filename = get_filename(filepath) + + matched = file_pattern.match(filename) + + if not matched: + message(log_to, "Skip : " + filepath) + continue + + content = file_read(filepath) + wrap = get_wrapper(filename) + + header_comment = wrap.wrapper(header_text, args) + + start = 0 + end = 0 + + # Look for an exact substr match + # + found = content.find(header_comment, 0, len(header_comment) + args.extra) + if found >= 0: + if not args.remove: + message(log_to, "OK : " + filepath) + continue + + start = found + end = found + len(header_comment) + else: + # Look for a fuzzy match in the first 60 chars + # + found = regex.search( + "(?be)(%s){e<=%d}" % (regex.escape(header_comment[0:60]), 6), + content[0 : 80 + args.extra], + ) + if found: + fuzzy = regex.compile( + "(?be)(%s){e<=%d}" % (regex.escape(header_comment), args.editdist) + ) + + # If the first 80 chars match - try harder for the rest of the header + # + found = fuzzy.search( + content[0 : len(header_comment) + args.extra], found.start() + ) + if found: + start = found.start() + end = found.end() + + if args.remove: + if start == 0 and end == 0: + if not args.inplace: + print(content, end="") + + message(log_to, "OK : " + filepath) + continue + + # If removing the header text, zero it out there. + header_comment = "" + + message(log_to, "Fix : " + filepath) + + if args.check: + fail = True + continue + + # Remove any partially matching header + # + content = content[0:start] + content[end:] + + if wrap.hashbang: + search = regex.search("^#!.*\n", content) + if search: + content = ( + content[search.start() : search.end()] + + header_comment + + content[search.end() :] + ) + else: + content = header_comment + content + else: + content = header_comment + content + + if args.inplace: + with open(filepath, "w") as file: + print(content, file=file, end="") + else: + print(content, end="") + + if fail: + return 1 + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/util.py b/scripts/util.py new file mode 100644 index 0000000..0af7812 --- /dev/null +++ b/scripts/util.py @@ -0,0 +1,90 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# Licensed 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. + +import gzip +import json +import os +import subprocess +import sys + +import regex + + +class attrdict(dict): + __getattr__ = dict.__getitem__ + __setattr__ = dict.__setitem__ + + +class string(str): + def extract(self, rexp): + return regex.match(rexp, self).group(1) + + def json(self): + return json.loads(self, object_hook=attrdict) + + +def run(command, compressed=False, **kwargs): + if "input" in kwargs: + input = kwargs["input"] + + if type(input) is list: + input = "\n".join(input) + "\n" + + kwargs["input"] = input.encode("utf-8") + + reply = subprocess.run( + command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, **kwargs + ) + + if compressed: + stdout = gzip.decompress(reply.stdout) + else: + stdout = reply.stdout + + stdout = ( + string(stdout.decode("utf-8", errors="ignore").strip()) + if stdout is not None + else "" + ) + stderr = ( + string(reply.stderr.decode("utf-8").strip()) if reply.stderr is not None else "" + ) + + if stderr != "": + print(stderr, file=sys.stderr) + + return reply.returncode, stdout, stderr + + +def get_filename(filename): + return os.path.basename(filename) + + +def get_fileextn(filename): + split = os.path.splitext(filename) + if len(split) <= 1: + return "" + + return split[-1] + + +def script_path(): + return os.path.dirname(os.path.realpath(sys.argv[0])) + + +def input_files(files): + if len(files) == 1 and files[0] == "-": + return [file.strip() for file in sys.stdin.readlines()] + else: + return files diff --git a/velox b/velox new file mode 160000 index 0000000..dadade0 --- /dev/null +++ b/velox @@ -0,0 +1 @@ +Subproject commit dadade0d423c1458052e84c00cd4fd3aa7d0b9bc