fix: add missing arguments to Producer.produce_substrait (#126)

Some producers had an additional parameter `validate` but the interface and the IbisProducer as well as one call site were missing it. This PR introduces that parameter in those places. In also changes the formatting in one of the producers to be consistent with the others. I have manually test: no test outcome changes. Several tests for the IbisProducer reported a wrong number of arguments before and now fail because the producer cannot handle the test cases. Signed-off-by: Ingo Müller <[email protected]>
substrait-io · Nov 5, 2024 · 4cc3091 · 4cc3091
1 parent 723bf0f
commit 4cc3091
Show file tree

Hide file tree

Showing 5 changed files with 5 additions and 5 deletions.
diff --git a/substrait_consumer/functional/common.py b/substrait_consumer/functional/common.py
@@ -135,7 +135,7 @@ def substrait_producer_sql_test(
     # Convert the SQL/Ibis expression to a substrait query plan
     if type(producer).__name__ == "IbisProducer":
         if ibis_expr:
-            substrait_plan = producer.produce_substrait(sql_query, ibis_expr(*args))
+            substrait_plan = producer.produce_substrait(sql_query, validate, ibis_expr(*args))
         else:
             pytest.xfail("ibis expression currently undefined")
     else:

diff --git a/substrait_consumer/producers/ibis_producer.py b/substrait_consumer/producers/ibis_producer.py
@@ -24,7 +24,7 @@ def __init__(self, db_connection=None):
     def set_db_connection(self, db_connection):
         self._db_connection = db_connection
 
-    def produce_substrait(self, sql_query: str, ibis_expr: str = None) -> str:
+    def produce_substrait(self, sql_query: str, validate = False, ibis_expr: str = None) -> str:
         """
         Produce the Ibis substrait plan using the given Ibis expression
 

diff --git a/substrait_consumer/producers/isthmus_producer.py b/substrait_consumer/producers/isthmus_producer.py
@@ -23,7 +23,7 @@ def __init__(self, db_connection=None):
     def set_db_connection(self, db_connection):
         self._db_connection = db_connection
 
-    def produce_substrait(self, sql_query: str, validate=False, ibis_expr: str=None) -> str:
+    def produce_substrait(self, sql_query: str, validate = False, ibis_expr: str = None) -> str:
         """
         Produce the Isthmus substrait plan using the given SQL query.
 

diff --git a/substrait_consumer/producers/producer.py b/substrait_consumer/producers/producer.py
@@ -12,7 +12,7 @@ def set_db_connection(self, db_connection):
         pass
 
     @abstractmethod
-    def produce_substrait(self, sql_query: str, ibis_expr: str = None) -> str:
+    def produce_substrait(self, sql_query: str, validate = False, ibis_expr: str = None) -> str:
         pass
 
     @abstractmethod

diff --git a/substrait_consumer/tests/functional/extension_functions/test_substrait_function_names.py b/substrait_consumer/tests/functional/extension_functions/test_substrait_function_names.py
@@ -191,7 +191,7 @@ def run_function_name_test(
         if type(producer).__name__ == "IbisProducer":
             if ibis_expr:
                 substrait_plan = producer.produce_substrait(
-                    sql_query, DuckDBConsumer, ibis_expr(*args)
+                    sql_query, validate=False, ibis_expr=ibis_expr(*args)
                 )
                 substrait_plan = json.loads(substrait_plan)
             else: