Add tools infer csv default behaviour (#20248)

Author: pnv1

ydb/public/lib/ydb_cli/commands/ydb_tools_infer.cpp

@@ -3,6 +3,8 @@
 #include <ydb/library/arrow_inference/arrow_inference.h>
 #include <ydb/public/lib/ydb_cli/common/interactive.h>
 #include <ydb/public/lib/ydb_cli/common/pretty_table.h>
+#include <ydb/public/lib/ydb_cli/common/print_utils.h>
+#include <ydb/public/lib/ydb_cli/common/csv_parser.h>
 #include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/query/client.h>
 
 #include <arrow/csv/options.h>
@@ -11,6 +13,8 @@
 #include <arrow/table.h>
 #include <util/string/builder.h>
 #include <library/cpp/string_utils/csv/csv.h>
+#include <util/stream/file.h>
+#include <regex>
 
 namespace NYdb::NConsoleClient {
 
@@ -22,8 +26,10 @@ TCommandToolsInfer::TCommandToolsInfer()
 
 TCommandToolsInferCsv::TCommandToolsInferCsv()
     : TYdbCommand("csv", {}, "Generate CREATE TABLE SQL query from CSV file"
-        "\n\nBy default, the command attempts to use the first row of the CSV as column names if possible."
-        " Use the \"--columns\", \"--gen-names\" or \"--header\" options to set the column names source explicitly.")
+        "\n\nBy default, if no options are specified, "
+        "the command uses the first row of the file as column names if possible"
+        " (i.e., if the values meet the requirements for column names and do not match data types in the other rows)."
+        " Otherwise, column names will be generated automatically.")
 {}
 
 void TCommandToolsInferCsv::Config(TConfig& config) {
@@ -33,19 +39,23 @@ void TCommandToolsInferCsv::Config(TConfig& config) {
             "One or more file paths to infer from. Or CSV data can be passed to stdin instead");
     config.Opts->AddLongOption('p', "path", "Database path to table that should be created")
         .RequiredArgument("STRING").DefaultValue("table").StoreResult(&Path);
-    config.Opts->AddLongOption("columns",
+    auto& columnsOption = config.Opts->AddLongOption("columns",
         "Explicitly specifies table column names, as a comma-separated list.")
         .RequiredArgument("NAMES").StoreResult(&ColumnNames);
-    config.Opts->AddLongOption("gen-columns",
+    auto& genColumnsOption = config.Opts->AddLongOption("gen-columns",
         "Explicitly indicates that table column names should be generated automatically.")
         .NoArgument().StoreTrue(&GenerateColumnNames);
-    config.Opts->AddLongOption("header", "Explicitly indicates that the first row in the CSV contains column names.")
+    auto& headerOption = config.Opts->AddLongOption("header", "Explicitly indicates that the first row in the CSV contains column names.")
         .NoArgument().StoreTrue(&HeaderHasColumnNames);
     config.Opts->AddLongOption("rows-to-analyze", "Number of rows to analyze. "
         "0 means unlimited. Reading will be stopped soon after this number of rows is read.")
         .DefaultValue(500000).StoreResult(&RowsToAnalyze);
     config.Opts->AddLongOption("execute", "Execute CREATE TABLE request right after generation.")
         .NoArgument().StoreTrue(&Execute);
+
+    config.Opts->MutuallyExclusiveOpt(columnsOption, genColumnsOption);
+    config.Opts->MutuallyExclusiveOpt(columnsOption, headerOption);
+    config.Opts->MutuallyExclusiveOpt(genColumnsOption, headerOption);
 }
 
 void TCommandToolsInferCsv::Parse(TConfig& config) {
@@ -67,11 +77,6 @@ void TCommandToolsInferCsv::Parse(TConfig& config) {
             ReadingFromStdin = true;
         }
     }
-
-    if (HeaderHasColumnNames && !ColumnNames.empty()) {
-        throw TMisuseException() << "Options --header and --columns are mutually exclusive."
-            " Use --header if first row in the file  containscolumn names. Use --columns to list column names manually.";
-    }
 }
 
 namespace {
@@ -93,6 +98,16 @@ namespace {
             builder << str;
         }
     }
+
+    bool IsValidColumnName(const std::string& name) {
+        if (name.empty()) {
+            return false;
+        }
+
+        // Column name must start with a letter or underscore and contain only letters, numbers and underscores
+        static const std::regex namePattern("^[a-zA-Z_][a-zA-Z0-9_]*$");
+        return std::regex_match(name, namePattern);
+    }
 }
 
 int TCommandToolsInferCsv::Run(TConfig& config) {
@@ -112,34 +127,174 @@ int TCommandToolsInferCsv::Run(TConfig& config) {
 
     auto formatConfig = std::make_shared<NArrowInference::TCsvConfig>();
     formatConfig->RowsToAnalyze = RowsToAnalyze;
+    formatConfig->Format = NArrowInference::EFileFormat::CsvWithNames;
+    formatConfig->ShouldMakeOptional = true;
+
+    // Configure CSV parsing options
+    formatConfig->ParseOpts.delimiter = ','; // Use comma as default delimiter
+    formatConfig->ParseOpts.quote_char = '"'; // Use double quotes as default quote character
+    formatConfig->ParseOpts.escape_char = '\\'; // Use backslash as default escape character
+
+    // Read the first line of the file if needed
+    std::vector<std::string> firstRowValues;
+    bool generateColumnNames = GenerateColumnNames;
     if (!ColumnNames.empty()) {
-        NCsvFormat::CsvSplitter splitter(ColumnNames);
-        auto tmp = static_cast<TVector<TString>>(splitter);
-        std::vector<std::string> columnNames;
+        // If --columns option is specified, use explicitly provided names
+        auto tmp = static_cast<TVector<TString>>(NCsvFormat::CsvSplitter(ColumnNames));
         for (const auto& columnName : tmp) {
-            columnNames.push_back(columnName.data());
+            firstRowValues.push_back(columnName.data());
         }
-        formatConfig->ReadOpts.column_names = columnNames;
-    } else if (!HeaderHasColumnNames) {
+        formatConfig->ReadOpts.column_names = firstRowValues;
+        formatConfig->ReadOpts.autogenerate_column_names = false;
+    } else if (HeaderHasColumnNames) {
+        // If --header option is specified, use first row as column names
+        formatConfig->ReadOpts.column_names = {};
+        formatConfig->ReadOpts.autogenerate_column_names = false;
+    } else if (GenerateColumnNames) {
+        // If --gen-columns option is specified, generate names automatically
+        formatConfig->ReadOpts.column_names = {};
         formatConfig->ReadOpts.autogenerate_column_names = true;
-    }
+    } else {
+        // If no option is specified:
+        // 1. Read the first line of the file
+        TFile file;
+        if (ReadingFromStdin) {
+            if (config.IsVerbose()) {
+                Cerr << "Reading first linefrom stdin" << Endl;
+            }
+            file = TFile(GetStdinFileno());
+        } else {
+            if (config.IsVerbose()) {
+                Cerr << "Reading first line from file " << FilePaths[0] << Endl;
+            }
+            file = TFile(FilePaths[0], RdOnly);
+        }
+        auto input = MakeHolder<TFileInput>(file);
+        NCsvFormat::TLinesSplitter csvSplitter(*input);
+        TString firstLine = csvSplitter.ConsumeLine();
 
-    formatConfig->Format = NArrowInference::EFileFormat::CsvWithNames;
+        // Check if the line contains newlines inside quotes
+        if (firstLine.find('\n') != TString::npos || firstLine.find('\r') != TString::npos) {
+            // If there are newlines, it's definitely not column names
+            formatConfig->ReadOpts.column_names = {};
+            formatConfig->ReadOpts.autogenerate_column_names = true;
+        } else {
+            // Split the line by delimiter
+            auto tmp = static_cast<TVector<TString>>(NCsvFormat::CsvSplitter(firstLine));
+            for (const auto& value : tmp) {
+                firstRowValues.push_back(value.data());
+            }
 
+            // 2. Tell the library to generate names automatically
+            // We will decide later if we would use generated names or first row as column names
+            formatConfig->ReadOpts.column_names = {};
+            formatConfig->ReadOpts.autogenerate_column_names = true;
+            formatConfig->ReadOpts.skip_rows = 1;
+        }
+    }
+
+    // Start file analysis
     auto result = NYdb::NArrowInference::InferTypes(inputs, formatConfig);
-    
     if (std::holds_alternative<TString>(result)) {
         throw TMisuseException() << "Failed to infer schema: " << std::get<TString>(result);
     }
 
     auto& arrowFields = std::get<NYdb::NArrowInference::ArrowFields>(result);
+    bool useFirstRowAsColumnNames = false;
+
+    // If no option is specified, check if the first row can be used as data
+    if (firstRowValues.size() > 0 && ColumnNames.empty() && !HeaderHasColumnNames && !GenerateColumnNames) {
+        bool canUseFirstRowAsColumnNames = true;
+        bool canUseFirstRowAsData = false;  // By default, assume we can't use it as data
+        if (firstRowValues.size() != arrowFields.size()) {
+            canUseFirstRowAsColumnNames = false;
+            if (config.IsVerbose()) {
+                Cerr << "First row size (" << firstRowValues.size() << ") doesn't match inferred fields count (" 
+                     << arrowFields.size() << "), can't use first row as header or data" << Endl;
+            }
+        } else {
+            // First check if all values in the first row can be column names
+            for (const auto& value : firstRowValues) {
+                if (!IsValidColumnName(value)) {
+                    canUseFirstRowAsColumnNames = false;
+                    generateColumnNames = true;
+                    if (config.IsVerbose()) {
+                        Cerr << "Value '" << value << "' is not a valid column name, can't use first row as header."
+                            " Column names will be generated automatically" << Endl;
+                    }
+                    break;
+                }
+            }
+
+            // Only if all values can be column names, check if they can be converted to types
+            if (canUseFirstRowAsColumnNames) {
+                if (config.IsVerbose()) {
+                    Cerr << "All values in first row are valid column names, checking if they can be used as data..." << Endl;
+                }
+                
+                canUseFirstRowAsData = true;  // Assume we can use as data until proven otherwise
+                for (size_t i = 0; i < arrowFields.size(); ++i) {
+                    auto field = arrowFields[i];
+                    auto value = firstRowValues[i];
+
+                    // Try to convert value to column type
+                    Ydb::Type inferredType;
+                    bool inferResult = NYdb::NArrowInference::ArrowToYdbType(inferredType, *field->type(), formatConfig);
+                    if (!inferResult) {
+                        canUseFirstRowAsData = false;
+                        if (config.IsVerbose()) {
+                            Cerr << "Failed to infer type for column " << i << ", assuming string type" << Endl;
+                        }
+                        break;
+                    }
+                    if (!NYdb::NConsoleClient::IsConvertibleToYdbValue(TString(value), inferredType)) {
+                        canUseFirstRowAsData = false;
+                        if (config.IsVerbose()) {
+                            Cerr << "Value '" << value << "' in column " << i << " cannot be converted to inferred type "
+                                 << TType(inferredType) << Endl;
+                        }
+                        break;
+                    }
+                }
+
+                if (canUseFirstRowAsData) {
+                    generateColumnNames = true;
+                    if (config.IsVerbose()) {
+                        Cerr << "All values in the first row can be used as data, "
+                            "so considering it as data and generating column names" << Endl;
+                    }
+                } else {
+                    useFirstRowAsColumnNames = true;
+                    if (config.IsVerbose()) {
+                        Cerr << "First row will be used as column names since values cannot be used as data" << Endl;
+                    }
+                }
+            }
+        }
+    }
+
+    // Generate SQL query for table creation
     TStringBuilder query;
     query << "CREATE TABLE ";
     PrintStringQuotedIfNeeded(query, GetRelativePath(Path, config));
     query << " (" << Endl;
+    int columnIndex = -1;
+    std::string firstColumnName;
     for (const auto& field : arrowFields) {
-        if (field->name().empty()) {
-            continue;
+        ++columnIndex;
+        std::string columnName;
+        if (useFirstRowAsColumnNames)  {
+            columnName = firstRowValues[columnIndex];
+        } else if (generateColumnNames) {
+            columnName = "column" + ToString(columnIndex);
+        } else {
+            columnName = field->name();
+            if (columnName.empty()) {
+                continue;
+            }
+        }
+        if (columnIndex == 0) {
+            firstColumnName = columnName;
         }
         Ydb::Type inferredType;
         bool inferResult = NYdb::NArrowInference::ArrowToYdbType(inferredType, *field->type(), formatConfig);
@@ -161,16 +316,18 @@ int TCommandToolsInferCsv::Run(TConfig& config) {
                     " Inferred type kind: " << parser.GetKind();
             }
         } else if (config.IsVerbose()) {
-            Cerr << "Failed to infer type for column " << field->name() << Endl;
+            Cerr << "Failed to infer type for column " << columnName << " with index " << columnIndex << Endl;
         }
         query << "    ";
-        PrintStringQuotedIfNeeded(query, field->name());
-        query << " " << resultType << ',' << Endl;
-        if (!field->nullable()) {
+        PrintStringQuotedIfNeeded(query, columnName);
+        query << " " << resultType;
+        // Only setting NOT NULL for the first column because we consider it a PRIMARY KEY
+        if (!columnIndex) {
             query << " NOT NULL";
         }
+        query << ',' << Endl;
     }
-    query << "    PRIMARY KEY (" << arrowFields[0]->name() << ") -- First column is chosen. Probably need to change this." << Endl;
+    query << "    PRIMARY KEY (`" << firstColumnName << "`) -- First column is chosen. Probably need to change this." << Endl;
     query <<
 R"()
 WITH (

ydb/public/lib/ydb_cli/common/csv_parser.cpp

@@ -659,6 +659,16 @@ const TString& TCsvParser::GetHeaderRow() const {
     return HeaderRow;
 }
 
+bool IsConvertibleToYdbValue(const TString& value, const Ydb::Type& type) {
+    try {
+        TTypeParser parser(type);
+        FieldToValue(parser, value, std::nullopt, {}, "columnName");
+        return true;
+    } catch (...) {
+        return false;
+    }
+}
+
 }
 }
 

ydb/public/lib/ydb_cli/common/csv_parser.h

@@ -2,6 +2,7 @@
 
 #include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/params/params.h>
 #include <google/protobuf/arena.h>
+#include <ydb/public/api/protos/ydb_value.pb.h>
 
 #include <library/cpp/string_utils/csv/csv.h>
 
@@ -82,5 +83,8 @@ class TCsvParser {
     ) const;
 };
 
+// Checks if a string value can be converted to a YDB type
+bool IsConvertibleToYdbValue(const TString& value, const Ydb::Type& type);
+
 }
 }

ydb/public/lib/ydb_cli/common/print_utils.cpp

@@ -138,5 +138,13 @@ int PrintProtoJsonBase64(const google::protobuf::Message& msg) {
     return EXIT_SUCCESS;
 }
 
+FHANDLE GetStdinFileno() {
+#if defined(_win32_)
+    return GetStdHandle(STD_INPUT_HANDLE);
+#elif defined(_unix_)
+    return STDIN_FILENO;
+#endif
+}
+
 }
 }

ydb/public/lib/ydb_cli/common/print_utils.h

@@ -2,6 +2,7 @@
 
 #include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/scheme/scheme.h>
 #include <library/cpp/colorizer/colors.h>
+#include <util/system/file.h>
 
 namespace NYdb {
 namespace NConsoleClient {
@@ -12,8 +13,8 @@ TString FormatDuration(TDuration duration);
 TString PrettySize(ui64 size);
 TString PrettyNumber(ui64 number);
 TString EntryTypeToString(NScheme::ESchemeEntryType entry);
-
 int PrintProtoJsonBase64(const google::protobuf::Message& msg);
+FHANDLE GetStdinFileno();
 
 }
 }

ydb/public/lib/ydb_cli/import/import.cpp

@@ -20,6 +20,7 @@
 #include <ydb/public/lib/ydb_cli/common/recursive_list.h>
 #include <ydb/public/lib/ydb_cli/common/interactive.h>
 #include <ydb/public/lib/ydb_cli/common/progress_bar.h>
+#include <ydb/public/lib/ydb_cli/common/print_utils.h>
 #include <ydb/public/lib/ydb_cli/commands/ydb_common.h>
 #include <ydb/public/lib/ydb_cli/dump/util/util.h>
 #include <ydb/public/lib/ydb_cli/import/cli_arrow_helpers.h>
@@ -163,14 +164,6 @@ void InitCsvParser(TCsvParser& parser,
     }
 }
 
-FHANDLE GetStdinFileno() {
-#if defined(_win32_)
-    return GetStdHandle(STD_INPUT_HANDLE);
-#elif defined(_unix_)
-    return STDIN_FILENO;
-#endif
-}
-
 class TMaxInflightGetter {
 public:
     TMaxInflightGetter(ui64 totalMaxInFlight, std::atomic<ui64>& currentFileCount)