[C++][Dataset] Add ORC predicate pushdown tests for all-null and filesystem sources

Author: cbb330

cpp/src/arrow/dataset/file_orc_test.cc

@@ -17,14 +17,18 @@
 
 #include "arrow/dataset/file_orc.h"
 
+#include <algorithm>
 #include <memory>
 #include <utility>
 
 #include "arrow/adapters/orc/adapter.h"
+#include "arrow/array/builder_primitive.h"
+#include "arrow/compute/api.h"
 #include "arrow/dataset/dataset_internal.h"
 #include "arrow/dataset/discovery.h"
 #include "arrow/dataset/file_base.h"
 #include "arrow/dataset/partition.h"
+#include "arrow/filesystem/mockfs.h"
 #include "arrow/dataset/test_util_internal.h"
 #include "arrow/io/memory.h"
 #include "arrow/record_batch.h"
@@ -88,6 +92,101 @@ TEST_P(TestOrcFileFormatScan, ScanRecordBatchReaderWithDuplicateColumnError) {
   TestScanWithDuplicateColumnError();
 }
 TEST_P(TestOrcFileFormatScan, ScanWithPushdownNulls) { TestScanWithPushdownNulls(); }
+
+TEST_P(TestOrcFileFormatScan, PredicatePushdownAllNullStripes) {
+  auto value_field = field("i64", int64());
+  const auto test_schema = schema({value_field});
+
+  const int64_t nrows = 2048;
+  Int64Builder null_builder;
+  ASSERT_OK(null_builder.AppendNulls(nrows));
+  ASSERT_OK_AND_ASSIGN(auto all_null_values, null_builder.Finish());
+
+  Int64Builder value_builder;
+  for (int64_t i = 0; i < nrows; ++i) {
+    ASSERT_OK(value_builder.Append(i));
+  }
+  ASSERT_OK_AND_ASSIGN(auto non_null_values, value_builder.Finish());
+
+  auto all_null_batch = RecordBatch::Make(test_schema, nrows, {all_null_values});
+  auto non_null_batch = RecordBatch::Make(test_schema, nrows, {non_null_values});
+
+  ASSERT_OK_AND_ASSIGN(auto sink, io::BufferOutputStream::Create());
+  adapters::orc::WriteOptions write_options;
+  write_options.stripe_size = 4096;
+  ASSERT_OK_AND_ASSIGN(auto writer,
+                       adapters::orc::ORCFileWriter::Open(sink.get(), write_options));
+  ASSERT_OK(writer->Write(*all_null_batch));
+  ASSERT_OK(writer->Write(*non_null_batch));
+  ASSERT_OK(writer->Close());
+  ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish());
+
+  FileSource source(buffer);
+  ASSERT_OK_AND_ASSIGN(auto fragment_base, format_->MakeFragment(source, literal(true)));
+  auto orc_fragment = checked_pointer_cast<OrcFileFragment>(fragment_base);
+
+  ASSERT_OK_AND_ASSIGN(auto input, source.Open());
+  ASSERT_OK_AND_ASSIGN(auto reader,
+                       adapters::orc::ORCFileReader::Open(std::move(input),
+                                                          default_memory_pool()));
+
+  std::vector<int> all_null_stripes;
+  std::vector<int> non_all_null_stripes;
+  const int64_t num_stripes = reader->NumberOfStripes();
+  for (int64_t stripe = 0; stripe < num_stripes; ++stripe) {
+    ASSERT_OK_AND_ASSIGN(auto stripe_stats, reader->GetStripeStatistics(stripe));
+    const auto* col_stats = stripe_stats->getColumnStatistics(1);
+    ASSERT_NE(col_stats, nullptr);
+
+    if (col_stats->hasNull() && col_stats->getNumberOfValues() == 0) {
+      all_null_stripes.push_back(static_cast<int>(stripe));
+    } else {
+      non_all_null_stripes.push_back(static_cast<int>(stripe));
+    }
+  }
+  ASSERT_FALSE(all_null_stripes.empty());
+  ASSERT_FALSE(non_all_null_stripes.empty());
+
+  ASSERT_OK_AND_ASSIGN(
+      auto is_null_selected,
+      orc_fragment->FilterStripes(compute::is_null(compute::field_ref("i64"))));
+  for (int stripe : all_null_stripes) {
+    EXPECT_NE(std::find(is_null_selected.begin(), is_null_selected.end(), stripe),
+              is_null_selected.end());
+  }
+
+  ASSERT_OK_AND_ASSIGN(
+      auto is_not_null_selected,
+      orc_fragment->FilterStripes(
+          compute::not_(compute::is_null(compute::field_ref("i64")))));
+  for (int stripe : all_null_stripes) {
+    EXPECT_EQ(std::find(is_not_null_selected.begin(), is_not_null_selected.end(), stripe),
+              is_not_null_selected.end());
+  }
+}
+
+TEST_P(TestOrcFileFormatScan, PredicatePushdownWithFileSystemSource) {
+  auto mock_fs = std::make_shared<fs::internal::MockFileSystem>(fs::kNoTime);
+  std::shared_ptr<Schema> test_schema = schema({field("x", int64())});
+  std::shared_ptr<RecordBatch> batch = RecordBatchFromJSON(test_schema, "[[0], [1], [2]]");
+
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<io::OutputStream> out_stream,
+                       mock_fs->OpenOutputStream("/foo.orc"));
+  ASSERT_OK_AND_ASSIGN(auto writer, adapters::orc::ORCFileWriter::Open(out_stream.get()));
+  ASSERT_OK(writer->Write(*batch));
+  ASSERT_OK(writer->Close());
+
+  FileSource source("/foo.orc", mock_fs);
+  ASSERT_OK_AND_ASSIGN(auto fragment_base, format_->MakeFragment(source, literal(true)));
+  auto orc_fragment = checked_pointer_cast<OrcFileFragment>(fragment_base);
+
+  ASSERT_OK_AND_ASSIGN(
+      auto stripes,
+      orc_fragment->FilterStripes(compute::greater(compute::field_ref("x"),
+                                                   compute::literal(int64_t{1}))));
+  ASSERT_FALSE(stripes.empty());
+}
+
 INSTANTIATE_TEST_SUITE_P(TestScan, TestOrcFileFormatScan,
                          ::testing::ValuesIn(TestFormatParams::Values()),
                          TestFormatParams::ToTestNameString);