validate that arrow IPC files include a valid embedded IPC stream wit…

…h EOS
bkietz · Aug 26, 2024 · 531ca4c · 531ca4c
1 parent 2328b6e
commit 531ca4c
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/cpp/src/arrow/integration/json_integration_test.cc b/cpp/src/arrow/integration/json_integration_test.cc
@@ -148,6 +148,16 @@ static Status ValidateFull(const RecordBatch& batch) {
   return Status::OK();
 }
 
+static Status ValidateEmbeddedStream(io::RandomAccessFile* arrow_file) {
+  RETURN_NOT_OK(arrow_file->Seek(8));
+  ARROW_ASSIGN_OR_RAISE(auto arrow_reader,
+                        ipc::RecordBatchStreamReader::Open(arrow_file));
+  for (auto maybe_batch : *arrow_reader) {
+    RETURN_NOT_OK(maybe_batch.status());
+  }
+  return arrow_file->Seek(0);
+}
+
 static Status ValidateArrowVsJson(const std::string& arrow_path,
                                   const std::string& json_path) {
   // Construct JSON reader
@@ -160,6 +170,7 @@ static Status ValidateArrowVsJson(const std::string& arrow_path,
 
   // Construct Arrow reader
   ARROW_ASSIGN_OR_RAISE(auto arrow_file, io::ReadableFile::Open(arrow_path));
+  RETURN_NOT_OK(ValidateEmbeddedStream(arrow_file.get()));
 
   std::shared_ptr<ipc::RecordBatchFileReader> arrow_reader;
   ARROW_ASSIGN_OR_RAISE(arrow_reader, ipc::RecordBatchFileReader::Open(arrow_file.get()));