Modify add_column() to optionally accept a pyarrow schema as param

Author: Varad Bhatnagar

src/datasets/arrow_dataset.py

@@ -5613,7 +5613,9 @@ def push_to_hub(
 
     @transmit_format
     @fingerprint_transform(inplace=False)
-    def add_column(self, name: str, column: Union[list, np.array], new_fingerprint: str):
+    def add_column(
+        self, name: str, column: Union[list, np.array], new_fingerprint: str, pyarrow_schema: pa.schema = None
+    ):
         """Add column to Dataset.
 
         <Added version="1.7"/>
@@ -5640,7 +5642,7 @@ def add_column(self, name: str, column: Union[list, np.array], new_fingerprint:
         })
         ```
         """
-        column_table = InMemoryTable.from_pydict({name: column})
+        column_table = InMemoryTable.from_pydict({name: column}, schema=pyarrow_schema)
         _check_column_names(self._data.column_names + column_table.column_names)
         dataset = self.flatten_indices() if self._indices is not None else self
         # Concatenate tables horizontally