From e270fb60d062307b34a6a342cefcc091ecd44a25 Mon Sep 17 00:00:00 2001
From: Xianjin YE <xianjin@apache.org>
Date: Sun, 25 Feb 2024 10:33:45 +0800
Subject: [PATCH] address comments

---
 .../execution/datafusion/shuffle_writer.rs    | 32 ++++++++++++-
 .../shuffle_writer/shuffle_writer_test.rs     | 48 -------------------
 2 files changed, 31 insertions(+), 49 deletions(-)
 delete mode 100644 core/src/execution/datafusion/shuffle_writer/shuffle_writer_test.rs

diff --git a/core/src/execution/datafusion/shuffle_writer.rs b/core/src/execution/datafusion/shuffle_writer.rs
index 480ba52d9..fc15facb4 100644
--- a/core/src/execution/datafusion/shuffle_writer.rs
+++ b/core/src/execution/datafusion/shuffle_writer.rs
@@ -1282,4 +1282,34 @@ impl RecordBatchStream for EmptyStream {
 }
 
 #[cfg(test)]
-mod shuffle_writer_test;
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_slot_size() {
+        let batch_size = 1usize;
+        // not inclusive of all supported types, but enough to test the function
+        let supported_primitive_types = [
+            DataType::Int32,
+            DataType::Int64,
+            DataType::UInt32,
+            DataType::UInt64,
+            DataType::Float32,
+            DataType::Float64,
+            DataType::Boolean,
+            DataType::Utf8,
+            DataType::LargeUtf8,
+            DataType::Binary,
+            DataType::LargeBinary,
+            DataType::FixedSizeBinary(16),
+        ];
+        let expected_slot_size = [4, 8, 4, 8, 4, 8, 1, 104, 108, 104, 108, 16];
+        supported_primitive_types
+            .iter()
+            .zip(expected_slot_size.iter())
+            .for_each(|(data_type, expected)| {
+                let slot_size = slot_size(batch_size, data_type);
+                assert_eq!(slot_size, *expected);
+            })
+    }
+}
diff --git a/core/src/execution/datafusion/shuffle_writer/shuffle_writer_test.rs b/core/src/execution/datafusion/shuffle_writer/shuffle_writer_test.rs
deleted file mode 100644
index e260f9e05..000000000
--- a/core/src/execution/datafusion/shuffle_writer/shuffle_writer_test.rs
+++ /dev/null
@@ -1,48 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements.  See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership.  The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License.  You may obtain a copy of the License at
-//
-//   http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied.  See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-//! unit tests for the shuffle writer
-
-use super::*;
-
-#[test]
-fn test_slot_size() {
-    let batch_size = 1usize;
-    // not inclusive of all supported types, but enough to test the function
-    let supported_primitive_types = [
-        DataType::Int32,
-        DataType::Int64,
-        DataType::UInt32,
-        DataType::UInt64,
-        DataType::Float32,
-        DataType::Float64,
-        DataType::Boolean,
-        DataType::Utf8,
-        DataType::LargeUtf8,
-        DataType::Binary,
-        DataType::LargeBinary,
-        DataType::FixedSizeBinary(16),
-    ];
-    let expected_slot_size = [4, 8, 4, 8, 4, 8, 1, 104, 108, 104, 108, 16];
-    supported_primitive_types
-        .iter()
-        .zip(expected_slot_size.iter())
-        .for_each(|(data_type, expected)| {
-            let slot_size = slot_size(batch_size, data_type);
-            assert_eq!(slot_size, *expected);
-        })
-}