Query local files (datafusion-contrib#253)

matthewmturner · Nov 24, 2024 · 7e0d76f · 7e0d76f
1 parent dec8788
commit 7e0d76f
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 0 deletions.
diff --git a/src/execution/local.rs b/src/execution/local.rs
@@ -111,6 +111,8 @@ impl ExecutionContext {
         }
 
         // Register Parquet Metadata Function
+        let session_ctx = session_ctx.enable_url_table();
+
         session_ctx.register_udtf(
             "parquet_metadata",
             Arc::new(datafusion_functions_parquet::ParquetMetadataFunc {}),

diff --git a/tests/cli_cases/basic.rs b/tests/cli_cases/basic.rs
@@ -361,3 +361,35 @@ fn test_write_file() {
 
     assert!(file.exists());
 }
+
+#[test]
+fn test_query_local_file() {
+    let sql = "SELECT c1 FROM 'data/aggregate_test_100.csv' LIMIT 1".to_string();
+    let assert = Command::cargo_bin("dft")
+        .unwrap()
+        .arg("-c")
+        .arg(sql)
+        .assert()
+        .success();
+
+    let expected = "
++----+
+| c1 |
++----+
+| c  |
++----+
+";
+
+    assert.stdout(contains_str(expected));
+}
+
+#[test]
+fn test_query_non_existent_local_file() {
+    let sql = "SELECT c1 FROM 'data/nofile.csv' LIMIT 1".to_string();
+    Command::cargo_bin("dft")
+        .unwrap()
+        .arg("-c")
+        .arg(sql)
+        .assert()
+        .failure();
+}