add test

Author: HeartLinked

test/avro_test.cc

@@ -28,11 +28,16 @@
 #include <avro/Generic.hh>
 #include <avro/GenericDatum.hh>
 #include <gtest/gtest.h>
+#include <iceberg/util/checked_cast.h>
 
 #include "iceberg/arrow/arrow_fs_file_io_internal.h"
+#include "iceberg/arrow_array_reader.h"
 #include "iceberg/avro/avro_register.h"
+#include "iceberg/file_format.h"
 #include "iceberg/file_reader.h"
+#include "iceberg/manifest_entry.h"
 #include "iceberg/schema.h"
+#include "iceberg/table_scan.h"
 #include "iceberg/type.h"
 #include "matchers.h"
 #include "temp_file_test_base.h"
@@ -77,7 +82,7 @@ class AvroReaderTest : public TempFileTestBase {
     writer.close();
   }
 
-  void VerifyNextBatch(Reader& reader, std::string_view expected_json) {
+  void VerifyNextBatch(ArrowArrayReader& reader, std::string_view expected_json) {
     // Boilerplate to get Arrow schema
     auto schema_result = reader.Schema();
     ASSERT_THAT(schema_result, IsOk());
@@ -99,7 +104,7 @@ class AvroReaderTest : public TempFileTestBase {
     ASSERT_TRUE(arrow_array->Equals(*expected_array));
   }
 
-  void VerifyExhausted(Reader& reader) {
+  void VerifyExhausted(ArrowArrayReader& reader) {
     auto data = reader.Next();
     ASSERT_THAT(data, IsOk());
     ASSERT_FALSE(data.value().has_value());
@@ -161,4 +166,30 @@ TEST_F(AvroReaderTest, ReadWithBatchSize) {
   ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader));
 }
 
+TEST_F(AvroReaderTest, ReadViaFileScanTask) {
+  CreateSimpleAvroFile();
+
+  auto data_file = std::make_shared<DataFile>();
+  data_file->file_path = temp_avro_file_;
+  data_file->file_format = FileFormatType::kAvro;
+  auto arrow_file_io = internal::checked_cast<arrow::ArrowFileSystemFileIO&>(*file_io_);
+  data_file->file_size_in_bytes =
+      arrow_file_io.fs()->GetFileInfo(temp_avro_file_).ValueOrDie().size();
+
+  TableScanContext context;
+  context.projected_schema = std::make_shared<Schema>(std::vector<SchemaField>{
+      SchemaField::MakeOptional(2, "name", std::make_shared<StringType>()),
+      SchemaField::MakeRequired(1, "id", std::make_shared<IntType>())});
+
+  FileScanTask task(data_file);
+
+  auto reader_result = task.ToArrowArrayReader(context, file_io_);
+  ASSERT_THAT(reader_result, IsOk());
+  auto reader = std::move(reader_result).value();
+
+  ASSERT_NO_FATAL_FAILURE(
+      VerifyNextBatch(*reader, R"([["Alice", 1], ["Bob", 2], ["Charlie", 3]])"));
+  ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader));
+}
+
 }  // namespace iceberg::avro

test/parquet_test.cc

@@ -28,9 +28,13 @@
 #include <parquet/metadata.h>
 
 #include "iceberg/arrow/arrow_fs_file_io_internal.h"
+#include "iceberg/arrow_array_reader.h"
+#include "iceberg/file_format.h"
+#include "iceberg/manifest_entry.h"
 #include "iceberg/parquet/parquet_reader.h"
 #include "iceberg/parquet/parquet_register.h"
 #include "iceberg/schema.h"
+#include "iceberg/table_scan.h"
 #include "iceberg/type.h"
 #include "iceberg/util/checked_cast.h"
 #include "matchers.h"
@@ -72,7 +76,7 @@ class ParquetReaderTest : public TempFileTestBase {
                     .ok());
   }
 
-  void VerifyNextBatch(Reader& reader, std::string_view expected_json) {
+  void VerifyNextBatch(ArrowArrayReader& reader, std::string_view expected_json) {
     // Boilerplate to get Arrow schema
     auto schema_result = reader.Schema();
     ASSERT_THAT(schema_result, IsOk());
@@ -94,7 +98,7 @@ class ParquetReaderTest : public TempFileTestBase {
     ASSERT_TRUE(arrow_array->Equals(*expected_array));
   }
 
-  void VerifyExhausted(Reader& reader) {
+  void VerifyExhausted(ArrowArrayReader& reader) {
     auto data = reader.Next();
     ASSERT_THAT(data, IsOk());
     ASSERT_FALSE(data.value().has_value());
@@ -204,4 +208,30 @@ TEST_F(ParquetReaderTest, ReadSplit) {
   }
 }
 
+TEST_F(ParquetReaderTest, ReadViaFileScanTask) {
+  CreateSimpleParquetFile();
+
+  auto data_file = std::make_shared<DataFile>();
+  data_file->file_path = temp_parquet_file_;
+  data_file->file_format = FileFormatType::kParquet;
+  auto io_internal = internal::checked_cast<arrow::ArrowFileSystemFileIO&>(*file_io_);
+  data_file->file_size_in_bytes =
+      io_internal.fs()->GetFileInfo(temp_parquet_file_).ValueOrDie().size();
+
+  TableScanContext context;
+  context.projected_schema = std::make_shared<Schema>(
+      std::vector<SchemaField>{SchemaField::MakeOptional(2, "name", string()),
+                               SchemaField::MakeRequired(1, "id", int32())});
+
+  FileScanTask task(data_file);
+
+  auto reader_result = task.ToArrowArrayReader(context, file_io_);
+  ASSERT_THAT(reader_result, IsOk());
+  auto reader = std::move(reader_result.value());
+
+  ASSERT_NO_FATAL_FAILURE(
+      VerifyNextBatch(*reader, R"([["Foo", 1], ["Bar", 2], ["Baz", 3]])"));
+  ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader));
+}
+
 }  // namespace iceberg::parquet