Merge pull request #259 from goodwanghan/fsql

FugueSQL implementation

.github/workflows/test-package.yml

@@ -104,6 +104,7 @@ jobs:
  - name: Install datacompy
  run: |
  python -m pip install --upgrade pip
+ python -m pip install "pandas<2.2"
  python -m pip install .[tests,duckdb,polars,dask,ray]
  - name: Test with pytest
  run: |

datacompy/_fsql_utils.py

@@ -0,0 +1,106 @@
+#
+# Copyright 2024 Capital One Services, LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Fugue SQL utils
+"""
+
+from contextlib import contextmanager
+from typing import Any, Callable, Dict, Iterator, List, Tuple, Iterable
+
+import duckdb
+import fugue.api as fa
+import pandas as pd
+from fsspec.implementations.local import LocalFileSystem
+from fugue import ExecutionEngine, NativeExecutionEngine
+from triad.utils.io import url_to_fs
+
+_CONF_GENERATORS: List[Tuple[Callable[[ExecutionEngine], bool], Any]] = []
+_DUMMY_DF = pd.DataFrame([[0]], columns=["a"])
+
+
+def _compare_conf_by_fugue_engine(
+ check: Callable[[ExecutionEngine], bool],
+) -> Any:
+ def wrapper(func: Any) -> Any:
+ _f = contextmanager(func)
+ _CONF_GENERATORS.append((check, _f))
+ return func
+
+ return wrapper
+
+
+@contextmanager
+def infer_fugue_engine(df1: Any, df2: Any) -> Iterator[Dict[str, Any]]:
+ infer_by = list(_get_infer_dfs(df1, df2))
+ with fa.engine_context(infer_by=infer_by) as engine:
+ for check, func in _CONF_GENERATORS:
+ if check(engine):
+ with func(engine) as conf:
+ yield conf
+ return
+ yield dict(
+ engine=engine,
+ persist_diff=False,
+ use_map=False,
+ num_buckets=1,
+ )
+
+
+@_compare_conf_by_fugue_engine(
+ lambda e: not e.is_distributed and isinstance(e, NativeExecutionEngine)
+)
+def _on_native_engine(_: ExecutionEngine) -> Iterator[Dict[str, Any]]:
+ with duckdb.connect() as con:
+ with fa.engine_context(con) as e:
+ yield dict(
+ engine=e,
+ persist_diff=False,
+ use_map=False,
+ num_buckets=1,
+ )
+
+
+try:
+ from fugue_ray import RayExecutionEngine
+
+ @_compare_conf_by_fugue_engine(lambda e: isinstance(e, RayExecutionEngine))
+ def _on_ray_engine(engine: ExecutionEngine) -> Iterator[Dict[str, Any]]:
+ yield dict(
+ engine=engine,
+ persist_diff=False,
+ use_map=True,
+ num_buckets=engine.get_current_parallelism() * 2,
+ )
+
+except (ImportError, ModuleNotFoundError):
+ pass
+
+
+def _get_infer_dfs(*dfs: Any) -> Iterable[Any]:
+ for df in dfs:
+ if isinstance(df, str):
+ if _is_local_path(df):
+ yield _DUMMY_DF
+ else:
+ yield df
+
+
+def _is_local_path(path: Any) -> bool:
+ try:
+ fs, _ = url_to_fs(path)
+ return isinstance(fs, LocalFileSystem)
+ except Exception:
+ return False