Fix: honour --head when --columns is provided. (#23)

This fix ensures that only the number of rows specified by the client
will be output when the --columns argument is supplied (fixes
#22).

parquet_tools/commands/csv.py

@@ -58,5 +58,5 @@ def _execute(df: pd.DataFrame, head: int, columns: list) -> None:
     # head
     df_head: pd.DataFrame = df.head(head) if head > 0 else df
     # select columns
-    df_select: pd.DataFrame = df[columns] if len(columns) else df_head
+    df_select: pd.DataFrame = df_head[columns] if len(columns) else df_head
     print(df_select.to_csv(index=None))

parquet_tools/commands/show.py

@@ -75,5 +75,5 @@ def _execute(df: pd.DataFrame, format: str, head: int, columns: list) -> None:
     # head
     df_head: pd.DataFrame = df.head(head) if head > 0 else df
     # select columns
-    df_select: pd.DataFrame = df[columns] if len(columns) else df_head
+    df_select: pd.DataFrame = df_head[columns] if len(columns) else df_head
     print(tabulate(df_select, df_select.columns, tablefmt=format, showindex=False))

tests/test_csv.py

@@ -75,9 +75,12 @@ def test_excute(capfd, parquet_file):
              'two': ['foo', 'bar', 'baz'],
              'three': [True, False, True]}
         ),
-        head=-1,
-        columns=[]
+        head=2,
+        columns=['one', 'three']
     )
     out, err = capfd.readouterr()
     assert out is not None
     assert err == ''
+    assert 'foo' not in out, 'Column two should not be output'
+    assert '2.5' not in out, 'Row 3 should not be output'
+

tests/test_show.py

@@ -79,9 +79,13 @@ def test_excute(capfd, parquet_file):
              'three': [True, False, True]}
         ),
         format='psql',
-        head=-1,
-        columns=[]
+        head=2,
+        columns=['one', 'three']
     )
     out, err = capfd.readouterr()
+
     assert out is not None
     assert err == ''
+    assert 'foo' not in out, 'Column two should not be output'
+    assert '2.5' not in out, 'Row 3 should not be output'
+