| # 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. |
| |
| # cython: profile=False |
| # distutils: language = c++ |
| |
| from cython.operator cimport dereference as deref |
| from libcpp.vector cimport vector as std_vector |
| from libcpp.utility cimport move |
| from pyarrow.includes.common cimport * |
| from pyarrow.includes.libarrow cimport * |
| from pyarrow.lib cimport (check_status, _Weakrefable, |
| MemoryPool, maybe_unbox_memory_pool, |
| pyarrow_wrap_schema, |
| pyarrow_wrap_batch, |
| Table, |
| pyarrow_wrap_table, |
| pyarrow_wrap_metadata, |
| pyarrow_unwrap_table, |
| get_reader, |
| get_writer) |
| from pyarrow.lib import frombytes, tobytes |
| from pyarrow.util import _stringify_path |
| |
| |
| cdef compression_type_from_enum(CCompressionType compression_type): |
| compression_map = { |
| CCompressionType_UNCOMPRESSED: 'UNCOMPRESSED', |
| CCompressionType_GZIP: 'ZLIB', |
| CCompressionType_SNAPPY: 'SNAPPY', |
| CCompressionType_LZ4: 'LZ4', |
| CCompressionType_ZSTD: 'ZSTD', |
| } |
| if compression_type in compression_map: |
| return compression_map[compression_type] |
| raise ValueError('Unsupported compression') |
| |
| |
| cdef CCompressionType compression_type_from_name(name) except *: |
| if not isinstance(name, str): |
| raise TypeError('compression must be a string') |
| name = name.upper() |
| if name == 'ZLIB': |
| return CCompressionType_GZIP |
| elif name == 'SNAPPY': |
| return CCompressionType_SNAPPY |
| elif name == 'LZ4': |
| return CCompressionType_LZ4 |
| elif name == 'ZSTD': |
| return CCompressionType_ZSTD |
| elif name == 'UNCOMPRESSED': |
| return CCompressionType_UNCOMPRESSED |
| raise ValueError(f'Unknown CompressionKind: {name}') |
| |
| |
| cdef compression_strategy_from_enum( |
| CompressionStrategy compression_strategy |
| ): |
| compression_strategy_map = { |
| _CompressionStrategy_SPEED: 'SPEED', |
| _CompressionStrategy_COMPRESSION: 'COMPRESSION', |
| } |
| if compression_strategy in compression_strategy_map: |
| return compression_strategy_map[compression_strategy] |
| raise ValueError('Unsupported compression strategy') |
| |
| |
| cdef CompressionStrategy compression_strategy_from_name(name) except *: |
| if not isinstance(name, str): |
| raise TypeError('compression strategy must be a string') |
| name = name.upper() |
| if name == 'COMPRESSION': |
| return _CompressionStrategy_COMPRESSION |
| elif name == 'SPEED': |
| return _CompressionStrategy_SPEED |
| raise ValueError(f'Unknown CompressionStrategy: {name}') |
| |
| |
| cdef file_version_from_class(FileVersion file_version): |
| return frombytes(file_version.ToString()) |
| |
| |
| cdef writer_id_from_enum(WriterId writer_id): |
| writer_id_map = { |
| _WriterId_ORC_JAVA_WRITER: 'ORC_JAVA', |
| _WriterId_ORC_CPP_WRITER: 'ORC_CPP', |
| _WriterId_PRESTO_WRITER: 'PRESTO', |
| _WriterId_SCRITCHLEY_GO: 'SCRITCHLEY_GO', |
| _WriterId_TRINO_WRITER: 'TRINO', |
| } |
| if writer_id in writer_id_map: |
| return writer_id_map[writer_id] |
| raise ValueError('Unsupported writer ID') |
| |
| |
| cdef writer_version_from_enum(WriterVersion writer_version): |
| writer_version_map = { |
| _WriterVersion_ORIGINAL: 'ORIGINAL', |
| _WriterVersion_HIVE_8732: 'HIVE_8732', |
| _WriterVersion_HIVE_4243: 'HIVE_4243', |
| _WriterVersion_HIVE_12055: 'HIVE_12055', |
| _WriterVersion_HIVE_13083: 'HIVE_13083', |
| _WriterVersion_ORC_101: 'ORC_101', |
| _WriterVersion_ORC_135: 'ORC_135', |
| _WriterVersion_ORC_517: 'ORC_517', |
| _WriterVersion_ORC_203: 'ORC_203', |
| _WriterVersion_ORC_14: 'ORC_14', |
| } |
| if writer_version in writer_version_map: |
| return writer_version_map[writer_version] |
| raise ValueError('Unsupported writer version') |
| |
| |
| cdef shared_ptr[WriteOptions] _create_write_options( |
| file_version=None, |
| batch_size=None, |
| stripe_size=None, |
| compression=None, |
| compression_block_size=None, |
| compression_strategy=None, |
| row_index_stride=None, |
| padding_tolerance=None, |
| dictionary_key_size_threshold=None, |
| bloom_filter_columns=None, |
| bloom_filter_fpp=None |
| ) except *: |
| """General writer options""" |
| cdef: |
| shared_ptr[WriteOptions] options |
| options = make_shared[WriteOptions]() |
| # batch_size |
| if batch_size is not None: |
| if isinstance(batch_size, int) and batch_size > 0: |
| deref(options).batch_size = batch_size |
| else: |
| raise ValueError(f"Invalid ORC writer batch size: {batch_size}") |
| # file_version |
| if file_version is not None: |
| if file_version == "0.12": |
| deref(options).file_version = FileVersion(0, 12) |
| elif file_version == "0.11": |
| deref(options).file_version = FileVersion(0, 11) |
| else: |
| raise ValueError(f"Unsupported ORC file version: {file_version}") |
| # stripe_size |
| if stripe_size is not None: |
| if isinstance(stripe_size, int) and stripe_size > 0: |
| deref(options).stripe_size = stripe_size |
| else: |
| raise ValueError(f"Invalid ORC stripe size: {stripe_size}") |
| # compression |
| if compression is not None: |
| if isinstance(compression, str): |
| deref(options).compression = compression_type_from_name( |
| compression) |
| else: |
| raise TypeError("Unsupported ORC compression type: " |
| f"{compression}") |
| # compression_block_size |
| if compression_block_size is not None: |
| if (isinstance(compression_block_size, int) and |
| compression_block_size > 0): |
| deref(options).compression_block_size = compression_block_size |
| else: |
| raise ValueError("Invalid ORC compression block size: " |
| f"{compression_block_size}") |
| # compression_strategy |
| if compression_strategy is not None: |
| if isinstance(compression, str): |
| deref(options).compression_strategy = \ |
| compression_strategy_from_name(compression_strategy) |
| else: |
| raise TypeError("Unsupported ORC compression strategy: " |
| f"{compression_strategy}") |
| # row_index_stride |
| if row_index_stride is not None: |
| if isinstance(row_index_stride, int) and row_index_stride > 0: |
| deref(options).row_index_stride = row_index_stride |
| else: |
| raise ValueError("Invalid ORC row index stride: " |
| f"{row_index_stride}") |
| # padding_tolerance |
| if padding_tolerance is not None: |
| try: |
| padding_tolerance = float(padding_tolerance) |
| deref(options).padding_tolerance = padding_tolerance |
| except Exception: |
| raise ValueError("Invalid ORC padding tolerance: " |
| f"{padding_tolerance}") |
| # dictionary_key_size_threshold |
| if dictionary_key_size_threshold is not None: |
| try: |
| dictionary_key_size_threshold = float( |
| dictionary_key_size_threshold) |
| assert 0 <= dictionary_key_size_threshold <= 1 |
| deref(options).dictionary_key_size_threshold = \ |
| dictionary_key_size_threshold |
| except Exception: |
| raise ValueError("Invalid ORC dictionary key size threshold: " |
| f"{dictionary_key_size_threshold}") |
| # bloom_filter_columns |
| if bloom_filter_columns is not None: |
| try: |
| bloom_filter_columns = list(bloom_filter_columns) |
| for col in bloom_filter_columns: |
| assert isinstance(col, int) and col >= 0 |
| deref(options).bloom_filter_columns = bloom_filter_columns |
| except Exception: |
| raise ValueError("Invalid ORC BloomFilter columns: " |
| f"{bloom_filter_columns}") |
| # Max false positive rate of the Bloom Filter |
| if bloom_filter_fpp is not None: |
| try: |
| bloom_filter_fpp = float(bloom_filter_fpp) |
| assert 0 <= bloom_filter_fpp <= 1 |
| deref(options).bloom_filter_fpp = bloom_filter_fpp |
| except Exception: |
| raise ValueError("Invalid ORC BloomFilter false positive rate: " |
| f"{bloom_filter_fpp}") |
| return options |
| |
| |
| cdef class ORCReader(_Weakrefable): |
| cdef: |
| object source |
| CMemoryPool* allocator |
| unique_ptr[ORCFileReader] reader |
| |
| def __cinit__(self, MemoryPool memory_pool=None): |
| self.allocator = maybe_unbox_memory_pool(memory_pool) |
| |
| def open(self, object source, c_bool use_memory_map=True): |
| cdef: |
| shared_ptr[CRandomAccessFile] rd_handle |
| |
| self.source = source |
| |
| get_reader(source, use_memory_map, &rd_handle) |
| with nogil: |
| self.reader = move(GetResultValue( |
| ORCFileReader.Open(rd_handle, self.allocator) |
| )) |
| |
| def metadata(self): |
| """ |
| The arrow metadata for this file. |
| |
| Returns |
| ------- |
| metadata : pyarrow.KeyValueMetadata |
| """ |
| cdef: |
| shared_ptr[const CKeyValueMetadata] sp_arrow_metadata |
| |
| with nogil: |
| sp_arrow_metadata = GetResultValue( |
| deref(self.reader).ReadMetadata() |
| ) |
| |
| return pyarrow_wrap_metadata(sp_arrow_metadata) |
| |
| def schema(self): |
| """ |
| The arrow schema for this file. |
| |
| Returns |
| ------- |
| schema : pyarrow.Schema |
| """ |
| cdef: |
| shared_ptr[CSchema] sp_arrow_schema |
| |
| with nogil: |
| sp_arrow_schema = GetResultValue(deref(self.reader).ReadSchema()) |
| |
| return pyarrow_wrap_schema(sp_arrow_schema) |
| |
| def nrows(self): |
| return deref(self.reader).NumberOfRows() |
| |
| def nstripes(self): |
| return deref(self.reader).NumberOfStripes() |
| |
| def file_version(self): |
| return file_version_from_class(deref(self.reader).GetFileVersion()) |
| |
| def software_version(self): |
| return frombytes(deref(self.reader).GetSoftwareVersion()) |
| |
| def compression(self): |
| return compression_type_from_enum( |
| GetResultValue(deref(self.reader).GetCompression())) |
| |
| def compression_size(self): |
| return deref(self.reader).GetCompressionSize() |
| |
| def row_index_stride(self): |
| return deref(self.reader).GetRowIndexStride() |
| |
| def writer(self): |
| writer_name = writer_id_from_enum(deref(self.reader).GetWriterId()) |
| if writer_name == 'UNKNOWN': |
| return deref(self.reader).GetWriterIdValue() |
| else: |
| return writer_name |
| |
| def writer_version(self): |
| return writer_version_from_enum(deref(self.reader).GetWriterVersion()) |
| |
| def nstripe_statistics(self): |
| return deref(self.reader).GetNumberOfStripeStatistics() |
| |
| def content_length(self): |
| return deref(self.reader).GetContentLength() |
| |
| def stripe_statistics_length(self): |
| return deref(self.reader).GetStripeStatisticsLength() |
| |
| def file_footer_length(self): |
| return deref(self.reader).GetFileFooterLength() |
| |
| def file_postscript_length(self): |
| return deref(self.reader).GetFilePostscriptLength() |
| |
| def file_length(self): |
| return deref(self.reader).GetFileLength() |
| |
| def serialized_file_tail(self): |
| return deref(self.reader).GetSerializedFileTail() |
| |
| def read_stripe(self, n, columns=None): |
| cdef: |
| shared_ptr[CRecordBatch] sp_record_batch |
| int64_t stripe |
| std_vector[c_string] c_names |
| |
| stripe = n |
| |
| if columns is None: |
| with nogil: |
| sp_record_batch = GetResultValue( |
| deref(self.reader).ReadStripe(stripe) |
| ) |
| else: |
| c_names = [tobytes(name) for name in columns] |
| with nogil: |
| sp_record_batch = GetResultValue( |
| deref(self.reader).ReadStripe(stripe, c_names) |
| ) |
| |
| return pyarrow_wrap_batch(sp_record_batch) |
| |
| def read(self, columns=None): |
| cdef: |
| shared_ptr[CTable] sp_table |
| std_vector[c_string] c_names |
| |
| if columns is None: |
| with nogil: |
| sp_table = GetResultValue(deref(self.reader).Read()) |
| else: |
| c_names = [tobytes(name) for name in columns] |
| with nogil: |
| sp_table = GetResultValue(deref(self.reader).Read(c_names)) |
| |
| return pyarrow_wrap_table(sp_table) |
| |
| |
| cdef class ORCWriter(_Weakrefable): |
| cdef: |
| unique_ptr[ORCFileWriter] writer |
| shared_ptr[COutputStream] sink |
| c_bool own_sink |
| |
| def open(self, object where, *, |
| file_version=None, |
| batch_size=None, |
| stripe_size=None, |
| compression=None, |
| compression_block_size=None, |
| compression_strategy=None, |
| row_index_stride=None, |
| padding_tolerance=None, |
| dictionary_key_size_threshold=None, |
| bloom_filter_columns=None, |
| bloom_filter_fpp=None): |
| cdef: |
| shared_ptr[WriteOptions] write_options |
| c_string c_where |
| try: |
| where = _stringify_path(where) |
| except TypeError: |
| get_writer(where, &self.sink) |
| self.own_sink = False |
| else: |
| c_where = tobytes(where) |
| with nogil: |
| self.sink = GetResultValue(FileOutputStream.Open(c_where)) |
| self.own_sink = True |
| |
| write_options = _create_write_options( |
| file_version=file_version, |
| batch_size=batch_size, |
| stripe_size=stripe_size, |
| compression=compression, |
| compression_block_size=compression_block_size, |
| compression_strategy=compression_strategy, |
| row_index_stride=row_index_stride, |
| padding_tolerance=padding_tolerance, |
| dictionary_key_size_threshold=dictionary_key_size_threshold, |
| bloom_filter_columns=bloom_filter_columns, |
| bloom_filter_fpp=bloom_filter_fpp |
| ) |
| |
| with nogil: |
| self.writer = move(GetResultValue( |
| ORCFileWriter.Open(self.sink.get(), |
| deref(write_options)))) |
| |
| def write(self, Table table): |
| cdef: |
| shared_ptr[CTable] sp_table |
| sp_table = pyarrow_unwrap_table(table) |
| with nogil: |
| check_status(deref(self.writer).Write(deref(sp_table))) |
| |
| def close(self): |
| with nogil: |
| check_status(deref(self.writer).Close()) |
| if self.own_sink: |
| check_status(deref(self.sink).Close()) |
