diff --git a/docs/cli.rst b/docs/cli.rst index 7a127fe..46ee335 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -665,6 +665,17 @@ The ``most_common`` and ``least_common`` columns will contain nested JSON arrays .. _cli_inserting_data: +Inserting Parquet data +====================== + +Parquet is a columnar storage format, frequently used in the Hadoop/Spark ecosystem as well as cloud providers. +Parquet files can be inserted via the ``--parquet`` flag. +Here's an example:: + + $ sqlite-utils insert --parquet data.db data ./data.parquet + +Parquet files, along with the data store data types too, thus making the ``--detect-types`` flag redundant, as data types are inferred automatically (using PyArrow) + Inserting JSON data =================== diff --git a/setup.py b/setup.py index 091c94a..0120e31 100644 --- a/setup.py +++ b/setup.py @@ -28,6 +28,7 @@ setup( "click-default-group", "tabulate", "dateutils", + "pyarrow" ], setup_requires=["pytest-runner"], extras_require={ diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index b29314e..cae1774 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -16,6 +16,7 @@ import os import sys import csv as csv_std import tabulate +import pyarrow.parquet as pq from .utils import ( file_progress, find_spatialite, @@ -25,6 +26,7 @@ from .utils import ( rows_from_file, Format, TypeTracker, + pyarrow_tuple_as_py ) CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"]) @@ -646,6 +648,7 @@ def insert_upsert_options(fn): click.option("--flatten", is_flag=True, help="Flatten nested JSON objects"), click.option("-c", "--csv", is_flag=True, help="Expect CSV"), click.option("--tsv", is_flag=True, help="Expect TSV"), + click.option("--parquet", is_flag=True, help="Expect Parquet"), click.option("--delimiter", help="Delimiter to use for CSV files"), click.option("--quotechar", help="Quote character to use for CSV/TSV"), click.option( @@ -682,7 +685,7 @@ def insert_upsert_options(fn): "--detect-types", is_flag=True, envvar="SQLITE_UTILS_DETECT_TYPES", - help="Detect types for columns in CSV/TSV data", + help="Detect types for columns in CSV/TSV. For Parquet, types are detected automatically", ), load_extension_option, click.option("--silent", is_flag=True, help="Do not show progress bar"), @@ -701,6 +704,7 @@ def insert_upsert_implementation( flatten, csv, tsv, + parquet, delimiter, quotechar, sniff, @@ -722,10 +726,10 @@ def insert_upsert_implementation( _load_extensions(db, load_extension) if (delimiter or quotechar or sniff or no_headers) and not tsv: csv = True - if (nl + csv + tsv) >= 2: - raise click.ClickException("Use just one of --nl, --csv or --tsv") - if (csv or tsv) and flatten: - raise click.ClickException("--flatten cannot be used with --csv or --tsv") + if (nl + csv + tsv + parquet) >= 2: + raise click.ClickException("Use just one of --nl, --csv, --tsv or --parquet") + if (csv or tsv or parquet) and flatten: + raise click.ClickException("--flatten cannot be used with --csv, --tsv or --parquet") if encoding and not (csv or tsv): raise click.ClickException("--encoding must be used with --csv or --tsv") if pk and len(pk) == 1: @@ -758,6 +762,28 @@ def insert_upsert_implementation( if detect_types: tracker = TypeTracker() docs = tracker.wrap(docs) + elif parquet: + """ + If path points to a file then we assume this file should go into the table. + It's often the case though were larger tables/dataframes are split into smaller + parquet files. In this case, if path points to a directory we read all the files in the directory and insert them. + The end-result should be the same as if multiple calls (with a single input file) were made. + """ + if pathlib.Path(json_file.name).is_file(): + # Maybe use buffered reader, or set buffer_size ? + pq_file = pq.ParquetFile(json_file) + columns = pq_file.schema.names + parquet_table = pq_file.read() # can pass columns=[...] + #TODO: should probably read in batches: see parquet_table.to_batches() + #TODO: the below: pyarrow_tuple_as_py(r) probably slows down things too much. + docs = (dict(zip(columns, pyarrow_tuple_as_py(r))) for r in zip(*parquet_table.columns)) + elif pathlib.Path(json_file.name).is_dir(): + for f in pathlib.Path(json_file.name).iterdir(): + raise NotImplementedError("Parquet directories are not supported yet") + else: + # Shouldn't happen + pass + else: try: if nl: @@ -768,7 +794,7 @@ def insert_upsert_implementation( docs = [docs] except json.decoder.JSONDecodeError: raise click.ClickException( - "Invalid JSON - use --csv for CSV or --tsv for TSV files" + "Invalid JSON - use --csv for CSV, --tsv for TSV or --parquet for Parquet files" ) if flatten: docs = (dict(_flatten(doc)) for doc in docs) @@ -782,6 +808,7 @@ def insert_upsert_implementation( extra_kwargs["upsert"] = upsert # Apply {"$base64": true, ...} decoding, if needed docs = (decode_base64_values(doc) for doc in docs) + try: db[table].insert_all( docs, pk=pk, batch_size=batch_size, alter=alter, **extra_kwargs @@ -856,6 +883,7 @@ def insert( flatten, csv, tsv, + parquet, delimiter, quotechar, sniff, @@ -876,7 +904,7 @@ def insert( Insert records from JSON file into a table, creating the table if it does not already exist. - Input should be a JSON array of objects, unless --nl or --csv is used. + Input should be a JSON array of objects, unless --nl, --csv or --parquet is used. """ try: insert_upsert_implementation( @@ -888,6 +916,7 @@ def insert( flatten, csv, tsv, + parquet, delimiter, quotechar, sniff, diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 00a3c02..6f55615 100644 --- a/sqlite_utils/utils.py +++ b/sqlite_utils/utils.py @@ -278,3 +278,7 @@ def progressbar(*args, **kwargs): else: with click.progressbar(*args, **kwargs) as bar: yield bar + + +def pyarrow_tuple_as_py(pyarrow_tuple): + return tuple(map(lambda f: f.as_py(), pyarrow_tuple)) \ No newline at end of file diff --git a/tests/test_cli.py b/tests/test_cli.py index 47920c6..cd2e8aa 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -7,6 +7,8 @@ import os import pytest from sqlite_utils.utils import sqlite3, find_spatialite import textwrap +from pyarrow import Table +from pyarrow.parquet import ParquetWriter from .utils import collapse_whitespace @@ -822,6 +824,27 @@ def test_insert_csv_tsv(content, options, db_path, tmpdir): assert [{"foo": "1", "bar": "2", "baz": "cat,dog"}] == list(db["data"].rows) +@pytest.mark.parametrize( + "content,options", + [ + ({"foo": [1, 11], "bar": [2, 22], "baz": ["cat,dog", "animal"]}, ["--parquet"]) + ], +) +def test_insert_parquet(content, options, db_path, tmpdir): + db = Database(db_path) + file_path = str(tmpdir / "insert.parquet") + arrow_table = Table.from_pydict(content) + ParquetWriter(file_path, arrow_table.schema).write_table(arrow_table) + result = CliRunner().invoke( + cli.cli, + ["insert", db_path, "data", file_path] + options, + catch_exceptions=False, + ) + assert 0 == result.exit_code + assert [{"foo": 1, "bar": 2, "baz": "cat,dog"}, + {"foo": 11, "bar": 22, "baz": "animal"}] == list(db["data"].rows) + + @pytest.mark.parametrize( "options", ( @@ -829,16 +852,18 @@ def test_insert_csv_tsv(content, options, db_path, tmpdir): ["--tsv", "--csv"], ["--csv", "--nl"], ["--csv", "--nl", "--tsv"], + ["--csv", "--nl", "--parquet"], + ["--csv", "--parquet"], ), ) -def test_only_allow_one_of_nl_tsv_csv(options, db_path, tmpdir): +def test_only_allow_one_of_nl_tsv_csv_parquet(options, db_path, tmpdir): file_path = str(tmpdir / "insert.csv-tsv") open(file_path, "w").write("foo") result = CliRunner().invoke( cli.cli, ["insert", db_path, "data", file_path] + options ) assert 0 != result.exit_code - assert "Error: Use just one of --nl, --csv or --tsv" == result.output.strip() + assert "Error: Use just one of --nl, --csv, --tsv or --parquet" == result.output.strip() def test_insert_replace(db_path, tmpdir):