ceph/src/arrow/cpp/src/arrow/dataset/file_orc_test.cc

   1 // Licensed to the Apache Software Foundation (ASF) under one
   2 // or more contributor license agreements.  See the NOTICE file
   3 // distributed with this work for additional information
   4 // regarding copyright ownership.  The ASF licenses this file
   5 // to you under the Apache License, Version 2.0 (the
   6 // "License"); you may not use this file except in compliance
   7 // with the License.  You may obtain a copy of the License at
   8 //
   9 //   http://www.apache.org/licenses/LICENSE-2.0
  10 //
  11 // Unless required by applicable law or agreed to in writing,
  12 // software distributed under the License is distributed on an
  13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  14 // KIND, either express or implied.  See the License for the
  15 // specific language governing permissions and limitations
  16 // under the License.
  17
  18 #include "arrow/dataset/file_orc.h"
  19
  20 #include <memory>
  21 #include <utility>
  22
  23 #include "arrow/adapters/orc/adapter.h"
  24 #include "arrow/dataset/dataset_internal.h"
  25 #include "arrow/dataset/discovery.h"
  26 #include "arrow/dataset/file_base.h"
  27 #include "arrow/dataset/partition.h"
  28 #include "arrow/dataset/scanner_internal.h"
  29 #include "arrow/dataset/test_util.h"
  30 #include "arrow/io/memory.h"
  31 #include "arrow/record_batch.h"
  32 #include "arrow/table.h"
  33 #include "arrow/testing/gtest_util.h"
  34 #include "arrow/testing/util.h"
  35
  36 namespace arrow {
  37 namespace dataset {
  38
  39 class OrcFormatHelper {
  40  public:
  41   using FormatType = OrcFileFormat;
  42   static Result<std::shared_ptr<Buffer>> Write(RecordBatchReader* reader) {
  43     ARROW_ASSIGN_OR_RAISE(auto sink, io::BufferOutputStream::Create());
  44     ARROW_ASSIGN_OR_RAISE(auto writer, adapters::orc::ORCFileWriter::Open(sink.get()));
  45     std::shared_ptr<Table> table;
  46     RETURN_NOT_OK(reader->ReadAll(&table));
  47     writer->Write(*table);
  48     RETURN_NOT_OK(writer->Close());
  49     return sink->Finish();
  50   }
  51
  52   static std::shared_ptr<OrcFileFormat> MakeFormat() {
  53     return std::make_shared<OrcFileFormat>();
  54   }
  55 };
  56
  57 class TestOrcFileFormat : public FileFormatFixtureMixin<OrcFormatHelper> {};
  58
  59 // TEST_F(TestOrcFileFormat, WriteRecordBatchReader) { TestWrite(); }
  60
  61 TEST_F(TestOrcFileFormat, InspectFailureWithRelevantError) {
  62   TestInspectFailureWithRelevantError(StatusCode::IOError, "ORC");
  63 }
  64 TEST_F(TestOrcFileFormat, Inspect) { TestInspect(); }
  65 TEST_F(TestOrcFileFormat, IsSupported) { TestIsSupported(); }
  66 TEST_F(TestOrcFileFormat, CountRows) { TestCountRows(); }
  67
  68 // TODO add TestOrcFileSystemDataset if write support is added
  69
  70 class TestOrcFileFormatScan : public FileFormatScanMixin<OrcFormatHelper> {};
  71
  72 TEST_P(TestOrcFileFormatScan, ScanRecordBatchReader) { TestScan(); }
  73 TEST_P(TestOrcFileFormatScan, ScanRecordBatchReaderWithVirtualColumn) {
  74   TestScanWithVirtualColumn();
  75 }
  76 TEST_P(TestOrcFileFormatScan, ScanRecordBatchReaderProjected) { TestScanProjected(); }
  77 TEST_P(TestOrcFileFormatScan, ScanRecordBatchReaderProjectedMissingCols) {
  78   TestScanProjectedMissingCols();
  79 }
  80 INSTANTIATE_TEST_SUITE_P(TestScan, TestOrcFileFormatScan,
  81                          ::testing::ValuesIn(TestFormatParams::Values()),
  82                          TestFormatParams::ToTestNameString);
  83
  84 }  // namespace dataset
  85 }  // namespace arrow