You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by ap...@apache.org on 2018/10/24 09:15:44 UTC
[arrow] branch master updated: ARROW-3381: [C++] Add bz2 codec
This is an automated email from the ASF dual-hosted git repository.
apitrou pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new 536e4f4 ARROW-3381: [C++] Add bz2 codec
536e4f4 is described below
commit 536e4f4c9fb861a9bdc1faf3a1d251e530b9923c
Author: Antoine Pitrou <an...@python.org>
AuthorDate: Wed Oct 24 11:15:35 2018 +0200
ARROW-3381: [C++] Add bz2 codec
Author: Antoine Pitrou <an...@python.org>
Closes #2812 from pitrou/ARROW-3381-bz2-codec and squashes the following commits:
1adbe8a2 <Antoine Pitrou> Pin flake8 (work around ARROW-3606)
9075dff5 <Antoine Pitrou> Fix checking return code
a54ee5b0 <Antoine Pitrou> ARROW-3381: Add bz2 codec
---
ci/appveyor-cpp-build.bat | 1 +
ci/cpp-msvc-build-main.bat | 7 +-
ci/travis_before_script_cpp.sh | 1 +
ci/travis_install_toolchain.sh | 1 +
ci/travis_lint.sh | 6 +-
cpp/CMakeLists.txt | 8 +
cpp/cmake_modules/FindBz2.cmake | 57 ++++++
cpp/cmake_modules/FindZSTD.cmake | 1 -
cpp/cmake_modules/ThirdpartyToolchain.cmake | 20 +++
cpp/src/arrow/CMakeLists.txt | 5 +
cpp/src/arrow/io/io-compressed-test.cc | 2 +
cpp/src/arrow/symbols.map | 2 +
cpp/src/arrow/util/CMakeLists.txt | 1 +
cpp/src/arrow/util/compression-test.cc | 17 ++
cpp/src/arrow/util/compression.cc | 11 ++
cpp/src/arrow/util/compression.h | 2 +-
cpp/src/arrow/util/compression_bz2.cc | 259 ++++++++++++++++++++++++++++
cpp/src/arrow/util/compression_bz2.h | 52 ++++++
python/pyarrow/includes/libarrow.pxd | 1 +
python/pyarrow/io.pxi | 2 +
python/pyarrow/tests/test_csv.py | 15 +-
21 files changed, 464 insertions(+), 7 deletions(-)
diff --git a/ci/appveyor-cpp-build.bat b/ci/appveyor-cpp-build.bat
index 3589871..bce722f 100644
--- a/ci/appveyor-cpp-build.bat
+++ b/ci/appveyor-cpp-build.bat
@@ -99,6 +99,7 @@ if "%JOB%" == "Toolchain" (
@rem Install pre-built "toolchain" packages for faster builds
conda install -q -y ^
brotli ^
+ bzip2 ^
cmake ^
flatbuffers ^
gflags ^
diff --git a/ci/cpp-msvc-build-main.bat b/ci/cpp-msvc-build-main.bat
index 5ddb648..e5eef1e 100644
--- a/ci/cpp-msvc-build-main.bat
+++ b/ci/cpp-msvc-build-main.bat
@@ -19,6 +19,11 @@
@rem (i.e. for usual configurations)
set ARROW_HOME=%CONDA_PREFIX%\Library
+set CMAKE_ARGS=
+
+if "%JOB%" == "Toolchain" (
+ set CMAKE_ARGS=%CMAKE_ARGS% -DARROW_WITH_BZ2=ON
+)
@rem Retrieve git submodules, configure env var for Parquet unit tests
git submodule update --init || exit /B
@@ -38,7 +43,7 @@ set CMAKE_CXX_FLAGS_RELEASE=/Od /UNDEBUG
mkdir cpp\build
pushd cpp\build
-cmake -G "%GENERATOR%" ^
+cmake -G "%GENERATOR%" %CMAKE_ARGS% ^
-DCMAKE_INSTALL_PREFIX=%CONDA_PREFIX%\Library ^
-DARROW_BOOST_USE_SHARED=OFF ^
-DCMAKE_BUILD_TYPE=%CONFIGURATION% ^
diff --git a/ci/travis_before_script_cpp.sh b/ci/travis_before_script_cpp.sh
index 1b389a2..2da4c4b 100755
--- a/ci/travis_before_script_cpp.sh
+++ b/ci/travis_before_script_cpp.sh
@@ -52,6 +52,7 @@ if [ "$ARROW_TRAVIS_USE_TOOLCHAIN" == "1" ]; then
# Set up C++ toolchain from conda-forge packages for faster builds
source $TRAVIS_BUILD_DIR/ci/travis_install_toolchain.sh
CMAKE_COMMON_FLAGS="${CMAKE_COMMON_FLAGS} -DARROW_JEMALLOC=ON"
+ CMAKE_COMMON_FLAGS="${CMAKE_COMMON_FLAGS} -DARROW_WITH_BZ2=ON"
fi
mkdir -p $ARROW_CPP_BUILD_DIR
diff --git a/ci/travis_install_toolchain.sh b/ci/travis_install_toolchain.sh
index eb4bdf2..fa3b710 100755
--- a/ci/travis_install_toolchain.sh
+++ b/ci/travis_install_toolchain.sh
@@ -26,6 +26,7 @@ if [ ! -e $CPP_TOOLCHAIN ]; then
conda create -y -q -p $CPP_TOOLCHAIN python=3.6 nomkl \
boost-cpp \
brotli \
+ bzip2 \
ccache \
cmake \
curl \
diff --git a/ci/travis_lint.sh b/ci/travis_lint.sh
index b8059f0..4a4578f 100755
--- a/ci/travis_lint.sh
+++ b/ci/travis_lint.sh
@@ -46,17 +46,17 @@ fi
if [ "$ARROW_CI_DEV_AFFECTED" != "0" ]; then
# crossbow requires python3
sudo apt-get install -y -q python3-pip
- sudo pip3 install -q flake8
+ sudo pip3 install -q flake8==3.5
python3 -m flake8 --count $ARROW_CROSSBOW_DIR
fi
if [ "$ARROW_CI_INTEGRATION_AFFECTED" != "0" ]; then
- sudo pip install -q flake8
+ sudo pip install -q flake8==3.5
python -m flake8 --count $ARROW_INTEGRATION_DIR
fi
if [ "$ARROW_CI_PYTHON_AFFECTED" != "0" ]; then
- sudo pip install -q flake8
+ sudo pip install -q flake8==3.5
python -m flake8 --count $ARROW_PYTHON_DIR
diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index e14f877..02499a1 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -228,6 +228,10 @@ Pass multiple labels by dividing with semicolons")
"Build with Brotli compression"
ON)
+ option(ARROW_WITH_BZ2
+ "Build with BZ2 compression"
+ OFF)
+
option(ARROW_WITH_LZ4
"Build with lz4 compression"
ON)
@@ -613,6 +617,10 @@ if (ARROW_WITH_BROTLI)
${ARROW_STATIC_LINK_LIBS})
endif()
+if (ARROW_WITH_BZ2)
+ SET(ARROW_STATIC_LINK_LIBS bz2_static ${ARROW_STATIC_LINK_LIBS})
+endif()
+
if (ARROW_WITH_LZ4)
SET(ARROW_STATIC_LINK_LIBS lz4_static ${ARROW_STATIC_LINK_LIBS})
endif()
diff --git a/cpp/cmake_modules/FindBz2.cmake b/cpp/cmake_modules/FindBz2.cmake
new file mode 100644
index 0000000..c26b8b7
--- /dev/null
+++ b/cpp/cmake_modules/FindBz2.cmake
@@ -0,0 +1,57 @@
+#
+# 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.
+#
+# Tries to find bz2 headers and libraries.
+#
+# Usage of this module as follows:
+#
+# find_package(Bz2)
+#
+# Variables used by this module, they can change the default behaviour and need
+# to be set before calling find_package:
+#
+# BZ2_HOME - When set, this path is inspected instead of standard library
+# locations as the root of the bz2 installation.
+#
+# This module defines
+# BZ2_INCLUDE_DIR, directory containing headers
+# BZ2_STATIC_LIB, path to libbz2 static library
+# BZ2_FOUND, whether bz2 has been found
+
+if( NOT "${BZ2_HOME}" STREQUAL "")
+ file( TO_CMAKE_PATH "${BZ2_HOME}" _native_path )
+ list( APPEND _bz2_roots ${_native_path} )
+elseif ( BZ2_HOME )
+ list( APPEND _bz2_roots ${BZ2_HOME} )
+endif()
+
+if (MSVC)
+ set(BZ2_STATIC_LIB_NAME
+ ${CMAKE_STATIC_LIBRARY_PREFIX}bz2_static${CMAKE_STATIC_LIBRARY_SUFFIX})
+else()
+ set(BZ2_STATIC_LIB_NAME
+ ${CMAKE_STATIC_LIBRARY_PREFIX}bz2${CMAKE_STATIC_LIBRARY_SUFFIX})
+endif()
+
+find_path(BZ2_INCLUDE_DIR NAMES bzlib.h
+ PATHS ${_bz2_roots}
+ NO_DEFAULT_PATH
+ PATH_SUFFIXES "include" )
+find_library(BZ2_STATIC_LIB NAMES ${BZ2_STATIC_LIB_NAME} lib${BZ2_STATIC_LIB_NAME}
+ PATHS ${_bz2_roots}
+ NO_DEFAULT_PATH
+ PATH_SUFFIXES "lib" )
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(BZ2 REQUIRED_VARS
+ BZ2_STATIC_LIB BZ2_INCLUDE_DIR)
diff --git a/cpp/cmake_modules/FindZSTD.cmake b/cpp/cmake_modules/FindZSTD.cmake
index 810e5c0..7f8d131 100644
--- a/cpp/cmake_modules/FindZSTD.cmake
+++ b/cpp/cmake_modules/FindZSTD.cmake
@@ -18,7 +18,6 @@
# - Find ZSTD (zstd.h, libzstd.a, libzstd.so, and libzstd.so.0)
# This module defines
# ZSTD_INCLUDE_DIR, directory containing headers
-# ZSTD_SHARED_LIB, path to libzstd shared library
# ZSTD_STATIC_LIB, path to libzstd static library
# ZSTD_FOUND, whether zstd has been found
diff --git a/cpp/cmake_modules/ThirdpartyToolchain.cmake b/cpp/cmake_modules/ThirdpartyToolchain.cmake
index b461764..c35dd4a 100644
--- a/cpp/cmake_modules/ThirdpartyToolchain.cmake
+++ b/cpp/cmake_modules/ThirdpartyToolchain.cmake
@@ -28,6 +28,7 @@ set(THIRDPARTY_DIR "${arrow_SOURCE_DIR}/thirdparty")
if (NOT "$ENV{ARROW_BUILD_TOOLCHAIN}" STREQUAL "")
set(BROTLI_HOME "$ENV{ARROW_BUILD_TOOLCHAIN}")
+ set(BZ2_HOME "$ENV{ARROW_BUILD_TOOLCHAIN}")
set(FLATBUFFERS_HOME "$ENV{ARROW_BUILD_TOOLCHAIN}")
set(GFLAGS_HOME "$ENV{ARROW_BUILD_TOOLCHAIN}")
set(GLOG_HOME "$ENV{ARROW_BUILD_TOOLCHAIN}")
@@ -61,6 +62,10 @@ if (DEFINED ENV{BROTLI_HOME})
set(BROTLI_HOME "$ENV{BROTLI_HOME}")
endif()
+if (DEFINED ENV{BZ2_HOME})
+ set(BZ2_HOME "$ENV{BZ2_HOME}")
+endif()
+
if (DEFINED ENV{DOUBLE_CONVERSION_HOME})
set(DOUBLE_CONVERSION_HOME "$ENV{DOUBLE_CONVERSION_HOME}")
endif()
@@ -975,6 +980,21 @@ if (ARROW_WITH_BROTLI)
endif()
endif()
+if (ARROW_WITH_BZ2)
+# ----------------------------------------------------------------------
+# BZ2
+
+ if("${BZ2_HOME}" STREQUAL "")
+ message(SEND_ERROR "a binary install of libbz2 must be present, please set the BZ2_HOME environment variable")
+ else()
+ find_package(Bz2 REQUIRED)
+ endif()
+
+ include_directories(SYSTEM ${BZ2_INCLUDE_DIR})
+ ADD_THIRDPARTY_LIB(bz2_static
+ STATIC_LIB ${BZ2_STATIC_LIB})
+endif()
+
if (ARROW_WITH_LZ4)
# ----------------------------------------------------------------------
# Lz4
diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt
index 9a2f3a4..63c1bf2 100644
--- a/cpp/src/arrow/CMakeLists.txt
+++ b/cpp/src/arrow/CMakeLists.txt
@@ -89,6 +89,11 @@ if (ARROW_WITH_BROTLI)
SET(ARROW_SRCS util/compression_brotli.cc ${ARROW_SRCS})
endif()
+if (ARROW_WITH_BZ2)
+ add_definitions(-DARROW_WITH_BZ2)
+ SET(ARROW_SRCS util/compression_bz2.cc ${ARROW_SRCS})
+endif()
+
if (ARROW_WITH_LZ4)
add_definitions(-DARROW_WITH_LZ4)
SET(ARROW_SRCS util/compression_lz4.cc ${ARROW_SRCS})
diff --git a/cpp/src/arrow/io/io-compressed-test.cc b/cpp/src/arrow/io/io-compressed-test.cc
index 78240da..4a3b323 100644
--- a/cpp/src/arrow/io/io-compressed-test.cc
+++ b/cpp/src/arrow/io/io-compressed-test.cc
@@ -191,6 +191,8 @@ TEST_P(CompressedInputStreamTest, InvalidData) {
// NOTE: Snappy doesn't support streaming decompression
+// NOTE: BZ2 doesn't support one-shot compression
+
// NOTE: LZ4 streaming decompression uses the LZ4 framing format,
// which must be tested against a streaming compressor
diff --git a/cpp/src/arrow/symbols.map b/cpp/src/arrow/symbols.map
index 18802ff..083b4b2 100644
--- a/cpp/src/arrow/symbols.map
+++ b/cpp/src/arrow/symbols.map
@@ -47,6 +47,8 @@
zlibCompileFlags;
zlibVersion;
_tr_*;
+ # bz2
+ BZ2_*;
# lz4
LZ4_*;
LZ4F_*;
diff --git a/cpp/src/arrow/util/CMakeLists.txt b/cpp/src/arrow/util/CMakeLists.txt
index 69e18a6..28c6116 100644
--- a/cpp/src/arrow/util/CMakeLists.txt
+++ b/cpp/src/arrow/util/CMakeLists.txt
@@ -28,6 +28,7 @@ install(FILES
compiler-util.h
compression.h
compression_brotli.h
+ compression_bz2.h
compression_lz4.h
compression_snappy.h
compression_zlib.h
diff --git a/cpp/src/arrow/util/compression-test.cc b/cpp/src/arrow/util/compression-test.cc
index 10447c6..899238b 100644
--- a/cpp/src/arrow/util/compression-test.cc
+++ b/cpp/src/arrow/util/compression-test.cc
@@ -314,6 +314,11 @@ class CodecTest : public ::testing::TestWithParam<Compression::type> {
};
TEST_P(CodecTest, CodecRoundtrip) {
+ if (GetCompression() == Compression::BZ2) {
+ // SKIP: BZ2 doesn't support one-shot compression
+ return;
+ }
+
int sizes[] = {0, 10000, 100000};
for (int data_size : sizes) {
vector<uint8_t> data = MakeRandomData(data_size);
@@ -329,6 +334,10 @@ TEST_P(CodecTest, StreamingCompressor) {
// SKIP: snappy doesn't support streaming compression
return;
}
+ if (GetCompression() == Compression::BZ2) {
+ // SKIP: BZ2 doesn't support one-shot decompression
+ return;
+ }
if (GetCompression() == Compression::LZ4) {
// SKIP: LZ4 streaming compression uses the LZ4 framing format,
// which must be tested against a streaming decompressor
@@ -352,6 +361,10 @@ TEST_P(CodecTest, StreamingDecompressor) {
// SKIP: snappy doesn't support streaming decompression
return;
}
+ if (GetCompression() == Compression::BZ2) {
+ // SKIP: BZ2 doesn't support one-shot compression
+ return;
+ }
if (GetCompression() == Compression::LZ4) {
// SKIP: LZ4 streaming decompression uses the LZ4 framing format,
// which must be tested against a streaming compressor
@@ -398,5 +411,9 @@ INSTANTIATE_TEST_CASE_P(TestLZ4, CodecTest, ::testing::Values(Compression::LZ4))
INSTANTIATE_TEST_CASE_P(TestBrotli, CodecTest, ::testing::Values(Compression::BROTLI));
+#if ARROW_WITH_BZ2
+INSTANTIATE_TEST_CASE_P(TestBZ2, CodecTest, ::testing::Values(Compression::BZ2));
+#endif
+
} // namespace util
} // namespace arrow
diff --git a/cpp/src/arrow/util/compression.cc b/cpp/src/arrow/util/compression.cc
index 8d4d838..7b13ade 100644
--- a/cpp/src/arrow/util/compression.cc
+++ b/cpp/src/arrow/util/compression.cc
@@ -39,6 +39,10 @@
#include "arrow/util/compression_zstd.h"
#endif
+#ifdef ARROW_WITH_BZ2
+#include "arrow/util/compression_bz2.h"
+#endif
+
#include "arrow/status.h"
namespace arrow {
@@ -91,6 +95,13 @@ Status Codec::Create(Compression::type codec_type, std::unique_ptr<Codec>* resul
return Status::NotImplemented("ZSTD codec support not built");
#endif
break;
+ case Compression::BZ2:
+#ifdef ARROW_WITH_BZ2
+ result->reset(new BZ2Codec());
+#else
+ return Status::NotImplemented("BZ2 codec support not built");
+#endif
+ break;
default:
return Status::Invalid("Unrecognized codec");
}
diff --git a/cpp/src/arrow/util/compression.h b/cpp/src/arrow/util/compression.h
index f34ae43..f935905 100644
--- a/cpp/src/arrow/util/compression.h
+++ b/cpp/src/arrow/util/compression.h
@@ -28,7 +28,7 @@ namespace arrow {
class Status;
struct Compression {
- enum type { UNCOMPRESSED, SNAPPY, GZIP, BROTLI, ZSTD, LZ4, LZO };
+ enum type { UNCOMPRESSED, SNAPPY, GZIP, BROTLI, ZSTD, LZ4, LZO, BZ2 };
};
namespace util {
diff --git a/cpp/src/arrow/util/compression_bz2.cc b/cpp/src/arrow/util/compression_bz2.cc
new file mode 100644
index 0000000..ae78f98
--- /dev/null
+++ b/cpp/src/arrow/util/compression_bz2.cc
@@ -0,0 +1,259 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/util/compression_bz2.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+#include <memory>
+#include <sstream>
+#include <string>
+
+// Avoid defining max() macro
+#include "arrow/util/windows_compatibility.h"
+
+#include <bzlib.h>
+
+#include "arrow/status.h"
+#include "arrow/util/logging.h"
+#include "arrow/util/macros.h"
+
+namespace arrow {
+namespace util {
+
+constexpr int kBZ2DefaultCompressionLevel = 9;
+
+// Max number of bytes the bz2 APIs accept at a time
+static constexpr auto kSizeLimit =
+ static_cast<int64_t>(std::numeric_limits<unsigned int>::max());
+
+Status BZ2Error(const char* prefix_msg, int bz_result) {
+ DCHECK(bz_result != BZ_OK && bz_result != BZ_RUN_OK && bz_result != BZ_FLUSH_OK &&
+ bz_result != BZ_FINISH_OK && bz_result != BZ_STREAM_END);
+ StatusCode code;
+ std::stringstream ss;
+ ss << prefix_msg;
+ switch (bz_result) {
+ case BZ_CONFIG_ERROR:
+ code = StatusCode::UnknownError;
+ ss << "bz2 library improperly configured (internal error)";
+ break;
+ case BZ_SEQUENCE_ERROR:
+ code = StatusCode::UnknownError;
+ ss << "wrong sequence of calls to bz2 library (internal error)";
+ break;
+ case BZ_PARAM_ERROR:
+ code = StatusCode::UnknownError;
+ ss << "wrong parameter to bz2 library (internal error)";
+ break;
+ case BZ_MEM_ERROR:
+ code = StatusCode::OutOfMemory;
+ ss << "could not allocate memory for bz2 library";
+ break;
+ case BZ_DATA_ERROR:
+ code = StatusCode::IOError;
+ ss << "invalid bz2 data";
+ break;
+ case BZ_DATA_ERROR_MAGIC:
+ code = StatusCode::IOError;
+ ss << "data is not bz2-compressed (no magic header)";
+ break;
+ default:
+ code = StatusCode::UnknownError;
+ ss << "unknown bz2 error " << bz_result;
+ break;
+ }
+ return Status(code, ss.str());
+}
+
+// ----------------------------------------------------------------------
+// bz2 decompressor implementation
+
+class BZ2Decompressor : public Decompressor {
+ public:
+ BZ2Decompressor() : initialized_(false) {}
+
+ ~BZ2Decompressor() override {
+ if (initialized_) {
+ ARROW_UNUSED(BZ2_bzDecompressEnd(&stream_));
+ }
+ }
+
+ Status Init() {
+ DCHECK(!initialized_);
+ memset(&stream_, 0, sizeof(stream_));
+ int ret;
+ ret = BZ2_bzDecompressInit(&stream_, 0, 0);
+ if (ret != BZ_OK) {
+ return BZ2Error("bz2 decompressor init failed: ", ret);
+ }
+ initialized_ = true;
+ finished_ = false;
+ return Status::OK();
+ }
+
+ Status Decompress(int64_t input_len, const uint8_t* input, int64_t output_len,
+ uint8_t* output, int64_t* bytes_read, int64_t* bytes_written,
+ bool* need_more_output) override {
+ stream_.next_in = const_cast<char*>(reinterpret_cast<const char*>(input));
+ stream_.avail_in = static_cast<unsigned int>(std::min(input_len, kSizeLimit));
+ stream_.next_out = reinterpret_cast<char*>(output);
+ stream_.avail_out = static_cast<unsigned int>(std::min(output_len, kSizeLimit));
+ int ret;
+
+ ret = BZ2_bzDecompress(&stream_);
+ if (ret == BZ_OK || ret == BZ_STREAM_END) {
+ *bytes_read = input_len - stream_.avail_in;
+ *bytes_written = output_len - stream_.avail_out;
+ finished_ = (ret == BZ_STREAM_END);
+ *need_more_output = (!finished_ && *bytes_read == 0 && *bytes_written == 0);
+ return Status::OK();
+ } else {
+ return BZ2Error("bz2 decompress failed: ", ret);
+ }
+ }
+
+ bool IsFinished() override { return finished_; }
+
+ protected:
+ bz_stream stream_;
+ bool initialized_;
+ bool finished_;
+};
+
+// ----------------------------------------------------------------------
+// bz2 compressor implementation
+
+class BZ2Compressor : public Compressor {
+ public:
+ BZ2Compressor() : initialized_(false) {}
+
+ ~BZ2Compressor() override {
+ if (initialized_) {
+ ARROW_UNUSED(BZ2_bzCompressEnd(&stream_));
+ }
+ }
+
+ Status Init() {
+ DCHECK(!initialized_);
+ memset(&stream_, 0, sizeof(stream_));
+ int ret;
+ ret = BZ2_bzCompressInit(&stream_, kBZ2DefaultCompressionLevel, 0, 0);
+ if (ret != BZ_OK) {
+ return BZ2Error("bz2 compressor init failed: ", ret);
+ }
+ initialized_ = true;
+ return Status::OK();
+ }
+
+ Status Compress(int64_t input_len, const uint8_t* input, int64_t output_len,
+ uint8_t* output, int64_t* bytes_read, int64_t* bytes_written) override {
+ stream_.next_in = const_cast<char*>(reinterpret_cast<const char*>(input));
+ stream_.avail_in = static_cast<unsigned int>(std::min(input_len, kSizeLimit));
+ stream_.next_out = reinterpret_cast<char*>(output);
+ stream_.avail_out = static_cast<unsigned int>(std::min(output_len, kSizeLimit));
+ int ret;
+
+ ret = BZ2_bzCompress(&stream_, BZ_RUN);
+ if (ret == BZ_RUN_OK) {
+ *bytes_read = input_len - stream_.avail_in;
+ *bytes_written = output_len - stream_.avail_out;
+ return Status::OK();
+ } else {
+ return BZ2Error("bz2 compress failed: ", ret);
+ }
+ }
+
+ Status Flush(int64_t output_len, uint8_t* output, int64_t* bytes_written,
+ bool* should_retry) override {
+ stream_.next_in = nullptr;
+ stream_.avail_in = 0;
+ stream_.next_out = reinterpret_cast<char*>(output);
+ stream_.avail_out = static_cast<unsigned int>(std::min(output_len, kSizeLimit));
+ int ret;
+
+ ret = BZ2_bzCompress(&stream_, BZ_FLUSH);
+ if (ret == BZ_RUN_OK || ret == BZ_FLUSH_OK) {
+ *bytes_written = output_len - stream_.avail_out;
+ *should_retry = (ret == BZ_FLUSH_OK);
+ return Status::OK();
+ } else {
+ return BZ2Error("bz2 compress failed: ", ret);
+ }
+ }
+
+ Status End(int64_t output_len, uint8_t* output, int64_t* bytes_written,
+ bool* should_retry) override {
+ stream_.next_in = nullptr;
+ stream_.avail_in = 0;
+ stream_.next_out = reinterpret_cast<char*>(output);
+ stream_.avail_out = static_cast<unsigned int>(std::min(output_len, kSizeLimit));
+ int ret;
+
+ ret = BZ2_bzCompress(&stream_, BZ_FINISH);
+ if (ret == BZ_STREAM_END || ret == BZ_FINISH_OK) {
+ *bytes_written = output_len - stream_.avail_out;
+ *should_retry = (ret == BZ_FINISH_OK);
+ return Status::OK();
+ } else {
+ return BZ2Error("bz2 compress failed: ", ret);
+ }
+ }
+
+ protected:
+ bz_stream stream_;
+ bool initialized_;
+};
+
+// ----------------------------------------------------------------------
+// bz2 codec implementation
+
+Status BZ2Codec::MakeCompressor(std::shared_ptr<Compressor>* out) {
+ auto ptr = std::make_shared<BZ2Compressor>();
+ RETURN_NOT_OK(ptr->Init());
+ *out = ptr;
+ return Status::OK();
+}
+
+Status BZ2Codec::MakeDecompressor(std::shared_ptr<Decompressor>* out) {
+ auto ptr = std::make_shared<BZ2Decompressor>();
+ RETURN_NOT_OK(ptr->Init());
+ *out = ptr;
+ return Status::OK();
+}
+
+Status BZ2Codec::Decompress(int64_t input_len, const uint8_t* input, int64_t output_len,
+ uint8_t* output_buffer) {
+ return Status::NotImplemented("One-shot bz2 decompression not supported");
+}
+
+int64_t BZ2Codec::MaxCompressedLen(int64_t input_len,
+ const uint8_t* ARROW_ARG_UNUSED(input)) {
+ // Cannot determine upper bound for bz2-compressed data
+ return 0;
+}
+
+Status BZ2Codec::Compress(int64_t input_len, const uint8_t* input,
+ int64_t output_buffer_len, uint8_t* output_buffer,
+ int64_t* output_length) {
+ return Status::NotImplemented("One-shot bz2 compression not supported");
+}
+
+} // namespace util
+} // namespace arrow
diff --git a/cpp/src/arrow/util/compression_bz2.h b/cpp/src/arrow/util/compression_bz2.h
new file mode 100644
index 0000000..304f4c0
--- /dev/null
+++ b/cpp/src/arrow/util/compression_bz2.h
@@ -0,0 +1,52 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#ifndef ARROW_UTIL_COMPRESSION_BZ2_H
+#define ARROW_UTIL_COMPRESSION_BZ2_H
+
+#include <cstdint>
+#include <memory>
+
+#include "arrow/status.h"
+#include "arrow/util/compression.h"
+#include "arrow/util/visibility.h"
+
+namespace arrow {
+namespace util {
+
+// BZ2 codec.
+class ARROW_EXPORT BZ2Codec : public Codec {
+ public:
+ Status Decompress(int64_t input_len, const uint8_t* input, int64_t output_len,
+ uint8_t* output_buffer) override;
+
+ Status Compress(int64_t input_len, const uint8_t* input, int64_t output_buffer_len,
+ uint8_t* output_buffer, int64_t* output_length) override;
+
+ int64_t MaxCompressedLen(int64_t input_len, const uint8_t* input) override;
+
+ Status MakeCompressor(std::shared_ptr<Compressor>* out) override;
+
+ Status MakeDecompressor(std::shared_ptr<Decompressor>* out) override;
+
+ const char* name() const override { return "bz2"; }
+};
+
+} // namespace util
+} // namespace arrow
+
+#endif // ARROW_UTIL_COMPRESSION_BZ2_H
diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd
index c53b4ad..9e3a879 100644
--- a/python/pyarrow/includes/libarrow.pxd
+++ b/python/pyarrow/includes/libarrow.pxd
@@ -1132,6 +1132,7 @@ cdef extern from 'arrow/util/compression.h' namespace 'arrow' nogil:
CompressionType_BROTLI" arrow::Compression::BROTLI"
CompressionType_ZSTD" arrow::Compression::ZSTD"
CompressionType_LZ4" arrow::Compression::LZ4"
+ CompressionType_BZ2" arrow::Compression::BZ2"
cdef cppclass CCodec" arrow::util::Codec":
@staticmethod
diff --git a/python/pyarrow/io.pxi b/python/pyarrow/io.pxi
index be75e7b..9881b8c 100644
--- a/python/pyarrow/io.pxi
+++ b/python/pyarrow/io.pxi
@@ -1135,6 +1135,8 @@ cdef CompressionType _get_compression_type(object name):
cdef CompressionType _get_compression_type_by_filename(filename):
if filename.endswith('.gz'):
return CompressionType_GZIP
+ elif filename.endswith('.bz2'):
+ return CompressionType_BZ2
elif filename.endswith('.lz4'):
return CompressionType_LZ4
elif filename.endswith('.zst'):
diff --git a/python/pyarrow/tests/test_csv.py b/python/pyarrow/tests/test_csv.py
index e3ea92a..dbfa4f3 100644
--- a/python/pyarrow/tests/test_csv.py
+++ b/python/pyarrow/tests/test_csv.py
@@ -15,6 +15,7 @@
# specific language governing permissions and limitations
# under the License.
+import bz2
import gzip
import io
import itertools
@@ -275,7 +276,11 @@ class BaseTestCompressedCSVRead:
csv, expected = make_random_csv(num_cols=2, num_rows=100)
csv_path = os.path.join(self.tmpdir, self.csv_filename)
self.write_file(csv_path, csv)
- table = read_csv(csv_path)
+ try:
+ table = read_csv(csv_path)
+ except pa.ArrowNotImplementedError as e:
+ pytest.skip(str(e))
+ return
table._validate()
assert table.schema == expected.schema
assert table.equals(expected)
@@ -288,3 +293,11 @@ class TestGZipCSVRead(BaseTestCompressedCSVRead, unittest.TestCase):
def write_file(self, path, contents):
with gzip.open(path, 'wb', 3) as f:
f.write(contents)
+
+
+class TestBZ2CSVRead(BaseTestCompressedCSVRead, unittest.TestCase):
+ csv_filename = "compressed.csv.bz2"
+
+ def write_file(self, path, contents):
+ with bz2.BZ2File(path, 'w') as f:
+ f.write(contents)