fix: handle input array for get_tranco_rank

Author: hatamiarash7

.gitignore

@@ -8,5 +8,4 @@ test/python/__pycache__/
 .Rhistory
 *.log
 *.csv
-!test/data/tranco.csv
-!test/data/examples.csv
+!test/data/*.csv

README.md

@@ -239,20 +239,20 @@ You can use this function to get the ranking of a domain:
 
 ```sql
 D SELECT get_tranco_rank('microsoft.com') as rank;
-┌───────┐
-│ rank  │
-│ int32 │
-├───────┤
-│     2 │
-└───────┘
+┌─────────┐
+│  rank   │
+│ varchar │
+├─────────┤
+│ 2       │
+└─────────┘
 
 D SELECT get_tranco_rank('cloudflare.com') as rank;
-┌───────┐
-│ rank  │
-│ int32 │
-├───────┤
-│    13 │
-└───────┘
+┌─────────┐
+│  rank   │
+│ varchar │
+├─────────┤
+│ 13      │
+└─────────┘
 ```
 
 ### Get Extension Version

src/functions/get_tranco.cpp

@@ -148,13 +148,28 @@ namespace duckdb
                 throw std::runtime_error ("Tranco table not found. Download it first using `SELECT update_tranco(true);`");
             }
 
-            auto &domain_vector = args.data[0];
-            auto domain         = domain_vector.GetValue (0).ToString ();
+            // Extract the input from the arguments
+            auto &input_vector = args.data[0];
+            auto result_data   = FlatVector::GetData<string_t> (result);
 
-            auto query        = "SELECT rank FROM tranco_list WHERE domain = '" + domain + "'";
-            auto query_result = con.Query (query);
+            for (idx_t i = 0; i < args.size (); i++)
+            {
+                auto input = input_vector.GetValue (i).ToString ();
+
+                try
+                {
+                    auto query = "SELECT rank FROM tranco_list WHERE domain = '" + input + "'";
+
+                    auto query_result = con.Query (query);
+                    auto rank         = query_result->RowCount () > 0 ? query_result->GetValue (0, 0) : Value ();
 
-            result.SetValue (0, query_result->RowCount () > 0 ? query_result->GetValue (0, 0) : Value ());
+                    result_data[i] = StringVector::AddString (result, rank.ToString ());
+                }
+                catch (const std::exception &e)
+                {
+                    result_data[i] = "Error extracting tranco rank: " + std::string (e.what ());
+                }
+            }
         }
     } // namespace netquack
 } // namespace duckdb

src/netquack_extension.cpp

@@ -95,7 +95,7 @@ namespace duckdb
         auto get_tranco_rank_function = ScalarFunction (
             "get_tranco_rank",
             { LogicalType::VARCHAR },
-            LogicalType::INTEGER,
+            LogicalType::VARCHAR,
             netquack::GetTrancoRankFunction);
         ExtensionUtil::RegisterFunction (instance, get_tranco_rank_function);
 

test/data/examples_tranco.csv

@@ -0,0 +1,4 @@
+microsoft.com
+googleapis.com
+gstatic.com
+apple.com

test/sql/get_tranco_rank.test

@@ -7,6 +7,9 @@ require netquack
 statement ok
 CREATE TABLE tranco_list AS SELECT * FROM read_csv('test/data/tranco.csv', header=false, columns={'rank': 'INTEGER', 'domain': 'VARCHAR'});
 
+statement ok
+CREATE TABLE uri_list AS SELECT * FROM read_csv('test/data/examples_tranco.csv', header=false, columns={'uri': 'VARCHAR'});
+
 query I
 SELECT COUNT(*) FROM tranco_list;
 ----
@@ -26,3 +29,11 @@ query I
 SELECT get_tranco_rank('notfound.com');
 ----
 NULL
+
+query I
+SELECT get_tranco_rank(uri) from uri_list;
+----
+2
+10
+19
+7