diff --git a/parquet_tools/commands/csv.py b/parquet_tools/commands/csv.py index 4a777f9..84c28f6 100644 --- a/parquet_tools/commands/csv.py +++ b/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)) diff --git a/parquet_tools/commands/show.py b/parquet_tools/commands/show.py index df2ca20..db3fce5 100644 --- a/parquet_tools/commands/show.py +++ b/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)) diff --git a/tests/test_csv.py b/tests/test_csv.py index 354d73a..516f733 100644 --- a/tests/test_csv.py +++ b/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' + diff --git a/tests/test_show.py b/tests/test_show.py index bc57578..5fd3e70 100644 --- a/tests/test_show.py +++ b/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' +