Added Dataframe output format

Author: nmreadelf

chdb/__init__.py

@@ -1,5 +1,7 @@
 import sys
 import os
+import pyarrow as pa
+
 
 chdb_version = (0, 1, 0)
 if sys.version_info[:2] >= (3, 7):
@@ -22,7 +24,19 @@
 except:  # pragma: no cover
     __version__ = "unknown"
 
-# wrap _chdb functions
 
+def _to_arrowTable(res):
+    """convert res to arrow table"""
+    return pa.RecordBatchFileReader(res.get_memview()).read_all()
+
+def to_df(r):
+    """"convert arrow table to Dataframe"""
+    t = _to_arrowTable(r)
+    return t.to_pandas(use_threads=True)
+
+# wrap _chdb functions
 def query(sql, output_format="CSV", **kwargs):
+    if output_format.lower() == "dataframe":
+        r = _chdb.query(sql, "Arrow", **kwargs)
+        return to_df(r)
     return _chdb.query(sql, output_format, **kwargs)

chdb/__main__.py

@@ -0,0 +1,16 @@
+import sys
+from .__init__ import query
+
+def main():
+    if len(sys.argv) < 2:
+        exit(0)
+    output_format = sys.argv[2] if len(sys.argv) > 2 else "CSV"
+    res = query(sys.argv[1], output_format)
+    if output_format.lower() == 'dataframe':
+        temp = res
+    else:
+        temp = res.data()
+    print(temp, end="")
+
+if __name__ == '__main__':
+    main()

chdb/test_smoke.sh

@@ -24,3 +24,5 @@ python3 -c \
 python3 -c \
     "import chdb; res = chdb.query('select version()', 'CSV'); print(str(res.get_memview().tobytes()))"
 
+# test cli
+python3 -m chdb "select 1112222222,555" Dataframe

setup.py

@@ -155,6 +155,7 @@ def build_extensions(self):
             exclude_package_data={'': ['*.pyc', 'src/**']},
             ext_modules=ext_modules,
             python_requires='>=3.7',
+            install_requires=['pyarrow', 'pandas'],
             cmdclass={'build_ext': BuildExt},
             test_suite="tests",
             zip_safe=False,