feat: add datafusion to parity tests (#1108)

* feat: add datafusion to parity tests

* fix: apply ruff linter rules

* fix: mark failing parity tests

* fix: reorganize parity tests

---------

Co-authored-by: tokoko <[email protected]>

.devcontainer/Dockerfile

@@ -0,0 +1,5 @@
+FROM mcr.microsoft.com/vscode/devcontainers/python:3.10-buster
+USER vscode
+RUN curl -sSL https://install.python-poetry.org | python -
+RUN /home/vscode/.local/bin/poetry config virtualenvs.in-project true
+USER root

.devcontainer/devcontainer.json

@@ -0,0 +1,24 @@
+{
+	"name": "ibis-substrait-devcontainer",
+	"build": {
+		"context": "..",
+		"dockerfile": "Dockerfile"
+	},
+
+	// Features to add to the dev container. More info: https://containers.dev/features.
+	"features": {
+		"ghcr.io/devcontainers/features/nix:1": {}
+	},
+
+	// Use 'forwardPorts' to make a list of ports inside the container available locally.
+	// "forwardPorts": [],
+
+	// Use 'postCreateCommand' to run commands after the container is created.
+	"postCreateCommand": "poetry install"
+
+	// Configure tool-specific properties.
+	// "customizations": {},
+
+	// Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root.
+	// "remoteUser": "root"
+}

ibis_substrait/compiler/translate.py

@@ -883,7 +883,7 @@ def _translate_join_type(join_kind: ops.JoinKind) -> stalg.JoinRel.JoinType.V:
         "left": stalg.JoinRel.JoinType.JOIN_TYPE_LEFT,
         "right": stalg.JoinRel.JoinType.JOIN_TYPE_RIGHT,
         "outer": stalg.JoinRel.JoinType.JOIN_TYPE_OUTER,
-        "semi": stalg.JoinRel.JoinType.JOIN_TYPE_SEMI,
+        "semi": stalg.JoinRel.JoinType.JOIN_TYPE_LEFT_SEMI,
         # "asof",
         # "anti",
         # "any_inner",

ibis_substrait/tests/compiler/parity_utils.py

@@ -0,0 +1,65 @@
+from abc import ABC, abstractmethod
+
+import datafusion
+import pyarrow as pa
+
+
+class SubstraitConsumer(ABC):
+    @abstractmethod
+    def with_tables(self, datasets: dict[str, pa.Table]):
+        pass
+
+    @abstractmethod
+    def execute(self, plan) -> pa.Table:
+        pass
+
+
+class AceroSubstraitConsumer(SubstraitConsumer):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def with_tables(self, datasets: dict[str, pa.Table]):
+        self.datasets = datasets
+        return self
+
+    def execute(self, plan) -> pa.Table:
+        import pyarrow.substrait as pa_substrait
+
+        def get_table_provider(datasets):
+            def table_provider(names, schema):
+                return datasets[names[0]]
+
+            return table_provider
+
+        query_bytes = plan.SerializeToString()
+        result = pa_substrait.run_query(
+            pa.py_buffer(query_bytes),
+            table_provider=get_table_provider(self.datasets),
+        )
+
+        return result.read_all()
+
+
+class DatafusionSubstraitConsumer(SubstraitConsumer):
+    def __init__(self) -> None:
+        self.connection = datafusion.SessionContext()
+
+    def with_tables(self, datasets: dict[str, pa.Table]):
+        for k, v in datasets.items():
+            self.connection.deregister_table(k)
+            self.connection.register_record_batches(k, [v.to_batches()])
+        return self
+
+    def execute(self, plan) -> pa.Table:
+        plan_data = plan.SerializeToString()
+        substrait_plan = datafusion.substrait.serde.deserialize_bytes(plan_data)
+        logical_plan = datafusion.substrait.consumer.from_substrait_plan(
+            self.connection, substrait_plan
+        )
+
+        df = self.connection.create_dataframe_from_logical_plan(logical_plan)
+        for column_number, column_name in enumerate(df.schema().names):
+            df = df.with_column_renamed(
+                column_name, plan.relations[0].root.names[column_number]
+            )
+        return df.to_arrow_table()

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h01/tpc_h01.json

@@ -694,7 +694,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h02/tpc_h02.json

@@ -1285,7 +1285,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h03/tpc_h03.json

@@ -728,7 +728,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h04/tpc_h04.json

@@ -506,7 +506,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h05/tpc_h05.json

@@ -1012,7 +1012,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h06/tpc_h06.json

@@ -439,7 +439,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h07/tpc_h07.json

@@ -1032,7 +1032,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h08/tpc_h08.json

@@ -1358,7 +1358,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h09/tpc_h09.json

@@ -878,7 +878,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h10/tpc_h10.json

@@ -843,7 +843,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h11/tpc_h11.json

@@ -863,7 +863,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h12/tpc_h12.json

@@ -774,7 +774,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h13/tpc_h13.json

@@ -421,7 +421,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h14/tpc_h14.json

@@ -775,7 +775,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h15/tpc_h15.json

@@ -1085,7 +1085,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h16/tpc_h16.json

@@ -704,7 +704,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h17/tpc_h17.json

@@ -773,7 +773,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h18/tpc_h18.json

@@ -774,7 +774,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h19/tpc_h19.json

@@ -1426,7 +1426,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h20/tpc_h20.json

@@ -1012,7 +1012,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h21/tpc_h21.json

@@ -1245,7 +1245,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }

ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h22/tpc_h22.json

@@ -879,7 +879,7 @@
     }
   ],
   "version": {
-    "minorNumber": 52,
+    "minorNumber": 54,
     "producer": "ibis-substrait"
   }
 }