futher fixing type inferring (#7456)

3 files changed, 24 insertions, 11 deletions

diff --git a/ydb/core/external_sources/object_storage/inference/arrow_inferencinator.cpp b/ydb/core/external_sources/object_storage/inference/arrow_inferencinator.cpp
index 33ae55c4a75..43e692cb5b9 100644
--- a/ydb/core/external_sources/object_storage/inference/arrow_inferencinator.cpp
+++ b/ydb/core/external_sources/object_storage/inference/arrow_inferencinator.cpp
@@ -14,8 +14,21 @@ namespace NKikimr::NExternalSource::NObjectStorage::NInference {
 
 namespace {
 
-bool ArrowToYdbType(Ydb::Type& optionalType, const arrow::DataType& type) {
-    auto& resType = *optionalType.mutable_optional_type()->mutable_item();
+bool ShouldBeOptional(const arrow::DataType& type) {
+    switch (type.id()) {
+    case arrow::Type::NA:
+    case arrow::Type::STRING:
+    case arrow::Type::BINARY:
+    case arrow::Type::LARGE_BINARY:
+    case arrow::Type::FIXED_SIZE_BINARY:
+        return false;
+    default:
+        return true;
+    }
+}
+
+bool ArrowToYdbType(Ydb::Type& maybeOptionalType, const arrow::DataType& type) {
+    auto& resType = ShouldBeOptional(type) ? *maybeOptionalType.mutable_optional_type()->mutable_item() : maybeOptionalType;
     switch (type.id()) {
     case arrow::Type::NA:
         resType.set_type_id(Ydb::Type::UTF8);
diff --git a/ydb/core/external_sources/object_storage/inference/ut/arrow_inference_ut.cpp b/ydb/core/external_sources/object_storage/inference/ut/arrow_inference_ut.cpp
index 9fe683f56c9..8dbb49f2834 100644
--- a/ydb/core/external_sources/object_storage/inference/ut/arrow_inference_ut.cpp
+++ b/ydb/core/external_sources/object_storage/inference/ut/arrow_inference_ut.cpp
@@ -97,8 +97,8 @@ TEST_F(ArrowInferenceTest, csv_simple) {
     ASSERT_EQ(fields[0].type().optional_type().item().type_id(), Ydb::Type::INT64);
     ASSERT_EQ(fields[0].name(), "A");
 
-    ASSERT_TRUE(fields[1].type().optional_type().item().has_type_id());
-    ASSERT_EQ(fields[1].type().optional_type().item().type_id(), Ydb::Type::UTF8);
+    ASSERT_TRUE(fields[1].type().has_type_id());
+    ASSERT_EQ(fields[1].type().type_id(), Ydb::Type::UTF8);
     ASSERT_EQ(fields[1].name(), "B");
 
     ASSERT_TRUE(fields[2].type().optional_type().item().has_type_id());
@@ -133,8 +133,8 @@ TEST_F(ArrowInferenceTest, tsv_simple) {
     ASSERT_EQ(fields[0].type().optional_type().item().type_id(), Ydb::Type::INT64);
     ASSERT_EQ(fields[0].name(), "A");
 
-    ASSERT_TRUE(fields[1].type().optional_type().item().has_type_id());
-    ASSERT_EQ(fields[1].type().optional_type().item().type_id(), Ydb::Type::UTF8);
+    ASSERT_TRUE(fields[1].type().has_type_id());
+    ASSERT_EQ(fields[1].type().type_id(), Ydb::Type::UTF8);
     ASSERT_EQ(fields[1].name(), "B");
 
     ASSERT_TRUE(fields[2].type().optional_type().item().has_type_id());
diff --git a/ydb/tests/fq/s3/test_s3_0.py b/ydb/tests/fq/s3/test_s3_0.py
index 76a27858a07..232513cc4d7 100644
--- a/ydb/tests/fq/s3/test_s3_0.py
+++ b/ydb/tests/fq/s3/test_s3_0.py
@@ -118,7 +118,7 @@ Pear,15,33,2024-05-06'''
         assert result_set.columns[0].name == "Date"
         assert result_set.columns[0].type.optional_type.item.type_id == ydb.Type.DATE
         assert result_set.columns[1].name == "Fruit"
-        assert result_set.columns[1].type.optional_type.item.type_id == ydb.Type.UTF8
+        assert result_set.columns[1].type.type_id == ydb.Type.UTF8
         assert result_set.columns[2].name == "Price"
         assert result_set.columns[2].type.optional_type.item.type_id == ydb.Type.INT64
         assert result_set.columns[3].name == "Weight"
@@ -176,9 +176,9 @@ Pear,15,'''
         logging.debug(str(result_set))
         assert len(result_set.columns) == 3
         assert result_set.columns[0].name == "Fruit"
-        assert result_set.columns[0].type.optional_type.item.type_id == ydb.Type.UTF8
+        assert result_set.columns[0].type.type_id == ydb.Type.UTF8
         assert result_set.columns[1].name == "Missing column"
-        assert result_set.columns[1].type.optional_type.item.type_id == ydb.Type.UTF8
+        assert result_set.columns[1].type.type_id == ydb.Type.UTF8
         assert result_set.columns[2].name == "Price"
         assert result_set.columns[2].type.optional_type.item.type_id == ydb.Type.INT64
         assert len(result_set.rows) == 3
@@ -233,7 +233,7 @@ Apple,2,22,
         assert result_set.columns[0].name == "Date"
         assert result_set.columns[0].type.optional_type.item.type_id == ydb.Type.DATE
         assert result_set.columns[1].name == "Fruit"
-        assert result_set.columns[1].type.optional_type.item.type_id == ydb.Type.UTF8
+        assert result_set.columns[1].type.type_id == ydb.Type.UTF8
         assert result_set.columns[2].name == "Price"
         assert result_set.columns[2].type.optional_type.item.type_id == ydb.Type.INT64
         assert result_set.columns[3].name == "Weight"
@@ -248,7 +248,7 @@ Apple,2,22,
         assert result_set.rows[1].items[2].int64_value == 2
         assert result_set.rows[1].items[3].int64_value == 22
         assert result_set.rows[2].items[0].uint32_value == 19849
-        assert result_set.rows[2].items[1].null_flag_value == NullValue.NULL_VALUE
+        assert result_set.rows[2].items[1].text_value == ""
         assert result_set.rows[2].items[2].int64_value == 15
         assert result_set.rows[2].items[3].int64_value == 33