Add implementation to read --parquet files.

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
 ===================
 

setup.py

@@ -28,6 +28,7 @@ setup(
         "click-default-group",
         "tabulate",
         "dateutils",
+        "pyarrow"
     ],
     setup_requires=["pytest-runner"],
     extras_require={

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,

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))

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):