Merge pull request #58 from xmnlab/master

Changed requirements for CI python35

ci/requirements-dev-3.5.yml

@@ -20,6 +20,7 @@ dependencies:
   - pyarrow>=0.6.0
   - pymapd>=0.3.2
   - pymysql
+  - pytables
   - pytest
   - pytest-xdist
   - python=3.5
@@ -30,6 +31,7 @@ dependencies:
   - ruamel.yaml
   - six
   - sqlalchemy>=1.0.0,<1.1.15
+  - thrift
   - toolz
   - xorg-libxpm
   - xorg-libxrender

ibis/bigquery/client.py

@@ -105,6 +105,11 @@ def columns(self):
         result = self.query.result()
         return [field.name for field in result.schema]
 
+    @property
+    def description(self):
+        result = self.query.result()
+        return [field for field in result.schema]
+
     def __enter__(self):
         # For compatibility when constructed from Query.execute()
         return self
@@ -379,6 +384,21 @@ def _get_table_schema(self, qualified_name):
         dataset, table = qualified_name.rsplit('.', 1)
         return self.get_schema(table, database=dataset)
 
+    def _get_schema_using_query(self, limited_query):
+        with self._execute(limited_query, results=True) as cur:
+            # resets the state of the cursor and closes operation
+            names, ibis_types = self._adapt_types(cur.description)
+        return sch.Schema(names, ibis_types)
+
+    def _adapt_types(self, descr):
+        names = []
+        adapted_types = []
+        for col in descr:
+            names.append(col.name)
+            typename = bigquery_field_to_ibis_dtype(col)
+            adapted_types.append(typename)
+        return names, adapted_types
+
     def _execute(self, stmt, results=True, query_parameters=None):
         job_config = bq.job.QueryJobConfig()
         job_config.query_parameters = query_parameters or []

ibis/bigquery/tests/test_client.py

@@ -527,3 +527,10 @@ def test_large_timestamp(client):
     expr = ibis.timestamp('4567-01-01 00:00:00')
     result = client.execute(expr)
     assert result == huge_timestamp
+
+
+def test_client_sql_query(client):
+    expr = client.sql('select * from testing.functional_alltypes limit 20')
+    result = expr.execute()
+    expected = client.table('functional_alltypes').head(20).execute()
+    tm.assert_frame_equal(result, expected)

ibis/bigquery/udf/api.py

@@ -19,7 +19,7 @@
 __all__ = 'udf',
 
 
-def udf(input_type, output_type, strict=True):
+def udf(input_type, output_type, strict=True, libraries=None):
     '''Define a UDF for BigQuery
 
     Parameters
@@ -128,6 +128,9 @@ class Rectangle {
     return my_rectangle(width, height);
     """;
     '''
+    if libraries is None:
+        libraries = []
+
     def wrapper(f):
         if not callable(f):
             raise TypeError('f must be callable, got {}'.format(f))
@@ -162,7 +165,7 @@ def compiles_udf_node(t, expr):
 LANGUAGE js AS """
 {strict}{source}
 return {name}({args});
-""";'''.format(
+"""{libraries};'''.format(
             name=f.__name__,
             return_type=ibis_type_to_bigquery_type(
                 dt.dtype(output_type), type_translation_context),
@@ -177,7 +180,12 @@ def compiles_udf_node(t, expr):
                 )
             ),
             strict=repr('use strict') + ';\n' if strict else '',
-            args=', '.join(inspect.signature(f).parameters.keys())
+            args=', '.join(inspect.signature(f).parameters.keys()),
+            libraries=(
+                '\nOPTIONS (\n    library={}\n)'.format(
+                    repr(list(libraries))
+                ) if libraries else ''
+            )
         )
 
         @functools.wraps(f)

ibis/bigquery/udf/tests/test_udf_execute.py

@@ -136,3 +136,22 @@ def my_add(x, y):
 SELECT my_add(1, 2) AS `tmp`'''
     result = client.execute(expr)
     assert result == 3
+
+
+def test_udf_libraries(client):
+    @udf(
+        [dt.Array(dt.string)],
+        dt.double,
+        # whatever symbols are exported in the library are visible inside the
+        # UDF, in this case lodash defines _ and we use that here
+        libraries=['gs://ibis-testing-libraries/lodash.min.js']
+    )
+    def string_length(strings):
+        return _.sum(_.map(strings, lambda x: x.length))  # noqa: F821
+
+    raw_data = ['aaa', 'bb', 'c']
+    data = ibis.literal(raw_data)
+    expr = string_length(data)
+    result = client.execute(expr)
+    expected = sum(map(len, raw_data))
+    assert result == expected