Repository: ibis-project/ibis

Issue:
FEAT: OmniScidb - Add "method" parameter to OmniSciDB load_table
Add to OmniSciDB load_data the paramter "method" (see https://pymapd.readthedocs.io/en/latest/api.html#pymapd.Connection.load_table)

method: string {‘infer’, ‘columnar’, ‘rows’, ‘arrow’}

Method to use for loading the data. Three options are available

- pyarrow and Apache Arrow loader
- columnar loader
- row-wise loader

The Arrow loader is typically the fastest, followed by the columnar loader, followed by the row-wise loader. If a DataFrame or pyarrow.Table is passed and pyarrow is installed, the Arrow-based loader will be used. If arrow isn’t available, the columnar loader is used. Finally, data is an iterable of tuples the row-wise loader is used.


Discussion:

Resolving patch:
diff --git a/ibis/backends/omniscidb/client.py b/ibis/backends/omniscidb/client.py
--- a/ibis/backends/omniscidb/client.py
+++ b/ibis/backends/omniscidb/client.py
@@ -4,6 +4,7 @@
 
 import pandas as pd
 import pkg_resources
+import pyarrow
 import pymapd
 import regex as re
 from pymapd._parsers import _extract_column_details
@@ -342,7 +343,7 @@ def truncate(self):
         """Delete all rows from, but do not drop, an existing table."""
         self._client.truncate_table(self._qualified_name)
 
-    def load_data(self, df):
+    def load_data(self, df: Union[pd.DataFrame, pyarrow.Table]):
         """
         Load a data frame into database.
 
@@ -1184,8 +1185,15 @@ def database(self, name=None):
             )
             return self.database_class(name, new_client)
 
-    def load_data(self, table_name, obj, database=None, **kwargs):
-        """Load data into a given table.
+    def load_data(
+        self,
+        table_name: str,
+        obj: Union[pd.DataFrame, pyarrow.Table],
+        database: Optional[str] = None,
+        method: str = 'rows',
+    ):
+        """
+        Load data into a given table.
 
         Wraps the LOAD DATA DDL statement. Loads data into an OmniSciDB table
         by physically moving data files.
@@ -1195,10 +1203,13 @@ def load_data(self, table_name, obj, database=None, **kwargs):
         table_name : string
         obj: pandas.DataFrame or pyarrow.Table
         database : string, optional
+        method : string, {‘infer’, ‘columnar’, ‘rows’, ‘arrow’}, default 'rows'
+            The Arrow loader is typically the fastest, followed by the columnar
+            loader, followed by the row-wise loader.
         """
         _database = self.db_name
         self.set_database(database)
-        self.con.load_table(table_name, obj)
+        self.con.load_table(table_name, obj, method=method)
         self.set_database(_database)
 
     @property
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -22,7 +22,7 @@
 postgres_requires = sqlite_requires + ['psycopg2']
 mysql_requires = sqlite_requires + ['pymysql']
 
-omniscidb_requires = ['pymapd==0.24']
+omniscidb_requires = ['pymapd==0.24', 'pyarrow']
 kerberos_requires = ['requests-kerberos']
 visualization_requires = ['graphviz']
 clickhouse_requires = [


Test patch:
diff --git a/ibis/backends/omniscidb/tests/test_client.py b/ibis/backends/omniscidb/tests/test_client.py
--- a/ibis/backends/omniscidb/tests/test_client.py
+++ b/ibis/backends/omniscidb/tests/test_client.py
@@ -3,6 +3,7 @@
 
 import mock
 import pandas as pd
+import pyarrow
 import pytest
 from pkg_resources import get_distribution, parse_version
 from pytest import param
@@ -14,6 +15,7 @@
 
 pymapd = pytest.importorskip('pymapd')
 
+
 pytestmark = pytest.mark.omniscidb
 
 
@@ -287,3 +289,45 @@ def test_cpu_execution_type(
 
     for mocked_method in mocked_methods:
         mocked_method.stop()
+
+
+@pytest.mark.parametrize(
+    'method,format',
+    [
+        ('rows', 'pandas'),
+        ('columnar', 'pandas'),
+        ('infer' 'pandas'),
+        ('infer', 'arrow'),
+        ('arrow', 'arrow'),
+    ],
+)
+def test_load_data(con, temp_table, method, format):
+    df_salary = pd.DataFrame(
+        {
+            'first_name': ['A', 'B', 'C'],
+            'last_name': ['D', 'E', 'F'],
+            'department_name': ['AA', 'BB', 'CC'],
+            'salary': [100.0, 200.0, 300.0],
+        }
+    )
+
+    sch_salary = ibis.schema(
+        [
+            ('first_name', 'string'),
+            ('last_name', 'string'),
+            ('department_name', 'string'),
+            ('salary', 'float64'),
+        ]
+    )
+
+    arrow_salary = pyarrow.Table.from_pandas(df_salary)
+
+    con.create_table(temp_table, schema=sch_salary)
+    con.load_data(
+        temp_table,
+        df_salary if format == 'pandas' else arrow_salary,
+        method=method,
+    )
+    result = con.table(temp_table).execute()
+
+    pd.testing.assert_frame_equal(result, df_salary)