duckdb 0.7.1-dev90.0 → 0.7.2-dev0.0

package/src/duckdb/extension/json/json_functions/json_transform.cpp

@@ -58,6 +58,9 @@ static LogicalType StructureToTypeObject(yyjson_val *obj, ClientContext &context
 		child_types.emplace_back(key_str, StructureStringToType(val, context));
 	}
 	D_ASSERT(yyjson_obj_size(obj) == names.size());
+	if (child_types.empty()) {
+		throw InvalidInputException("Empty object in JSON structure");
+	}
 	return LogicalType::STRUCT(child_types);
 }
 
@@ -87,7 +90,7 @@ static unique_ptr<FunctionData> JSONTransformBind(ClientContext &context, Scalar
 	} else {
 		auto structure_val = ExpressionExecutor::EvaluateScalar(context, *arguments[1]);
 		if (!structure_val.DefaultTryCastAs(JSONCommon::JSONType())) {
-			throw InvalidInputException("cannot cast JSON structure to string");
+			throw InvalidInputException("Cannot cast JSON structure to string");
 		}
 		auto structure_string = structure_val.GetValueUnsafe<string_t>();
 		JSONAllocator json_allocator(Allocator::DefaultAllocator());
@@ -251,7 +254,10 @@ static bool TransformDecimal(yyjson_val *vals[], Vector &result, const idx_t cou
 
 bool JSONTransform::GetStringVector(yyjson_val *vals[], const idx_t count, const LogicalType &target,
                                     Vector &string_vector, JSONTransformOptions &options) {
-	auto data = (string_t *)FlatVector::GetData(string_vector);
+	if (count > STANDARD_VECTOR_SIZE) {
+		string_vector.Initialize(false, count);
+	}
+	auto data = FlatVector::GetData<string_t>(string_vector);
 	auto &validity = FlatVector::Validity(string_vector);
 	validity.SetAllValid(count);
 
@@ -380,12 +386,20 @@ bool JSONTransform::TransformObject(yyjson_val *objects[], yyjson_alc *alc, cons
 	size_t idx, max;
 	yyjson_val *key, *val;
 	for (idx_t i = 0; i < count; i++) {
-		if (objects[i]) {
+		if (objects[i] && !unsafe_yyjson_is_null(objects[i])) {
+			if (!unsafe_yyjson_is_obj(objects[i]) && options.strict_cast) {
+				options.error_message =
+				    StringUtil::Format("Expected OBJECT, but got %s: %s", JSONCommon::ValTypeToString(objects[i]),
+				                       JSONCommon::ValToString(objects[i], 50));
+				options.object_index = i;
+				success = false;
+				break;
+			}
 			found_key_count = 0;
 			memset(found_keys, false, column_count);
 			yyjson_obj_foreach(objects[i], idx, max, key, val) {
-				auto key_ptr = yyjson_get_str(key);
-				auto key_len = yyjson_get_len(key);
+				auto key_ptr = unsafe_yyjson_get_str(key);
+				auto key_len = unsafe_yyjson_get_len(key);
 				auto it = key_map.find({key_ptr, key_len});
 				if (it != key_map.end()) {
 					const auto &col_idx = it->second;
@@ -476,13 +490,24 @@ static bool TransformArray(yyjson_val *arrays[], yyjson_alc *alc, Vector &result
 	auto &list_validity = FlatVector::Validity(result);
 	idx_t offset = 0;
 	for (idx_t i = 0; i < count; i++) {
-		if (!arrays[i] || yyjson_is_null(arrays[i])) {
+		if (!arrays[i] || unsafe_yyjson_is_null(arrays[i])) {
 			list_validity.SetInvalid(i);
+		} else if (!unsafe_yyjson_is_arr(arrays[i])) {
+			if (options.strict_cast) {
+				options.error_message =
+				    StringUtil::Format("Expected ARRAY, but got %s: %s", JSONCommon::ValTypeToString(arrays[i]),
+				                       JSONCommon::ValToString(arrays[i], 50));
+				options.object_index = i;
+				return false;
+			} else {
+				list_validity.SetInvalid(i);
+			}
+		} else {
+			auto &entry = list_entries[i];
+			entry.offset = offset;
+			entry.length = unsafe_yyjson_get_len(arrays[i]);
+			offset += entry.length;
 		}
-		auto &entry = list_entries[i];
-		entry.offset = offset;
-		entry.length = yyjson_arr_size(arrays[i]);
-		offset += entry.length;
 	}
 	ListVector::SetListSize(result, offset);
 	ListVector::Reserve(result, offset);
@@ -523,6 +548,21 @@ static bool TransformArray(yyjson_val *arrays[], yyjson_alc *alc, Vector &result
 	return success;
 }
 
+bool TransformToJSON(yyjson_val *vals[], yyjson_alc *alc, Vector &result, const idx_t count) {
+	auto data = (string_t *)FlatVector::GetData(result);
+	auto &validity = FlatVector::Validity(result);
+	for (idx_t i = 0; i < count; i++) {
+		const auto &val = vals[i];
+		if (!val) {
+			validity.SetInvalid(i);
+		} else {
+			data[i] = JSONCommon::WriteVal(val, alc);
+		}
+	}
+	// Can always transform to JSON
+	return true;
+}
+
 bool JSONTransform::Transform(yyjson_val *vals[], yyjson_alc *alc, Vector &result, const idx_t count,
                               JSONTransformOptions &options) {
 	auto result_type = result.GetType();
@@ -531,6 +571,10 @@ bool JSONTransform::Transform(yyjson_val *vals[], yyjson_alc *alc, Vector &resul
 		return TransformFromStringWithFormat(vals, result, count, options);
 	}
 
+	if (JSONCommon::LogicalTypeIsJSON(result_type)) {
+		return TransformToJSON(vals, alc, result, count);
+	}
+
 	switch (result_type.id()) {
 	case LogicalTypeId::SQLNULL:
 		return true;

package/src/duckdb/extension/json/json_functions/read_json.cpp

@@ -13,63 +13,88 @@ void JSONScan::AutoDetect(ClientContext &context, JSONScanData &bind_data, vecto
 	JSONScanLocalState lstate(context, gstate);
 	ArenaAllocator allocator(BufferAllocator::Get(context));
 
-	static const unordered_map<LogicalTypeId, vector<const char *>, LogicalTypeIdHash> FORMAT_TEMPLATES = {
-	    {LogicalTypeId::DATE, {"%m-%d-%Y", "%m-%d-%y", "%d-%m-%Y", "%d-%m-%y", "%Y-%m-%d", "%y-%m-%d"}},
-	    {LogicalTypeId::TIMESTAMP,
-	     {"%Y-%m-%d %H:%M:%S.%f", "%m-%d-%Y %I:%M:%S %p", "%m-%d-%y %I:%M:%S %p", "%d-%m-%Y %H:%M:%S",
-	      "%d-%m-%y %H:%M:%S", "%Y-%m-%d %H:%M:%S", "%y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%SZ"}},
-	};
-
-	// Populate possible date/timestamp formats, assume this is consistent across columns
-	for (auto &kv : FORMAT_TEMPLATES) {
-		const auto &type = kv.first;
-		if (bind_data.date_format_map.HasFormats(type)) {
-			continue; // Already populated
-		}
-		const auto &format_strings = kv.second;
-		for (auto &format_string : format_strings) {
-			bind_data.date_format_map.AddFormat(type, format_string);
-		}
-	}
-
 	// Read for the specified sample size
 	JSONStructureNode node;
+	bool more_than_one = false;
 	Vector string_vector(LogicalType::VARCHAR);
 	idx_t remaining = bind_data.sample_size;
 	while (remaining != 0) {
 		allocator.Reset();
 		auto read_count = lstate.ReadNext(gstate);
+		if (lstate.scan_count > 1) {
+			more_than_one = true;
+		}
 		if (read_count == 0) {
 			break;
 		}
 		idx_t next = MinValue<idx_t>(read_count, remaining);
+		yyjson_val **values;
+		if (bind_data.record_type == JSONRecordType::ARRAY_OF_RECORDS ||
+		    bind_data.record_type == JSONRecordType::ARRAY_OF_JSON) {
+			values = lstate.array_values;
+		} else {
+			values = lstate.values;
+		}
 		for (idx_t i = 0; i < next; i++) {
-			if (lstate.objects[i]) {
-				JSONStructure::ExtractStructure(lstate.objects[i], node);
+			if (values[i]) {
+				JSONStructure::ExtractStructure(values[i], node);
 			}
 		}
 		if (!node.ContainsVarchar()) { // Can't refine non-VARCHAR types
 			continue;
 		}
 		node.InitializeCandidateTypes(bind_data.max_depth);
-		node.RefineCandidateTypes(lstate.objects, next, string_vector, allocator, bind_data.date_format_map);
+		node.RefineCandidateTypes(values, next, string_vector, allocator, bind_data.date_format_map);
 		remaining -= next;
+
+		if (gstate.file_index == 10) {
+			// We really shouldn't open more than 10 files when sampling
+			break;
+		}
 	}
 	bind_data.type = original_scan_type;
-	bind_data.transform_options.date_format_map = &bind_data.date_format_map;
 
-	const auto type = JSONStructure::StructureToType(context, node, bind_data.max_depth);
-	if (type.id() != LogicalTypeId::STRUCT) {
-		return_types.emplace_back(type);
-		names.emplace_back("json");
-		bind_data.objects = false;
-	} else {
-		const auto &child_types = StructType::GetChildTypes(type);
-		return_types.reserve(child_types.size());
-		names.reserve(child_types.size());
-		for (auto &child_type : child_types) {
-			return_types.emplace_back(child_type.second);
-			names.emplace_back(child_type.first);
+	// Convert structure to logical type
+	auto type = JSONStructure::StructureToType(context, node, bind_data.max_depth);
+
+	// Detect record type
+	if (bind_data.record_type == JSONRecordType::AUTO) {
+		switch (type.id()) {
+		case LogicalTypeId::STRUCT:
+			bind_data.record_type = JSONRecordType::RECORDS;
+			break;
+		case LogicalTypeId::LIST: {
+			if (more_than_one) {
+				bind_data.record_type = JSONRecordType::JSON;
+			} else {
+				type = ListType::GetChildType(type);
+				if (type.id() == LogicalTypeId::STRUCT) {
+					bind_data.record_type = JSONRecordType::ARRAY_OF_RECORDS;
+				} else {
+					bind_data.record_type = JSONRecordType::ARRAY_OF_JSON;
+				}
+			}
+			break;
+		}
+		default:
+			bind_data.record_type = JSONRecordType::JSON;
+		}
+	}
+
+	// Detect return type
+	if (bind_data.auto_detect) {
+		bind_data.transform_options.date_format_map = &bind_data.date_format_map;
+		if (type.id() != LogicalTypeId::STRUCT) {
+			return_types.emplace_back(type);
+			names.emplace_back("json");
+		} else {
+			const auto &child_types = StructType::GetChildTypes(type);
+			return_types.reserve(child_types.size());
+			names.reserve(child_types.size());
+			for (auto &child_type : child_types) {
+				return_types.emplace_back(child_type.second);
+				names.emplace_back(child_type.first);
+			}
 		}
 	}
 
@@ -150,6 +175,22 @@ void JSONScan::InitializeBindData(ClientContext &context, JSONScanData &bind_dat
 			if (!error.empty()) {
 				throw InvalidInputException("Could not parse TIMESTAMPFORMAT: %s", error.c_str());
 			}
+		} else if (loption == "json_format") {
+			auto arg = StringValue::Get(kv.second);
+			if (arg == "records") {
+				bind_data.record_type = JSONRecordType::RECORDS;
+			} else if (arg == "array_of_records") {
+				bind_data.record_type = JSONRecordType::ARRAY_OF_RECORDS;
+			} else if (arg == "values") {
+				bind_data.record_type = JSONRecordType::JSON;
+			} else if (arg == "array_of_values") {
+				bind_data.record_type = JSONRecordType::ARRAY_OF_JSON;
+			} else if (arg == "auto") {
+				bind_data.record_type = JSONRecordType::AUTO;
+			} else {
+				throw InvalidInputException("\"json_format\" must be one of ['records', 'array_of_records', 'json', "
+				                            "'array_of_json', 'auto']");
+			}
 		}
 	}
 }
@@ -170,7 +211,7 @@ unique_ptr<FunctionData> ReadJSONBind(ClientContext &context, TableFunctionBindI
 
 	bind_data.InitializeFormats();
 
-	if (bind_data.auto_detect) {
+	if (bind_data.auto_detect || bind_data.record_type == JSONRecordType::AUTO) {
 		JSONScan::AutoDetect(context, bind_data, return_types, names);
 		bind_data.names = names;
 	}
@@ -189,9 +230,16 @@ static void ReadJSONFunction(ClientContext &context, TableFunctionInput &data_p,
 	auto &gstate = ((JSONGlobalTableFunctionState &)*data_p.global_state).state;
 	auto &lstate = ((JSONLocalTableFunctionState &)*data_p.local_state).state;
 
-	// Fetch next lines
 	const auto count = lstate.ReadNext(gstate);
-	const auto objects = lstate.objects;
+	yyjson_val **values;
+	if (gstate.bind_data.record_type == JSONRecordType::ARRAY_OF_RECORDS ||
+	    gstate.bind_data.record_type == JSONRecordType::ARRAY_OF_JSON) {
+		values = lstate.array_values;
+	} else {
+		D_ASSERT(gstate.bind_data.record_type != JSONRecordType::AUTO);
+		values = lstate.values;
+	}
+	output.SetCardinality(count);
 
 	vector<Vector *> result_vectors;
 	result_vectors.reserve(output.ColumnCount());
@@ -202,22 +250,23 @@ static void ReadJSONFunction(ClientContext &context, TableFunctionInput &data_p,
 
 	// Pass current reader to transform options so we can get line number information if an error occurs
 	bool success;
-	if (gstate.bind_data.objects) {
-		success = JSONTransform::TransformObject(objects, lstate.GetAllocator(), count, gstate.bind_data.names,
+	if (gstate.bind_data.record_type == JSONRecordType::RECORDS ||
+	    gstate.bind_data.record_type == JSONRecordType::ARRAY_OF_RECORDS) {
+		success = JSONTransform::TransformObject(values, lstate.GetAllocator(), count, gstate.bind_data.names,
 		                                         result_vectors, lstate.transform_options);
 	} else {
-		success = JSONTransform::Transform(objects, lstate.GetAllocator(), *result_vectors[0], count,
+		success = JSONTransform::Transform(values, lstate.GetAllocator(), *result_vectors[0], count,
 		                                   lstate.transform_options);
 	}
+
 	if (!success) {
 		string hint = gstate.bind_data.auto_detect
 		                  ? "\nTry increasing 'sample_size', reducing 'maximum_depth', specifying 'columns' manually, "
-		                    "or setting 'ignore_errors' to true."
-		                  : "";
-		lstate.ThrowTransformError(count, lstate.transform_options.object_index,
+		                    "specifying 'lines' or 'json_format' manually, or setting 'ignore_errors' to true."
+		                  : "\n Try specifying 'lines' or 'json_format' manually, or setting 'ignore_errors' to true.";
+		lstate.ThrowTransformError(lstate.transform_options.object_index,
 		                           lstate.transform_options.error_message + hint);
 	}
-	output.SetCardinality(count);
 }
 
 TableFunction JSONFunctions::GetReadJSONTableFunction(bool list_parameter, shared_ptr<JSONScanInfo> function_info) {
@@ -233,8 +282,10 @@ TableFunction JSONFunctions::GetReadJSONTableFunction(bool list_parameter, share
 	table_function.named_parameters["date_format"] = LogicalType::VARCHAR;
 	table_function.named_parameters["timestampformat"] = LogicalType::VARCHAR;
 	table_function.named_parameters["timestamp_format"] = LogicalType::VARCHAR;
+	table_function.named_parameters["json_format"] = LogicalType::VARCHAR;
 
 	table_function.projection_pushdown = true;
+	// TODO: might be able to do filter pushdown/prune too
 
 	table_function.function_info = std::move(function_info);
 
@@ -249,7 +300,8 @@ TableFunction GetReadJSONAutoTableFunction(bool list_parameter, shared_ptr<JSONS
 
 CreateTableFunctionInfo JSONFunctions::GetReadJSONFunction() {
 	TableFunctionSet function_set("read_json");
-	auto function_info = make_shared<JSONScanInfo>(JSONScanType::READ_JSON, JSONFormat::UNSTRUCTURED, false);
+	auto function_info =
+	    make_shared<JSONScanInfo>(JSONScanType::READ_JSON, JSONFormat::UNSTRUCTURED, JSONRecordType::RECORDS, false);
 	function_set.AddFunction(JSONFunctions::GetReadJSONTableFunction(false, function_info));
 	function_set.AddFunction(JSONFunctions::GetReadJSONTableFunction(true, function_info));
 	return CreateTableFunctionInfo(function_set);
@@ -257,7 +309,8 @@ CreateTableFunctionInfo JSONFunctions::GetReadJSONFunction() {
 
 CreateTableFunctionInfo JSONFunctions::GetReadNDJSONFunction() {
 	TableFunctionSet function_set("read_ndjson");
-	auto function_info = make_shared<JSONScanInfo>(JSONScanType::READ_JSON, JSONFormat::NEWLINE_DELIMITED, false);
+	auto function_info = make_shared<JSONScanInfo>(JSONScanType::READ_JSON, JSONFormat::NEWLINE_DELIMITED,
+	                                               JSONRecordType::RECORDS, false);
 	function_set.AddFunction(JSONFunctions::GetReadJSONTableFunction(false, function_info));
 	function_set.AddFunction(JSONFunctions::GetReadJSONTableFunction(true, function_info));
 	return CreateTableFunctionInfo(function_set);
@@ -265,7 +318,8 @@ CreateTableFunctionInfo JSONFunctions::GetReadNDJSONFunction() {
 
 CreateTableFunctionInfo JSONFunctions::GetReadJSONAutoFunction() {
 	TableFunctionSet function_set("read_json_auto");
-	auto function_info = make_shared<JSONScanInfo>(JSONScanType::READ_JSON, JSONFormat::AUTO_DETECT, true);
+	auto function_info =
+	    make_shared<JSONScanInfo>(JSONScanType::READ_JSON, JSONFormat::AUTO_DETECT, JSONRecordType::AUTO, true);
 	function_set.AddFunction(GetReadJSONAutoTableFunction(false, function_info));
 	function_set.AddFunction(GetReadJSONAutoTableFunction(true, function_info));
 	return CreateTableFunctionInfo(function_set);
@@ -273,7 +327,8 @@ CreateTableFunctionInfo JSONFunctions::GetReadJSONAutoFunction() {
 
 CreateTableFunctionInfo JSONFunctions::GetReadNDJSONAutoFunction() {
 	TableFunctionSet function_set("read_ndjson_auto");
-	auto function_info = make_shared<JSONScanInfo>(JSONScanType::READ_JSON, JSONFormat::NEWLINE_DELIMITED, true);
+	auto function_info =
+	    make_shared<JSONScanInfo>(JSONScanType::READ_JSON, JSONFormat::NEWLINE_DELIMITED, JSONRecordType::AUTO, true);
 	function_set.AddFunction(GetReadJSONAutoTableFunction(false, function_info));
 	function_set.AddFunction(GetReadJSONAutoTableFunction(true, function_info));
 	return CreateTableFunctionInfo(function_set);

package/src/duckdb/extension/json/json_functions/read_json_objects.cpp

@@ -20,7 +20,7 @@ static void ReadJSONObjectsFunction(ClientContext &context, TableFunctionInput &
 	// Fetch next lines
 	const auto count = lstate.ReadNext(gstate);
 	const auto lines = lstate.lines;
-	const auto objects = lstate.objects;
+	const auto objects = lstate.values;
 
 	// Create the strings without copying them
 	auto strings = FlatVector::GetData<string_t>(output.data[0]);
@@ -48,7 +48,8 @@ TableFunction GetReadJSONObjectsTableFunction(bool list_parameter, shared_ptr<JS
 
 CreateTableFunctionInfo JSONFunctions::GetReadJSONObjectsFunction() {
 	TableFunctionSet function_set("read_json_objects");
-	auto function_info = make_shared<JSONScanInfo>(JSONScanType::READ_JSON_OBJECTS, JSONFormat::UNSTRUCTURED);
+	auto function_info =
+	    make_shared<JSONScanInfo>(JSONScanType::READ_JSON_OBJECTS, JSONFormat::UNSTRUCTURED, JSONRecordType::JSON);
 	function_set.AddFunction(GetReadJSONObjectsTableFunction(false, function_info));
 	function_set.AddFunction(GetReadJSONObjectsTableFunction(true, function_info));
 	return CreateTableFunctionInfo(function_set);
@@ -56,7 +57,8 @@ CreateTableFunctionInfo JSONFunctions::GetReadJSONObjectsFunction() {
 
 CreateTableFunctionInfo JSONFunctions::GetReadNDJSONObjectsFunction() {
 	TableFunctionSet function_set("read_ndjson_objects");
-	auto function_info = make_shared<JSONScanInfo>(JSONScanType::READ_JSON_OBJECTS, JSONFormat::NEWLINE_DELIMITED);
+	auto function_info =
+	    make_shared<JSONScanInfo>(JSONScanType::READ_JSON_OBJECTS, JSONFormat::NEWLINE_DELIMITED, JSONRecordType::JSON);
 	function_set.AddFunction(GetReadJSONObjectsTableFunction(false, function_info));
 	function_set.AddFunction(GetReadJSONObjectsTableFunction(true, function_info));
 	return CreateTableFunctionInfo(function_set);

package/src/duckdb/extension/json/json_functions.cpp

@@ -166,7 +166,14 @@ vector<CreateTableFunctionInfo> JSONFunctions::GetTableFunctions() {
 unique_ptr<TableRef> JSONFunctions::ReadJSONReplacement(ClientContext &context, const string &table_name,
                                                         ReplacementScanData *data) {
 	auto lower_name = StringUtil::Lower(table_name);
+	// remove any compression
+	if (StringUtil::EndsWith(lower_name, ".gz")) {
+		lower_name = lower_name.substr(0, lower_name.size() - 3);
+	} else if (StringUtil::EndsWith(lower_name, ".zst")) {
+		lower_name = lower_name.substr(0, lower_name.size() - 4);
+	}
 	if (!StringUtil::EndsWith(lower_name, ".json") && !StringUtil::Contains(lower_name, ".json?") &&
+	    !StringUtil::EndsWith(lower_name, ".jsonl") && !StringUtil::Contains(lower_name, ".jsonl?") &&
 	    !StringUtil::EndsWith(lower_name, ".ndjson") && !StringUtil::Contains(lower_name, ".ndjson?")) {
 		return nullptr;
 	}