From b185d860c4a8c5611002e59073bd1a1b43f32bad Mon Sep 17 00:00:00 2001
From: ZENOTME <43447882+ZENOTME@users.noreply.github.com>
Date: Mon, 31 Jul 2023 16:15:31 +0800
Subject: [PATCH] fix: return Vec<Datafile> instead of HashMap in TaskWriter
 (#126)

---
 icelake/src/io/task_writer.rs | 13 ++++---------
 1 file changed, 4 insertions(+), 9 deletions(-)

diff --git a/icelake/src/io/task_writer.rs b/icelake/src/io/task_writer.rs
index 114dece..9caa12a 100644
--- a/icelake/src/io/task_writer.rs
+++ b/icelake/src/io/task_writer.rs
@@ -1,8 +1,6 @@
 //! task_writer module provide a task writer for writing data in a table.
 //! table writer used directly by the compute engine.
 
-use std::collections::HashMap;
-
 use arrow_array::RecordBatch;
 use arrow_schema::Schema as ArrowSchema;
 use opendal::Operator;
@@ -11,7 +9,7 @@ use super::data_file_writer::DataFileWriter;
 use super::location_generator;
 use crate::error::Result;
 use crate::io::location_generator::DataFileLocationGenerator;
-use crate::types::{DataFile, StructValue, TableMetadata};
+use crate::types::{DataFile, TableMetadata};
 
 /// `TaskWriter` used to write data for a table.
 ///
@@ -91,7 +89,7 @@ impl TaskWriter {
     }
 
     /// Close the writer and return the data files.
-    pub async fn close(self) -> Result<HashMap<StructValue, Vec<DataFile>>> {
+    pub async fn close(self) -> Result<Vec<DataFile>> {
         match self {
             Self::Unpartitioned(writer) => writer.close().await,
         }
@@ -137,10 +135,7 @@ impl UnpartitionedWriter {
     /// # Note
     ///
     /// For unpartitioned table, the key of the result map is default partition key.
-    pub async fn close(self) -> Result<HashMap<StructValue, Vec<DataFile>>> {
-        let datafiles = self.data_file_writer.close().await?;
-        let mut result = HashMap::new();
-        result.insert(StructValue::default(), datafiles);
-        Ok(result)
+    pub async fn close(self) -> Result<Vec<DataFile>> {
+        self.data_file_writer.close().await
     }
 }