From c2131a3c9a7d0c99675834cfec570aa3dccd0a01 Mon Sep 17 00:00:00 2001 From: Vaibhav Rabber Date: Sun, 14 Jan 2024 07:08:55 +0530 Subject: [PATCH] fix: Infer and scan BSON files from file path (#2410) --- crates/datafusion_ext/src/planner/relation/mod.rs | 4 ++++ tests/tests/test_bson.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/crates/datafusion_ext/src/planner/relation/mod.rs b/crates/datafusion_ext/src/planner/relation/mod.rs index 034c076d7..bf1eac29d 100644 --- a/crates/datafusion_ext/src/planner/relation/mod.rs +++ b/crates/datafusion_ext/src/planner/relation/mod.rs @@ -255,6 +255,10 @@ fn infer_func_for_file(path: &str) -> Result { schema: "public".into(), table: "ndjson_scan".into(), }, + "bson" => OwnedTableReference::Partial { + schema: "public".into(), + table: "read_bson".into(), + }, ext => { return Err(DataFusionError::Plan(format!( "unable to infer how to handle file extension: {ext}" diff --git a/tests/tests/test_bson.py b/tests/tests/test_bson.py index 700f73ffa..14da0f380 100644 --- a/tests/tests/test_bson.py +++ b/tests/tests/test_bson.py @@ -74,7 +74,7 @@ def test_read_bson( f"create external table bson_beatles from bson options ( location='{data_path}', file_type='bson')" ) - for from_clause in ["bson_beatles", f"read_bson('{data_path}')"]: + for from_clause in ["bson_beatles", f"read_bson('{data_path}')", f"'{data_path}'"]: with glaredb_connection.cursor() as curr: curr.execute(f"select count(*) from {from_clause}") r = curr.fetchone()