feat: add head, tail methods

Author: ion-elgreco

python/datafusion/dataframe.py

@@ -223,6 +223,30 @@ def limit(self, count: int, offset: int = 0) -> DataFrame:
         """
         return DataFrame(self.df.limit(count, offset))
 
+    def head(self, n: int) -> DataFrame:
+        """Return a new :py:class:`DataFrame` with a limited number of rows.
+
+        Args:
+            n: Number of rows to take from the head of the DataFrame.
+
+        Returns:
+            DataFrame after limiting.
+        """
+        return DataFrame(self.df.limit(n, 0))
+
+    def tail(self, n: int) -> DataFrame:
+        """Return a new :py:class:`DataFrame` with a limited number of rows.
+
+        Be aware this could be potentially expensive due to the size of the frame.
+
+        Args:
+            n: Number of rows to take from the tail of the DataFrame.
+
+        Returns:
+            DataFrame after limiting.
+        """
+        return DataFrame(self.df.limit(n, max(0, self.count() - n)))
+
     def collect(self) -> list[pa.RecordBatch]:
         """Execute this :py:class:`DataFrame` and collect results into memory.
 

python/tests/test_dataframe.py

@@ -190,6 +190,28 @@ def test_limit_with_offset(df):
     assert len(result.column(1)) == 1
 
 
+def test_head(df):
+    df = df.head(1)
+
+    # execute and collect the first (and only) batch
+    result = df.collect()[0]
+
+    assert result.column(0) == pa.array([1])
+    assert result.column(1) == pa.array([4])
+    assert result.column(2) == pa.array([8])
+
+
+def test_tail(df):
+    df = df.tail(1)
+
+    # execute and collect the first (and only) batch
+    result = df.collect()[0]
+
+    assert result.column(0) == pa.array([3])
+    assert result.column(1) == pa.array([6])
+    assert result.column(2) == pa.array([8])
+
+
 def test_with_column(df):
     df = df.with_column("c", column("a") + column("b"))