Adding a "Discover" object for finding candidate tables that may be joined on the main table

.gitignore

@@ -20,6 +20,9 @@ var/
 .installed.cfg
 *.egg
 *.pkl
+data/
+
+data/
 
 .pytest_cache/
 

10_discovery.py

@@ -0,0 +1,24 @@
+# %%
+import polars as pl
+
+from skrub._discover import Discover, find_unique_values
+
+# %%
+# working with binary to debug
+data_lake_path = "data/binary_update/*.parquet"
+base_table_path = "data/source_tables/company_employees-yadl-depleted.parquet"
+query_column = "col_to_embed"
+
+
+base_table = pl.read_parquet(base_table_path)
+# %%
+find_unique_values(base_table, ["col_to_embed"])
+# %%
+discover = Discover(data_lake_path, [query_column])
+print("fitting")
+discover.fit(base_table)
+print("transforming")
+ranking = discover.transform(base_table)
+print(ranking)
+
+# %%

main_discover.py

@@ -0,0 +1,18 @@
+import polars as pl
+
+from skrub._discover import Discover
+
+if __name__ == "__main__":
+    # working with binary to debug
+    data_lake_path = "data/binary_update/*.parquet"
+    base_table_path = "data/source_tables/yadl/movies_large-yadl-depleted.parquet"
+    query_column = "col_to_embed"
+
+    base_table = pl.read_parquet(base_table_path)
+
+    discover = Discover(data_lake_path, [query_column])
+    print("fitting")
+    discover.fit(base_table)
+    print("transforming")
+    joined_table = discover.transform(base_table)
+    print(joined_table)

skrub/_dataframe/_common.py

@@ -43,6 +43,11 @@
     "col_by_idx",
     "collect",
     #
+    # Loading data
+    #
+    "read_parquet",
+    "read_csv",
+    #
     # Querying and modifying metadata
     #
     "shape",
@@ -397,6 +402,43 @@ def _collect_polars_lazyframe(df):
     return df.collect()
 
 
+#
+# Loading data
+# ============
+#
+
+
+# TODO: Adding X here as a placeholder to get around the type check,
+@dispatch
+def read_parquet(X, input_path):
+    raise NotImplementedError()
+
+
+@read_parquet.specialize("pandas", argument_type=["DataFrame"])
+def _read_parquet_pandas(X, input_path):
+    return pd.read_parquet(input_path)
+
+
+@read_parquet.specialize("polars", argument_type=["DataFrame"])
+def _read_parquet_polars(X, input_path):
+    return pl.read_parquet(input_path)
+
+
+@dispatch
+def read_csv(X, input_path):
+    raise NotImplementedError()
+
+
+@read_csv.specialize("pandas", argument_type=["DataFrame"])
+def _read_csv_pandas(X, input_path):
+    return pd.read_csv(input_path)
+
+
+@read_csv.specialize("polars", argument_type=["DataFrame"])
+def _read_csv_polars(X, input_path):
+    return pl.read_csv(input_path)
+
+
 #
 # Querying and modifying metadata
 # ===============================