feat: add top level cast

Author: ion-elgreco

python/datafusion/dataframe.py

@@ -245,6 +245,19 @@ def sort(self, *exprs: Expr | SortExpr) -> DataFrame:
         exprs_raw = [sort_or_default(expr) for expr in exprs]
         return DataFrame(self.df.sort(*exprs_raw))
 
+    def cast(self, mapping: dict[str, pa.DataType[Any]]) -> DataFrame:
+        """Cast all or a subset of columns to new dtype.
+
+        Args:
+            mapping (dict[str, pa.DataType[Any]]):  Mapped with column as key and column
+                dtype as value.
+
+        Returns:
+            DataFrame after casting columns
+        """
+        exprs = [Expr.column(col).cast(dtype) for col, dtype in mapping.items()]
+        return self.with_columns(exprs)
+
     def limit(self, count: int, offset: int = 0) -> DataFrame:
         """Return a new :py:class:`DataFrame` with a limited number of rows.
 

python/tests/test_dataframe.py

@@ -236,6 +236,15 @@ def test_with_columns(df):
     assert result.column(6) == pa.array([5, 7, 9])
 
 
+def test_cast(df):
+    df = df.cast({"a": pa.float16(), "b": pa.list_(pa.uint32())})
+    expected = pa.schema(
+        [("a", pa.float16()), ("b", pa.list_(pa.uint32())), ("c", pa.int64())]
+    )
+
+    assert df.schema() == expected
+
+
 def test_with_column_renamed(df):
     df = df.with_column("c", column("a") + column("b")).with_column_renamed("c", "sum")